supabase/apps/studio/lib/ai/generate-assistant-response.ts

import * as ai from 'ai'
import {
  convertToModelMessages,
  isToolUIPart,
  type LanguageModel,
  type ModelMessage,
  stepCountIs,
  type ToolSet,
  type UIMessage,
} from 'ai'
import { wrapAISDK } from 'braintrust'
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'

const { streamText } = wrapAISDK(ai)

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 }),
  })
}