update

Author: pashidlos

package-lock.json

@@ -52,6 +52,7 @@
     "ldapts": "^7.1.0",
     "looks-same": "^9.0.0",
     "odiff-bin": "^2.6.1",
+    "ollama": "^0.6.3",
     "passport": "^0.6.0",
     "passport-jwt": "^4.0.1",
     "passport-local": "^1.0.0",

package.json

@@ -50,9 +50,14 @@ export class OllamaController {
 
     return this.ollamaService.generate({
       model,
-      prompt,
+      messages: [
+        {
+          role: 'user',
+          content: prompt,
+          images: files.map((f) => new Uint8Array(f.buffer)),
+        },
+      ],
       format: 'json',
-      images: files.map((f) => f.buffer.toString('base64')),
       options: { temperature: Number(temperature) },
     });
   }

src/compare/libs/vlm/ollama.controller.ts

@@ -2,10 +2,28 @@ import { Test, TestingModule } from '@nestjs/testing';
 import { ConfigService } from '@nestjs/config';
 import { OllamaService } from './ollama.service';
 
+// Mock the ollama module
+const mockChat = jest.fn();
+const mockList = jest.fn();
+
+jest.mock('ollama', () => {
+  const MockOllama = jest.fn().mockImplementation(() => ({
+    chat: mockChat,
+    list: mockList,
+  }));
+  return {
+    Ollama: MockOllama,
+  };
+});
+
+
 describe('OllamaService', () => {
   let service: OllamaService;
 
   beforeEach(async () => {
+    // Reset mocks
+    jest.clearAllMocks();
+
     const module: TestingModule = await Test.createTestingModule({
       providers: [
         OllamaService,
@@ -22,39 +40,106 @@ describe('OllamaService', () => {
   });
 
   describe('generate', () => {
-    it('should call Ollama API with correct parameters', async () => {
-      const mockResponse = { response: 'YES', done: true };
-      globalThis.fetch = jest.fn().mockResolvedValue({
-        ok: true,
-        json: () => Promise.resolve(mockResponse),
-      });
+    it('should call Ollama SDK with correct parameters for Uint8Array', async () => {
+      const mockResponse = {
+        model: 'llava',
+        created_at: new Date(),
+        message: { content: 'YES', role: 'assistant' },
+        done: true,
+        done_reason: 'stop',
+        total_duration: 1000,
+        load_duration: 100,
+        prompt_eval_count: 10,
+        prompt_eval_duration: 200,
+        eval_count: 5,
+        eval_duration: 300,
+      };
+      mockChat.mockResolvedValue(mockResponse);
 
+      const testBytes = new Uint8Array([1, 2, 3, 4]);
       const result = await service.generate({
         model: 'llava',
-        prompt: 'Test prompt',
-        images: ['base64img'],
+        messages: [
+          {
+            role: 'user',
+            content: 'Test prompt',
+            images: [testBytes],
+          },
+        ],
       });
 
-      expect(fetch).toHaveBeenCalledWith(
-        'http://localhost:11434/api/generate',
-        expect.objectContaining({
-          method: 'POST',
-          headers: { 'Content-Type': 'application/json' },
-        })
-      );
-      expect(result).toEqual(mockResponse);
+      expect(mockChat).toHaveBeenCalledWith({
+        model: 'llava',
+        messages: [
+          {
+            role: 'user',
+            content: 'Test prompt',
+            images: [testBytes],
+          },
+        ],
+        stream: false,
+        format: undefined,
+        options: undefined,
+      });
+      expect(result.message.content).toBe('YES');
+      expect(result.done).toBe(true);
     });
 
-    it('should throw error when API returns non-ok status', async () => {
-      globalThis.fetch = jest.fn().mockResolvedValue({
-        ok: false,
-        status: 500,
-        text: () => Promise.resolve('Internal Server Error'),
+    it('should call Ollama SDK with correct parameters for base64 strings', async () => {
+      const mockResponse = {
+        model: 'llava',
+        created_at: new Date(),
+        message: { content: 'YES', role: 'assistant' },
+        done: true,
+        done_reason: 'stop',
+        total_duration: 1000,
+        load_duration: 100,
+        prompt_eval_count: 10,
+        prompt_eval_duration: 200,
+        eval_count: 5,
+        eval_duration: 300,
+      };
+      mockChat.mockResolvedValue(mockResponse);
+
+      // Use a longer base64 string
+      const longBase64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==';
+      const result = await service.generate({
+        model: 'llava',
+        messages: [
+          {
+            role: 'user',
+            content: 'Test prompt',
+            images: [longBase64], // base64 string - passed through as-is
+          },
+        ],
       });
 
-      await expect(service.generate({ model: 'llava', prompt: 'Test' })).rejects.toThrow(
-        'Ollama API returned status 500'
-      );
+      expect(mockChat).toHaveBeenCalledWith({
+        model: 'llava',
+        messages: [
+          {
+            role: 'user',
+            content: 'Test prompt',
+            images: [longBase64],
+          },
+        ],
+        stream: false,
+        format: undefined,
+        options: undefined,
+      });
+      expect(result.message.content).toBe('YES');
+      expect(result.done).toBe(true);
+    });
+
+    it('should throw error when SDK call fails', async () => {
+      mockChat.mockRejectedValue(new Error('Connection refused'));
+
+      await expect(
+        service.generate({
+          model: 'llava',
+          messages: [{ role: 'user', content: 'Test' }],
+        })
+      ).rejects.toThrow('Connection refused');
     });
 
     it('should throw error when OLLAMA_BASE_URL is not configured', async () => {
@@ -65,32 +150,68 @@ describe('OllamaService', () => {
       } as any;
       const newService = new OllamaService(mockConfigService);
 
-      await expect(newService.generate({ model: 'llava', prompt: 'Test' })).rejects.toThrow('OLLAMA_BASE_URL');
+      await expect(
+        newService.generate({
+          model: 'llava',
+          messages: [{ role: 'user', content: 'Test' }],
+        })
+      ).rejects.toThrow('OLLAMA_BASE_URL');
     });
   });
 
   describe('listModels', () => {
     it('should return list of models', async () => {
-      const mockModels = { models: [{ name: 'llava:7b' }, { name: 'moondream' }] };
-      globalThis.fetch = jest.fn().mockResolvedValue({
-        ok: true,
-        json: () => Promise.resolve(mockModels),
-      });
+      const mockDate = new Date('2024-01-01');
+      const mockResponse = {
+        models: [
+          {
+            name: 'llava:7b',
+            model: 'llava:7b',
+            size: 1000,
+            digest: 'abc123',
+            modified_at: mockDate,
+            expires_at: mockDate,
+            size_vram: 500,
+            details: {
+              parent_model: '',
+              format: 'gguf',
+              family: 'llama',
+              families: ['llama'],
+              parameter_size: '7B',
+              quantization_level: 'Q4_0',
+            },
+          },
+          {
+            name: 'moondream',
+            model: 'moondream',
+            size: 2000,
+            digest: 'def456',
+            modified_at: mockDate,
+            expires_at: mockDate,
+            size_vram: 1000,
+            details: {
+              parent_model: '',
+              format: 'gguf',
+              family: 'moondream',
+              families: ['moondream'],
+              parameter_size: '1.6B',
+              quantization_level: 'Q4_0',
+            },
+          },
+        ],
+      };
+      mockList.mockResolvedValue(mockResponse);
 
       const result = await service.listModels();
 
-      expect(fetch).toHaveBeenCalledWith('http://localhost:11434/api/tags');
-      expect(result).toEqual(mockModels.models);
+      expect(mockList).toHaveBeenCalled();
+      expect(result).toEqual(mockResponse.models);
     });
 
     it('should throw error when API fails', async () => {
-      globalThis.fetch = jest.fn().mockResolvedValue({
-        ok: false,
-        status: 503,
-        text: () => Promise.resolve('Service Unavailable'),
-      });
+      mockList.mockRejectedValue(new Error('Service Unavailable'));
 
-      await expect(service.listModels()).rejects.toThrow('Failed to list models');
+      await expect(service.listModels()).rejects.toThrow('Service Unavailable');
     });
   });
 });

src/compare/libs/vlm/ollama.service.spec.ts

@@ -1,54 +1,43 @@
 import { Injectable, Logger } from '@nestjs/common';
 import { ConfigService } from '@nestjs/config';
-import { OllamaGenerateRequest, OllamaGenerateResponse, OllamaModel, OllamaModelsResponse } from './ollama.types';
+import { Ollama, ChatRequest, ChatResponse, ListResponse, ModelResponse } from 'ollama';
 
 @Injectable()
 export class OllamaService {
   private readonly logger: Logger = new Logger(OllamaService.name);
-  private baseUrl: string | null = null;
+  private ollamaClient: Ollama | null = null;
 
   constructor(private readonly configService: ConfigService) {}
 
-  private getBaseUrl(): string {
-    if (!this.baseUrl) {
-      this.baseUrl = this.configService.getOrThrow<string>('OLLAMA_BASE_URL');
+  private getOllamaClient(): Ollama {
+    if (!this.ollamaClient) {
+      const baseUrl = this.configService.getOrThrow<string>('OLLAMA_BASE_URL');
+      this.ollamaClient = new Ollama({ host: baseUrl });
     }
-    return this.baseUrl;
+    return this.ollamaClient;
   }
 
-  async generate(request: OllamaGenerateRequest): Promise<OllamaGenerateResponse> {
-    const baseUrl = this.getBaseUrl();
+  async generate(request: ChatRequest): Promise<ChatResponse> {
+    const client = this.getOllamaClient();
+
     try {
-      const response = await fetch(`${baseUrl}/api/generate`, {
-        method: 'POST',
-        headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({ ...request, stream: request.stream ?? false }),
+      const response = await client.chat({
+        ...request,
+        stream: false,
       });
 
-      if (!response.ok) {
-        const errorText = await response.text();
-        throw new Error(`Ollama API returned status ${response.status}: ${errorText}`);
-      }
-
-      return await response.json();
+      return response;
     } catch (error) {
       this.logger.error(`Ollama generate request failed: ${error.message}`);
       throw error;
     }
   }
 
-  async listModels(): Promise<OllamaModel[]> {
-    const baseUrl = this.getBaseUrl();
+  async listModels(): Promise<ModelResponse[]> {
+    const client = this.getOllamaClient();
     try {
-      const response = await fetch(`${baseUrl}/api/tags`);
-
-      if (!response.ok) {
-        const errorText = await response.text();
-        throw new Error(`Failed to list models: ${response.status} ${errorText}`);
-      }
-
-      const data: OllamaModelsResponse = await response.json();
-      return data.models;
+      const response: ListResponse = await client.list();
+      return response.models;
     } catch (error) {
       this.logger.error(`Failed to list models: ${error.message}`);
       throw error;