dima/supabase
1
0 Fork 0
mirror of https://github.com/supabase/supabase.git synced 2026-10-18 15:49:09 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat(db-playground): use native tool calls

Browse Source
This commit is contained in:
Greg Richardson committed 2024-06-07 18:52:43 -06:00
1 parent 49b246089a
commit 91ca121aa3
4 files changed
+189 -131
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

apps/database-playground/app/api/chat/route.ts
+68
View File
@@ -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()
}
apps/database-playground/app/page.tsx
+92 -131
View File
@@ -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<string, FunctionBinding> = useMemo(() => {
return {
runSql: {
description: codeBlock`
Runs Postgres SQL.
type Row<T = {
[key: string]: any;
}> = T;
type Results<T = {
[key: string]: any;
}> = {
rows: Row<T>[];
affectedRows?: number;
fields: {
name: string;
dataTypeID: number;
}[];
};
`,
typeDef: '(sql: string): Promise<Result>',
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 (
<div className="w-full h-full flex p-6">
<div className="h-full w-[50rem] py-4 rounded-md bg-[#1e1e1e]">
<Editor
language="pgsql"
value={sql}
theme="vs-dark"
options={{
tabSize: 2,
minimap: {
enabled: false,
},
fontSize: 13,
}}
onMount={async (editor, monaco) => {
// 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 && (
<div className="h-full w-[50rem] py-4 rounded-md bg-[#1e1e1e]">
<Editor
language="pgsql"
value={sql}
theme="vs-dark"
options={{
tabSize: 2,
minimap: {
enabled: false,
},
})
fontSize: 13,
}}
onMount={async (editor, monaco) => {
// 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()
}}
/>
</div>
}}
/>
</div>
)}
<div className="flex-1 max-h-full flex flex-col items-stretch">
<div className="flex-1 flex flex-col items-center overflow-y-auto">
<div className="flex flex-col gap-4 w-full max-w-4xl p-10">
@@ -173,31 +143,21 @@ export default function Page() {
>
{message.content}
</ReactMarkdown>
{message.toolInvocations &&
message.toolInvocations.map((invocation) => (
<div key={invocation.toolCallId}>
<div className="flex items-center px-3 py-2 bg-neutral-600 text-white text-xs rounded-t-md">
<h3>JavaScript (Interpreter)</h3>
</div>
<div className="">
<div className="py-4 bg-neutral-100">
<CodeBlock className="language-js border-none p-0">
{invocation.args.source}
</CodeBlock>
</div>
{'result' in invocation && (
<div className="bg-white p-4 border border-neutral-100 rounded-b-md">
<CodeBlock
className="language-js border-none p-0"
hideLineNumbers
>
{JSON.stringify(invocation.result, undefined, 2)}
</CodeBlock>
</div>
)}
</div>
</div>
))}
{message.toolInvocations?.map((toolInvocation) => {
switch (toolInvocation.toolName) {
case 'executeSql': {
return <div>Executing SQL...</div>
}
case 'generateChart': {
return (
<Chart
key={toolInvocation.toolCallId}
{...toolInvocation.args.config}
/>
)
}
}
})}
</div>
)
}
@@ -215,6 +175,7 @@ export default function Page() {
value={input}
onChange={handleInputChange}
placeholder="Message Supabase AI"
autoFocus
id="input"
/>
<Button className="rounded-full w-8 h-8 p-1.5 text-white bg-neutral-800" type="submit">
apps/database-playground/package.json
+1
View File
@@ -19,6 +19,7 @@
"monaco-editor": "^0.49.0",
"next": "14.2.3",
"react": "^18",
"react-chartjs-2": "^5.2.0",
"react-dom": "^18",
"react-markdown": "^9.0.1",
"rehype-katex": "^7.0.0",
package-lock.json
Generated
+28
View File
@@ -234,6 +234,7 @@
"monaco-editor": "^0.49.0",
"next": "14.2.3",
"react": "^18",
"react-chartjs-2": "^5.2.0",
"react-dom": "^18",
"react-markdown": "^9.0.1",
"rehype-katex": "^7.0.0",
@@ -9798,6 +9799,12 @@
"dev": true,
"license": "Apache-2.0"
},
"node_modules/@kurkle/color": {
"version": "0.3.2",
"resolved": "https://registry.npmjs.org/@kurkle/color/-/color-0.3.2.tgz",
"integrity": "sha512-fuscdXJ9G1qb7W8VdHi+IwRqij3lBkosAm4ydQtEmbY58OzHXqQhvlxqEkoz0yssNVn38bcpRWgA9PP+OGoisw==",
"peer": true
},
"node_modules/@lezer/common": {
"version": "1.1.0",
"license": "MIT",
@@ -23906,6 +23913,18 @@
"url": "https://github.com/sponsors/wooorm"
}
},
"node_modules/chart.js": {
"version": "4.4.3",
"resolved": "https://registry.npmjs.org/chart.js/-/chart.js-4.4.3.tgz",
"integrity": "sha512-qK1gkGSRYcJzqrrzdR6a+I0vQ4/R+SoODXyAjscQ/4mzuNzySaMCd+hyVxitSY1+L2fjPD1Gbn+ibNqRmwQeLw==",
"peer": true,
"dependencies": {
"@kurkle/color": "^0.3.0"
},
"engines": {
"pnpm": ">=8"
}
},
"node_modules/check-error": {
"version": "1.0.3",
"resolved": "https://registry.npmjs.org/check-error/-/check-error-1.0.3.tgz",
@@ -43229,6 +43248,15 @@
"react-dom": "^16.8.5 || ^17.0.0 || ^18.0.0"
}
},
"node_modules/react-chartjs-2": {
"version": "5.2.0",
"resolved": "https://registry.npmjs.org/react-chartjs-2/-/react-chartjs-2-5.2.0.tgz",
"integrity": "sha512-98iN5aguJyVSxp5U3CblRLH67J8gkfyGNbiK3c+l1QI/G4irHMPQw44aEPmjVag+YKTyQ260NcF82GTQ3bdscA==",
"peerDependencies": {
"chart.js": "^4.1.1",
"react": "^16.8.0 || ^17.0.0 || ^18.0.0"
}
},
"node_modules/react-colorful": {
"version": "5.6.1",
"dev": true,