Feta/ai upgrade

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"
@@ -17,7 +17,7 @@ export const AIMessage = ({
   msgIndex,
 }: {
   msgIndex: number
-  message: string
+  message?: string
   onRun: (props: {
     code: string
     lang: string
@@ -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.")) {
@@ -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>
   )

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>
+  )
+}

app/[database]/ai-chat.tsx

@@ -1,79 +1,123 @@
 "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, useSize } from "ahooks"
-import { Bot, Loader2, Paintbrush, User } from "lucide-react"
+import { useCallback, useRef, useState } from "react"
 
-import { useAI } from "@/hooks/use-ai"
-import { useAutoRunCode } from "@/hooks/use-auto-run-code"
 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 { AIMessage } from "./ai-chat-message-prisma"
-import { useTableChange } from "./hook"
+import { AIChatMessage } from "./ai-chat-message"
 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 { database } = 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"
-      content: string
-    }[]
-  >([])
+  const { aiMessages: messages, setAiMessages: setMessages } =
+    useDatabaseAppStore()
   const { allTables } = useSqliteStore()
+
   const cleanMessages = useCallback(() => {
     setMessages([])
-  }, [])
+  }, [setMessages])
 
   const textInputRef = useRef<HTMLTextAreaElement>()
 
   useKeyPress("ctrl.forwardslash", () => {
     textInputRef.current?.focus()
   })
 
-  useTableChange(cleanMessages)
+  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 sendMessages = async (_messages: any) => {
+    setLoading(true)
+    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)
+      }
+    } catch (error: any) {
+      toast({
+        title: "Error",
+        description: 'oops, something went wrong. please try again later.',
+      })
+    } finally {
+      setLoading(false)
+    }
+  }
 
   const handleSend = async () => {
     if (loading) return
     if (!input.trim().length) return
-    setLoading(true)
     const _messages: any = [...messages, { role: "user", content: input }]
     setMessages(_messages)
     setInput("")
-    const response = await askAI(_messages, {
-      tableSchema: currentTableSchema,
-      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)
-    }
-    setLoading(false)
+    await sendMessages(_messages)
   }
 
   const handleEnter = async (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
@@ -97,28 +141,20 @@ 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}
-                />
-              </>
-            ) : (
-              <>
-                <User className="h-4 w-4 shrink-0" />
-                <p className="grow">{message.content}</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>

app/[database]/base-layout.tsx

@@ -72,13 +72,13 @@ 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>
       <div
         className={cn(
-          " h-full  lg:border-l",
+          " h-screen  lg:border-l",
           isAiOpen ? "col-span-3 hidden md:block" : "hidden"
         )}
       >

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()
 

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 }),