From 7978c9e94f0652b8df6597cb414c913f4f9624e9 Mon Sep 17 00:00:00 2001
From: Mayne <i@gine.me>
Date: Wed, 14 Jun 2023 08:12:53 +0800
Subject: [PATCH 1/7] feat(ai): openai function call

---
 app/[database]/ai-chat-message-prisma.tsx |   2 +-
 app/[database]/ai-chat.tsx                | 108 ++++++++++++------
 app/settings/ai/ai-form.tsx               |  38 +++++--
 components/sidebar/create-table.tsx       |  15 ++-
 components/sidebar/helper.ts              |  17 ++-
 components/ui/progress.tsx                |  28 +++++
 hooks/use-ai.ts                           |  77 +------------
 hooks/use-auto-run-code.ts                |  19 +++-
 hooks/use-sqlite.ts                       |  18 +++
 hooks/use-table.ts                        |  22 +++-
 lib/ai/functions/index.ts                 |  24 ++++
 lib/ai/functions/sql-query.ts             |  15 +++
 lib/ai/openai.ts                          | 128 ++++++++++++++++++++++
 package.json                              |   4 +-
 pnpm-lock.yaml                            |  42 ++++++-
 15 files changed, 422 insertions(+), 135 deletions(-)
 create mode 100644 components/ui/progress.tsx
 create mode 100644 lib/ai/functions/index.ts
 create mode 100644 lib/ai/functions/sql-query.ts
 create mode 100644 lib/ai/openai.ts

diff --git a/app/[database]/ai-chat-message-prisma.tsx b/app/[database]/ai-chat-message-prisma.tsx
index d150963f..08e3bfd9 100644
--- a/app/[database]/ai-chat-message-prisma.tsx
+++ b/app/[database]/ai-chat-message-prisma.tsx
@@ -68,7 +68,7 @@ export const AIMessage = ({
   }
   return (
     <div className="grow">
-      <Markdown markdown={message} renderers={renderers} />
+      {message && <Markdown markdown={message} renderers={renderers} />}
       <div id={`chart-${msgIndex}`} />
     </div>
   )
diff --git a/app/[database]/ai-chat.tsx b/app/[database]/ai-chat.tsx
index 8f3c94af..b3c3ae19 100644
--- a/app/[database]/ai-chat.tsx
+++ b/app/[database]/ai-chat.tsx
@@ -7,8 +7,10 @@ import { useParams } from "next/navigation"
 import { useKeyPress, useSize } from "ahooks"
 import { Bot, Loader2, Paintbrush, User } from "lucide-react"
 
+import { handleOpenAIFunctionCall } from "@/lib/ai/openai"
 import { useAI } from "@/hooks/use-ai"
 import { useAutoRunCode } from "@/hooks/use-auto-run-code"
+import { useSqliteStore } from "@/hooks/use-sqlite"
 import { Button } from "@/components/ui/button"
 import { Textarea } from "@/components/ui/textarea"
 
@@ -16,21 +18,20 @@ import { useConfigStore } from "../settings/store"
 import { AIMessage } from "./ai-chat-message-prisma"
 import { useTableChange } from "./hook"
 import { useDatabaseAppStore } from "./store"
-import { useSqliteStore } from "@/hooks/use-sqlite"
 
 export const AIChat = () => {
-  const { currentTableSchema, setCurrentQuery } = useDatabaseAppStore()
+  const { currentTableSchema } = useDatabaseAppStore()
   const { askAI } = useAI()
   const { database, table } = useParams()
   const { aiConfig } = useConfigStore()
   const [input, setInput] = useState("")
   const [loading, setLoading] = useState(false)
-  const { autoRun: runCode, handleRunCode } = useAutoRunCode()
+  const { handleFunctionCall, handleRunCode } = useAutoRunCode()
   const divRef = useRef<HTMLDivElement>()
   const size = useSize(divRef)
   const [messages, setMessages] = useState<
     {
-      role: "user" | "assistant"
+      role: "user" | "assistant" | "function"
       content: string
     }[]
   >([])
@@ -59,20 +60,46 @@ export const AIChat = () => {
       allTables,
       databaseName: database,
     })
-    const newMessages = [
-      ..._messages,
-      { role: "assistant", content: response?.content! },
-    ]
-    const thisMsgIndex = newMessages.length - 1
-    setMessages(newMessages)
-    if (response?.content && aiConfig.autoRunScope) {
-      setTimeout(() => {
-        runCode(response.content, {
-          msgIndex: thisMsgIndex,
-          width: size?.width ?? 300,
-        })
-      }, 1000)
+
+    if (response?.finish_reason == "function_call") {
+      if (aiConfig.autoRunScope) {
+        const res = await handleOpenAIFunctionCall(
+          response.message!,
+          handleFunctionCall
+        )
+        if (res) {
+          const { name, resp } = res
+          const newMessages = [
+            ..._messages,
+            response.message,
+            {
+              role: "function",
+              name,
+              content: JSON.stringify(resp),
+            },
+          ]
+          const newResponse = await askAI(newMessages, {
+            tableSchema: currentTableSchema,
+            allTables,
+            databaseName: database,
+          })
+          const _newMessages = [
+            ...newMessages,
+            { role: "assistant", content: newResponse?.message?.content },
+          ]
+          console.log({ _newMessages })
+          setMessages(_newMessages as any)
+        }
+      }
+    } else if (response?.message) {
+      const newMessages = [
+        ..._messages,
+        { role: "assistant", content: response?.message?.content },
+      ]
+      const thisMsgIndex = newMessages.length - 1
+      setMessages(newMessages)
     }
+
     setLoading(false)
   }
 
@@ -97,28 +124,37 @@ export const AIChat = () => {
             first
           </p>
         )}
-        {messages.map((message, i) => (
-          <div
-            className="flex w-full items-start gap-2 rounded-lg bg-gray-200 p-2 dark:bg-gray-700"
-            key={i}
-          >
-            {message.role === "assistant" ? (
-              <>
-                <Bot className="h-4 w-4 shrink-0" />
-                <AIMessage
-                  message={message.content}
-                  onRun={handleRunCode}
-                  msgIndex={i}
-                />
-              </>
-            ) : (
+        {messages
+          .filter((m) => (m.role === "user" || m.role == "assistant") && m.content)
+          .map((message, i) => (
+            <div
+              className="flex w-full items-start gap-2 rounded-lg bg-gray-200 p-2 dark:bg-gray-700"
+              key={i}
+            >
+              {message.role === "assistant" && (
+                <>
+                  <Bot className="h-4 w-4 shrink-0" />
+                  <AIMessage
+                    message={message.content}
+                    onRun={handleRunCode}
+                    msgIndex={i}
+                  />
+                </>
+              )}
+              {message.role === "user" && (
+                <>
+                  <User className="h-4 w-4 shrink-0" />
+                  <p className="grow">{message.content}</p>
+                </>
+              )}
+              {/* {message.role === "function" && (
               <>
                 <User className="h-4 w-4 shrink-0" />
-                <p className="grow">{message.content}</p>
+                <p className="grow">run function </p>
               </>
-            )}
-          </div>
-        ))}
+            )} */}
+            </div>
+          ))}
         <div className="flex w-full justify-center">
           {loading && <Loader2 className="h-5 w-5 animate-spin" />}
         </div>
diff --git a/app/settings/ai/ai-form.tsx b/app/settings/ai/ai-form.tsx
index 3803c7cc..742b7a4a 100644
--- a/app/settings/ai/ai-form.tsx
+++ b/app/settings/ai/ai-form.tsx
@@ -28,7 +28,27 @@ const AIConfigFormSchema = z.object({
 })
 
 export type AIConfigFormValues = z.infer<typeof AIConfigFormSchema>
-export const AutoRunScopes = ["SQL.SELECT", "SQL.INSERT", "D3.CHART"]
+export const AutoRunScopesWithDesc = [
+  {
+    value: "SQL.SELECT",
+    description: "Selects data from a SQL table.",
+  },
+  {
+    value: "UI.REFRESH",
+    description: "Refresh the UI after querying data.",
+  },
+  {
+    value: "SQL.INSERT",
+    description: "Inserts data into a SQL table.",
+  },
+  {
+    value: "D3.CHART",
+    description: "Creates a D3 chart.",
+  },
+]
+
+export const AutoRunScopes = AutoRunScopesWithDesc.map((item) => item.value)
+
 // This can come from your database or API.
 const defaultValues: Partial<AIConfigFormValues> = {
   // name: "Your name",
@@ -89,32 +109,34 @@ export function AIConfigForm() {
                   run.
                 </FormDescription>
               </div>
-              {AutoRunScopes.map((_item) => (
+              {AutoRunScopesWithDesc.map(({ value, description }) => (
                 <FormField
-                  key={_item}
+                  key={value}
                   control={form.control}
                   name="autoRunScope"
                   render={({ field }) => {
                     return (
                       <FormItem
-                        key={_item}
+                        key={value}
                         className="flex flex-row items-start space-x-3 space-y-0"
                       >
                         <FormControl>
                           <Checkbox
-                            checked={field.value?.includes(_item)}
+                            checked={field.value?.includes(value)}
                             onCheckedChange={(checked: any) => {
                               return checked
-                                ? field.onChange([...field.value, _item])
+                                ? field.onChange([...field.value, value])
                                 : field.onChange(
                                     field.value?.filter(
-                                      (value) => value !== _item
+                                      (value) => value !== value
                                     )
                                   )
                             }}
                           />
                         </FormControl>
-                        <FormLabel className="font-normal">{_item}</FormLabel>
+                        <FormLabel className="font-normal">
+                          {description}
+                        </FormLabel>
                       </FormItem>
                     )
                   }}
diff --git a/components/sidebar/create-table.tsx b/components/sidebar/create-table.tsx
index 7d9f3130..7847d7b4 100644
--- a/components/sidebar/create-table.tsx
+++ b/components/sidebar/create-table.tsx
@@ -15,22 +15,31 @@ import {
 } from "@/components/ui/dialog"
 import { Input } from "@/components/ui/input"
 import { Label } from "@/components/ui/label"
+import { Progress } from "@/components/ui/progress"
 
 import { csvFile2Sql } from "./helper"
 
 export function CreateTableDialog() {
   const [open, setOpen] = useState(false)
   const [tableName, setTableName] = useState("")
+  const [importing, setImporting] = useState(false)
+  const [progress, setProgress] = useState(0)
   const [file, setFile] = useState<File | null>(null)
   const params = useParams()
   const router = useRouter()
   const { database } = params
-  const { createTable, createTableWithSql } = useSqlite(database)
+  const { createTable, createTableWithSqlAndInsertSqls } = useSqlite(database)
 
   const handleCreateTable = async () => {
     if (file) {
       const res = await csvFile2Sql(file, tableName.trim())
-      await createTableWithSql(res.createTableSql, res.insertSql)
+      setImporting(true)
+      await createTableWithSqlAndInsertSqls(
+        res.createTableSql,
+        res.sqls,
+        setProgress
+      )
+      setImporting(false)
       // await createTableWithSql(res.createTableSql, res.insertSql)
       router.push(`/${database}/${tableName}`)
       setOpen(false)
@@ -49,6 +58,7 @@ export function CreateTableDialog() {
           <Plus size={16} className="mr-2" />
         </Button>
       </DialogTrigger>
+
       <DialogContent className="sm:max-w-[425px]">
         <DialogHeader>
           <DialogTitle>Create Table</DialogTitle>
@@ -86,6 +96,7 @@ export function CreateTableDialog() {
             />
           </div>
         </div>
+        {importing && <Progress value={progress} />}
         <DialogFooter>
           <Button type="submit" onClick={handleCreateTable}>
             Create
diff --git a/components/sidebar/helper.ts b/components/sidebar/helper.ts
index a9c03074..9c4acc9a 100644
--- a/components/sidebar/helper.ts
+++ b/components/sidebar/helper.ts
@@ -3,16 +3,20 @@ import { v4 as uuidv4 } from "uuid"
 
 import { logger } from "@/lib/log"
 
+export type ISqls = {
+  sql: string
+  bind: any[]
+}[]
 export const csvFile2Sql = async (
   file: File,
   tableName: string
 ): Promise<{
   createTableSql: string
-  insertSql: string
+  sqls: ISqls
 }> => {
   let hasCreateTable = false
   let createTableSql = ""
-  let insertSql = ""
+  const sqls: ISqls = []
 
   return new Promise((resolve, reject) => {
     Papa.parse(file, {
@@ -22,7 +26,7 @@ export const csvFile2Sql = async (
         logger.info("Parsing complete:", results, file)
         resolve({
           createTableSql,
-          insertSql,
+          sqls,
         })
       },
       step: (results: any, parser: any) => {
@@ -42,8 +46,11 @@ export const csvFile2Sql = async (
           const _columns = ["_id", ...columns]
           const sql = `INSERT INTO ${tableName} (${_columns.join(
             ","
-          )}) VALUES (${_values.join(",")});\n`
-          insertSql += sql
+          )}) VALUES (${Array.from(_values).fill("?").join(",")});`
+          sqls.push({
+            sql,
+            bind: [uuidv4(), ...values],
+          })
         }
       },
     })
diff --git a/components/ui/progress.tsx b/components/ui/progress.tsx
new file mode 100644
index 00000000..5c87ea48
--- /dev/null
+++ b/components/ui/progress.tsx
@@ -0,0 +1,28 @@
+"use client"
+
+import * as React from "react"
+import * as ProgressPrimitive from "@radix-ui/react-progress"
+
+import { cn } from "@/lib/utils"
+
+const Progress = React.forwardRef<
+  React.ElementRef<typeof ProgressPrimitive.Root>,
+  React.ComponentPropsWithoutRef<typeof ProgressPrimitive.Root>
+>(({ className, value, ...props }, ref) => (
+  <ProgressPrimitive.Root
+    ref={ref}
+    className={cn(
+      "relative h-4 w-full overflow-hidden rounded-full bg-secondary",
+      className
+    )}
+    {...props}
+  >
+    <ProgressPrimitive.Indicator
+      className="h-full w-full flex-1 bg-primary transition-all"
+      style={{ transform: `translateX(-${100 - (value || 0)}%)` }}
+    />
+  </ProgressPrimitive.Root>
+))
+Progress.displayName = ProgressPrimitive.Root.displayName
+
+export { Progress }
diff --git a/hooks/use-ai.ts b/hooks/use-ai.ts
index 6a9bc5c8..807fc8ba 100644
--- a/hooks/use-ai.ts
+++ b/hooks/use-ai.ts
@@ -1,38 +1,8 @@
-import { useCallback, useEffect, useState } from "react"
-import { Configuration, OpenAIApi } from "openai"
+import { OpenAIApi } from "openai"
+import { useEffect, useState } from "react"
 
 import { useConfigStore } from "@/app/settings/store"
-
-const getOpenAI = (token: string) => {
-  const configuration = new Configuration({
-    apiKey: token ?? process.env.OPENAI_API_KEY,
-  })
-  const openai = new OpenAIApi(configuration)
-  return openai
-}
-
-// TODO: expose to user config
-const baseSysPrompt = `you're a sql generator and a d3.js master. must abide by the following rules:
-*important: if user doesn't have intention to generate a d3.js chart, you can just return sql.*
-
-when you act as a sql generator:
-1. your engine is sqlite, what you return is *pure sql* that can be executed in sqlite
-2. you return markdown, sql you return must be wrapped in \`\`\`sql\`\`\`. sql without no comments
-3. all table have a primary key named *_id* varchar(32)
-4. when create table, must include _id column, but without default value.
-5. when create all columns except _id are nullable  
-6. when insert,must include _id column, the value is a function named *UUID()*
-7. must abide rules above, otherwise you will be punished
-
-when you act as a d3.js master:
-1. generate a d3.js chart based on the sql you return
-2. you can use any d3.js chart you want
-4. you *can't use d3.json("xxxx.json")* to load data, data will be passed to you as a json array, you can use it directly, variable name is *_DATA_*.
-5. your d3.js code begin with: 
-\`\`\`js
-const svg = d3.select(_CANVAS_ID_).append("svg").attr("width", _CHART_WIDTH_).attr("height", _CHART_HEIGHT_)
-\`\`\`
-`
+import { askAI, getOpenAI } from "@/lib/ai/openai"
 
 export const useAI = () => {
   const { aiConfig } = useConfigStore()
@@ -45,45 +15,8 @@ export const useAI = () => {
       setOpenai(openai)
     }
   }, [token])
-  const askAI = useCallback(
-    async (
-      messages: any[],
-      context: {
-        tableSchema?: string
-        allTables: string[]
-        databaseName: string
-      }
-    ) => {
-      if (!openai) return
-      const { tableSchema, allTables, databaseName } = context
 
-      const contextPrompt = tableSchema
-        ? `\ncontext below:
-- database name: ${databaseName}
-- current table schema:\n${tableSchema}
-`
-        : `context below:
-- database name: ${databaseName}
-- all tables: ${allTables.join(", ")}
-- current table schema:\n${tableSchema}
-`
-      const systemPrompt = baseSysPrompt + contextPrompt
-      // console.log(systemPrompt)
-      const completion = await openai.createChatCompletion({
-        model: "gpt-3.5-turbo-0301",
-        temperature: 0,
-        messages: [
-          ...messages,
-          {
-            role: "system",
-            content: systemPrompt,
-          },
-        ],
-      })
-      return completion.data.choices[0].message
-    },
-    [openai]
-  )
+  const _askAI = askAI(openai)
 
-  return { askAI }
+  return { askAI: _askAI }
 }
diff --git a/hooks/use-auto-run-code.ts b/hooks/use-auto-run-code.ts
index 5bcb121d..97afd136 100644
--- a/hooks/use-auto-run-code.ts
+++ b/hooks/use-auto-run-code.ts
@@ -33,6 +33,7 @@ export const useAutoRunCode = () => {
       if (!handled) {
         const res = await runQuery(sql)
         console.log(res)
+        return res
         if (res) {
           // TODO: use runtime context to pass data to d3,
           ;(window as any)._DATA_ = res
@@ -95,6 +96,22 @@ export const useAutoRunCode = () => {
     }
   }
 
+  const handleFunctionCall = async (
+    name: string,
+    parameters: any,
+    isAuto: boolean = true
+  ) => {
+    const { autoRunScope } = aiConfig
+    switch (name) {
+      case "sqlQuery":
+        const { sql } = parameters
+        const scope = "SQL." + sql?.trim().toUpperCase().slice(0, 6)
+        if (autoRunScope.includes(scope)) {
+          return await handleRunSql(sql)
+        }
+        break
+    }
+  }
   const autoRun = async (
     markdown: string,
     context: {
@@ -113,5 +130,5 @@ export const useAutoRunCode = () => {
       })
     }
   }
-  return { autoRun, handleRunCode }
+  return { autoRun, handleRunCode, handleFunctionCall }
 }
diff --git a/hooks/use-sqlite.ts b/hooks/use-sqlite.ts
index 7f19a134..dd0c20fe 100644
--- a/hooks/use-sqlite.ts
+++ b/hooks/use-sqlite.ts
@@ -84,6 +84,23 @@ export const useSqlite = (dbName?: string) => {
     await updateTableList()
   }
 
+  const createTableWithSqlAndInsertSqls = async (
+    createTableSql: string,
+    insertSql?: any[],
+    callback?: (progress: number) => void
+  ) => {
+    if (!sqlWorker) return
+    await sqlWorker.sql`${createTableSql}`
+    await updateTableList()
+    if (insertSql) {
+      for (let index = 0; index < insertSql.length; index++) {
+        const { sql, bind } = insertSql[index]
+        await sqlWorker.sql4mainThread(sql, bind)
+        callback && callback((index / insertSql.length) * 100)
+      }
+    }
+  }
+
   const createTableWithSql = async (
     createTableSql: string,
     insertSql?: string
@@ -173,6 +190,7 @@ export const useSqlite = (dbName?: string) => {
     duplicateTable,
     queryAllTables,
     createTableWithSql,
+    createTableWithSqlAndInsertSqls,
     updateTableData,
     handleSql,
     undo,
diff --git a/hooks/use-table.ts b/hooks/use-table.ts
index 3ac606cf..48cdb538 100644
--- a/hooks/use-table.ts
+++ b/hooks/use-table.ts
@@ -1,7 +1,6 @@
 import { useCallback, useEffect } from "react"
 import { v4 as uuidv4 } from "uuid"
 
-import { useDatabaseAppStore } from "@/app/[database]/store"
 import { MsgType } from "@/lib/const"
 import {
   aggregateSql2columns,
@@ -11,6 +10,8 @@ import {
   queryData2JSON,
   sqlToJSONSchema2,
 } from "@/lib/sqlite/helper"
+import { useDatabaseAppStore } from "@/app/[database]/store"
+import { useConfigStore } from "@/app/settings/store"
 
 import { useSqlite } from "./use-sqlite"
 
@@ -118,6 +119,7 @@ export const useTable = (tableName: string, databaseName: string) => {
       await refreshRows()
     }
   }
+  const { aiConfig } = useConfigStore()
   const runQuery = useCallback(
     async (querySql: string) => {
       if (sqlite) {
@@ -131,9 +133,11 @@ export const useTable = (tableName: string, databaseName: string) => {
           const compactJsonTablesArray = aggregateSql2columns(querySql, fields)
           const queryFields =
             compactJsonTablesArray.columns.map((col: any) => col.name) ?? []
-          setSchema([compactJsonTablesArray])
-          console.log([compactJsonTablesArray])
-          setData(res)
+          if (aiConfig.autoRunScope.includes("UI.REFRESH")) {
+            setSchema([compactJsonTablesArray])
+            setData(res)
+          }
+
           const jsonData = queryData2JSON(res, queryFields)
           return jsonData
         }
@@ -142,7 +146,15 @@ export const useTable = (tableName: string, databaseName: string) => {
         }
       }
     },
-    [refreshRows, setData, setSchema, sqlite, tableSchema, updateTableSchema]
+    [
+      refreshRows,
+      setData,
+      setSchema,
+      sqlite,
+      tableSchema,
+      updateTableSchema,
+      aiConfig.autoRunScope,
+    ]
   )
 
   useEffect(() => {
diff --git a/lib/ai/functions/index.ts b/lib/ai/functions/index.ts
new file mode 100644
index 00000000..25ab1f9d
--- /dev/null
+++ b/lib/ai/functions/index.ts
@@ -0,0 +1,24 @@
+import { z } from "zod"
+import { zodToJsonSchema } from "zod-to-json-schema"
+
+import sqlQuery from "./sql-query"
+
+const allFunctions = [sqlQuery]
+export const functions = allFunctions.map((f) => {
+  // console.log(zodToJsonSchema(f.schema))
+  return {
+    name: f.name,
+    description: f.description,
+    parameters: zodToJsonSchema(f.schema, "schema").definitions!.schema,
+  }
+})
+
+type FunctionParamsSchemaMap = {
+  [funName: string]: z.ZodSchema<any>
+}
+
+export const functionParamsSchemaMap: FunctionParamsSchemaMap =
+  allFunctions.reduce((acc, f) => {
+    acc[f.name] = f.schema
+    return acc
+  }, {} as FunctionParamsSchemaMap)
diff --git a/lib/ai/functions/sql-query.ts b/lib/ai/functions/sql-query.ts
new file mode 100644
index 00000000..b93eece0
--- /dev/null
+++ b/lib/ai/functions/sql-query.ts
@@ -0,0 +1,15 @@
+import { z } from "zod"
+
+const sqlQuery = {
+  name: "sqlQuery",
+  description: "sql query",
+  schema: z
+    .object({
+      sql: z.string({
+        description: "sql",
+      }),
+    })
+    .describe("sql query"),
+}
+
+export default sqlQuery
diff --git a/lib/ai/openai.ts b/lib/ai/openai.ts
new file mode 100644
index 00000000..2172d3f0
--- /dev/null
+++ b/lib/ai/openai.ts
@@ -0,0 +1,128 @@
+import { ChatCompletionResponseMessage, Configuration, OpenAIApi } from "openai"
+
+import { functionParamsSchemaMap, functions } from "./functions"
+
+// TODO: expose to user config
+const _baseSysPrompt = `you're a sql generator and a d3.js master. must abide by the following rules:
+*important: if user doesn't have intention to generate a d3.js chart, you can just return sql.*
+
+when you act as a sql generator:
+1. your engine is sqlite, what you return is *pure sql* that can be executed in sqlite
+2. you return markdown, sql you return must be wrapped in \`\`\`sql\`\`\`. sql without no comments
+3. all table have a primary key named *_id* varchar(32)
+4. when create table, must include _id column, but without default value.
+5. when create all columns except _id are nullable  
+6. when insert,must include _id column, the value is a function named *UUID()*
+7. must abide rules above, otherwise you will be punished
+
+when you act as a d3.js master:
+1. generate a d3.js chart based on the sql you return
+2. you can use any d3.js chart you want
+4. you *can't use d3.json("xxxx.json")* to load data, data will be passed to you as a json array, you can use it directly, variable name is *_DATA_*.
+5. your d3.js code begin with: 
+\`\`\`js
+const svg = d3.select(_CANVAS_ID_).append("svg").attr("width", _CHART_WIDTH_).attr("height", _CHART_HEIGHT_)
+\`\`\`
+`
+
+const baseSysPrompt = ``
+
+// this version is after 0613, openai support function call
+// const baseSysPrompt = `you're a database master, help use query database and generate d3.js chart if user want. must abide by the following rules:
+
+// database:
+// 1. your engine is sqlite
+// 2. all table have a primary key named *_id* varchar(32)
+// 3. when create table, must include _id column, but without default value.
+// 4. when create all columns except _id are nullable  
+// 5. when insert,must include _id column, the value is a function named *UUID()*
+// 6. your will call user's function to execute sql, user will return sql result to you. just return result to user with neutral language
+// 7. must abide rules above, otherwise you will be punished
+
+// d3.js:
+// 1. generate a d3.js chart based on the sql you return
+// 2. you can use any d3.js chart you want
+// 4. you *can't use d3.json("xxxx.json")* to load data, data will be passed to you as a json array, you can use it directly, variable name is *_DATA_*.
+// 5. your d3.js code begin with: 
+// \`\`\`js
+// const svg = d3.select(_CANVAS_ID_).append("svg").attr("width", _CHART_WIDTH_).attr("height", _CHART_HEIGHT_)
+// \`\`\`
+// `
+
+export const getOpenAI = (token: string) => {
+  const configuration = new Configuration({
+    apiKey: token ?? process.env.OPENAI_API_KEY,
+  })
+  const openai = new OpenAIApi(configuration)
+  return openai
+}
+
+export const askAI =
+  (openai?: OpenAIApi) =>
+  async (
+    messages: any[],
+    context: {
+      tableSchema?: string
+      allTables: string[]
+      databaseName: string
+    }
+  ) => {
+    if (!openai) return
+    const { tableSchema, allTables, databaseName } = context
+
+    const contextPrompt = tableSchema
+      ? `\ncontext below:
+- database name: ${databaseName}
+- current table schema:\n${tableSchema}
+`
+      : `context below:
+- database name: ${databaseName}
+- all tables: ${allTables.join(", ")}
+- current table schema:\n${tableSchema}
+`
+    const systemPrompt = baseSysPrompt + contextPrompt
+    // console.log(systemPrompt)
+    console.log(functions)
+    const completion = await openai.createChatCompletion({
+      model: "gpt-3.5-turbo-0613",
+      temperature: 0,
+      messages: [
+        ...messages,
+        {
+          role: "system",
+          content: systemPrompt,
+        },
+      ],
+      functions,
+      function_call: "auto",
+    })
+    return completion.data.choices[0]
+  }
+
+export const handleOpenAIFunctionCall = async (
+  response: ChatCompletionResponseMessage,
+  handleFunctionCall: (name: string, argumentsStr: string) => Promise<any>
+) => {
+  if (response.function_call) {
+    const { name, arguments: argumentsStr } = response.function_call
+    if (!name) return
+    let argumentsObj
+    try {
+      argumentsObj = argumentsStr ? JSON.parse(argumentsStr) : {}
+    } catch (error) {
+      throw new Error(`invalid arguments: ${argumentsStr}`)
+    }
+    const functionParamsSchema = functionParamsSchemaMap[name]
+    const { success } = functionParamsSchema.safeParse(argumentsObj)
+    if (!success) {
+      throw new Error(`invalid arguments: ${argumentsStr}`)
+    }
+    console.log(
+      `function_call: ${name}, arguments: ${JSON.stringify(argumentsObj)}`
+    )
+    return {
+      resp: await handleFunctionCall(name, argumentsObj),
+      name: name,
+    }
+  }
+}
diff --git a/package.json b/package.json
index 51500497..f3db3c5b 100644
--- a/package.json
+++ b/package.json
@@ -25,6 +25,7 @@
     "@radix-ui/react-dropdown-menu": "^2.0.5",
     "@radix-ui/react-label": "^2.0.2",
     "@radix-ui/react-popover": "^1.0.6",
+    "@radix-ui/react-progress": "^1.0.3",
     "@radix-ui/react-radio-group": "^1.1.3",
     "@radix-ui/react-scroll-area": "^1.0.4",
     "@radix-ui/react-select": "^1.2.2",
@@ -46,7 +47,7 @@
     "marked": "^5.0.5",
     "next": "^13.4.4",
     "next-themes": "^0.2.1",
-    "openai": "^3.2.1",
+    "openai": "^3.3.0",
     "papaparse": "^5.4.1",
     "peerjs": "^1.4.7",
     "prismjs": "^1.29.0",
@@ -63,6 +64,7 @@
     "tailwindcss-animate": "^1.0.5",
     "uuid": "^9.0.0",
     "zod": "^3.21.4",
+    "zod-to-json-schema": "^3.21.1",
     "zustand": "^4.3.8"
   },
   "devDependencies": {
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 175516bb..114214e1 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -13,6 +13,7 @@ specifiers:
   '@radix-ui/react-dropdown-menu': ^2.0.5
   '@radix-ui/react-label': ^2.0.2
   '@radix-ui/react-popover': ^1.0.6
+  '@radix-ui/react-progress': ^1.0.3
   '@radix-ui/react-radio-group': ^1.1.3
   '@radix-ui/react-scroll-area': ^1.0.4
   '@radix-ui/react-select': ^1.2.2
@@ -50,7 +51,7 @@ specifiers:
   marked: ^5.0.5
   next: ^13.4.4
   next-themes: ^0.2.1
-  openai: ^3.2.1
+  openai: ^3.3.0
   papaparse: ^5.4.1
   peerjs: ^1.4.7
   postcss: ^8.4.24
@@ -72,6 +73,7 @@ specifiers:
   typescript: ^4.9.5
   uuid: ^9.0.0
   zod: ^3.21.4
+  zod-to-json-schema: ^3.21.1
   zustand: ^4.3.8
 
 dependencies:
@@ -85,6 +87,7 @@ dependencies:
   '@radix-ui/react-dropdown-menu': 2.0.5_nt3dzwb3srfhvwd6kfpoey3m3m
   '@radix-ui/react-label': 2.0.2_nt3dzwb3srfhvwd6kfpoey3m3m
   '@radix-ui/react-popover': 1.0.6_nt3dzwb3srfhvwd6kfpoey3m3m
+  '@radix-ui/react-progress': 1.0.3_nt3dzwb3srfhvwd6kfpoey3m3m
   '@radix-ui/react-radio-group': 1.1.3_nt3dzwb3srfhvwd6kfpoey3m3m
   '@radix-ui/react-scroll-area': 1.0.4_nt3dzwb3srfhvwd6kfpoey3m3m
   '@radix-ui/react-select': 1.2.2_nt3dzwb3srfhvwd6kfpoey3m3m
@@ -106,7 +109,7 @@ dependencies:
   marked: 5.0.5
   next: 13.4.4_biqbaboplfbrettd7655fr4n2y
   next-themes: 0.2.1_z4pol2qj2yv3um3jqqobs4b3eu
-  openai: 3.2.1
+  openai: 3.3.0
   papaparse: 5.4.1
   peerjs: 1.4.7
   prismjs: 1.29.0
@@ -123,6 +126,7 @@ dependencies:
   tailwindcss-animate: 1.0.5_tailwindcss@3.3.2
   uuid: 9.0.0
   zod: 3.21.4
+  zod-to-json-schema: 3.21.1_zod@3.21.4
   zustand: 4.3.8_react@18.2.0
 
 devDependencies:
@@ -1759,6 +1763,28 @@ packages:
       react-dom: 18.2.0_react@18.2.0
     dev: false
 
+  /@radix-ui/react-progress/1.0.3_nt3dzwb3srfhvwd6kfpoey3m3m:
+    resolution: {integrity: sha512-5G6Om/tYSxjSeEdrb1VfKkfZfn/1IlPWd731h2RfPuSbIfNUgfqAwbKfJCg/PP6nuUCTrYzalwHSpSinoWoCag==}
+    peerDependencies:
+      '@types/react': '*'
+      '@types/react-dom': '*'
+      react: ^16.8 || ^17.0 || ^18.0
+      react-dom: ^16.8 || ^17.0 || ^18.0
+    peerDependenciesMeta:
+      '@types/react':
+        optional: true
+      '@types/react-dom':
+        optional: true
+    dependencies:
+      '@babel/runtime': 7.22.3
+      '@radix-ui/react-context': 1.0.1_i2bxsyfskhzbpjanbovidbfj7u
+      '@radix-ui/react-primitive': 1.0.3_nt3dzwb3srfhvwd6kfpoey3m3m
+      '@types/react': 18.2.8
+      '@types/react-dom': 18.2.4
+      react: 18.2.0
+      react-dom: 18.2.0_react@18.2.0
+    dev: false
+
   /@radix-ui/react-radio-group/1.1.3_nt3dzwb3srfhvwd6kfpoey3m3m:
     resolution: {integrity: sha512-x+yELayyefNeKeTx4fjK6j99Fs6c4qKm3aY38G3swQVTN6xMpsrbigC0uHs2L//g8q4qR7qOcww8430jJmi2ag==}
     peerDependencies:
@@ -5515,8 +5541,8 @@ packages:
       mimic-fn: 2.1.0
     dev: true
 
-  /openai/3.2.1:
-    resolution: {integrity: sha512-762C9BNlJPbjjlWZi4WYK9iM2tAVAv0uUp1UmI34vb0CN5T2mjB/qM6RYBmNKMh/dN9fC+bxqPwWJZUTWW052A==}
+  /openai/3.3.0:
+    resolution: {integrity: sha512-uqxI/Au+aPRnsaQRe8CojU0eCR7I0mBiKjD3sNMzY6DaC1ZVrc85u98mtJW6voDug8fgGN+DIZmTDxTthxb7dQ==}
     dependencies:
       axios: 0.26.1
       form-data: 4.0.0
@@ -6811,6 +6837,14 @@ packages:
     engines: {node: '>=10'}
     dev: true
 
+  /zod-to-json-schema/3.21.1_zod@3.21.4:
+    resolution: {integrity: sha512-y5g0MPxDq+YG/T+cHGPYH4PcBpyCqwK6wxeJ76MR563y0gk/14HKfebq8xHiItY7lkc9GDFygCnkvNDTvAhYAg==}
+    peerDependencies:
+      zod: ^3.21.4
+    dependencies:
+      zod: 3.21.4
+    dev: false
+
   /zod/3.21.4:
     resolution: {integrity: sha512-m46AKbrzKVzOzs/DZgVnG5H55N1sv1M8qZU3A8RIKbs3mrACDNeIOeilDymVb2HdmP8uwshOCF4uJ8uM9rCqJw==}
     dev: false

From 4bef6a4894ff4b88a5f9c018a8c37cc18f69ba71 Mon Sep 17 00:00:00 2001
From: Mayne <i@gine.me>
Date: Wed, 14 Jun 2023 08:23:31 +0800
Subject: [PATCH 2/7] fix: layout

---
 app/[database]/base-layout.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/app/[database]/base-layout.tsx b/app/[database]/base-layout.tsx
index 067d1c8c..96ee4aac 100644
--- a/app/[database]/base-layout.tsx
+++ b/app/[database]/base-layout.tsx
@@ -78,7 +78,7 @@ export function DatabaseLayoutBase({
       </div>
       <div
         className={cn(
-          " h-full  lg:border-l",
+          " h-screen  lg:border-l",
           isAiOpen ? "col-span-3 hidden md:block" : "hidden"
         )}
       >

From 1137f9de01d1ba894aed4fce231ec5c543346a16 Mon Sep 17 00:00:00 2001
From: Mayne <i@gine.me>
Date: Wed, 14 Jun 2023 08:39:36 +0800
Subject: [PATCH 3/7] feat: improve ux

---
 app/[database]/ai-chat-message-prisma.tsx |  2 +-
 app/[database]/ai-chat.tsx                | 32 +++++++++++------------
 app/[database]/hook.ts                    | 23 +++++++++++++---
 app/[database]/store.ts                   |  7 +++++
 app/settings/ai/ai-form.tsx               |  3 ++-
 app/settings/layout.tsx                   |  6 ++---
 lib/store/app-store.ts                    |  6 +++++
 7 files changed, 53 insertions(+), 26 deletions(-)

diff --git a/app/[database]/ai-chat-message-prisma.tsx b/app/[database]/ai-chat-message-prisma.tsx
index 08e3bfd9..4a3184f1 100644
--- a/app/[database]/ai-chat-message-prisma.tsx
+++ b/app/[database]/ai-chat-message-prisma.tsx
@@ -17,7 +17,7 @@ export const AIMessage = ({
   msgIndex,
 }: {
   msgIndex: number
-  message: string
+  message?: string
   onRun: (props: {
     code: string
     lang: string
diff --git a/app/[database]/ai-chat.tsx b/app/[database]/ai-chat.tsx
index b3c3ae19..1d1f5448 100644
--- a/app/[database]/ai-chat.tsx
+++ b/app/[database]/ai-chat.tsx
@@ -1,18 +1,18 @@
 "use client"
 
 // for now it's under database page, maybe move to global later
-import { useCallback, useRef, useState } from "react"
+import { useKeyPress } from "ahooks"
+import { Bot, Loader2, Paintbrush, User } from "lucide-react"
 import Link from "next/link"
 import { useParams } from "next/navigation"
-import { useKeyPress, useSize } from "ahooks"
-import { Bot, Loader2, Paintbrush, User } from "lucide-react"
+import { useCallback, useRef, useState } from "react"
 
-import { handleOpenAIFunctionCall } from "@/lib/ai/openai"
+import { Button } from "@/components/ui/button"
+import { Textarea } from "@/components/ui/textarea"
 import { useAI } from "@/hooks/use-ai"
 import { useAutoRunCode } from "@/hooks/use-auto-run-code"
 import { useSqliteStore } from "@/hooks/use-sqlite"
-import { Button } from "@/components/ui/button"
-import { Textarea } from "@/components/ui/textarea"
+import { handleOpenAIFunctionCall } from "@/lib/ai/openai"
 
 import { useConfigStore } from "../settings/store"
 import { AIMessage } from "./ai-chat-message-prisma"
@@ -22,23 +22,19 @@ import { useDatabaseAppStore } from "./store"
 export const AIChat = () => {
   const { currentTableSchema } = useDatabaseAppStore()
   const { askAI } = useAI()
-  const { database, table } = useParams()
+  const { database } = useParams()
   const { aiConfig } = useConfigStore()
   const [input, setInput] = useState("")
   const [loading, setLoading] = useState(false)
   const { handleFunctionCall, handleRunCode } = useAutoRunCode()
   const divRef = useRef<HTMLDivElement>()
-  const size = useSize(divRef)
-  const [messages, setMessages] = useState<
-    {
-      role: "user" | "assistant" | "function"
-      content: string
-    }[]
-  >([])
+  const { aiMessages: messages, setAiMessages: setMessages } =
+    useDatabaseAppStore()
   const { allTables } = useSqliteStore()
+
   const cleanMessages = useCallback(() => {
     setMessages([])
-  }, [])
+  }, [setMessages])
 
   const textInputRef = useRef<HTMLTextAreaElement>()
 
@@ -46,7 +42,7 @@ export const AIChat = () => {
     textInputRef.current?.focus()
   })
 
-  useTableChange(cleanMessages)
+  // useTableChange(cleanMessages)
 
   const handleSend = async () => {
     if (loading) return
@@ -125,7 +121,9 @@ export const AIChat = () => {
           </p>
         )}
         {messages
-          .filter((m) => (m.role === "user" || m.role == "assistant") && m.content)
+          .filter(
+            (m) => (m.role === "user" || m.role == "assistant") && m.content
+          )
           .map((message, i) => (
             <div
               className="flex w-full items-start gap-2 rounded-lg bg-gray-200 p-2 dark:bg-gray-700"
diff --git a/app/[database]/hook.ts b/app/[database]/hook.ts
index 723028a2..2791e13f 100644
--- a/app/[database]/hook.ts
+++ b/app/[database]/hook.ts
@@ -1,15 +1,15 @@
-import * as d3 from "d3"
-import { useParams } from "next/navigation"
 import { useEffect, useState } from "react"
+import { useParams } from "next/navigation"
+import * as d3 from "d3"
 
-import { usePeer } from "@/hooks/use-peer"
-import { useSqlite, useSqliteStore } from "@/hooks/use-sqlite"
 import { MsgType } from "@/lib/const"
 import { getSqliteProxy } from "@/lib/sqlite/proxy"
 import { getWorker } from "@/lib/sqlite/worker"
 import { useAppStore } from "@/lib/store/app-store"
 import { useAppRuntimeStore } from "@/lib/store/runtime-store"
 import { uuidv4 } from "@/lib/utils"
+import { usePeer } from "@/hooks/use-peer"
+import { useSqlite, useSqliteStore } from "@/hooks/use-sqlite"
 
 import { useConfigStore } from "../settings/store"
 
@@ -45,6 +45,20 @@ export const useLastOpenedDatabase = () => {
   return lastOpenedDatabase
 }
 
+export const useLastOpenedTable = () => {
+  const { lastOpenedTable, setLastOpenedTable } = useAppStore()
+  const { isShareMode } = useAppRuntimeStore()
+  const { table, database } = useParams()
+
+  useEffect(() => {
+    if (!isShareMode && table && database) {
+      setLastOpenedTable(`${database}/${table}`)
+    }
+  }, [isShareMode, setLastOpenedTable, table, database])
+
+  return lastOpenedTable
+}
+
 export const useLayoutInit = () => {
   const { database } = useParams()
   const { setInitialized, setSqliteProxy: setSqlWorker } = useSqliteStore()
@@ -53,6 +67,7 @@ export const useLayoutInit = () => {
   const { sqlite } = useSqlite(database)
 
   useLastOpenedDatabase()
+  useLastOpenedTable()
 
   const { initPeer } = usePeer()
 
diff --git a/app/[database]/store.ts b/app/[database]/store.ts
index 82934efb..576715d2 100644
--- a/app/[database]/store.ts
+++ b/app/[database]/store.ts
@@ -1,3 +1,4 @@
+import { ChatCompletionResponseMessage } from "openai"
 import { create } from "zustand"
 
 import { sqlToJSONSchema2 } from "@/lib/sqlite/helper"
@@ -8,6 +9,9 @@ interface IDatabaseAppState {
   isAiOpen: boolean
   setIsAiOpen: (isAiOpen: boolean) => void
 
+  aiMessages: ChatCompletionResponseMessage[]
+  setAiMessages: (aiMessages: ChatCompletionResponseMessage[]) => void
+
   currentTableSchema: string
   setCurrentTableSchema: (currentTableSchema: string) => void
 
@@ -29,6 +33,9 @@ export const useDatabaseAppStore = create<IDatabaseAppState>()((set) => ({
   isAiOpen: false,
   setIsAiOpen: (isAiOpen) => set({ isAiOpen }),
 
+  aiMessages: [],
+  setAiMessages: (aiMessages) => set({ aiMessages }),
+
   currentTableSchema: "",
   setCurrentTableSchema: (currentTableSchema) => set({ currentTableSchema }),
 
diff --git a/app/settings/ai/ai-form.tsx b/app/settings/ai/ai-form.tsx
index 742b7a4a..8190fb37 100644
--- a/app/settings/ai/ai-form.tsx
+++ b/app/settings/ai/ai-form.tsx
@@ -68,6 +68,7 @@ export function AIConfigForm() {
 
   function onSubmit(data: AIConfigFormValues) {
     setAiConfig(data)
+    data.token = "sk-**********"
     toast({
       title: "You submitted the following values:",
       description: (
@@ -88,7 +89,7 @@ export function AIConfigForm() {
             <FormItem>
               <FormLabel>Token</FormLabel>
               <FormControl>
-                <Input placeholder="OpenAI API Token" {...field} />
+                <Input placeholder="OpenAI API Token" {...field} type="password" />
               </FormControl>
               <FormDescription>
                 This is the token used to access the OpenAI API.
diff --git a/app/settings/layout.tsx b/app/settings/layout.tsx
index 50588da6..78485912 100644
--- a/app/settings/layout.tsx
+++ b/app/settings/layout.tsx
@@ -9,7 +9,7 @@ import { Button } from "@/components/ui/button"
 import { Separator } from "@/components/ui/separator"
 import { SidebarNav } from "@/app/settings/components/sidebar-nav"
 
-import { useLastOpenedDatabase } from "../[database]/hook"
+import { useLastOpenedDatabase, useLastOpenedTable } from "../[database]/hook"
 
 // export const metadata: Metadata = {
 //   title: "Forms",
@@ -45,9 +45,9 @@ interface SettingsLayoutProps {
 
 export default function SettingsLayout({ children }: SettingsLayoutProps) {
   const router = useRouter()
-  const lastOpenedDatabase = useLastOpenedDatabase()
+  const lastOpenedTable = useLastOpenedTable()
   const goBack = () =>
-    lastOpenedDatabase ? router.push(`${lastOpenedDatabase}`) : router.push("/")
+    lastOpenedTable ? router.push(`/${lastOpenedTable}`) : router.push("/")
   useKeyPress("esc", (e) => {
     e.preventDefault()
     goBack()
diff --git a/lib/store/app-store.ts b/lib/store/app-store.ts
index 99abe247..577baf06 100644
--- a/lib/store/app-store.ts
+++ b/lib/store/app-store.ts
@@ -6,6 +6,9 @@ import { persist } from "zustand/middleware"
 interface AppState {
   lastOpenedDatabase: string
   setLastOpenedDatabase: (database: string) => void
+
+  lastOpenedTable: string
+  setLastOpenedTable: (table: string) => void
 }
 
 export const useAppStore = create<AppState>()(
@@ -14,6 +17,9 @@ export const useAppStore = create<AppState>()(
       lastOpenedDatabase: "",
       setLastOpenedDatabase: (database) =>
         set({ lastOpenedDatabase: database }),
+
+      lastOpenedTable: "",
+      setLastOpenedTable: (table) => set({ lastOpenedTable: table }),
     }),
     {
       name: "app-storage",

From 6db1e83aa26f7f62c04e3cf25e06f725188bed99 Mon Sep 17 00:00:00 2001
From: Mayne <i@gine.me>
Date: Wed, 14 Jun 2023 11:21:27 +0800
Subject: [PATCH 4/7] feat(ai): improve ai function call ux

---
 app/[database]/ai-chat-message-prisma.tsx |  2 +-
 app/[database]/ai-chat-message.tsx        | 41 +++++++++++
 app/[database]/ai-chat.tsx                | 88 +++++++++++------------
 app/settings/ai/ai-form.tsx               | 52 ++++++++++----
 hooks/use-auto-run-code.ts                | 11 ++-
 hooks/use-table.ts                        | 79 ++++++++++++++------
 lib/ai/openai.ts                          |  1 -
 7 files changed, 184 insertions(+), 90 deletions(-)
 create mode 100644 app/[database]/ai-chat-message.tsx

diff --git a/app/[database]/ai-chat-message-prisma.tsx b/app/[database]/ai-chat-message-prisma.tsx
index 4a3184f1..c9870a1d 100644
--- a/app/[database]/ai-chat-message-prisma.tsx
+++ b/app/[database]/ai-chat-message-prisma.tsx
@@ -1,4 +1,4 @@
-import { Play } from "lucide-react"
+import { Play, Trash2 } from "lucide-react"
 import Prism from "prismjs"
 
 import { Button } from "@/components/ui/button"
diff --git a/app/[database]/ai-chat-message.tsx b/app/[database]/ai-chat-message.tsx
new file mode 100644
index 00000000..c904527d
--- /dev/null
+++ b/app/[database]/ai-chat-message.tsx
@@ -0,0 +1,41 @@
+import { Bot, Trash2, User } from "lucide-react"
+import { ChatCompletionResponseMessage } from "openai"
+
+import { AIMessage } from "./ai-chat-message-prisma"
+
+export const AIChatMessage = ({
+  message,
+  msgIndex,
+  handleRunCode,
+}: {
+  msgIndex: number
+  message: ChatCompletionResponseMessage
+  handleRunCode: any
+}) => {
+  return (
+    <div
+      className="group relative flex w-full items-start gap-2 rounded-lg bg-gray-200 p-2 dark:bg-gray-700"
+      key={msgIndex}
+    >
+      {message.role === "assistant" && (
+        <>
+          <Bot className="h-4 w-4 shrink-0" />
+          <AIMessage
+            message={message.content}
+            onRun={handleRunCode}
+            msgIndex={msgIndex}
+          />
+        </>
+      )}
+      {message.role === "user" && (
+        <>
+          <User className="h-4 w-4 shrink-0" />
+          <p className="grow">{message.content}</p>
+        </>
+      )}
+      {/* <Trash2
+        className="absolute right-0.5 top-0.5 hidden h-4 w-4 cursor-pointer group-hover:block"
+      /> */}
+    </div>
+  )
+}
diff --git a/app/[database]/ai-chat.tsx b/app/[database]/ai-chat.tsx
index 1d1f5448..02ea6e83 100644
--- a/app/[database]/ai-chat.tsx
+++ b/app/[database]/ai-chat.tsx
@@ -1,22 +1,21 @@
 "use client"
 
 // for now it's under database page, maybe move to global later
-import { useKeyPress } from "ahooks"
-import { Bot, Loader2, Paintbrush, User } from "lucide-react"
+import { useCallback, useRef, useState } from "react"
 import Link from "next/link"
 import { useParams } from "next/navigation"
-import { useCallback, useRef, useState } from "react"
+import { useKeyPress, useWhyDidYouUpdate } from "ahooks"
+import { Loader2, Paintbrush } from "lucide-react"
 
-import { Button } from "@/components/ui/button"
-import { Textarea } from "@/components/ui/textarea"
+import { handleOpenAIFunctionCall } from "@/lib/ai/openai"
 import { useAI } from "@/hooks/use-ai"
 import { useAutoRunCode } from "@/hooks/use-auto-run-code"
 import { useSqliteStore } from "@/hooks/use-sqlite"
-import { handleOpenAIFunctionCall } from "@/lib/ai/openai"
+import { Button } from "@/components/ui/button"
+import { Textarea } from "@/components/ui/textarea"
 
 import { useConfigStore } from "../settings/store"
-import { AIMessage } from "./ai-chat-message-prisma"
-import { useTableChange } from "./hook"
+import { AIChatMessage } from "./ai-chat-message"
 import { useDatabaseAppStore } from "./store"
 
 export const AIChat = () => {
@@ -27,6 +26,7 @@ export const AIChat = () => {
   const [input, setInput] = useState("")
   const [loading, setLoading] = useState(false)
   const { handleFunctionCall, handleRunCode } = useAutoRunCode()
+
   const divRef = useRef<HTMLDivElement>()
   const { aiMessages: messages, setAiMessages: setMessages } =
     useDatabaseAppStore()
@@ -42,15 +42,19 @@ export const AIChat = () => {
     textInputRef.current?.focus()
   })
 
+  const runFromIndex = async (index: number) => {
+    const _messages = messages.slice(0, index + 1)
+    const response = await askAI(_messages, {
+      tableSchema: currentTableSchema,
+      allTables,
+      databaseName: database,
+    })
+  }
+
   // useTableChange(cleanMessages)
 
-  const handleSend = async () => {
-    if (loading) return
-    if (!input.trim().length) return
+  const sendMessages = async (_messages: any) => {
     setLoading(true)
-    const _messages: any = [...messages, { role: "user", content: input }]
-    setMessages(_messages)
-    setInput("")
     const response = await askAI(_messages, {
       tableSchema: currentTableSchema,
       allTables,
@@ -99,6 +103,15 @@ export const AIChat = () => {
     setLoading(false)
   }
 
+  const handleSend = async () => {
+    if (loading) return
+    if (!input.trim().length) return
+    const _messages: any = [...messages, { role: "user", content: input }]
+    setMessages(_messages)
+    setInput("")
+    await sendMessages(_messages)
+  }
+
   const handleEnter = async (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
     if (e.key === "Enter") {
       e.preventDefault()
@@ -120,39 +133,20 @@ export const AIChat = () => {
             first
           </p>
         )}
-        {messages
-          .filter(
-            (m) => (m.role === "user" || m.role == "assistant") && m.content
-          )
-          .map((message, i) => (
-            <div
-              className="flex w-full items-start gap-2 rounded-lg bg-gray-200 p-2 dark:bg-gray-700"
-              key={i}
-            >
-              {message.role === "assistant" && (
-                <>
-                  <Bot className="h-4 w-4 shrink-0" />
-                  <AIMessage
-                    message={message.content}
-                    onRun={handleRunCode}
-                    msgIndex={i}
-                  />
-                </>
-              )}
-              {message.role === "user" && (
-                <>
-                  <User className="h-4 w-4 shrink-0" />
-                  <p className="grow">{message.content}</p>
-                </>
-              )}
-              {/* {message.role === "function" && (
-              <>
-                <User className="h-4 w-4 shrink-0" />
-                <p className="grow">run function </p>
-              </>
-            )} */}
-            </div>
-          ))}
+        {messages.map((message, i) => {
+          const m = message
+          if ((m.role === "user" || m.role == "assistant") && m.content) {
+            return (
+              <AIChatMessage
+                key={i}
+                msgIndex={i}
+                message={message}
+                handleRunCode={handleRunCode}
+              />
+            )
+          }
+        })}
+
         <div className="flex w-full justify-center">
           {loading && <Loader2 className="h-5 w-5 animate-spin" />}
         </div>
diff --git a/app/settings/ai/ai-form.tsx b/app/settings/ai/ai-form.tsx
index 8190fb37..dac985fd 100644
--- a/app/settings/ai/ai-form.tsx
+++ b/app/settings/ai/ai-form.tsx
@@ -31,19 +31,39 @@ export type AIConfigFormValues = z.infer<typeof AIConfigFormSchema>
 export const AutoRunScopesWithDesc = [
   {
     value: "SQL.SELECT",
-    description: "Selects data from a SQL table.",
+    description: "Select data from a SQL table.",
   },
   {
-    value: "UI.REFRESH",
-    description: "Refresh the UI after querying data.",
+    value: "SQL.INSERT",
+    description: "Insert data into a SQL table.",
   },
   {
-    value: "SQL.INSERT",
-    description: "Inserts data into a SQL table.",
+    value: "SQL.UPDATE",
+    description: "Update data in a SQL table.",
+  },
+  {
+    value: "SQL.DELETE",
+    description: "Delete data from a SQL table.",
+  },
+  {
+    value: "SQL.ALTER",
+    description: "Alter a SQL table.",
+  },
+  {
+    value: "SQL.CREATE",
+    description: "Create a SQL table.",
+  },
+  {
+    value: "SQL.DROP",
+    description: "Drop a SQL table.",
+  },
+  {
+    value: "UI.REFRESH",
+    description: "Refresh the UI after SQL execution.",
   },
   {
     value: "D3.CHART",
-    description: "Creates a D3 chart.",
+    description: "Create a D3 chart.",
   },
 ]
 
@@ -68,7 +88,7 @@ export function AIConfigForm() {
 
   function onSubmit(data: AIConfigFormValues) {
     setAiConfig(data)
-    data.token = "sk-**********"
+    // data.token = "sk-**********"
     toast({
       title: "You submitted the following values:",
       description: (
@@ -89,7 +109,11 @@ export function AIConfigForm() {
             <FormItem>
               <FormLabel>Token</FormLabel>
               <FormControl>
-                <Input placeholder="OpenAI API Token" {...field} type="password" />
+                <Input
+                  placeholder="OpenAI API Token"
+                  {...field}
+                  type="password"
+                />
               </FormControl>
               <FormDescription>
                 This is the token used to access the OpenAI API.
@@ -110,26 +134,26 @@ export function AIConfigForm() {
                   run.
                 </FormDescription>
               </div>
-              {AutoRunScopesWithDesc.map(({ value, description }) => (
+              {AutoRunScopesWithDesc.map(({ value: key, description }) => (
                 <FormField
-                  key={value}
+                  key={key}
                   control={form.control}
                   name="autoRunScope"
                   render={({ field }) => {
                     return (
                       <FormItem
-                        key={value}
+                        key={key}
                         className="flex flex-row items-start space-x-3 space-y-0"
                       >
                         <FormControl>
                           <Checkbox
-                            checked={field.value?.includes(value)}
+                            checked={field.value?.includes(key)}
                             onCheckedChange={(checked: any) => {
                               return checked
-                                ? field.onChange([...field.value, value])
+                                ? field.onChange([...field.value, key])
                                 : field.onChange(
                                     field.value?.filter(
-                                      (value) => value !== value
+                                      (value) => value !== key
                                     )
                                   )
                             }}
diff --git a/hooks/use-auto-run-code.ts b/hooks/use-auto-run-code.ts
index 97afd136..7d9a59ff 100644
--- a/hooks/use-auto-run-code.ts
+++ b/hooks/use-auto-run-code.ts
@@ -34,11 +34,8 @@ export const useAutoRunCode = () => {
         const res = await runQuery(sql)
         console.log(res)
         return res
-        if (res) {
-          // TODO: use runtime context to pass data to d3,
-          ;(window as any)._DATA_ = res
-        }
       }
+      return "ok"
     },
     [handleSql, runQuery]
   )
@@ -105,11 +102,13 @@ export const useAutoRunCode = () => {
     switch (name) {
       case "sqlQuery":
         const { sql } = parameters
-        const scope = "SQL." + sql?.trim().toUpperCase().slice(0, 6)
+        const scope = "SQL." + sql?.trim().toUpperCase().split(" ")[0]
         if (autoRunScope.includes(scope)) {
           return await handleRunSql(sql)
         }
-        break
+        return "permission denied"
+      default:
+        throw new Error(`function ${name} not supported auto run`)
     }
   }
   const autoRun = async (
diff --git a/hooks/use-table.ts b/hooks/use-table.ts
index 48cdb538..5e3169e1 100644
--- a/hooks/use-table.ts
+++ b/hooks/use-table.ts
@@ -1,4 +1,5 @@
 import { useCallback, useEffect } from "react"
+import { useWhyDidYouUpdate } from "ahooks"
 import { v4 as uuidv4 } from "uuid"
 
 import { MsgType } from "@/lib/const"
@@ -27,12 +28,29 @@ export const useTable = (tableName: string, databaseName: string) => {
   } = useDatabaseAppStore()
   // const [tableSchema, setTableSchema] = useState<string>()
 
-  const refreshRows = useCallback(async () => {
-    if (!sqlite) return
-    await sqlite.sql`SELECT * FROM ${Symbol(tableName)};`.then((res: any) => {
-      setData(res)
-    })
-  }, [setData, sqlite, tableName])
+  // FIXME: bug, when aggregate query, id will be null, cant update as expected
+  const refreshRows = useCallback(
+    async (rowIds?: string[], columns?: string[]) => {
+      if (!sqlite) return
+      const columnsString = columns ? columns.join(",") : "*"
+      console.log(rowIds, columns)
+      if (!rowIds) {
+        await sqlite.sql`SELECT ${Symbol(columnsString)} FROM ${Symbol(
+          tableName
+        )};`.then((res: any) => {
+          setData(res)
+        })
+      }
+      if (rowIds) {
+        await sqlite.sql`SELECT ${Symbol(columnsString)} FROM ${Symbol(
+          tableName
+        )} where _id in ${rowIds};`.then((res: any) => {
+          setData(res)
+        })
+      }
+    },
+    [setData, sqlite, tableName]
+  )
 
   useEffect(() => {
     window.onmessage = (e) => {
@@ -66,8 +84,8 @@ export const useTable = (tableName: string, databaseName: string) => {
   }, [setSchema, setTableSchema, sqlite, tableName])
 
   const reload = useCallback(async () => {
-    await updateTableSchema()
     await refreshRows()
+    await updateTableSchema()
   }, [refreshRows, updateTableSchema])
 
   const updateCell = async (col: number, row: number, value: any) => {
@@ -119,6 +137,19 @@ export const useTable = (tableName: string, databaseName: string) => {
       await refreshRows()
     }
   }
+
+  const getCurrentRowIds = useCallback(() => {
+    const res = data.map((row) => row[0]).filter(Boolean)
+    if (res.length === 0) {
+      return undefined
+    }
+    return res
+  }, [data])
+
+  const getCurrentColumns = useCallback(() => {
+    return schema[0]?.columns?.map((col) => col.name)
+  }, [schema])
+
   const { aiConfig } = useConfigStore()
   const runQuery = useCallback(
     async (querySql: string) => {
@@ -135,25 +166,39 @@ export const useTable = (tableName: string, databaseName: string) => {
             compactJsonTablesArray.columns.map((col: any) => col.name) ?? []
           if (aiConfig.autoRunScope.includes("UI.REFRESH")) {
             setSchema([compactJsonTablesArray])
-            setData(res)
+            if (res.length === 1 && !res[0]) {
+              setData([])
+            } else {
+              setData(res)
+            }
           }
 
           const jsonData = queryData2JSON(res, queryFields)
           return jsonData
         }
         if (checkSqlIsModifyTableData(querySql)) {
+          const columns = getCurrentColumns()
+          if (querySql.includes("UPDATE")) {
+            const rowIds = getCurrentRowIds()
+            refreshRows(rowIds, columns)
+          } else {
+            refreshRows(undefined, columns)
+          }
           refreshRows()
         }
+        return res
       }
     },
     [
-      refreshRows,
-      setData,
-      setSchema,
       sqlite,
-      tableSchema,
       updateTableSchema,
+      tableSchema,
       aiConfig.autoRunScope,
+      setSchema,
+      setData,
+      refreshRows,
+      getCurrentRowIds,
+      getCurrentColumns,
     ]
   )
 
@@ -164,15 +209,7 @@ export const useTable = (tableName: string, databaseName: string) => {
         updateTableSchema()
       })
     }
-  }, [
-    sqlite,
-    tableName,
-    updateTableSchema,
-    refreshRows,
-    tableSchema,
-    runQuery,
-    setData,
-  ])
+  }, [setData, sqlite, tableName, updateTableSchema])
 
   return {
     data,
diff --git a/lib/ai/openai.ts b/lib/ai/openai.ts
index 2172d3f0..0cc977dc 100644
--- a/lib/ai/openai.ts
+++ b/lib/ai/openai.ts
@@ -26,7 +26,6 @@ const svg = d3.select(_CANVAS_ID_).append("svg").attr("width", _CHART_WIDTH_).at
 `
 
 const baseSysPrompt = ``
-
 // this version is after 0613, openai support function call
 // const baseSysPrompt = `you're a database master, help use query database and generate d3.js chart if user want. must abide by the following rules:
 

From 2a0a57d4edd939242680b7146dddb394d4c99482 Mon Sep 17 00:00:00 2001
From: Mayne <i@gine.me>
Date: Wed, 14 Jun 2023 11:33:47 +0800
Subject: [PATCH 5/7] fix: md render set default lang

---
 app/[database]/ai-chat-message-prisma.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/app/[database]/ai-chat-message-prisma.tsx b/app/[database]/ai-chat-message-prisma.tsx
index c9870a1d..21f3f426 100644
--- a/app/[database]/ai-chat-message-prisma.tsx
+++ b/app/[database]/ai-chat-message-prisma.tsx
@@ -35,7 +35,7 @@ export const AIMessage = ({
       return <code className="language-none">{children}</code>
     },
     codeblock: function Code(props) {
-      const { lang, text } = props
+      const { lang='sql', text } = props
       const codeHtml = Prism.highlight(text, Prism.languages[lang], lang)
       // if it's a d3 codeblock, we need to render it differently
       if (lang === "js" && text.includes("d3.")) {

From 6230453b9e7614e148e390f33e8ba6083a55c78c Mon Sep 17 00:00:00 2001
From: Mayne <i@gine.me>
Date: Wed, 14 Jun 2023 11:46:28 +0800
Subject: [PATCH 6/7] feat(ai): handle api error

---
 app/[database]/ai-chat.tsx | 94 +++++++++++++++++++++-----------------
 1 file changed, 51 insertions(+), 43 deletions(-)

diff --git a/app/[database]/ai-chat.tsx b/app/[database]/ai-chat.tsx
index 02ea6e83..1432e55a 100644
--- a/app/[database]/ai-chat.tsx
+++ b/app/[database]/ai-chat.tsx
@@ -13,6 +13,7 @@ import { useAutoRunCode } from "@/hooks/use-auto-run-code"
 import { useSqliteStore } from "@/hooks/use-sqlite"
 import { Button } from "@/components/ui/button"
 import { Textarea } from "@/components/ui/textarea"
+import { toast } from "@/components/ui/use-toast"
 
 import { useConfigStore } from "../settings/store"
 import { AIChatMessage } from "./ai-chat-message"
@@ -55,52 +56,59 @@ export const AIChat = () => {
 
   const sendMessages = async (_messages: any) => {
     setLoading(true)
-    const response = await askAI(_messages, {
-      tableSchema: currentTableSchema,
-      allTables,
-      databaseName: database,
-    })
-
-    if (response?.finish_reason == "function_call") {
-      if (aiConfig.autoRunScope) {
-        const res = await handleOpenAIFunctionCall(
-          response.message!,
-          handleFunctionCall
-        )
-        if (res) {
-          const { name, resp } = res
-          const newMessages = [
-            ..._messages,
-            response.message,
-            {
-              role: "function",
-              name,
-              content: JSON.stringify(resp),
-            },
-          ]
-          const newResponse = await askAI(newMessages, {
-            tableSchema: currentTableSchema,
-            allTables,
-            databaseName: database,
-          })
-          const _newMessages = [
-            ...newMessages,
-            { role: "assistant", content: newResponse?.message?.content },
-          ]
-          console.log({ _newMessages })
-          setMessages(_newMessages as any)
+    try {
+      const response = await askAI(_messages, {
+        tableSchema: currentTableSchema,
+        allTables,
+        databaseName: database,
+      })
+
+      if (response?.finish_reason == "function_call") {
+        if (aiConfig.autoRunScope) {
+          const res = await handleOpenAIFunctionCall(
+            response.message!,
+            handleFunctionCall
+          )
+          if (res) {
+            const { name, resp } = res
+            const newMessages = [
+              ..._messages,
+              response.message,
+              {
+                role: "function",
+                name,
+                content: JSON.stringify(resp),
+              },
+            ]
+            const newResponse = await askAI(newMessages, {
+              tableSchema: currentTableSchema,
+              allTables,
+              databaseName: database,
+            })
+            const _newMessages = [
+              ...newMessages,
+              { role: "assistant", content: newResponse?.message?.content },
+            ]
+            console.log({ _newMessages })
+            setMessages(_newMessages as any)
+          }
         }
+      } else if (response?.message) {
+        const newMessages = [
+          ..._messages,
+          { role: "assistant", content: response?.message?.content },
+        ]
+        const thisMsgIndex = newMessages.length - 1
+        setMessages(newMessages)
       }
-    } else if (response?.message) {
-      const newMessages = [
-        ..._messages,
-        { role: "assistant", content: response?.message?.content },
-      ]
-      const thisMsgIndex = newMessages.length - 1
-      setMessages(newMessages)
+    } catch (error: any) {
+      toast({
+        title: "Error",
+        description: 'oops, something went wrong. please try again later.',
+      })
+    } finally {
+      setLoading(false)
     }
-
-    setLoading(false)
   }
 
   const handleSend = async () => {

From 2979f179be825ac65e582737ca8c4493b59b40a9 Mon Sep 17 00:00:00 2001
From: Mayne <i@gine.me>
Date: Wed, 14 Jun 2023 21:51:11 +0800
Subject: [PATCH 7/7] fix: z-index

---
 app/[database]/ai-chat.tsx     | 14 +++++++-------
 app/[database]/base-layout.tsx |  2 +-
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/app/[database]/ai-chat.tsx b/app/[database]/ai-chat.tsx
index 1432e55a..79055eec 100644
--- a/app/[database]/ai-chat.tsx
+++ b/app/[database]/ai-chat.tsx
@@ -1,19 +1,19 @@
 "use client"
 
 // for now it's under database page, maybe move to global later
-import { useCallback, useRef, useState } from "react"
+import { useKeyPress } from "ahooks"
+import { Loader2, Paintbrush } from "lucide-react"
 import Link from "next/link"
 import { useParams } from "next/navigation"
-import { useKeyPress, useWhyDidYouUpdate } from "ahooks"
-import { Loader2, Paintbrush } from "lucide-react"
+import { useCallback, useRef, useState } from "react"
 
-import { handleOpenAIFunctionCall } from "@/lib/ai/openai"
-import { useAI } from "@/hooks/use-ai"
-import { useAutoRunCode } from "@/hooks/use-auto-run-code"
-import { useSqliteStore } from "@/hooks/use-sqlite"
 import { Button } from "@/components/ui/button"
 import { Textarea } from "@/components/ui/textarea"
 import { toast } from "@/components/ui/use-toast"
+import { useAI } from "@/hooks/use-ai"
+import { useAutoRunCode } from "@/hooks/use-auto-run-code"
+import { useSqliteStore } from "@/hooks/use-sqlite"
+import { handleOpenAIFunctionCall } from "@/lib/ai/openai"
 
 import { useConfigStore } from "../settings/store"
 import { AIChatMessage } from "./ai-chat-message"
diff --git a/app/[database]/base-layout.tsx b/app/[database]/base-layout.tsx
index 96ee4aac..cb636144 100644
--- a/app/[database]/base-layout.tsx
+++ b/app/[database]/base-layout.tsx
@@ -72,7 +72,7 @@ export function DatabaseLayoutBase({
           <MobileSideBar />
           <Nav />
         </div>
-        <div className="flex h-[calc(100vh-4rem)] grow overflow-auto">
+        <div className="z-[1] flex h-[calc(100vh-4rem)] grow overflow-auto">
           <div className="grow">{children}</div>
         </div>
       </div>