Add semantic cache support using ChromaDB

tests/semantic-cache.test.ts

@@ -0,0 +1,267 @@
+// 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().mockResolvedValue({
+        query: jest.fn(),
+        add: jest.fn(),
+      }),
+    };
+  }),
+}));
+
+// Now import SemanticCacheService after mocking
+import { SemanticCacheService } from '../src/semantic-cache';
+
+jest.spyOn(global, 'fetch').mockImplementation(jest.fn());
+
+const mockChromaClient = {
+  getOrCreateCollection: jest.fn().mockResolvedValue({
+    query: jest.fn(),
+    add: 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',
+    };
+
+    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] }),
+      });
+
+      // crypto.randomUUID mock
+      const mockUuid = 'mock-uuid-123' as any;
+      jest.spyOn(crypto, 'randomUUID').mockReturnValue(mockUuid);
+
+      await service.setCache('test prompt', 'test response');
+
+      const mockCollection = mockChromaClient.getOrCreateCollection();
+      expect(mockCollection.add).toHaveBeenCalledWith({
+        ids: [mockUuid],
+        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().add).not.toHaveBeenCalled();
+    });
+
+    it('should not add entry when response is empty', async () => {
+      await service.setCache('test prompt', '   ');
+
+      expect(mockFetch).not.toHaveBeenCalled();
+      expect(mockChromaClient.getOrCreateCollection().add).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] }),
+      });
+
+      const mockUuid = 'mock-uuid-456' as any;
+      jest.spyOn(crypto, 'randomUUID').mockReturnValue(mockUuid);
+
+      mockChromaClient.getOrCreateCollection().add.mockRejectedValueOnce(new Error('Add failed'));
+
+      // Should not throw
+      await expect(service.setCache('test prompt', 'test response')).resolves.toBeUndefined();
+    });
+  });
+});