From 8a12ffdfc0c62acbaf09ccffa975d79e21a657d3 Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Mon, 8 Sep 2025 10:57:27 +1000 Subject: [PATCH] add tools --- .../ui/AIAssistantPanel/AIAssistant.tsx | 8 + .../ui/AIAssistantPanel/Message.tsx | 198 ++++++++++++++++++ apps/studio/lib/ai/tool-filter.ts | 4 + apps/studio/lib/ai/tools/rendering-tools.ts | 17 ++ 4 files changed, 227 insertions(+) diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx index 84db19d38db..91cf64e47a6 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx @@ -272,6 +272,13 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { message={message} status={chatStatus} onResults={updateMessage} + addToolResult={ + addToolResult as unknown as (args: { + tool: string + toolCallId: string + output: unknown + }) => Promise + } onDelete={deleteMessageFromHere} onEdit={editMessage} isAfterEditedMessage={isAfterEditedMessage} @@ -288,6 +295,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => { cancelEdit, editingMessageId, chatStatus, + addToolResult, ] ) diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx index a692bade1e0..e43b9848007 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx @@ -7,6 +7,11 @@ import remarkGfm from 'remark-gfm' import { toast } from 'sonner' import { ProfileImage } from 'components/ui/ProfileImage' +import { useParams } from 'common' +import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject' +import { constructHeaders } from 'data/fetchers' +import { executeSql } from 'data/sql/execute-sql-query' +import { useEdgeFunctionDeployMutation } from 'data/edge-functions/edge-functions-deploy-mutation' import { useProfile } from 'lib/profile' import { cn, markdownComponents, WarningIcon } from 'ui' import { ButtonTooltip } from '../ButtonTooltip' @@ -55,6 +60,7 @@ interface MessageProps { resultId?: string results: any[] }) => void + addToolResult?: (args: { tool: string; toolCallId: string; output: unknown }) => Promise onDelete: (id: string) => void onEdit: (id: string) => void isAfterEditedMessage: boolean @@ -70,6 +76,7 @@ const Message = function Message({ action = null, variant = 'default', onResults, + addToolResult, onDelete, onEdit, isAfterEditedMessage = false, @@ -78,6 +85,9 @@ const Message = function Message({ onCancelEdit, }: PropsWithChildren) { const { profile } = useProfile() + const { ref } = useParams() + const { data: project } = useSelectedProjectQuery() + const { mutateAsync: deployFunction } = useEdgeFunctionDeployMutation() const [showDeleteConfirmModal, setShowDeleteConfirmModal] = useState(false) const allMarkdownComponents: Partial = useMemo( () => ({ @@ -220,6 +230,191 @@ const Message = function Message({ } return null } + case 'tool-execute_sql': { + const { toolCallId, state, input } = part + if (state === 'input-streaming') { + return ( +
+ + {`Preparing SQL execution...`} +
+ ) + } + if (state === 'input-available') { + return ( +
+
+ About to run: +
+ 
+                                {(input as any).sql}
+
+
+ + +
+
+ ) + } + if (state === 'output-available') { + return ( +
+ SQL executed. +
+ ) + } + if (state === 'output-error') { + return ( +
+ Failed to execute SQL. +
+ ) + } + return null + } + case 'tool-deploy_edge_function': { + const { toolCallId, state, input } = part + if (state === 'input-streaming') { + return ( +
+ + {`Preparing Edge Function deployment...`} +
+ ) + } + if (state === 'input-available' || state === 'output-available') { + return ( +
+ + + +
+ ) : undefined + } + /> + + ) + } + if (state === 'output-error') { + return ( +
+ Failed to deploy Edge Function. +
+ ) + } + return null + } case 'tool-display_edge_function': { const { toolCallId, state, input } = part if (state === 'input-streaming' || state === 'input-available') { @@ -329,6 +524,7 @@ export const MemoizedMessage = memo( message, status, onResults, + addToolResult, onDelete, onEdit, isAfterEditedMessage, @@ -346,6 +542,7 @@ export const MemoizedMessage = memo( resultId?: string results: any[] }) => void + addToolResult?: (args: { tool: string; toolCallId: string; output: unknown }) => Promise onDelete: (id: string) => void onEdit: (id: string) => void isAfterEditedMessage: boolean @@ -360,6 +557,7 @@ export const MemoizedMessage = memo( isLoading={status === 'submitted' || status === 'streaming'} status={status} onResults={onResults} + addToolResult={addToolResult} onDelete={onDelete} onEdit={onEdit} isAfterEditedMessage={isAfterEditedMessage} diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index 2a32b8638b3..4c89b30fc24 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/apps/studio/lib/ai/tool-filter.ts @@ -27,6 +27,8 @@ export const toolSetValidationSchema = z.record( 'get_advisors', // Local tools + 'execute_sql', + 'deploy_edge_function', 'display_query', 'display_edge_function', 'rename_chat', @@ -87,6 +89,8 @@ type ToolCategory = (typeof TOOL_CATEGORIES)[keyof typeof TOOL_CATEGORIES] */ export const TOOL_CATEGORY_MAP: Record = { // UI tools - always available + execute_sql: TOOL_CATEGORIES.UI, + deploy_edge_function: TOOL_CATEGORIES.UI, display_query: TOOL_CATEGORIES.UI, display_edge_function: TOOL_CATEGORIES.UI, rename_chat: TOOL_CATEGORIES.UI, diff --git a/apps/studio/lib/ai/tools/rendering-tools.ts b/apps/studio/lib/ai/tools/rendering-tools.ts index c94a981dafb..0096b3e0038 100644 --- a/apps/studio/lib/ai/tools/rendering-tools.ts +++ b/apps/studio/lib/ai/tools/rendering-tools.ts @@ -2,6 +2,23 @@ import { tool } from 'ai' import { z } from 'zod' export const getRenderingTools = () => ({ + execute_sql: tool({ + description: + 'Execute a SQL statement on the client. Use for SELECT/DDL/DML. Client will confirm before running and return results.', + inputSchema: z.object({ + sql: z.string().describe('The SQL statement to execute.'), + }), + // No execute – runs client-side in the UI + }), + deploy_edge_function: tool({ + description: + 'Deploy a Supabase Edge Function from provided code on the client. Client will confirm before deploying and return the result.', + inputSchema: z.object({ + name: z.string().describe('The URL-friendly name/slug of the Edge Function.'), + code: z.string().describe('The TypeScript code for the Edge Function.'), + }), + // No execute – runs client-side in the UI + }), display_query: tool({ description: 'Displays SQL query results (table or chart) or renders SQL for write/DDL operations. Use this for all query display needs. Optionally references a previous execute_sql call via manualToolCallId for displaying SELECT results.',