obsidian_ollama/tests/ollama-client-cache.test.ts

// tests/ollama-client-cache.test.ts

import { OllamaMessage, OllamaTool, CacheConfig } from '../src/types';

// Create mock functions for the cache service
const mockInitialize = jest.fn().mockResolvedValue(undefined);
const mockGetCache = jest.fn().mockResolvedValue(null);
const mockSetCache = jest.fn().mockResolvedValue(undefined);
const mockClearCache = jest.fn().mockResolvedValue(undefined);

// Mock the semantic cache service BEFORE importing OllamaClient
jest.mock('../src/semantic-cache', () => ({
  SemanticCacheService: jest.fn().mockImplementation(() => ({
    initialize: mockInitialize,
    getCache: mockGetCache,
    setCache: mockSetCache,
    clearCache: mockClearCache,
  })),
}));

// Mock fetch globally
jest.spyOn(global, 'fetch').mockImplementation(jest.fn());

// Import OllamaClient AFTER mocking
import { OllamaClient } from '../src/ollama-client';
import { SemanticCacheService } from '../src/semantic-cache';

describe('OllamaClient with Semantic Cache', () => {
  let client: OllamaClient;
  let mockFetch: jest.Mock;

  function createMockReader(data: string) {
    const encoder = new TextEncoder();
    const encoded = encoder.encode(data);
    let called = false;
    return {
      read: () => {
        if (!called) {
          called = true;
          return Promise.resolve({ done: false, value: encoded });
        }
        return Promise.resolve({ done: true, value: new Uint8Array(0) });
      },
      releaseLock: jest.fn(),
    };
  }

  const mockMessages: OllamaMessage[] = [
    { role: 'system', content: 'You are helpful.' },
    { role: 'user', content: 'What is AI?' },
  ];

  const mockTools: OllamaTool[] = [
    {
      type: 'function',
      function: {
        name: 'test_tool',
        description: 'A test tool',
        parameters: {
          type: 'object',
          properties: { input: { type: 'string' } },
          required: ['input'],
        },
      },
    },
  ];

  const cacheConfig: CacheConfig = {
    enabled: true,
    similarityThreshold: 0.85,
    collectionName: 'test_cache',
    embeddingModel: 'nomic-embed-text',
    chromaURL: 'http://localhost:8000',
  };

  beforeEach(() => {
    jest.clearAllMocks();
    mockFetch = global.fetch as jest.Mock;

    client = new OllamaClient('http://localhost:11434', 'llama3', mockFetch, cacheConfig);
  });

  afterEach(() => {
    jest.clearAllMocks();
  });

  describe('initializeCache', () => {
    it('should call initialize on cache service when enabled', async () => {
      await client.initializeCache();

      // The SemanticCacheService mock was instantiated in constructor
      expect(SemanticCacheService).toHaveBeenCalledWith('http://localhost:11434', cacheConfig);
    });

    it('should not create cache service when disabled', () => {
      const disabledConfig: CacheConfig = { ...cacheConfig, enabled: false };
      new OllamaClient('http://localhost:11434', 'llama3', mockFetch, disabledConfig);

      // Constructor should not have created a cache service
      expect(SemanticCacheService).not.toHaveBeenCalledWith(
        'http://localhost:11434',
        disabledConfig
      );
    });

    it('should not create cache service when no config provided', () => {
      jest.clearAllMocks(); // Reset the call recorded by beforeEach before checking
      new OllamaClient('http://localhost:11434', 'llama3', mockFetch);

      expect(SemanticCacheService).not.toHaveBeenCalled();
    });
  });

  describe('chat (non-streaming) with cache', () => {
    it('should return cached response when available', async () => {
      // Setup cache hit
      const cachedResponse = 'Cached AI definition';
      mockGetCache.mockResolvedValueOnce(cachedResponse);

      const result = await client.chat(mockMessages);

      expect(result.content).toBe(cachedResponse);
      expect(mockFetch).not.toHaveBeenCalled();
    });

    it('should call LLM and cache response on cache miss', async () => {
      // Setup cache miss
      mockGetCache.mockResolvedValueOnce(null);

      const mockResponse = {
        ok: true,
        json: () =>
          Promise.resolve({
            message: {
              content: 'AI is the simulation of intelligence.',
            },
          }),
      };
      mockFetch.mockResolvedValueOnce(mockResponse);

      const result = await client.chat(mockMessages);

      expect(result.content).toBe('AI is the simulation of intelligence.');
      expect(mockFetch).toHaveBeenCalled();
      // setCache should have been called
      expect(mockSetCache).toHaveBeenCalled();
    });

    it('should bypass cache when tools are present', async () => {
      const mockResponse = {
        ok: true,
        json: () =>
          Promise.resolve({
            message: {
              content: 'Tool response',
            },
          }),
      };
      mockFetch.mockResolvedValueOnce(mockResponse);

      const result = await client.chat(mockMessages, mockTools);

      expect(mockFetch).toHaveBeenCalled();
      expect(mockGetCache).not.toHaveBeenCalled();
      expect(mockSetCache).not.toHaveBeenCalled();
    });

    it('should not cache failed responses', async () => {
      mockGetCache.mockResolvedValueOnce(null);
      mockFetch.mockResolvedValueOnce({ ok: false, status: 500 });

      await expect(client.chat(mockMessages)).rejects.toThrow();
      expect(mockSetCache).not.toHaveBeenCalled();
    });

    it('should cache successful responses after LLM call', async () => {
      mockGetCache.mockResolvedValueOnce(null);

      const mockResponse = {
        ok: true,
        json: () =>
          Promise.resolve({
            message: {
              content: 'Fresh response from LLM',
            },
          }),
      };
      mockFetch.mockResolvedValueOnce(mockResponse);

      await client.chat(mockMessages);

      expect(mockSetCache).toHaveBeenCalledWith('What is AI?', 'Fresh response from LLM');
    });
  });

  describe('streamChat with cache', () => {
    it('should return cached response when available', async () => {
      const cachedResponse = 'Cached streaming response';
      mockGetCache.mockResolvedValueOnce(cachedResponse);

      const chunks: OllamaMessage[] = [];
      for await (const chunk of client.streamChat(mockMessages)) {
        chunks.push(chunk);
      }

      expect(chunks.length).toBe(1);
      expect(chunks[0].content).toBe(cachedResponse);
      expect(mockFetch).not.toHaveBeenCalled();
    });

    it('should stream from LLM and cache on cache miss', async () => {
      mockGetCache.mockResolvedValueOnce(null);

      const streamData = [
        JSON.stringify({ message: { content: 'AI' } }),
        '\n',
        JSON.stringify({ message: { content: ' is' } }),
        '\n',
        JSON.stringify({ message: { content: ' cool' } }),
        '\n',
      ].join('');

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValueOnce({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: () => 'application/x-ndjson',
        },
      });

      const chunks: OllamaMessage[] = [];
      for await (const chunk of client.streamChat(mockMessages)) {
        chunks.push(chunk);
      }

      expect(chunks.length).toBe(3);
      expect(mockFetch).toHaveBeenCalled();
      expect(mockSetCache).toHaveBeenCalled();
    });

    it('should cache combined stream content on miss', async () => {
      mockGetCache.mockResolvedValueOnce(null);

      const streamData = [
        JSON.stringify({ message: { content: 'Hello' } }),
        '\n',
        JSON.stringify({ message: { content: ' world' } }),
        '\n',
      ].join('');

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValueOnce({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: () => 'application/x-ndjson',
        },
      });

      const chunks: OllamaMessage[] = [];
      for await (const chunk of client.streamChat(mockMessages)) {
        chunks.push(chunk);
      }

      expect(mockSetCache).toHaveBeenCalledWith('What is AI?', 'Hello world');
    });

    it('should bypass cache when tools are present for streaming', async () => {
      const streamData = [
        JSON.stringify({ message: { content: 'Tool call response' } }),
        '\n',
      ].join('');

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValueOnce({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: () => 'application/x-ndjson',
        },
      });

      const chunks: OllamaMessage[] = [];
      for await (const chunk of client.streamChat(mockMessages, mockTools)) {
        chunks.push(chunk);
      }

      expect(mockFetch).toHaveBeenCalled();
      expect(mockGetCache).not.toHaveBeenCalled();
    });

    it('should not cache when stream fails', async () => {
      mockGetCache.mockResolvedValueOnce(null);
      mockFetch.mockResolvedValueOnce({ ok: false, status: 500 });

      try {
        for await (const _ of client.streamChat(mockMessages)) {
          // Should throw
        }
      } catch (e) {
        // Expected to throw
      }

      expect(mockSetCache).not.toHaveBeenCalled();
    });
  });

  describe('streamChatAsPromise with cache', () => {
    it('should return cached response when available', async () => {
      const cachedResponse = 'Cached response via promise';
      mockGetCache.mockResolvedValueOnce(cachedResponse);

      const chunks = await client.streamChatAsPromise(mockMessages);

      expect(chunks.length).toBe(1);
      expect(chunks[0].content).toBe(cachedResponse);
      expect(mockFetch).not.toHaveBeenCalled();
    });

    it('should stream and cache on miss', async () => {
      mockGetCache.mockResolvedValueOnce(null);

      const streamData = [
        JSON.stringify({ message: { content: 'Full' } }),
        '\n',
        JSON.stringify({ message: { content: ' response' } }),
        '\n',
      ].join('');

      const mockReader = createMockReader(streamData);

      mockFetch.mockResolvedValueOnce({
        ok: true,
        body: { getReader: () => mockReader },
        headers: {
          get: () => 'application/x-ndjson',
        },
      });

      const chunks = await client.streamChatAsPromise(mockMessages);

      expect(chunks.length).toBe(2);
      expect(mockFetch).toHaveBeenCalled();
      expect(mockSetCache).toHaveBeenCalled();
    });
  });

  describe('edge cases', () => {
    it('should handle cache service errors gracefully during chat', async () => {
      // Mock cache service to throw an error
      mockGetCache.mockRejectedValueOnce(new Error('Cache error'));

      // Note: fetch is never reached because the cache throws first.
      // The chat method does not handle cache errors, so it should propagate.
      await expect(client.chat(mockMessages)).rejects.toThrow('Cache error');
    });

    it('should handle cache service errors gracefully during streaming', async () => {
      mockGetCache.mockRejectedValueOnce(new Error('Cache error'));

      await expect(async () => {
        for await (const _ of client.streamChat(mockMessages)) {
          // Should throw
        }
      }).rejects.toThrow('Cache error');
    });

    it('should work without cache when no config provided', async () => {
      const noCacheClient = new OllamaClient('http://localhost:11434', 'llama3', mockFetch);

      const mockResponse = {
        ok: true,
        json: () =>
          Promise.resolve({
            message: {
              content: 'Response without cache',
            },
          }),
      };
      mockFetch.mockResolvedValueOnce(mockResponse);

      const result = await noCacheClient.chat(mockMessages);

      expect(result.content).toBe('Response without cache');
      expect(mockFetch).toHaveBeenCalled();
    });

    it('should use last user message for cache lookup', async () => {
      const multiMessageList: OllamaMessage[] = [
        { role: 'system', content: 'You are helpful.' },
        { role: 'user', content: 'First question' },
        { role: 'assistant', content: 'First answer' },
        { role: 'user', content: 'Second question' },
      ];

      const cachedResponse = 'Cached second answer';
      mockGetCache.mockResolvedValueOnce(cachedResponse);

      const result = await client.chat(multiMessageList);

      expect(result.content).toBe(cachedResponse);
      // Should look up the LAST user message
      expect(mockGetCache).toHaveBeenCalledWith('Second question');
    });

    it('should call clearCache on the cache service', async () => {
      await client.clearCache();
      expect(mockClearCache).toHaveBeenCalledTimes(1);
    });

    it('should not throw when clearCache is called without a cache service', async () => {
      const noCacheClient = new OllamaClient('http://localhost:11434', 'llama3', mockFetch);
      jest.clearAllMocks();
      await expect(noCacheClient.clearCache()).resolves.toBeUndefined();
      expect(mockClearCache).not.toHaveBeenCalled();
    });

    it('should skip cache when no user message found', async () => {
      const onlyAssistantMessages: OllamaMessage[] = [
        { role: 'system', content: 'You are helpful.' },
        { role: 'assistant', content: 'Hello!' },
      ];

      const mockResponse = {
        ok: true,
        json: () =>
          Promise.resolve({
            message: {
              content: 'Response for assistant-only messages',
            },
          }),
      };
      mockFetch.mockResolvedValueOnce(mockResponse);

      const result = await client.chat(onlyAssistantMessages);

      expect(result.content).toBe('Response for assistant-only messages');
      expect(mockGetCache).not.toHaveBeenCalled();
      expect(mockSetCache).not.toHaveBeenCalled();
    });
  });
});