toolcall typing & ingestion improvements

Author: Veikkosuhonen

src/client/components/ChatV2/Conversation.tsx

@@ -18,6 +18,7 @@ import { t } from 'i18next'
 import FormatQuoteIcon from '@mui/icons-material/FormatQuote'
 import useLocalStorageState from '../../hooks/useLocalStorageState'
 import { BlueButton } from './general/Buttons'
+import { ToolCallStatusEvent } from '../../../shared/chat'
 
 const UserMessage = ({ content, attachements }: { content: string; attachements?: string }) => (
   <Box
@@ -324,7 +325,7 @@ export const Conversation = ({
   expandedNodeHeight: number
   messages: Message[]
   completion: string
-  toolCalls: { id: string; name?: string }[]
+  toolCalls: { [callId: string]: ToolCallStatusEvent }
   isStreaming: boolean
   setActiveFileSearchResult: (data: FileSearchCompletedData) => void
   setShowFileSearchResults: (show: boolean) => void
@@ -371,7 +372,7 @@ export const Conversation = ({
               setShowFileSearchResults={setShowFileSearchResults}
             />
           ) : (
-            <LoadingMessage expandedNodeHeight={expandedNodeHeight} isFileSearching={toolCalls.length > 0} />
+            <LoadingMessage expandedNodeHeight={expandedNodeHeight} isFileSearching={Object.values(toolCalls).some((call) => !!call.result)} />
           ))}
       </Box>
       {!reminderSeen && !isStreaming && messages.length > 15 && (

src/client/components/ChatV2/useChatStream.ts

@@ -1,22 +1,9 @@
 import { useState } from 'react'
-import type { FileSearchCompletedData, ResponseStreamEventData } from '../../../shared/types'
+import type { FileSearchCompletedData } from '../../../shared/types'
 import type { Message } from '../../types'
-import { ChatToolDef } from '../../../shared/tools'
+import { ChatEvent, ToolCallStatusEvent } from '../../../shared/chat'
 
-type ToolCallState = { id: string; name: ChatToolDef['name'] } & (
-  | {
-      status: 'starting'
-    }
-  | {
-      input: ChatToolDef['input']
-      status: 'started'
-    }
-  | {
-      status: 'completed'
-      input?: ChatToolDef['input'] // Have to allow optional because it is possible "started" state was not reached before completed.
-      artifacts: ChatToolDef['artifacts']
-    }
-)
+type ToolCallState = ToolCallStatusEvent
 
 export const useChatStream = ({
   onFileSearchComplete,
@@ -31,7 +18,7 @@ export const useChatStream = ({
 }) => {
   const [completion, setCompletion] = useState('')
   const [isStreaming, setIsStreaming] = useState(false)
-  const [toolCalls, setToolCalls] = useState<ToolCallState[]>([])
+  const [toolCalls, setToolCalls] = useState<{ [callId: string]: ToolCallState }>({})
   const [streamController, setStreamController] = useState<AbortController>()
 
   const decoder = new TextDecoder()
@@ -55,7 +42,7 @@ export const useChatStream = ({
         for (const chunk of data.split('\n')) {
           if (!chunk || chunk.trim().length === 0) continue
 
-          let parsedChunk: ResponseStreamEventData | undefined
+          let parsedChunk: ChatEvent | undefined
           try {
             parsedChunk = JSON.parse(chunk)
           } catch (e: any) {
@@ -81,29 +68,8 @@ export const useChatStream = ({
               content += parsedChunk.text
               break
 
-            case 'toolCallStarting':
-              setToolCalls((prev) => [
-                ...prev,
-                {
-                  id: parsedChunk.id,
-                  name: parsedChunk.name,
-                  status: 'starting',
-                },
-              ])
-              break
-
-            case 'toolCallStarted':
-              setToolCalls((prev) =>
-                prev.map((call) =>
-                  call.id === parsedChunk.id && call.status === 'starting' ? { ...call, status: 'started', input: parsedChunk.input } : call,
-                ),
-              )
-              break
-
-            case 'toolCallCompleted':
-              setToolCalls((prev) =>
-                prev.map((call) => (call.id === parsedChunk.id ? { ...call, status: 'completed', artifacts: parsedChunk.artifacts } : call)),
-              )
+            case 'toolCallStatus':
+              setToolCalls((prev) => ({ ...prev, [parsedChunk.callId]: parsedChunk }))
               break
 
             case 'error':
@@ -125,7 +91,7 @@ export const useChatStream = ({
     } finally {
       setStreamController(undefined)
       setCompletion('')
-      setToolCalls([])
+      setToolCalls({})
       setIsStreaming(false)
 
       onComplete({

src/server/routes/ai/v3.ts

@@ -1,7 +1,6 @@
 import express from 'express'
 import { DEFAULT_TOKEN_LIMIT, FREE_MODEL, inProduction } from '../../../config'
 import type { ChatMessage } from '../../../shared/llmTypes'
-import type { ResponseStreamEventData } from '../../../shared/types'
 import { ChatInstance, Discussion, RagIndex, UserChatInstanceUsage } from '../../db/models'
 import { calculateUsage, checkCourseUsage, checkUsage, incrementCourseUsage, incrementUsage } from '../../services/chatInstances/usage'
 import { streamChat } from '../../services/langchain/chat'
@@ -15,6 +14,7 @@ import { upload } from './multer'
 import { PostStreamSchemaV3 } from './types'
 import { StructuredTool } from '@langchain/core/tools'
 import { getRagIndexSearchTool } from '../../services/rag/searchTool'
+import { ChatEvent } from '../../../shared/chat'
 
 const router = express.Router()
 
@@ -137,7 +137,7 @@ router.post('/stream', upload.single('file'), async (r, res) => {
     systemMessage: options.systemMessage,
     model: options.model,
     tools,
-    writeEvent: async (event: ResponseStreamEventData) => {
+    writeEvent: async (event: ChatEvent) => {
       await new Promise((resolve) => {
         const success = res.write(`${JSON.stringify(event)}\n`, (err) => {
           if (err) {

src/server/routes/rag/ragIndex.ts

@@ -181,7 +181,7 @@ const upload = multer({
     bucket: S3_BUCKET,
     acl: 'private',
     metadata: (req, file, cb) => {
-      cb(null, { fieldName: file.fieldName })
+      cb(null, { fieldName: file.fieldname })
     },
     key: (req, file, cb) => {
       const { ragIndex } = req as RagIndexRequest

src/server/services/langchain/chat.ts

@@ -1,22 +1,14 @@
-import { AzureChatOpenAI } from '@langchain/openai'
-import { AZURE_API_KEY, AZURE_RESOURCE } from '../../util/config'
-import { validModels } from '../../../config'
-import { ChatMessage, Message } from '../../../shared/llmTypes'
-import { Response } from 'express'
-import { AIMessageChunk, BaseMessage, BaseMessageLike } from '@langchain/core/messages'
-import { IterableReadableStream } from '@langchain/core/utils/stream'
-import { ResponseStreamEventData } from '../../../shared/types'
-import { Tiktoken } from '@dqbd/tiktoken'
-import { FakeStreamingChatModel } from '@langchain/core/utils/testing'
-import { MockModel } from './MockModel'
+import { BaseChatModel } from '@langchain/core/language_models/chat_models'
+import { AIMessageChunk, BaseMessageLike } from '@langchain/core/messages'
 import { StructuredTool } from '@langchain/core/tools'
 import { concat } from '@langchain/core/utils/stream'
-import { BaseChatModel } from '@langchain/core/language_models/chat_models'
-import { Runnable } from '@langchain/core/runnables'
-import logger from '../../util/logger'
-import { BaseLanguageModelInput } from '@langchain/core/language_models/base'
-import { ToolCall } from '@langchain/core/messages/tool'
+import { AzureChatOpenAI } from '@langchain/openai'
+import { validModels } from '../../../config'
+import { ChatEvent } from '../../../shared/chat'
+import { ChatMessage } from '../../../shared/llmTypes'
 import { ChatToolDef } from '../../../shared/tools'
+import { AZURE_API_KEY, AZURE_RESOURCE } from '../../util/config'
+import { MockModel } from './MockModel'
 
 const getChatModel = (model: string, tools: StructuredTool[]): BaseChatModel => {
   const deploymentName = validModels.find((m) => m.name === model)?.deployment
@@ -37,7 +29,7 @@ const getChatModel = (model: string, tools: StructuredTool[]): BaseChatModel =>
   }).bindTools(tools) as BaseChatModel
 }
 
-type WriteEventFunction = (data: ResponseStreamEventData) => Promise<void>
+type WriteEventFunction = (data: ChatEvent) => Promise<void>
 
 type ChatTool = StructuredTool<any, any, any, string>
 
@@ -115,9 +107,10 @@ const chatTurn = async (model: BaseChatModel, messages: BaseMessageLike[], tools
           status: 'pending',
         }
         await writeEvent({
-          type: 'toolCallStarting',
-          name: toolCall.name as ChatToolDef['name'],
-          id,
+          type: 'toolCallStatus',
+          toolName: toolCall.name as ChatToolDef['name'],
+          callId: id,
+          text: 'Starting',
         })
       }
     }
@@ -143,21 +136,25 @@ const chatTurn = async (model: BaseChatModel, messages: BaseMessageLike[], tools
     const input = toolCall.args as ChatToolDef['input']
     if (id && tool) {
       await writeEvent({
-        type: 'toolCallStarted',
-        name,
-        input,
-        id,
+        type: 'toolCallStatus',
+        toolName: name,
+        callId: id,
+        text: `Searching for '${input.query}'`,
       })
       const result = await tool.invoke(toolCall)
       messages.push(result)
       toolCallStatuses[id] = {
         status: 'completed',
       }
       await writeEvent({
-        type: 'toolCallCompleted',
-        name,
-        artifacts: result.artifact,
-        id,
+        type: 'toolCallStatus',
+        toolName: name,
+        callId: id,
+        text: 'Completed search',
+        result: {
+          artifacts: result.artifact,
+          input,
+        },
       })
     }
   }

src/server/services/rag/ingestion.ts

@@ -26,6 +26,7 @@ export const ingestRagFiles = async (ragIndex: RagIndex) => {
 
   const vectorStore = getRedisVectorStore(ragFiles[0].ragIndexId, ragIndex.metadata.language)
   const allDocuments: Document[] = []
+  const allEmbeddings: number[][] = []
 
   for (const ragFile of ragFiles) {
     console.time(`Ingestion ${ragFile.filename}`)
@@ -40,24 +41,27 @@ export const ingestRagFiles = async (ragIndex: RagIndex) => {
 
     const chunkDocuments = await splitter.splitDocuments([document])
 
-    chunkDocuments.forEach((chunkDocument, idx) => {
+    let idx = 0
+    for (const chunkDocument of chunkDocuments) {
       chunkDocument.id = `ragIndex-${ragFile.ragIndexId}-${ragFile.filename}-${idx}`
       chunkDocument.metadata = {
         ...chunkDocument.metadata,
         ragFileName: ragFile.filename,
       }
-    })
+      idx++
+    }
+
+    const embeddings = await vectorStore.embeddings.embedDocuments(chunkDocuments.map((d) => d.pageContent))
 
-    // console.log(await redisClient.ft.info(vectorStore.indexName))
     allDocuments.push(...chunkDocuments)
+    allEmbeddings.push(...embeddings)
 
     console.timeEnd(`Ingestion ${ragFile.filename}`)
-    // console.log(await redisClient.ft.info(vectorStore.indexName))
-    //
+
     ragFile.pipelineStage = 'completed'
     await ragFile.save()
   }
 
   // @todo we can only call this once. How to handle new documents?
-  await vectorStore.addDocuments(allDocuments)
+  await vectorStore.addVectors(allEmbeddings, allDocuments)
 }

src/shared/types.ts

@@ -52,39 +52,6 @@ export type FileSearchCompletedData = {
 
 export type FileSearchResultData = Record<string, unknown>
 
-export type ResponseStreamEventData =
-  | {
-      type: 'writing'
-      text: string
-    }
-  | {
-      type: 'complete'
-      prevResponseId: string // Not yet sure what to put here in v3
-    }
-  | ({
-      type: 'toolCallStarting'
-      id: string
-    } & Pick<ChatToolDef, 'name'>)
-  | ({
-      type: 'toolCallStarted'
-      id: string
-    } & Pick<ChatToolDef, 'name' | 'input'>)
-  | ({
-      type: 'toolCallCompleted'
-      id: string
-    } & Pick<ChatToolDef, 'name' | 'artifacts'>)
-  | {
-      type: 'error'
-      error: string
-    }
-
-export interface CourseAssistant {
-  course_id: string | null
-  name: string
-  assistant_instruction: string
-  vector_store_id: string | null
-}
-
 export type Locale = {
   fi?: string
   en?: string