Merge branch 'modularize-indexing-pipeline'
This commit is contained in:
@@ -16,431 +16,84 @@ jest.mock('../src/semantic-cache', () => ({
|
||||
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',
|
||||
};
|
||||
describe('OllamaClient', () => {
|
||||
const mockBaseUrl = 'http://localhost:11434';
|
||||
const mockModel = 'llama3';
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockFetch = global.fetch as jest.Mock;
|
||||
|
||||
client = new OllamaClient('http://localhost:11434', 'llama3', mockFetch, cacheConfig);
|
||||
mockInitialize.mockClear();
|
||||
mockGetCache.mockClear();
|
||||
mockSetCache.mockClear();
|
||||
mockClearCache.mockClear();
|
||||
});
|
||||
|
||||
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.',
|
||||
},
|
||||
}),
|
||||
describe('constructor', () => {
|
||||
it('should initialize cache service when enabled', () => {
|
||||
const cacheConfig: CacheConfig = {
|
||||
enabled: true,
|
||||
similarityThreshold: 0.85,
|
||||
collectionName: 'test_cache',
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
chromaURL: 'http://localhost:8000'
|
||||
};
|
||||
mockFetch.mockResolvedValueOnce(mockResponse);
|
||||
|
||||
const result = await client.chat(mockMessages);
|
||||
const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig);
|
||||
|
||||
expect(result.content).toBe('AI is the simulation of intelligence.');
|
||||
expect(mockFetch).toHaveBeenCalled();
|
||||
// setCache should have been called
|
||||
expect(mockSetCache).toHaveBeenCalled();
|
||||
expect(client).toBeInstanceOf(OllamaClient);
|
||||
expect(mockInitialize).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should bypass cache when tools are present', async () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: () =>
|
||||
Promise.resolve({
|
||||
message: {
|
||||
content: 'Tool response',
|
||||
},
|
||||
}),
|
||||
it('should not initialize cache service when disabled', () => {
|
||||
const cacheConfig: CacheConfig = {
|
||||
enabled: false,
|
||||
similarityThreshold: 0.85,
|
||||
collectionName: 'test_cache',
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
chromaURL: 'http://localhost:8000'
|
||||
};
|
||||
mockFetch.mockResolvedValueOnce(mockResponse);
|
||||
|
||||
const result = await client.chat(mockMessages, mockTools);
|
||||
const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig);
|
||||
|
||||
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');
|
||||
expect(client).toBeInstanceOf(OllamaClient);
|
||||
expect(mockInitialize).toHaveBeenCalledTimes(0);
|
||||
});
|
||||
});
|
||||
|
||||
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',
|
||||
},
|
||||
}),
|
||||
describe('clearCache', () => {
|
||||
it('should clear the cache when enabled', async () => {
|
||||
const cacheConfig: CacheConfig = {
|
||||
enabled: true,
|
||||
similarityThreshold: 0.85,
|
||||
collectionName: 'test_cache',
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
chromaURL: 'http://localhost:8000'
|
||||
};
|
||||
mockFetch.mockResolvedValueOnce(mockResponse);
|
||||
|
||||
const result = await noCacheClient.chat(mockMessages);
|
||||
const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig);
|
||||
|
||||
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',
|
||||
},
|
||||
}),
|
||||
it('should not clear cache when disabled', async () => {
|
||||
const cacheConfig: CacheConfig = {
|
||||
enabled: false,
|
||||
similarityThreshold: 0.85,
|
||||
collectionName: 'test_cache',
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
chromaURL: 'http://localhost:8000'
|
||||
};
|
||||
mockFetch.mockResolvedValueOnce(mockResponse);
|
||||
|
||||
const result = await client.chat(onlyAssistantMessages);
|
||||
const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig);
|
||||
|
||||
expect(result.content).toBe('Response for assistant-only messages');
|
||||
expect(mockGetCache).not.toHaveBeenCalled();
|
||||
expect(mockSetCache).not.toHaveBeenCalled();
|
||||
await client.clearCache();
|
||||
|
||||
expect(mockClearCache).toHaveBeenCalledTimes(0);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user