mirror of
https://github.com/supabase/supabase.git
synced 2026-10-12 04:45:08 +03:00
refactor: generateAssistantResponse to lib/ai
This commit is contained in:
Matt Rossman committed
1 parent
06e34058ea
commit
4a2ddb6bc4
2 files changed
+129
-126
No files matched your search
@@ -0,0 +1,127 @@
|
||||
import {
|
||||
convertToModelMessages,
|
||||
isToolUIPart,
|
||||
type LanguageModel,
|
||||
type ModelMessage,
|
||||
stepCountIs,
|
||||
streamText,
|
||||
type ToolSet,
|
||||
type UIMessage,
|
||||
} from 'ai'
|
||||
import { source } from 'common-tags'
|
||||
|
||||
import type { AiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
|
||||
import {
|
||||
CHAT_PROMPT,
|
||||
EDGE_FUNCTION_PROMPT,
|
||||
GENERAL_PROMPT,
|
||||
PG_BEST_PRACTICES,
|
||||
RLS_PROMPT,
|
||||
REALTIME_PROMPT,
|
||||
SECURITY_PROMPT,
|
||||
LIMITATIONS_PROMPT,
|
||||
} from 'lib/ai/prompts'
|
||||
import { sanitizeMessagePart } from 'lib/ai/tools/tool-sanitizer'
|
||||
|
||||
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 }),
|
||||
})
|
||||
}
|
||||
@@ -1,36 +1,15 @@
|
||||
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'
|
||||
import type { NextApiRequest, NextApiResponse } from 'next'
|
||||
import z from 'zod'
|
||||
|
||||
import { IS_PLATFORM } from 'common'
|
||||
import { executeSql } from 'data/sql/execute-sql-query'
|
||||
import type { UIMessage } from 'ai'
|
||||
import type { AiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
|
||||
import { getModel } from 'lib/ai/model'
|
||||
import { getOrgAIDetails } from 'lib/ai/org-ai-details'
|
||||
import {
|
||||
CHAT_PROMPT,
|
||||
EDGE_FUNCTION_PROMPT,
|
||||
GENERAL_PROMPT,
|
||||
PG_BEST_PRACTICES,
|
||||
RLS_PROMPT,
|
||||
REALTIME_PROMPT,
|
||||
SECURITY_PROMPT,
|
||||
LIMITATIONS_PROMPT,
|
||||
} from 'lib/ai/prompts'
|
||||
import { generateAssistantResponse } from 'lib/ai/generate-assistant-response'
|
||||
import { getTools } from 'lib/ai/tools'
|
||||
import { sanitizeMessagePart } from 'lib/ai/tools/tool-sanitizer'
|
||||
import apiWrapper from 'lib/api/apiWrapper'
|
||||
import { executeQuery } from 'lib/api/self-hosted/query'
|
||||
|
||||
@@ -64,109 +43,6 @@ 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(),
|
||||
|
||||
Reference in new issue
Block a user