Merge branch 'modularize-indexing-pipeline'
This commit is contained in:
+91
-216
@@ -7,10 +7,10 @@ import { CacheConfig } from '../src/types';
|
||||
jest.mock('chromadb', () => ({
|
||||
ChromaClient: jest.fn().mockImplementation(() => {
|
||||
return {
|
||||
getOrCreateCollection: jest.fn().mockReturnValue({
|
||||
getOrCreateCollection: jest.fn().mockResolvedValue({
|
||||
query: jest.fn(),
|
||||
add: jest.fn(),
|
||||
upsert: jest.fn(),
|
||||
reset: jest.fn(),
|
||||
}),
|
||||
deleteCollection: jest.fn(),
|
||||
};
|
||||
@@ -23,270 +23,145 @@ jest.mock('chromadb', () => ({
|
||||
},
|
||||
}));
|
||||
|
||||
// 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;
|
||||
const mockOllamaUrl = 'http://localhost:11434';
|
||||
const mockCacheConfig: CacheConfig = {
|
||||
enabled: true,
|
||||
similarityThreshold: 0.85,
|
||||
collectionName: 'test_cache',
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
chromaURL: 'http://localhost:8000',
|
||||
};
|
||||
|
||||
let cacheService: SemanticCacheService;
|
||||
let mockChromaClient: any;
|
||||
let mockCollection: any;
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset all mocks
|
||||
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',
|
||||
};
|
||||
// Create a fresh instance for each test
|
||||
cacheService = new SemanticCacheService(mockOllamaUrl, mockCacheConfig);
|
||||
|
||||
service = new SemanticCacheService('http://localhost:11434', config);
|
||||
// Access the internal mocks
|
||||
mockChromaClient = (ChromaClient as jest.Mock).mock.instances[0];
|
||||
mockCollection = mockChromaClient.getOrCreateCollection.mock.results[0].value;
|
||||
});
|
||||
|
||||
describe('constructor', () => {
|
||||
it('should initialize with correct configuration', () => {
|
||||
expect(mockChromaClient).toBeDefined();
|
||||
expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({
|
||||
name: mockCacheConfig.collectionName,
|
||||
metadata: { 'hnsw:space': 'cosine' },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('initialize', () => {
|
||||
it('should create or get the collection on initialize', async () => {
|
||||
await service.initialize();
|
||||
it('should initialize the cache collection', async () => {
|
||||
await cacheService.initialize();
|
||||
|
||||
expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({
|
||||
name: 'test_cache',
|
||||
name: mockCacheConfig.collectionName,
|
||||
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();
|
||||
it('should not initialize when cache is disabled', async () => {
|
||||
const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
|
||||
const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);
|
||||
|
||||
await disabledCacheService.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 null when cache is disabled', async () => {
|
||||
const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
|
||||
const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);
|
||||
|
||||
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');
|
||||
const result = await disabledCacheService.getCache('test query');
|
||||
|
||||
expect(result).toBeNull();
|
||||
expect(mockCollection.query).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return null when no results found', async () => {
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
|
||||
it('should return null when no cache hit', async () => {
|
||||
mockCollection.query.mockResolvedValue({
|
||||
ids: [[]],
|
||||
documents: [[]],
|
||||
distances: [[]],
|
||||
});
|
||||
|
||||
mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce({
|
||||
distances: [],
|
||||
metadatas: [],
|
||||
const result = await cacheService.getCache('test query');
|
||||
|
||||
expect(result).toBeNull();
|
||||
expect(mockCollection.query).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return cached content when hit', async () => {
|
||||
const cachedContent = 'cached response';
|
||||
mockCollection.query.mockResolvedValue({
|
||||
ids: [['test-id']],
|
||||
documents: [[cachedContent]],
|
||||
distances: [[0.9]], // Above threshold
|
||||
});
|
||||
|
||||
const result = await service.getCache('test prompt');
|
||||
const result = await cacheService.getCache('test query');
|
||||
|
||||
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();
|
||||
expect(result).toBe(cachedContent);
|
||||
expect(mockCollection.query).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('setCache', () => {
|
||||
beforeEach(async () => {
|
||||
await service.initialize();
|
||||
it('should not set cache when disabled', async () => {
|
||||
const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
|
||||
const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);
|
||||
|
||||
await disabledCacheService.setCache('test query', 'test response');
|
||||
|
||||
expect(mockCollection.add).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should add entry to collection', async () => {
|
||||
mockFetch.mockResolvedValue({
|
||||
it('should add content to cache', async () => {
|
||||
const mockEmbedding = [0.1, 0.2, 0.3];
|
||||
|
||||
// Mock the fetch function for embedding generation
|
||||
global.fetch = jest.fn().mockResolvedValue({
|
||||
ok: true,
|
||||
json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
|
||||
json: () => Promise.resolve({ embedding: mockEmbedding }),
|
||||
});
|
||||
|
||||
await service.setCache('test prompt', 'test response');
|
||||
await cacheService.setCache('test query', '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' }],
|
||||
})
|
||||
);
|
||||
});
|
||||
expect(mockCollection.add).toHaveBeenCalled();
|
||||
|
||||
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();
|
||||
// Clean up
|
||||
global.fetch = undefined as any;
|
||||
});
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await service.initialize();
|
||||
mockChromaClient.deleteCollection.mockResolvedValue(undefined);
|
||||
describe('clearCache', () => {
|
||||
it('should not clear when cache is disabled', async () => {
|
||||
const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
|
||||
const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);
|
||||
|
||||
await disabledCacheService.clearCache();
|
||||
|
||||
expect(mockCollection.reset).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should delete the collection and re-initialize', async () => {
|
||||
|
||||
expect(mockChromaClient.deleteCollection).toHaveBeenCalledWith({ name: 'test_cache' });
|
||||
expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should propagate errors from deleteCollection', async () => {
|
||||
mockChromaClient.deleteCollection.mockRejectedValueOnce(new Error('Delete failed'));
|
||||
it('should clear the cache collection', async () => {
|
||||
await cacheService.clearCache();
|
||||
|
||||
expect(mockCollection.reset).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user