implement mock model tool calling so rag tests work again

Author: Veikkosuhonen

e2e/courseChatRag.spec.ts

@@ -20,7 +20,7 @@ test.describe('Course Chat v2', () => {
     await expect(page.getByTestId('assistant-message')).toContainText('You are calling mock endpoint for streaming mock data')
   })
 
-  /* test('Course chat RAG feature', async ({ page }) => {
+  test.only('Course chat RAG feature', async ({ page }) => {
     const ragName = `rag-${test.info().workerIndex}`
     await page.locator('#rag-index-selector').first().click()
     await page.getByRole('menuitem', { name: ragName }).click()
@@ -30,10 +30,10 @@ test.describe('Course Chat v2', () => {
     await chatInput.press('Shift+Enter')
 
     // Shows file search loading indicator
-    await expect(page.getByTestId('file-searching-message')).toBeVisible()
+    await expect(page.getByTestId('tool-call-message')).toBeVisible()
 
-    // Responds with RAG mock text
-    await expect(page.getByTestId('assistant-message')).toContainText('This is a mock response for file search stream.')
+    // Responds with RAG mock document text
+    await expect(page.getByTestId('assistant-message')).toContainText('This is the first mock document')
 
     // Source button is visible
     await expect(page.getByTestId('file-search-sources')).toBeVisible()
@@ -43,5 +43,5 @@ test.describe('Course Chat v2', () => {
 
     // Three source items should be visible
     await expect(page.getByTestId('sources-truncated-item')).toHaveCount(3)
-    }) */
+  })
 })

src/server/routes/ai/v3.ts

@@ -15,6 +15,7 @@ import { PostStreamSchemaV3 } from './types'
 import { StructuredTool } from '@langchain/core/tools'
 import { getRagIndexSearchTool } from '../../services/rag/searchTool'
 import { ChatEvent } from '../../../shared/chat'
+import { getMockRagIndexSearchTool } from '../../services/rag/mockSearchTool'
 
 const router = express.Router()
 
@@ -126,7 +127,11 @@ router.post('/stream', upload.single('file'), async (r, res) => {
       return
     }
 
-    tools.push(getRagIndexSearchTool(ragIndex))
+    const searchTool = model === 'mock' ? getMockRagIndexSearchTool(ragIndex) : getRagIndexSearchTool(ragIndex)
+
+    console.log('Tool given: ' + searchTool.name)
+
+    tools.push(searchTool)
   }
 
   // Prepare for streaming response

src/server/routes/testUtils.ts

@@ -3,11 +3,9 @@ import { inProduction } from '../../config'
 import { getTestUserHeaders, TEST_COURSES } from '../../shared/testData'
 import { ChatInstanceRagIndex, Enrolment, Prompt, RagIndex, User, UserChatInstanceUsage } from '../db/models'
 import { headersToUser } from '../middleware/user'
-import type { RequestWithUser } from '../types'
 import { ApplicationError } from '../util/ApplicationError'
 import { getCompletionEvents } from '../util/azure/client'
 import logger from '../util/logger'
-import getEncoding from '../util/tiktoken'
 
 const router = Router()
 

src/server/services/langchain/MockModel.ts

@@ -1,7 +1,6 @@
-import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager'
-import { isSystemMessage } from '@langchain/core/messages'
-import { BaseMessage } from '@langchain/core/messages'
-import { ChatGenerationChunk, ChatResult } from '@langchain/core/outputs'
+import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager'
+import { AIMessage, AIMessageChunk, type BaseMessage, isHumanMessage, isSystemMessage, isToolMessage } from '@langchain/core/messages'
+import type { ChatGenerationChunk, ChatResult } from '@langchain/core/outputs'
 import { FakeStreamingChatModel } from '@langchain/core/utils/testing'
 import { basicTestContent } from '../../util/azure/mocks/mockContent'
 
@@ -17,13 +16,26 @@ export class MockModel extends FakeStreamingChatModel {
     })
   }
 
-  async _generate(messages: BaseMessage[], _options: this['ParsedCallOptions'], _runManager?: CallbackManagerForLLMRun): Promise<ChatResult> {
-    const firstMessage = messages[0]
-    if (isSystemMessage(firstMessage) && (firstMessage.content as string).startsWith('mocktest')) {
+  setupTestResponse(messages: BaseMessage[]) {
+    const firstSystemMessage = messages.find(isSystemMessage)
+    const lastHumanMessage = messages.findLast(isHumanMessage)
+    const toolMessage = isToolMessage(messages[messages.length - 1]) ? messages[messages.length - 1] : null
+
+    if (toolMessage) {
+      this.chunks = [new AIMessageChunk(`Ok! Got some great results from that mock tool call!: "${toolMessage.content}"`)]
+    } else if (firstSystemMessage && (firstSystemMessage.content as string).startsWith('mocktest')) {
+      // testing a system message
       // Do nothing. FakeStreamingChatModel echoes the first message.
+    } else if (((lastHumanMessage?.content ?? '') as string).startsWith('rag')) {
+      // Do a tool call
+      this.chunks = toolCallChunks
     } else {
-      firstMessage.content = basicTestContent
+      this.responses = defaultResponse
     }
+  }
+
+  async _generate(messages: BaseMessage[], _options: this['ParsedCallOptions'], _runManager?: CallbackManagerForLLMRun): Promise<ChatResult> {
+    this.setupTestResponse(messages)
     return super._generate(messages, _options, _runManager)
   }
 
@@ -32,12 +44,22 @@ export class MockModel extends FakeStreamingChatModel {
     _options: this['ParsedCallOptions'],
     runManager?: CallbackManagerForLLMRun,
   ): AsyncGenerator<ChatGenerationChunk> {
-    const firstMessage = messages[0]
-    if (isSystemMessage(firstMessage) && (firstMessage.content as string).startsWith('mocktest')) {
-      // Do nothing. FakeStreamingChatModel echoes the first message.
-    } else {
-      firstMessage.content = basicTestContent
-    }
+    this.setupTestResponse(messages)
     yield* super._streamResponseChunks(messages, _options, runManager)
   }
 }
+
+const defaultResponse = [new AIMessage(basicTestContent)]
+
+const toolCallChunks = [
+  new AIMessageChunk({
+    content: '',
+    tool_call_chunks: [
+      {
+        name: 'mock_document_search',
+        args: JSON.stringify({ query: 'mock test query' }),
+        id: 'mock_document_search_id',
+      },
+    ],
+  }),
+]

src/server/services/langchain/chat.ts

@@ -18,17 +18,20 @@ const getChatModel = (model: string, tools: StructuredTool[]): BaseChatModel =>
     throw new Error(`Invalid model: ${model}`)
   }
 
-  if (deploymentName === 'mock') {
-    return new MockModel()
-  }
+  const chatModel =
+    deploymentName === 'mock'
+      ? new MockModel()
+      : new AzureChatOpenAI({
+          model,
+          azureOpenAIApiKey: AZURE_API_KEY,
+          azureOpenAIApiVersion: '2023-05-15',
+          azureOpenAIApiDeploymentName: deploymentName,
+          azureOpenAIApiInstanceName: AZURE_RESOURCE,
+        })
+
+  chatModel.bindTools(tools)
 
-  return new AzureChatOpenAI({
-    model,
-    azureOpenAIApiKey: AZURE_API_KEY,
-    azureOpenAIApiVersion: '2023-05-15',
-    azureOpenAIApiDeploymentName: deploymentName,
-    azureOpenAIApiInstanceName: AZURE_RESOURCE,
-  }).bindTools(tools) as BaseChatModel
+  return chatModel
 }
 
 type WriteEventFunction = (data: ChatEvent) => Promise<void>

src/server/services/rag/mockSearchTool.ts

@@ -0,0 +1,44 @@
+import { tool } from '@langchain/core/tools'
+import { z } from 'zod/v4'
+import type { RagIndex } from '../../db/models'
+import { SearchSchema } from '../../../shared/rag'
+import { Document } from '@langchain/core/documents'
+import type { search } from './search'
+import type { getRagIndexSearchTool } from './searchTool'
+
+const mockDocuments = [
+  new Document({ pageContent: 'This is the first mock document.', metadata: { ragFileName: 'mock_document1.pdf' } }),
+  new Document({ pageContent: 'This is the second mock document.', metadata: { ragFileName: 'mock_document2.pdf' } }),
+  new Document({ pageContent: 'This is the third mock document.', metadata: { ragFileName: 'mock_document3.pdf' } }),
+]
+
+const mockSearch: typeof search = async (_index: RagIndex, _params: { query: string }) => {
+  await new Promise((resolve) => setTimeout(resolve, 300))
+
+  return {
+    results: mockDocuments.map((doc) => ({
+      id: doc.id,
+      content: doc.pageContent,
+      metadata: doc.metadata,
+    })),
+    timings: { search: 1000 },
+  }
+}
+
+export const getMockRagIndexSearchTool: typeof getRagIndexSearchTool = (ragIndex: RagIndex) =>
+  tool(
+    async ({ query }: { query: string }) => {
+      console.log('Mock search tool invoked with query:', query)
+      const { results: documents } = await mockSearch(ragIndex, SearchSchema.parse({ query }))
+      // With responseFormat: content_and_artifact, return content and artifact like this:
+      return [documents.map((doc) => doc.content).join('\n\n'), documents]
+    },
+    {
+      name: `mock_document_search`, // Gotcha: function name must match '^[a-zA-Z0-9_\.-]+$' at least in AzureOpenAI. This name must satisfy the name in ChatToolDef type
+      description: `Search documents in the materials (titled '${ragIndex.metadata.name}'). Prefer ${ragIndex.metadata.language}, which is the language used in the documents.`,
+      schema: z.object({
+        query: z.string().describe('the query to search for'),
+      }),
+      responseFormat: 'content_and_artifact',
+    },
+  )

src/shared/tools.ts

@@ -1,7 +1,7 @@
 import type { RagChunk } from './rag'
 
 export type ChatToolDef = {
-  name: 'document_search'
+  name: 'document_search' | 'mock_document_search'
   input: { query: string }
   result: { files: { fileName: string; score?: number }[] }
   output: RagChunk[]