From 91ca121aa3c3283d7e95213bc58ab75c86f0914e Mon Sep 17 00:00:00 2001 From: Greg Richardson Date: Fri, 7 Jun 2024 18:52:43 -0600 Subject: [PATCH] feat(db-playground): use native tool calls --- .../database-playground/app/api/chat/route.ts | 68 ++++++ apps/database-playground/app/page.tsx | 223 ++++++++---------- apps/database-playground/package.json | 1 + package-lock.json | 28 +++ 4 files changed, 189 insertions(+), 131 deletions(-) create mode 100644 apps/database-playground/app/api/chat/route.ts diff --git a/apps/database-playground/app/api/chat/route.ts b/apps/database-playground/app/api/chat/route.ts new file mode 100644 index 00000000000..733357cfb8b --- /dev/null +++ b/apps/database-playground/app/api/chat/route.ts @@ -0,0 +1,68 @@ +import { openai } from '@ai-sdk/openai' +import { convertToCoreMessages, streamText } from 'ai' +import { codeBlock } from 'common-tags' +import { z } from 'zod' + +// Allow streaming responses up to 30 seconds +export const maxDuration = 30 + +export async function POST(req: Request) { + const { messages } = await req.json() + + const result = await streamText({ + system: codeBlock` + You are a helpful database assistant. Under the hood you have access to a Postgres database. + + You also know math. All math equations and expressions must be written in KaTex and must be wrapped in double dollar \`$$\`: + - Inline: $$\\sqrt{26}$$ + - Multiline: + $$ + \\sqrt{26} + $$ + + No images are allowed. Do not try to link images. + + Err on the side of caution. Ask the user to confirm before any mutating operations. + + If you're just querying schema, data, or showing charts, go ahead and do it without asking. + `, + model: openai('gpt-4o-2024-05-13'), + messages: convertToCoreMessages(messages), + tools: { + getDatabaseSchema: { + description: + 'Executes Postgres SQL using information_schema and other meta tables to retrieve the requested schema info. When requesting tables, go ahead and also fetch basic column info via join. Present it in a table.', + parameters: z.object({ sql: z.string() }), + }, + executeSql: { + description: "Executes Postgres SQL against the user's database", + parameters: z.object({ sql: z.string() }), + }, + brainstormReports: { + description: + 'Brainstorms some interesting reports to show to the user. Call `getDatabaseSchema` first.', + parameters: z.object({ + reports: z.array(z.object({ name: z.string() })), + }), + }, + generateChart: { + description: + 'Generates a chart using Chart.js for a given SQL query. Call `executeSql` first.', + parameters: z.object({ + config: z + .any() + .describe( + 'The `config` passed to `new Chart(ctx, config). Includes `type`, `data`, `options`, etc.' + ), + }), + }, + appendSqlToMigration: { + description: + 'Appends schema changes to a SQL migration file. Always call this after modifying database structure.', + parameters: z.object({ sql: z.string() }), + }, + }, + }) + + return result.toAIStreamResponse() +} diff --git a/apps/database-playground/app/page.tsx b/apps/database-playground/app/page.tsx index ab426d06540..285e3c0ebce 100644 --- a/apps/database-playground/app/page.tsx +++ b/apps/database-playground/app/page.tsx @@ -1,18 +1,19 @@ 'use client' +import 'chart.js/auto' + import { PGlite } from '@electric-sql/pglite' import { Editor } from '@monaco-editor/react' -import { FunctionBinding, executeJS } from 'ai-sandbox' import { useChat } from 'ai/react' -import { codeBlock } from 'common-tags' import { ArrowUp } from 'lucide-react' -import { useMemo, useState } from 'react' +import { useState } from 'react' +import { Chart } from 'react-chartjs-2' import ReactMarkdown from 'react-markdown' import rehypeKatex from 'rehype-katex' import remarkGfm from 'remark-gfm' import remarkMath from 'remark-math' import { format } from 'sql-formatter' -import { CodeBlock, markdownComponents } from 'ui' +import { markdownComponents } from 'ui' import { Button } from 'ui/src/components/shadcn/ui/button' import { Input } from 'ui/src/components/shadcn/ui/input' @@ -22,29 +23,37 @@ let db: PGlite = new PGlite('idb://local') export default function Page() { const [sql, setSql] = useState('') + const [isEditorVisible, setIsEditorVisible] = useState(false) - const functionBindings: Record = useMemo(() => { - return { - runSql: { - description: codeBlock` - Runs Postgres SQL. - - type Row = T; - type Results = { - rows: Row[]; - affectedRows?: number; - fields: { - name: string; - dataTypeID: number; - }[]; - }; - `, - typeDef: '(sql: string): Promise', - fn: async (sql: string) => { + const { messages, input, handleInputChange, handleSubmit } = useChat({ + api: 'api/chat', + maxToolRoundtrips: 5, + async onToolCall({ toolCall }) { + console.log('tool call', toolCall) + switch (toolCall.toolName) { + case 'getDatabaseSchema': { + const { sql } = toolCall.args as any + console.log(sql) + const results = await db.exec(sql) + console.log(results) + return results + } + case 'brainstormReports': { + return 'Reports have been brainstormed. Relay this info to the user.' + } + case 'executeSql': { + const { sql } = toolCall.args as any + console.log(sql) + const results = await db.exec(sql) + console.log(results) + return results + } + case 'generateChart': { + const { config } = toolCall.args as any + return 'Chart has been generated and shown to the user. Say something like, "Above is a chart ...".' + } + case 'appendSqlToMigration': { + const { sql } = toolCall.args as any setSql((s) => { const newSql = (s + '\n' + sql).trim() return format(newSql, { @@ -55,96 +64,57 @@ export default function Page() { functionCase: 'lower', }) }) - - console.log('about to exec') - const results = await db.exec(sql) - console.log('after exec') - return results - }, - }, - } - }, []) - - const expose = useMemo(() => { - return Object.entries(functionBindings).reduce( - (merged, [name, { fn }]) => ({ - ...merged, - [name]: fn, - }), - {} - ) - }, [functionBindings]) - - const { messages, input, handleInputChange, handleSubmit } = useChat({ - api: 'api/plan', - body: { - functionBindings, - }, - maxToolRoundtrips: 5, - async onToolCall({ toolCall }) { - console.log('tool', toolCall) - if (toolCall.toolName === 'interpret') { - const args: any = toolCall.args - const { exports, error } = await executeJS(args.source, { - expose, - }) - - if (error) { - // Make `message` enumerable so that it gets serialized - Object.defineProperty(error, 'message', { - enumerable: true, - }) + setIsEditorVisible(true) + return 'SQL has successfully been appended to the migration file.' } - - const result = error ?? exports - - return result } }, }) return (
-
- { - // Register pgsql formatter - monaco.languages.registerDocumentFormattingEditProvider('pgsql', { - async provideDocumentFormattingEdits(model) { - const currentCode = editor.getValue() - const formattedCode = format(currentCode, { - language: 'postgresql', - keywordCase: 'lower', - }) - return [ - { - range: model.getFullModelRange(), - text: formattedCode, - }, - ] + {isEditorVisible && ( +
+ { + // Register pgsql formatter + monaco.languages.registerDocumentFormattingEditProvider('pgsql', { + async provideDocumentFormattingEdits(model) { + const currentCode = editor.getValue() + const formattedCode = format(currentCode, { + language: 'postgresql', + keywordCase: 'lower', + }) + return [ + { + range: model.getFullModelRange(), + text: formattedCode, + }, + ] + }, + }) - // Format on cmd+s - editor.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyS, async () => { + // Format on cmd+s + editor.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyS, async () => { + await editor.getAction('editor.action.formatDocument').run() + }) + + // Run format on the initial value await editor.getAction('editor.action.formatDocument').run() - }) - - // Run format on the initial value - await editor.getAction('editor.action.formatDocument').run() - }} - /> -
+ }} + /> +
+ )}
@@ -173,31 +143,21 @@ export default function Page() { > {message.content} - {message.toolInvocations && - message.toolInvocations.map((invocation) => ( -
-
-

JavaScript (Interpreter)

-
-
-
- - {invocation.args.source} - -
- {'result' in invocation && ( -
- - {JSON.stringify(invocation.result, undefined, 2)} - -
- )} -
-
- ))} + {message.toolInvocations?.map((toolInvocation) => { + switch (toolInvocation.toolName) { + case 'executeSql': { + return
Executing SQL...
+ } + case 'generateChart': { + return ( + + ) + } + } + })}
) } @@ -215,6 +175,7 @@ export default function Page() { value={input} onChange={handleInputChange} placeholder="Message Supabase AI" + autoFocus id="input" />