konnosDashboard
generate-assistant-response.ts204 lines · main
1import * as ai from 'ai'
2import {
3 convertToModelMessages,
4 isToolUIPart,
5 stepCountIs,
6 type LanguageModel,
7 type ModelMessage,
8 type SystemModelMessage,
9 type ToolSet,
10 type UIMessage,
11} from 'ai'
12import { startSpan, traced, withCurrent, wrapAISDK, type Span } from 'braintrust'
13import { source } from 'common-tags'
14
15import type { AssistantEvalInput } from '@/evals/scorer'
16import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
17import { IS_TRACING_ENABLED } from '@/lib/ai/braintrust-logger'
18import { CHAT_PROMPT, GENERAL_PROMPT, LIMITATIONS_PROMPT, SECURITY_PROMPT } from '@/lib/ai/prompts'
19import { sanitizeMessagePart } from '@/lib/ai/tools/tool-sanitizer'
20
21const { streamText: tracedStreamText } = wrapAISDK(ai)
22
23export async function generateAssistantResponse({
24 messages: rawMessages,
25 model,
26 tools,
27 aiOptInLevel = 'schema',
28 getSchemas,
29 projectRef,
30 chatId,
31 chatName,
32 allowTracing,
33 userId,
34 orgId,
35 planId,
36 systemProviderOptions,
37 providerOptions,
38 requestedModel,
39 abortSignal,
40 onSpanCreated,
41}: {
42 messages: UIMessage[]
43 model: LanguageModel
44 tools: ToolSet
45 aiOptInLevel?: AiOptInLevel
46 getSchemas?: () => Promise<string>
47 projectRef?: string
48 chatId?: string
49 chatName?: string
50 allowTracing?: boolean
51 userId?: string
52 orgId?: number
53 planId?: string
54 requestedModel?: string
55 systemProviderOptions?: Record<string, any>
56 providerOptions?: Record<string, any>
57 abortSignal?: AbortSignal
58 onSpanCreated?: (spanId: string) => void
59}) {
60 const shouldTrace = allowTracing ?? IS_TRACING_ENABLED
61
62 const run = async (span?: Span) => {
63 // Only returns last 7 messages
64 // Filters out tools with invalid states
65 // Filters out tool outputs based on opt-in level
66 const messages = (rawMessages || []).slice(-7).map((msg) => {
67 if (msg && msg.role === 'assistant' && 'results' in msg) {
68 const cleanedMsg = { ...msg }
69 delete cleanedMsg.results
70 return cleanedMsg
71 }
72 if (msg && msg.role === 'assistant' && msg.parts) {
73 const cleanedParts = msg.parts
74 .filter((part) => {
75 if (isToolUIPart(part)) {
76 const invalidStates = [
77 'input-streaming',
78 'input-available',
79 'approval-requested',
80 'output-error',
81 ]
82 return !invalidStates.includes(part.state)
83 }
84 return true
85 })
86 .map((part) => {
87 return sanitizeMessagePart(part, aiOptInLevel)
88 })
89 return { ...msg, parts: cleanedParts }
90 }
91 return msg
92 })
93
94 const schemasString =
95 aiOptInLevel !== 'disabled' && getSchemas
96 ? shouldTrace
97 ? await traced(async () => getSchemas(), { name: 'getSchemas', type: 'function' })
98 : await getSchemas()
99 : "You don't have access to any schemas."
100
101 // Important: do not use dynamic content in the system prompt or Bedrock will not cache it
102 const system = source`
103 ${GENERAL_PROMPT}
104 ${CHAT_PROMPT}
105 ${SECURITY_PROMPT}
106 ${LIMITATIONS_PROMPT}
107
108 ## Available Knowledge
109
110 Before writing SQL or answering questions about the following topics, call \`load_knowledge\` to load detailed knowledge:
111 - \`pg_best_practices\` — PostgreSQL best practices. Always load before writing any SQL, even simple queries.
112 - \`rls\` — Row Level Security policies
113 - \`edge_functions\` — Briven Edge Functions
114 - \`realtime\` — Briven Realtime
115 `
116
117 const hasProjectContext =
118 projectRef || chatName || schemasString !== "You don't have access to any schemas."
119
120 const assistantContent = hasProjectContext
121 ? `The user's current project is ${projectRef || 'unknown'}. Their available schemas are: ${schemasString}. The current chat name is: ${chatName || 'unnamed'}.`
122 : undefined
123
124 const systemMessage: SystemModelMessage = {
125 role: 'system',
126 content: system,
127 ...(systemProviderOptions && { providerOptions: systemProviderOptions }),
128 }
129
130 const coreMessages: ModelMessage[] = [
131 ...(assistantContent
132 ? [
133 {
134 role: 'assistant' as const,
135 content: assistantContent,
136 },
137 ]
138 : []),
139 ...(await convertToModelMessages(messages)),
140 ]
141
142 const streamTextFn = shouldTrace ? tracedStreamText : ai.streamText
143
144 return streamTextFn({
145 model,
146 system: systemMessage,
147 stopWhen: stepCountIs(5),
148 messages: coreMessages,
149 ...(providerOptions && { providerOptions }),
150 tools,
151 ...(abortSignal && { abortSignal }),
152 ...(span && {
153 onFinish: ({ steps, finishReason }) => {
154 const metadata: Record<string, unknown> = {
155 isFinalStep: finishReason === 'stop',
156 }
157 for (const step of steps) {
158 for (const toolCall of step.toolCalls) {
159 if (toolCall.toolName === 'rename_chat') {
160 const { newName } = toolCall.input as { newName: string }
161 metadata.chatName = newName
162 }
163 }
164 }
165 span.log({ metadata })
166 span.end()
167 },
168 }),
169 } satisfies Parameters<typeof ai.streamText>[0])
170 }
171
172 if (shouldTrace) {
173 // startSpan instead of traced() so we control when the span closes via onFinish.
174 // Scorers read from child spans (LLM + tool) in the trace rather than a root span output field.
175 const span = startSpan({ name: 'generateAssistantResponse', type: 'function' })
176 onSpanCreated?.(span.id)
177
178 const lastUserMessage = rawMessages.findLast((m) => m.role === 'user')
179 const lastUserText = lastUserMessage?.parts
180 ?.filter((p): p is { type: 'text'; text: string } => p.type === 'text')
181 .map((p) => p.text)
182 .join('\n')
183
184 span.log({
185 input: { prompt: lastUserText ?? '' } satisfies AssistantEvalInput,
186 metadata: {
187 projectRef,
188 chatId,
189 chatName,
190 aiOptInLevel,
191 userId,
192 orgId,
193 planId,
194 requestedModel,
195 gitBranch: process.env.VERCEL_GIT_COMMIT_REF,
196 environment: process.env.NEXT_PUBLIC_ENVIRONMENT,
197 },
198 })
199
200 return withCurrent(span, () => run(span))
201 }
202
203 return run()
204}