Add semantic cache support using ChromaDB
This commit is contained in:
@@ -0,0 +1,441 @@
|
||||
// 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);
|
||||
|
||||
// Mock the semantic cache service BEFORE importing OllamaClient
|
||||
jest.mock('../src/semantic-cache', () => ({
|
||||
SemanticCacheService: jest.fn().mockImplementation(() => ({
|
||||
initialize: mockInitialize,
|
||||
getCache: mockGetCache,
|
||||
setCache: mockSetCache,
|
||||
})),
|
||||
}));
|
||||
|
||||
// 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',
|
||||
};
|
||||
|
||||
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', () => {
|
||||
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'));
|
||||
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
message: {
|
||||
content: 'LLM response after cache failure',
|
||||
},
|
||||
}),
|
||||
};
|
||||
mockFetch.mockResolvedValueOnce(mockResponse);
|
||||
|
||||
// The chat method does not handle cache errors, so it should propagate
|
||||
// However, the client should still be usable
|
||||
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 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();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user