konnosDashboard
studio-tools.ts127 lines · main
1import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta'
2import { tool } from 'ai'
3import { z } from 'zod'
4
5import { deployEdgeFunction } from '@/data/edge-functions/edge-functions-deploy-mutation'
6import { executeSql } from '@/data/sql/execute-sql-query'
7import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
8import {
9 EDGE_FUNCTION_PROMPT,
10 PG_BEST_PRACTICES,
11 REALTIME_PROMPT,
12 RLS_PROMPT,
13} from '@/lib/ai/prompts'
14import { NO_DATA_PERMISSIONS } from '@/lib/ai/tools/tool-sanitizer'
15import { fixSqlBackslashEscapes } from '@/lib/ai/util'
16
17const KNOWLEDGE = {
18 pg_best_practices: PG_BEST_PRACTICES,
19 rls: RLS_PROMPT,
20 edge_functions: EDGE_FUNCTION_PROMPT,
21 realtime: REALTIME_PROMPT,
22} as const
23
24type KnowledgeName = keyof typeof KNOWLEDGE
25
26export const executeSqlInputSchema = z.object({
27 // Transform at parse time so the corrected SQL is what gets stored in
28 // toolCall.input — ensuring evals and logs reflect what actually runs.
29 sql: z.string().describe('The SQL statement to execute.').transform(fixSqlBackslashEscapes),
30 label: z.string().describe('A short 2-4 word label for the SQL statement.'),
31 chartConfig: z
32 .object({
33 view: z.enum(['table', 'chart']).describe('How to render the results after execution'),
34 xAxis: z.string().optional().describe('The column to use for the x-axis of the chart.'),
35 yAxis: z.string().optional().describe('The column to use for the y-axis of the chart.'),
36 })
37 .describe('Chart configuration for rendering the results'),
38 isWriteQuery: z
39 .boolean()
40 .default(false)
41 .describe(
42 'Whether the SQL statement performs a write operation or has side effects. Set true for INSERT/UPDATE/DELETE/DDL and for SELECT statements that call side-effecting functions, such as select cron.schedule(...), cron.unschedule(...), or functions that create, modify, schedule, enqueue, notify, or trigger work.'
43 ),
44})
45
46export const loadKnowledgeInputSchema = z.object({
47 name: z
48 .enum(Object.keys(KNOWLEDGE) as [KnowledgeName, ...KnowledgeName[]])
49 .describe('The knowledge to load'),
50})
51
52export type StudioToolsContext = {
53 projectRef?: string
54 connectionString?: string
55 authorization?: string
56 aiOptInLevel?: AiOptInLevel
57}
58
59export const getStudioTools = (ctx: StudioToolsContext = {}) => {
60 const { projectRef, connectionString, authorization, aiOptInLevel = 'schema' } = ctx
61 const authHeaders = authorization
62 ? { 'Content-Type': 'application/json', Authorization: authorization }
63 : undefined
64
65 return {
66 execute_sql: tool({
67 description:
68 'Asks the user to execute a SQL statement and return the results. Requires user approval before executing.',
69 inputSchema: executeSqlInputSchema,
70 needsApproval: true,
71 execute: async ({ sql }) => {
72 // The `needsApproval: true` gate on this tool means the user has
73 // explicitly approved this AI-generated SQL before execute runs —
74 // that approval is the user gesture that promotes untrusted to safe.
75 const { result } = await executeSql(
76 { projectRef, connectionString, sql: acceptUntrustedSql(untrustedSql(sql)) },
77 undefined,
78 authHeaders
79 )
80 return result
81 },
82 toModelOutput: ({ output }) => {
83 return aiOptInLevel === 'schema_and_log_and_data'
84 ? { type: 'json', value: output }
85 : { type: 'text', value: NO_DATA_PERMISSIONS }
86 },
87 }),
88 deploy_edge_function: tool({
89 description:
90 'Asks the user to deploy a Briven Edge Function from provided code. Requires user approval before deploying.',
91 inputSchema: z.object({
92 name: z.string().describe('The URL-friendly name/slug of the Edge Function.'),
93 code: z.string().describe('The TypeScript code for the Edge Function.'),
94 }),
95 needsApproval: true,
96 execute: async ({ name, code }) => {
97 await deployEdgeFunction({
98 projectRef: projectRef ?? '',
99 slug: name,
100 metadata: {
101 entrypoint_path: 'index.ts',
102 name,
103 verify_jwt: true,
104 },
105 files: [{ name: 'index.ts', content: code }],
106 authorization,
107 })
108 return { success: true }
109 },
110 }),
111 rename_chat: tool({
112 description: `Rename the current chat session when the current chat name doesn't describe the conversation topic.`,
113 inputSchema: z.object({
114 newName: z.string().describe('The new name for the chat session. Five words or less.'),
115 }),
116 execute: async () => {
117 return { status: 'Chat request sent to client' }
118 },
119 }),
120 load_knowledge: tool({
121 description:
122 'Load detailed knowledge about a Briven topic before answering questions about it.',
123 inputSchema: loadKnowledgeInputSchema,
124 execute: ({ name }) => KNOWLEDGE[name],
125 }),
126 }
127}