mirror of
https://github.com/supabase/supabase.git
synced 2026-10-12 04:45:08 +03:00
refactor: generateAssistantResponse out of handlePost
This commit is contained in:
Matt Rossman committed
1 parent
775ee5a98c
commit
06e34058ea
1 file changed
+138
-87
@@ -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<string>
|
||||
projectRef?: string
|
||||
chatName?: string
|
||||
promptProviderOptions?: Record<string, any>
|
||||
providerOptions?: Record<string, any>
|
||||
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<UIMessage>()),
|
||||
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<Schemas>(
|
||||
{
|
||||
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<string> => {
|
||||
const pgMetaSchemasList = pgMeta.schemas.list()
|
||||
type Schemas = z.infer<(typeof pgMetaSchemasList)['zod']>
|
||||
|
||||
const { result: schemas } = await executeSql<Schemas>(
|
||||
{
|
||||
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,
|
||||
})
|
||||
|
||||
|
||||
Reference in new issue
Block a user