fix ollama chat dropping tool data

Author: grypez

packages/kernel-language-model-service/src/ollama/base.test.ts

@@ -1,4 +1,4 @@
-import type { GenerateResponse, ListResponse } from 'ollama';
+import type { ChatResponse, GenerateResponse, ListResponse } from 'ollama';
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 
 import type { OllamaClient, OllamaModelOptions } from './types.ts';
@@ -15,6 +15,7 @@ describe('OllamaBaseService', () => {
     mockClient = {
       list: vi.fn(),
       generate: vi.fn(),
+      chat: vi.fn(),
     };
 
     mockMakeClient = vi.fn().mockResolvedValue(mockClient);
@@ -134,6 +135,111 @@ describe('OllamaBaseService', () => {
     });
   });
 
+  describe('chat', () => {
+    const makeChatResponse = (
+      overrides: Partial<ChatResponse> = {},
+    ): ChatResponse =>
+      ({
+        model: 'llama3',
+        message: { role: 'assistant', content: 'hello' },
+        done: true,
+        done_reason: 'stop',
+        prompt_eval_count: 5,
+        eval_count: 3,
+        ...overrides,
+      }) as ChatResponse;
+
+    it('forwards tools to the Ollama client', async () => {
+      vi.mocked(mockClient.chat).mockResolvedValue(makeChatResponse());
+
+      await service.chat({
+        model: 'llama3',
+        messages: [{ role: 'user', content: 'hi' }],
+        tools: [
+          {
+            type: 'function',
+            function: { name: 'get_time', description: 'Returns the time' },
+          },
+        ],
+      });
+
+      expect(mockClient.chat).toHaveBeenCalledWith(
+        expect.objectContaining({
+          tools: [
+            {
+              type: 'function',
+              function: { name: 'get_time', description: 'Returns the time' },
+            },
+          ],
+        }),
+      );
+    });
+
+    it('forwards message tool_calls with parsed arguments', async () => {
+      vi.mocked(mockClient.chat).mockResolvedValue(makeChatResponse());
+
+      await service.chat({
+        model: 'llama3',
+        messages: [
+          {
+            role: 'assistant',
+            content: null,
+            tool_calls: [
+              {
+                id: 'tc-1',
+                type: 'function',
+                function: { name: 'get_time', arguments: '{"tz":"UTC"}' },
+              },
+            ],
+          },
+        ],
+      });
+
+      expect(mockClient.chat).toHaveBeenCalledWith(
+        expect.objectContaining({
+          messages: [
+            expect.objectContaining({
+              tool_calls: [
+                { function: { name: 'get_time', arguments: { tz: 'UTC' } } },
+              ],
+            }),
+          ],
+        }),
+      );
+    });
+
+    it('maps response tool_calls back with stringified arguments', async () => {
+      vi.mocked(mockClient.chat).mockResolvedValue(
+        makeChatResponse({
+          message: {
+            role: 'assistant',
+            content: '',
+            tool_calls: [
+              { function: { name: 'get_time', arguments: { tz: 'UTC' } } },
+            ],
+          },
+        }),
+      );
+
+      const result = await service.chat({
+        model: 'llama3',
+        messages: [{ role: 'user', content: 'what time is it?' }],
+      });
+
+      expect(result.choices[0]?.message).toStrictEqual({
+        role: 'assistant',
+        content: '',
+        tool_calls: [
+          {
+            id: 'tool-0',
+            type: 'function',
+            function: { name: 'get_time', arguments: '{"tz":"UTC"}' },
+          },
+        ],
+      });
+    });
+  });
+
   describe('makeInstance', () => {
     it('should create instance with model', async () => {
       const config: InstanceConfig<OllamaModelOptions> = {

packages/kernel-language-model-service/src/ollama/base.ts

@@ -19,6 +19,7 @@ import type {
   OllamaModel,
   OllamaClient,
   OllamaModelOptions,
+  OllamaTool,
 } from './types.ts';
 
 /**
@@ -73,18 +74,30 @@ export class OllamaBaseService<Ollama extends OllamaClient>
    * @returns A hardened chat result.
    */
   async chat(params: ChatParams): Promise<ChatResult> {
-    const { model, messages, temperature, seed, stop } = params;
+    const { model, messages, tools, temperature, seed, stop } = params;
     const ollama = await this.#makeClient();
     let stopArr: string[] | undefined;
     if (stop !== undefined) {
       stopArr = Array.isArray(stop) ? stop : [stop];
     }
     const response = await ollama.chat({
       model,
-      messages: messages.map(({ role, content }) => ({
+      // eslint-disable-next-line camelcase
+      messages: messages.map(({ role, content, tool_calls }) => ({
         role,
         content: content ?? '',
+        // eslint-disable-next-line camelcase
+        ...(tool_calls && {
+          // eslint-disable-next-line camelcase
+          tool_calls: tool_calls.map(({ function: fn }) => ({
+            function: {
+              name: fn.name,
+              arguments: JSON.parse(fn.arguments) as Record<string, unknown>,
+            },
+          })),
+        }),
       })),
+      ...(tools && { tools: tools as OllamaTool[] }),
       stream: false,
       options: ifDefined({
         temperature,
@@ -96,6 +109,7 @@ export class OllamaBaseService<Ollama extends OllamaClient>
     });
     const promptTokens = response.prompt_eval_count ?? 0;
     const completionTokens = response.eval_count ?? 0;
+    const { tool_calls: responseToolCalls } = response.message;
     return harden({
       id: 'ollama-chat',
       model: response.model,
@@ -104,6 +118,16 @@ export class OllamaBaseService<Ollama extends OllamaClient>
           message: {
             role: response.message.role as ChatRole,
             content: response.message.content,
+            ...(responseToolCalls && {
+              tool_calls: responseToolCalls.map((tc, index) => ({
+                id: `tool-${index}`,
+                type: 'function' as const,
+                function: {
+                  name: tc.function.name,
+                  arguments: JSON.stringify(tc.function.arguments),
+                },
+              })),
+            }),
           },
           index: 0,
           finish_reason: response.done_reason ?? 'stop',

packages/kernel-language-model-service/src/ollama/types.ts

@@ -8,6 +8,7 @@ import type {
   Config,
   ChatRequest,
   ChatResponse,
+  Tool as OllamaTool,
 } from 'ollama';
 
 import type { LanguageModel } from '../types.ts';
@@ -30,6 +31,7 @@ export type {
   GenerateRequest,
   GenerateResponse,
   OllamaClient,
+  OllamaTool,
   ChatRequest,
   ChatResponse,
 };