diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx index 8707919ba29..67a0013fa42 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx @@ -236,7 +236,8 @@ export function MessagePartSwitcher({ } case 'tool-list_policies': case 'tool-search_docs': - case 'tool-get_active_incidents': { + case 'tool-get_active_incidents': + case 'tool-load_knowledge': { return } case 'reasoning': diff --git a/apps/studio/evals/assistant.eval.ts b/apps/studio/evals/assistant.eval.ts index f3aaae14db7..9e4babf6e1f 100644 --- a/apps/studio/evals/assistant.eval.ts +++ b/apps/studio/evals/assistant.eval.ts @@ -13,6 +13,7 @@ import { correctnessScorer, docsFaithfulnessScorer, goalCompletionScorer, + knowledgeUsageScorer, toolUsageScorer, urlValidityScorer, } from './scorer' @@ -49,6 +50,7 @@ Eval('Assistant', { }, scores: [ toolUsageScorer, + knowledgeUsageScorer, sqlSyntaxScorer, sqlIdentifierQuotingScorer, goalCompletionScorer, diff --git a/apps/studio/evals/dataset.ts b/apps/studio/evals/dataset.ts index 84da4f65264..bed3c4f1fc0 100644 --- a/apps/studio/evals/dataset.ts +++ b/apps/studio/evals/dataset.ts @@ -19,6 +19,7 @@ export const dataset: AssistantEvalCase[] = [ input: { prompt: 'Create a new table "foods" with columns for "name" and "color"' }, expected: { requiredTools: ['execute_sql'], + requiredKnowledge: ['pg_best_practices'], }, metadata: { category: ['sql_generation', 'schema_design'] }, }, @@ -29,6 +30,7 @@ export const dataset: AssistantEvalCase[] = [ }, expected: { requiredTools: ['execute_sql'], + requiredKnowledge: ['pg_best_practices'], }, metadata: { category: ['sql_generation'] }, }, @@ -36,6 +38,7 @@ export const dataset: AssistantEvalCase[] = [ input: { prompt: 'Create an index on the projects table for the name column' }, expected: { requiredTools: ['execute_sql'], + requiredKnowledge: ['pg_best_practices'], }, metadata: { category: ['sql_generation', 'database_optimization'] }, }, @@ -87,6 +90,7 @@ export const dataset: AssistantEvalCase[] = [ }, expected: { requiredTools: ['execute_sql'], + requiredKnowledge: ['pg_best_practices'], }, metadata: { category: ['sql_generation'], @@ -100,6 +104,7 @@ export const dataset: AssistantEvalCase[] = [ }, expected: { requiredTools: ['execute_sql'], + requiredKnowledge: ['pg_best_practices'], }, metadata: { category: ['sql_generation', 'schema_design'], @@ -128,6 +133,34 @@ export const dataset: AssistantEvalCase[] = [ 'Verifies template URLs like https://.supabase.co/auth/v1/callback are excluded from URL validity scoring', }, }, + { + input: { prompt: "How do I write an RLS policy to restrict access to a user's own rows?" }, + expected: { + requiredTools: ['list_tables', 'list_policies', 'execute_sql'], + requiredKnowledge: ['rls'], + }, + metadata: { category: ['rls_policies'] }, + }, + { + input: { prompt: 'Write an edge function that sends a welcome email when a user signs up' }, + expected: { + requiredTools: ['deploy_edge_function'], + requiredKnowledge: ['edge_functions'], + }, + metadata: { category: ['edge_functions'] }, + }, + { + input: { prompt: 'What indexes should I add to improve query performance?' }, + expected: { requiredKnowledge: ['pg_best_practices'] }, + metadata: { category: ['database_optimization'] }, + }, + { + input: { prompt: 'How do I subscribe to realtime changes on a table?' }, + expected: { + requiredKnowledge: ['realtime'], + }, + metadata: { category: ['general_help'] }, + }, { input: { prompt: diff --git a/apps/studio/evals/output.ts b/apps/studio/evals/output.ts index f9b51f3486d..6c2518b83d3 100644 --- a/apps/studio/evals/output.ts +++ b/apps/studio/evals/output.ts @@ -27,7 +27,16 @@ function parseToolCall( case 'search_docs': { const content = toolResult.output?.content if (!content || !Array.isArray(content)) return {} - const docs = content.map((item) => item?.text).filter((text) => typeof text === 'string') + const docs = content + .map((item) => item?.text) + .filter((text) => { + if (typeof text !== 'string') return false + try { + return !JSON.parse(text)?.error + } catch { + return true + } + }) if (docs.length === 0) return {} return { docs } } diff --git a/apps/studio/evals/scorer.ts b/apps/studio/evals/scorer.ts index 5ce66cf68fb..7f4d8d30e51 100644 --- a/apps/studio/evals/scorer.ts +++ b/apps/studio/evals/scorer.ts @@ -28,6 +28,7 @@ export type AssistantEvalOutput = { export type Expected = { requiredTools?: string[] + requiredKnowledge?: string[] correctAnswer?: string } @@ -92,6 +93,40 @@ export const toolUsageScorer: EvalScorer< } } +export const knowledgeUsageScorer: EvalScorer< + AssistantEvalInput, + AssistantEvalOutput, + Expected +> = async ({ output, expected }) => { + if (!expected.requiredKnowledge) return null + + const loadedKnowledge = output.steps + .flatMap((step) => step.toolCalls) + .filter((call) => call.toolName === 'load_knowledge') + .flatMap((call) => { + const input = call.input + if ( + typeof input !== 'object' || + input === null || + !('name' in input) || + typeof input.name !== 'string' + ) + return [] + return [input.name] + }) + + const presentCount = expected.requiredKnowledge.filter((knowledge) => + loadedKnowledge.includes(knowledge) + ).length + const totalCount = expected.requiredKnowledge.length + const ratio = totalCount === 0 ? 1 : presentCount / totalCount + + return { + name: 'Knowledge Usage', + score: ratio, + } +} + const concisenessEvaluator = LLMClassifierFromTemplate<{ input: string }>({ name: 'Conciseness', promptTemplate: stripIndent` diff --git a/apps/studio/lib/ai/generate-assistant-response.ts b/apps/studio/lib/ai/generate-assistant-response.ts index 0a740b37acc..4229fdccd5c 100644 --- a/apps/studio/lib/ai/generate-assistant-response.ts +++ b/apps/studio/lib/ai/generate-assistant-response.ts @@ -14,16 +14,7 @@ import { buildAssistantEvalOutput } from 'evals/output' import type { AssistantEvalInput, AssistantEvalOutput } from 'evals/scorer' import type { AiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi' import { IS_TRACING_ENABLED } from 'lib/ai/braintrust-logger' -import { - CHAT_PROMPT, - EDGE_FUNCTION_PROMPT, - GENERAL_PROMPT, - LIMITATIONS_PROMPT, - PG_BEST_PRACTICES, - REALTIME_PROMPT, - RLS_PROMPT, - SECURITY_PROMPT, -} from 'lib/ai/prompts' +import { CHAT_PROMPT, GENERAL_PROMPT, LIMITATIONS_PROMPT, SECURITY_PROMPT } from 'lib/ai/prompts' import { sanitizeMessagePart } from 'lib/ai/tools/tool-sanitizer' const { streamText: tracedStreamText } = wrapAISDK(ai) @@ -70,7 +61,7 @@ export async function generateAssistantResponse({ const run = async (span?: Span) => { // Only returns last 7 messages // Filters out tools with invalid states - // Filters out tool outputs based on opt-in level using renderingToolOutputParser + // Filters out tool outputs based on opt-in level const messages = (rawMessages || []).slice(-7).map((msg) => { if (msg && msg.role === 'assistant' && 'results' in msg) { const cleanedMsg = { ...msg } @@ -103,12 +94,16 @@ export async function generateAssistantResponse({ const system = source` ${GENERAL_PROMPT} ${CHAT_PROMPT} - ${PG_BEST_PRACTICES} - ${RLS_PROMPT} - ${EDGE_FUNCTION_PROMPT} - ${REALTIME_PROMPT} ${SECURITY_PROMPT} ${LIMITATIONS_PROMPT} + + ## Available Knowledge + + Before writing SQL or answering questions about the following topics, call \`load_knowledge\` to load detailed knowledge: + - \`pg_best_practices\` — PostgreSQL best practices. Always load before writing any SQL, even simple queries. + - \`rls\` — Row Level Security policies + - \`edge_functions\` — Supabase Edge Functions + - \`realtime\` — Supabase Realtime ` // Note: these must be of type `CoreMessage` to prevent AI SDK from stripping `providerOptions` diff --git a/apps/studio/lib/ai/prompts.ts b/apps/studio/lib/ai/prompts.ts index 4e0fbccea75..8df426e80ad 100644 --- a/apps/studio/lib/ai/prompts.ts +++ b/apps/studio/lib/ai/prompts.ts @@ -577,7 +577,7 @@ Support the user by: Before using tools, determine the task type (not exhaustive): **For questions about Supabase features/capabilities/limitations, or tasks** -- Use \`search_docs\` FIRST before making claims or gathering database context +- Use \`search_docs\` and/or \`load_knowledge\` FIRST before making claims or gathering database context - Examples: "How do I...", "Can Supabase...", "Is it possible to..." **For database interactions:** @@ -598,7 +598,7 @@ Before using tools, determine the task type (not exhaustive): - Never use tables in responses and use emojis minimally. If a tool output should be summarized, integrate the information clearly into the Markdown response. When a tool call returns an error, provide a concise inline explanation or summary of the error. Quote large error messages only if essential to user action. Upon each tool call or code edit, validate the result in 1–2 lines and proceed or self-correct if validation fails. ## Documentation Search -- When users ask about Supabase features, limitations, or capabilities, use \`search_docs\` BEFORE attempting database operations or making claims +- When users ask about Supabase features, limitations, or capabilities, use \`search_docs\` BEFORE attempting database operations or making claims. This DOES NOT replace the need for \`load_knowledge\`. - If \`search_docs\` reveals a limitation, inform the user immediately without gathering database context - Do not make claims unsupported by documentation ` @@ -653,7 +653,7 @@ export const OUTPUT_ONLY_PROMPT = ` - **CRITICAL: Final message must be only raw code needed to fulfill the request.** - **If you lack privelages to use a tool, do your best to generate the code without it. No need to explain why you couldn't use the tool.** - **No explanations, no commentary, no markdown**. Do not wrap output in backticks. -- **Do not call UI display tools** (no \`display_query\`, no \`display_edge_function\"). +- **Do not call UI display tools** (no \`execute_sql\`, no \`deploy_edge_function\`). ` export const SECURITY_PROMPT = ` diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index 53f5a905fcf..bcef7b6d75b 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/apps/studio/lib/ai/tool-filter.ts @@ -40,6 +40,8 @@ export const toolSetValidationSchema = z.record( 'getRlsKnowledge', 'getFunctions', 'getEdgeFunctionKnowledge', + + 'load_knowledge', ]), basicToolSchema ) @@ -71,6 +73,7 @@ export const TOOL_CATEGORY_MAP: Record = { rename_chat: TOOL_CATEGORIES.UI, search_docs: TOOL_CATEGORIES.UI, get_active_incidents: TOOL_CATEGORIES.UI, + load_knowledge: TOOL_CATEGORIES.UI, // Schema tools - MCP list_tables: TOOL_CATEGORIES.SCHEMA, diff --git a/apps/studio/lib/ai/tools/index.ts b/apps/studio/lib/ai/tools/index.ts index 00290964526..2a8d3059940 100644 --- a/apps/studio/lib/ai/tools/index.ts +++ b/apps/studio/lib/ai/tools/index.ts @@ -6,7 +6,7 @@ import { IS_PLATFORM } from 'common' import { getIncidentTools } from './incident-tools' import { getMcpTools } from './mcp-tools' import { getSchemaTools } from './schema-tools' -import { getRenderingTools } from './rendering-tools' +import { getStudioTools } from './studio-tools' export const getTools = async ({ projectRef, @@ -23,8 +23,8 @@ export const getTools = async ({ accessToken?: string baseUrl?: string }) => { - // Always include rendering tools - let tools: ToolSet = getRenderingTools() + // Always include studio tools + let tools: ToolSet = getStudioTools() // If self-hosted, only add fallback tools if (!IS_PLATFORM) { diff --git a/apps/studio/lib/ai/tools/mock-tools.ts b/apps/studio/lib/ai/tools/mock-tools.ts index 68460d05182..c8195777987 100644 --- a/apps/studio/lib/ai/tools/mock-tools.ts +++ b/apps/studio/lib/ai/tools/mock-tools.ts @@ -1,5 +1,5 @@ import { tool, type ToolSet } from 'ai' -import { getRenderingTools } from '../tools/rendering-tools' +import { getStudioTools } from '../tools/studio-tools' import { z } from 'zod' import { getMcpTools } from 'lib/ai/tools/mcp-tools' import assert from 'node:assert' @@ -142,11 +142,11 @@ const MOCK_LOGS_DATA = [ }, ] -function createMockedRenderingTools() { - const renderingTools = getRenderingTools() +function createMockedStudioTools() { + const studioTools = getStudioTools() return Object.fromEntries( - Object.entries(renderingTools).map(([name, baseTool]) => { + Object.entries(studioTools).map(([name, baseTool]) => { if (typeof baseTool.execute === 'function') { return [name, baseTool] } @@ -166,7 +166,7 @@ function createMockedRenderingTools() { }, ] }) - ) as typeof renderingTools + ) as typeof studioTools } function createMockListTablesTool(overrideData?: Record) { @@ -305,7 +305,7 @@ export type MockToolOverrides = { * Note: search_docs uses the real implementation */ export async function getMockTools(overrides?: MockToolOverrides) { - const mockedRenderingTools = createMockedRenderingTools() + const mockedStudioTools = createMockedStudioTools() const { search_docs } = await getMcpTools({ accessToken: 'mock-access-token', @@ -316,7 +316,7 @@ export async function getMockTools(overrides?: MockToolOverrides) { assert(search_docs, 'search_docs tool not available from MCP server') return { - ...mockedRenderingTools, + ...mockedStudioTools, search_docs, list_tables: createMockListTablesTool(overrides?.list_tables), list_extensions: createMockListExtensionsTool(), diff --git a/apps/studio/lib/ai/tools/rendering-tools.test.ts b/apps/studio/lib/ai/tools/studio-tools.test.ts similarity index 86% rename from apps/studio/lib/ai/tools/rendering-tools.test.ts rename to apps/studio/lib/ai/tools/studio-tools.test.ts index 5bfcab87b4c..56dff53f758 100644 --- a/apps/studio/lib/ai/tools/rendering-tools.test.ts +++ b/apps/studio/lib/ai/tools/studio-tools.test.ts @@ -1,49 +1,50 @@ import { describe, expect, it } from 'vitest' -import { getRenderingTools } from './rendering-tools' +import { getStudioTools } from './studio-tools' -describe('ai/tools/rendering-tools', () => { - describe('getRenderingTools', () => { +describe('ai/tools/studio-tools', () => { + describe('getStudioTools', () => { it('should return an object with tool definitions', () => { - const tools = getRenderingTools() + const tools = getStudioTools() expect(tools).toBeDefined() expect(typeof tools).toBe('object') }) it('should include execute_sql tool', () => { - const tools = getRenderingTools() + const tools = getStudioTools() expect(tools.execute_sql).toBeDefined() expect(tools.execute_sql.description).toContain('execute a SQL statement') }) it('should include deploy_edge_function tool', () => { - const tools = getRenderingTools() + const tools = getStudioTools() expect(tools.deploy_edge_function).toBeDefined() expect(tools.deploy_edge_function.description).toContain('deploy a Supabase Edge Function') }) it('should include rename_chat tool', () => { - const tools = getRenderingTools() + const tools = getStudioTools() expect(tools.rename_chat).toBeDefined() expect(tools.rename_chat.description).toContain('Rename the current chat session') }) - it('should have exactly 3 tools', () => { - const tools = getRenderingTools() + it('should have exactly 4 tools', () => { + const tools = getStudioTools() const toolNames = Object.keys(tools) - expect(toolNames).toHaveLength(3) + expect(toolNames).toHaveLength(4) + expect(toolNames).toContain('load_knowledge') expect(toolNames).toContain('execute_sql') expect(toolNames).toContain('deploy_edge_function') expect(toolNames).toContain('rename_chat') }) it('should have execute_sql with correct input schema fields', () => { - const tools = getRenderingTools() + const tools = getStudioTools() const executeSqlTool = tools.execute_sql // Check that the tool has an input schema @@ -56,7 +57,7 @@ describe('ai/tools/rendering-tools', () => { }) it('should have deploy_edge_function with input schema', () => { - const tools = getRenderingTools() + const tools = getStudioTools() const deployTool = tools.deploy_edge_function expect(deployTool.inputSchema).toBeDefined() @@ -67,7 +68,7 @@ describe('ai/tools/rendering-tools', () => { }) it('should have rename_chat with execute function', async () => { - const tools = getRenderingTools() + const tools = getStudioTools() const renameTool = tools.rename_chat expect(renameTool.execute).toBeDefined() @@ -83,7 +84,7 @@ describe('ai/tools/rendering-tools', () => { }) it('should validate execute_sql input schema correctly', () => { - const tools = getRenderingTools() + const tools = getStudioTools() const schema = tools.execute_sql.inputSchema // Check if schema is a Zod schema with safeParse @@ -119,7 +120,7 @@ describe('ai/tools/rendering-tools', () => { }) it('should validate rename_chat input schema correctly', () => { - const tools = getRenderingTools() + const tools = getStudioTools() const schema = tools.rename_chat.inputSchema // Check if schema is a Zod schema with safeParse diff --git a/apps/studio/lib/ai/tools/rendering-tools.ts b/apps/studio/lib/ai/tools/studio-tools.ts similarity index 73% rename from apps/studio/lib/ai/tools/rendering-tools.ts rename to apps/studio/lib/ai/tools/studio-tools.ts index b9036cc336f..ce97d011301 100644 --- a/apps/studio/lib/ai/tools/rendering-tools.ts +++ b/apps/studio/lib/ai/tools/studio-tools.ts @@ -1,8 +1,23 @@ import { tool } from 'ai' +import { + EDGE_FUNCTION_PROMPT, + PG_BEST_PRACTICES, + REALTIME_PROMPT, + RLS_PROMPT, +} from 'lib/ai/prompts' import { fixSqlBackslashEscapes } from 'lib/ai/util' import { z } from 'zod' -export const getRenderingTools = () => ({ +const KNOWLEDGE = { + pg_best_practices: PG_BEST_PRACTICES, + rls: RLS_PROMPT, + edge_functions: EDGE_FUNCTION_PROMPT, + realtime: REALTIME_PROMPT, +} as const + +type KnowledgeName = keyof typeof KNOWLEDGE + +export const getStudioTools = () => ({ execute_sql: tool({ description: 'Asks the user to execute a SQL statement and return the results', inputSchema: z.object({ @@ -41,4 +56,14 @@ export const getRenderingTools = () => ({ return { status: 'Chat request sent to client' } }, }), + load_knowledge: tool({ + description: + 'Load detailed knowledge about a Supabase topic before answering questions about it.', + inputSchema: z.object({ + name: z + .enum(Object.keys(KNOWLEDGE) as [KnowledgeName, ...KnowledgeName[]]) + .describe('The knowledge to load'), + }), + execute: ({ name }) => KNOWLEDGE[name], + }), })