obsidian_ollama/tests/semantic-cache.test.ts

// tests/semantic-cache.test.ts

import { ChromaClient } from 'chromadb';
import { CacheConfig } from '../src/types';

// Mock ChromaDB module
jest.mock('chromadb', () => ({
  ChromaClient: jest.fn().mockImplementation(() => {
    return {
      getOrCreateCollection: jest.fn().mockReturnValue({
        query: jest.fn(),
        add: jest.fn(),
        upsert: jest.fn(),
      }),
      deleteCollection: jest.fn(),
    };
  }),
  IncludeEnum: {
    Documents: 'documents',
    Embeddings: 'embeddings',
    Metadatas: 'metadatas',
    Distances: 'distances',
  },
}));

// Now import SemanticCacheService after mocking
import { SemanticCacheService } from '../src/semantic-cache';

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

const mockChromaClient = {
  getOrCreateCollection: jest.fn().mockReturnValue({
    query: jest.fn(),
    add: jest.fn(),
    upsert: jest.fn(),
  }),
  deleteCollection: jest.fn(),
};

// Set up mock instance
(ChromaClient as jest.Mock).mockImplementation(() => mockChromaClient as any);

describe('SemanticCacheService', () => {
  let service: SemanticCacheService;
  let config: CacheConfig;
  let mockFetch: jest.Mock;

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

    config = {
      enabled: true,
      similarityThreshold: 0.85,
      collectionName: 'test_cache',
      embeddingModel: 'nomic-embed-text',
      chromaUrl: 'http://localhost:8000',
    };

    service = new SemanticCacheService('http://localhost:11434', config);
  });

  describe('initialize', () => {
    it('should create or get the collection on initialize', async () => {
      await service.initialize();

      expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({
        name: 'test_cache',
        metadata: { 'hnsw:space': 'cosine' },
      });
    });

    it('should not initialize if cache is disabled', async () => {
      const disabledConfig = { ...config, enabled: false };
      service = new SemanticCacheService('http://localhost:11434', disabledConfig);
      await service.initialize();

      expect(mockChromaClient.getOrCreateCollection).not.toHaveBeenCalled();
    });
  });

  describe('getEmbedding', () => {
    it('should call Ollama embeddings API correctly', async () => {
      mockFetch.mockResolvedValueOnce({
        ok: true,
        json: () =>
          Promise.resolve({
            embedding: [0.1, 0.2, 0.3],
          }),
      });

      // Call getCache to trigger embedding generation
      const mockQueryResult = {
        distances: [[0.1]],
        metadatas: [[{ fullResponse: 'Cached response' }]],
      };
      mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce(mockQueryResult);

      await service.initialize();
      await service.getCache('test prompt');

      expect(mockFetch).toHaveBeenCalledWith(
        'http://localhost:11434/api/embeddings',
        expect.objectContaining({
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({
            model: 'nomic-embed-text',
            prompt: 'test prompt',
          }),
        })
      );
    });

    it('should return empty array on embedding failure', async () => {
      mockFetch.mockResolvedValueOnce({
        ok: false,
        status: 500,
      });

      await service.initialize();
      // We need to test the private method indirectly via getCache
      const result = await service.getCache('test prompt');
      // Embedding failed, so getCache should return null
      expect(result).toBeNull();
    });
  });

  describe('getCache', () => {
    beforeEach(async () => {
      await service.initialize();
    });

    it('should return cached response when similarity is above threshold', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
      });

      const mockQueryResult = {
        distances: [[0.1]], // distance < 0.15 means similarity > 0.85
        metadatas: [[{ fullResponse: 'Cached answer' }]],
      };
      mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce(mockQueryResult);

      const result = await service.getCache('test prompt');

      expect(result).toBe('Cached answer');
    });

    it('should return null when similarity is below threshold', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
      });

      const mockQueryResult = {
        distances: [[0.2]], // distance > 0.15 means similarity < 0.85
        metadatas: [[{ fullResponse: 'Cached answer' }]],
      };
      mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce(mockQueryResult);

      const result = await service.getCache('test prompt');

      expect(result).toBeNull();
    });

    it('should return null when no results found', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
      });

      mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce({
        distances: [],
        metadatas: [],
      });

      const result = await service.getCache('test prompt');

      expect(result).toBeNull();
    });

    it('should return null when prompt is empty', async () => {
      const result = await service.getCache('   ');

      expect(result).toBeNull();
      expect(mockFetch).not.toHaveBeenCalled();
    });

    it('should return null when cache is not initialized', async () => {
      // Don't call initialize
      const result = await service.getCache('test prompt');

      expect(result).toBeNull();
    });

    it('should handle query errors gracefully', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
      });

      mockChromaClient
        .getOrCreateCollection()
        .query.mockRejectedValueOnce(new Error('Query failed'));

      const result = await service.getCache('test prompt');

      expect(result).toBeNull();
    });
  });

  describe('setCache', () => {
    beforeEach(async () => {
      await service.initialize();
    });

    it('should add entry to collection', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
      });

      await service.setCache('test prompt', 'test response');

      const mockCollection = mockChromaClient.getOrCreateCollection();
      expect(mockCollection.upsert).toHaveBeenCalledWith(
        expect.objectContaining({
          ids: [expect.any(String)],
          embeddings: [[0.1, 0.2, 0.3]],
          metadatas: [{ fullResponse: 'test response' }],
        })
      );
    });

    it('should not add entry when prompt is empty', async () => {
      await service.setCache('   ', 'test response');

      expect(mockFetch).not.toHaveBeenCalled();
      expect(mockChromaClient.getOrCreateCollection().upsert).not.toHaveBeenCalled();
    });

    it('should not add entry when response is empty', async () => {
      await service.setCache('test prompt', '   ');

      expect(mockFetch).not.toHaveBeenCalled();
      expect(mockChromaClient.getOrCreateCollection().upsert).not.toHaveBeenCalled();
    });

    it('should not add entry when cache is disabled', async () => {
      const disabledConfig = { ...config, enabled: false };
      service = new SemanticCacheService('http://localhost:11434', disabledConfig);
      await service.initialize();

      await service.setCache('test prompt', 'test response');

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

    it('should handle add errors gracefully', async () => {
      mockFetch.mockResolvedValue({
        ok: true,
        json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
      });

      mockChromaClient
        .getOrCreateCollection()
        .upsert.mockRejectedValueOnce(new Error('Add failed'));

      // Should not throw
      await expect(service.setCache('test prompt', 'test response')).resolves.toBeUndefined();
    });
  });

  describe('clearCache', () => {
    beforeEach(async () => {
      await service.initialize();
      mockChromaClient.deleteCollection.mockResolvedValue(undefined);
    });

    it('should delete the collection and re-initialize', async () => {
      await service.clearCache();

      expect(mockChromaClient.deleteCollection).toHaveBeenCalledWith({ name: 'test_cache' });
      // getOrCreateCollection called once in beforeEach initialize, once in clearCache re-init
      expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledTimes(2);
    });

    it('should propagate errors from deleteCollection', async () => {
      mockChromaClient.deleteCollection.mockRejectedValueOnce(new Error('Delete failed'));

      await expect(service.clearCache()).rejects.toThrow('Delete failed');
    });
  });
});