From 06e34058ea59792bd4ac1ffcc3cd4659dd66b2f8 Mon Sep 17 00:00:00 2001 From: Matt Rossman <22670878+mattrossman@users.noreply.github.com> Date: Fri, 5 Dec 2025 17:07:01 -0500 Subject: [PATCH] refactor: `generateAssistantResponse` out of `handlePost` --- apps/studio/pages/api/ai/sql/generate-v4.ts | 225 ++++++++++++-------- 1 file changed, 138 insertions(+), 87 deletions(-) diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index 68cec6d4059..de6672ea412 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -1,10 +1,13 @@ import pgMeta from '@supabase/pg-meta' import { convertToModelMessages, + generateText, isToolUIPart, + type LanguageModel, type ModelMessage, stepCountIs, streamText, + type ToolSet, type UIMessage, } from 'ai' import { source } from 'common-tags' @@ -61,6 +64,109 @@ const wrapper = (req: NextApiRequest, res: NextApiResponse) => export default wrapper +export async function generateAssistantResponse({ + messages: rawMessages, + model, + tools, + aiOptInLevel = 'schema', + getSchemas, + projectRef, + chatName, + promptProviderOptions, + providerOptions, + abortSignal, +}: { + messages: UIMessage[] + model: LanguageModel + tools: ToolSet + aiOptInLevel?: AiOptInLevel + getSchemas?: () => Promise + projectRef?: string + chatName?: string + promptProviderOptions?: Record + providerOptions?: Record + abortSignal?: AbortSignal +}) { + // Only returns last 7 messages + // Filters out tools with invalid states + // Filters out tool outputs based on opt-in level using renderingToolOutputParser + const messages = (rawMessages || []).slice(-7).map((msg) => { + if (msg && msg.role === 'assistant' && 'results' in msg) { + const cleanedMsg = { ...msg } + delete cleanedMsg.results + return cleanedMsg + } + if (msg && msg.role === 'assistant' && msg.parts) { + const cleanedParts = msg.parts + .filter((part) => { + if (isToolUIPart(part)) { + const invalidStates = ['input-streaming', 'input-available', 'output-error'] + return !invalidStates.includes(part.state) + } + return true + }) + .map((part) => { + return sanitizeMessagePart(part, aiOptInLevel) + }) + return { ...msg, parts: cleanedParts } + } + return msg + }) + + const schemasString = + aiOptInLevel !== 'disabled' && getSchemas + ? await getSchemas() + : "You don't have access to any schemas." + + // Important: do not use dynamic content in the system prompt or Bedrock will not cache it + const system = source` + ${GENERAL_PROMPT} + ${CHAT_PROMPT} + ${PG_BEST_PRACTICES} + ${RLS_PROMPT} + ${EDGE_FUNCTION_PROMPT} + ${REALTIME_PROMPT} + ${SECURITY_PROMPT} + ${LIMITATIONS_PROMPT} + ` + + // Note: these must be of type `CoreMessage` to prevent AI SDK from stripping `providerOptions` + // https://github.com/vercel/ai/blob/81ef2511311e8af34d75e37fc8204a82e775e8c3/packages/ai/core/prompt/standardize-prompt.ts#L83-L88 + const assistantContent = + projectRef || chatName || schemasString !== "You don't have access to any schemas." + ? `The user's current project is ${projectRef || 'unknown'}. Their available schemas are: ${schemasString}. The current chat name is: ${chatName || 'unnamed'}` + : undefined + + const coreMessages: ModelMessage[] = [ + { + role: 'system', + content: system, + ...(promptProviderOptions && { + providerOptions: promptProviderOptions, + }), + }, + ...(assistantContent + ? [ + { + role: 'assistant' as const, + // Add any dynamic context here + content: assistantContent, + }, + ] + : []), + ...convertToModelMessages(messages), + ] + + return streamText({ + model, + stopWhen: stepCountIs(5), + messages: coreMessages, + ...(providerOptions && { providerOptions }), + tools, + ...(abortSignal && { abortSignal }), + }) +} + const requestBodySchema = z.object({ messages: z.array(z.custom()), projectRef: z.string(), @@ -121,32 +227,6 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { } } - // Only returns last 7 messages - // Filters out tools with invalid states - // Filters out tool outputs based on opt-in level using renderingToolOutputParser - const messages = (rawMessages || []).slice(-7).map((msg) => { - if (msg && msg.role === 'assistant' && 'results' in msg) { - const cleanedMsg = { ...msg } - delete cleanedMsg.results - return cleanedMsg - } - if (msg && msg.role === 'assistant' && msg.parts) { - const cleanedParts = msg.parts - .filter((part) => { - if (isToolUIPart(part)) { - const invalidStates = ['input-streaming', 'input-available', 'output-error'] - return !invalidStates.includes(part.state) - } - return true - }) - .map((part) => { - return sanitizeMessagePart(part, aiOptInLevel) - }) - return { ...msg, parts: cleanedParts } - } - return msg - }) - const { model, error: modelError, @@ -164,67 +244,10 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { } try { - // Get a list of all schemas to add to context - const pgMetaSchemasList = pgMeta.schemas.list() - type Schemas = z.infer<(typeof pgMetaSchemasList)['zod']> - - const { result: schemas } = - aiOptInLevel !== 'disabled' - ? await executeSql( - { - projectRef, - connectionString, - sql: pgMetaSchemasList.sql, - }, - undefined, - { - 'Content-Type': 'application/json', - ...(authorization && { Authorization: authorization }), - }, - IS_PLATFORM ? undefined : executeQuery - ) - : { result: [] } - - const schemasString = - schemas?.length > 0 - ? `The available database schema names are: ${JSON.stringify(schemas)}` - : "You don't have access to any schemas." - - // Important: do not use dynamic content in the system prompt or Bedrock will not cache it - const system = source` - ${GENERAL_PROMPT} - ${CHAT_PROMPT} - ${PG_BEST_PRACTICES} - ${RLS_PROMPT} - ${EDGE_FUNCTION_PROMPT} - ${REALTIME_PROMPT} - ${SECURITY_PROMPT} - ${LIMITATIONS_PROMPT} - ` - - // Note: these must be of type `CoreMessage` to prevent AI SDK from stripping `providerOptions` - // https://github.com/vercel/ai/blob/81ef2511311e8af34d75e37fc8204a82e775e8c3/packages/ai/core/prompt/standardize-prompt.ts#L83-L88 - const coreMessages: ModelMessage[] = [ - { - role: 'system', - content: system, - ...(promptProviderOptions && { - providerOptions: promptProviderOptions, - }), - }, - { - role: 'assistant', - // Add any dynamic context here - content: `The user's current project is ${projectRef}. Their available schemas are: ${schemasString}. The current chat name is: ${chatName}`, - }, - ...convertToModelMessages(messages), - ] - const abortController = new AbortController() req.on('close', () => abortController.abort()) req.on('aborted', () => abortController.abort()) - // Get tools const tools = await getTools({ projectRef, connectionString, @@ -233,12 +256,40 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { accessToken, }) - const result = streamText({ + // Get a list of all schemas to add to context + const getSchemas = async (): Promise => { + const pgMetaSchemasList = pgMeta.schemas.list() + type Schemas = z.infer<(typeof pgMetaSchemasList)['zod']> + + const { result: schemas } = await executeSql( + { + projectRef, + connectionString, + sql: pgMetaSchemasList.sql, + }, + undefined, + { + 'Content-Type': 'application/json', + ...(authorization && { Authorization: authorization }), + }, + IS_PLATFORM ? undefined : executeQuery + ) + + return schemas?.length > 0 + ? `The available database schema names are: ${JSON.stringify(schemas)}` + : "You don't have access to any schemas." + } + + const result = await generateAssistantResponse({ + messages: rawMessages, model, - stopWhen: stepCountIs(5), - messages: coreMessages, - ...(providerOptions && { providerOptions }), tools, + aiOptInLevel, + getSchemas: aiOptInLevel !== 'disabled' ? getSchemas : undefined, + projectRef, + chatName, + promptProviderOptions, + providerOptions, abortSignal: abortController.signal, })