Refactor code to handle chat completion and update chat messages and …

…storage in real-time

src/components/Sidebar/chat/chatInput.tsx

@@ -3,11 +3,10 @@ import TextareaAutosize from 'react-textarea-autosize'
 import { GiMagicBroom } from 'react-icons/gi'
 import { IoSend } from 'react-icons/io5'
 import { HiHand } from 'react-icons/hi'
-import { ChatMessageParams, ChatRole } from '../../../hooks/useOpenAI'
 
 interface SidebarInputProps {
   loading: boolean
-  submitMessage: (messages: ChatMessageParams[]) => void
+  submitMessage: (prompt: string) => void
   clearMessages: () => void
   chatIsEmpty: boolean
   cancelRequest: () => void
@@ -33,7 +32,7 @@ export function SidebarInput({
   }, [loading])
 
   const handleSubmit = () => {
-    submitMessage([{ content: text, role: ChatRole.USER }])
+    submitMessage(text)
     setText('')
   }
 

src/components/Sidebar/chat/index.tsx

@@ -1,21 +1,23 @@
-import React, { useEffect } from 'react'
+import { useEffect } from 'react'
 import ChatList from './chatList'
 import { SidebarInput } from './chatInput'
-import { useChatCompletion } from '../../../hooks/useOpenAI'
+import { useChatCompletion } from '../../../hooks/useChatCompletion'
 import { SYSTEM_PROMPT } from '../../../config/prompts'
 import { Settings } from '../../../config/settings'
 
 interface ChatProps {
   settings: Settings
+  chatId: string
 }
 
-const Chat = ({ settings }: ChatProps) => {
-  const { messages, submitQuery, clearMessages, loading, cancelRequest } =
+const Chat = ({ settings, chatId }: ChatProps) => {
+  const { messages, submitQuery, clearMessages, generating, cancelRequest } =
     useChatCompletion({
       model: settings.chat.modal,
       apiKey: settings.chat.openAIKey!,
       mode: settings.chat.mode,
       systemPrompt: SYSTEM_PROMPT,
+      chatId
     })
 
   useEffect(() => {
@@ -25,7 +27,7 @@ const Chat = ({ settings }: ChatProps) => {
         prompt: string
       }
       if (action === 'generate') {
-        submitQuery([{ content: prompt, role: 'user' }])
+        submitQuery(prompt)
       }
     }
     window.addEventListener('message', handleWindowMessage)
@@ -39,7 +41,7 @@ const Chat = ({ settings }: ChatProps) => {
     <>
       <ChatList messages={messages} />
       <SidebarInput
-        loading={loading}
+        loading={generating}
         submitMessage={submitQuery}
         chatIsEmpty={messages.length <= 1}
         clearMessages={clearMessages}

src/hooks/useChatCompletion.tsx

@@ -1,33 +1,68 @@
 import { AvailableModels, Mode } from "../config/settings"
 import { ChatOpenAI } from "langchain/chat_models/openai"
 import { useStorage } from "./useStorage"
+import { useCurrentChat, Role } from "./useCurrentChat"
+import { useMemo } from "react"
+import { AIMessage, HumanMessage, SystemMessage } from "langchain/dist/schema"
+import { useState } from "react"
 
 interface UseChatCompletionProps {
   model: AvailableModels
   apiKey: string
   mode: Mode
   systemPrompt: string
+  chatId: string
 }
 
 
-type Role = "user" | "assistant" | "system"
-
-type Message = {
-  role: Role
-  message: string
-  timestamp: number
-}
-
 export const useChatCompletion = ({
   model,
   apiKey,
   mode,
   systemPrompt,
+  chatId,
 }: UseChatCompletionProps) => {
-  const chat = new ChatOpenAI({
+  const { messages, updateAssistantMessage, updateStoredMessages, clearMessages } = useCurrentChat(chatId)
+  const [generating, setGenerating] = useState(false)
+
+  const controller = new AbortController();
+
+  const chat = useMemo(() => new ChatOpenAI({
     streaming: true,
-  });
+  }), [])
+
+
+
+  const submitQuery = async (query: string) => {
+    const previousMessages = messages.map((msg) => {
+      switch (msg.role) {
+        case Role.ASSISTANT:
+          return new AIMessage(msg.message)
+        case Role.SYSTEM:
+          return new SystemMessage(msg.message)
+        case Role.USER:
+          return new HumanMessage(msg.message)
+      }
+    })
+    setGenerating(true)
+    const response = await chat.call([...previousMessages, new HumanMessage(query)], {
+      signal: controller.signal,
+      callbacks: [
+        {
+          handleLLMNewToken: updateAssistantMessage
+        },
+      ],
+    });
+    setGenerating(false)
+    updateStoredMessages()
+    return response.content
+  }
+
+  const cancelRequest = () => {
+    controller.abort()
+    setGenerating(false)
+  }
 
-  const [messages, setMessages] = useStorage()
 
+  return { messages, submitQuery, generating, cancelRequest, clearMessages }
 }

src/hooks/useCurrentChat.tsx

@@ -0,0 +1,55 @@
+import { useEffect } from "react"
+import { useState } from "react"
+import { useStorage } from "./useStorage"
+
+
+export enum Role {
+  "USER",
+  "ASSISTANT",
+  "SYSTEM"
+}
+
+export type Message = {
+  role: Role
+  message: string
+  timestamp: number
+}
+
+export const useCurrentChat = (chatId: string) => {
+  const [storedMessages, setStoredMessages] = useStorage<Message[]>(`CHAT-${chatId}`, [])
+  const [messages, setMessages] = useState<Message[]>(storedMessages) // we don't directly update storedMessages for performance reasons
+
+  const updateAssistantMessage = (chunk: string) => {
+    setMessages(messages => {
+      const lastMessage = messages[messages.length - 1]
+      lastMessage.message += chunk
+      return [...messages]
+    })
+  }
+
+  const addNewMessage = (role: Role, message: string) => {
+    const newMessage: Message = {
+      role,
+      message,
+      timestamp: Date.now(),
+    }
+    setMessages([...messages, newMessage])
+  }
+
+  const updateStoredMessages = () => {
+    setStoredMessages(messages)
+  }
+
+  const clearMessages = () => {
+    setMessages([])
+    updateStoredMessages()
+  }
+
+  return {
+    messages,
+    updateAssistantMessage,
+    addNewMessage,
+    updateStoredMessages,
+    clearMessages
+  }
+}

src/hooks/useHistory.tsx

@@ -11,12 +11,12 @@ interface ChatHistory {
 export const useHistory = () => {
   const [history, setHistory] = useStorage<ChatHistory[]>("HISTORY", [])
 
-  const createChatHistory = () => {
+  const createChatHistory = (name: string) => {
     const newId = randomUUID()
 
     setHistory(prev => [...prev, {
       id: newId,
-      name: "",
+      name,
       createdAt: new Date().toISOString(),
       updatedAt: new Date().toISOString(),
     }])