Merge branch 'modularize-indexing-pipeline'
This commit is contained in:
@@ -0,0 +1,308 @@
|
||||
import { ContentExtractor } from '../src/indexing-pipeline/extraction';
|
||||
import { ContentNormalizer } from '../src/indexing-pipeline/normalization';
|
||||
import { ContentVectorizer } from '../src/indexing-pipeline/vectorization';
|
||||
import { IndexingPipeline } from '../src/indexing-pipeline/pipeline';
|
||||
|
||||
// Mock VaultFile interface for testing
|
||||
interface MockVaultFile {
|
||||
basename: string;
|
||||
path: string;
|
||||
}
|
||||
|
||||
describe('Indexing Pipeline Components', () => {
|
||||
describe('ContentExtractor', () => {
|
||||
let extractor: ContentExtractor;
|
||||
|
||||
beforeEach(() => {
|
||||
extractor = new ContentExtractor();
|
||||
});
|
||||
|
||||
it('should extract frontmatter correctly', () => {
|
||||
const content = `---
|
||||
title: Test Title
|
||||
tags: algorithm, programming
|
||||
date: 2023-01-01
|
||||
---
|
||||
|
||||
# Heading
|
||||
|
||||
Content here`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
|
||||
expect(extracted.frontmatter.title).toBe('Test Title');
|
||||
expect(extracted.frontmatter.tags).toBe('algorithm, programming');
|
||||
expect(extracted.frontmatter.date).toBe('2023-01-01');
|
||||
expect(extracted.headings).toContain('Heading');
|
||||
});
|
||||
|
||||
it('should extract headings correctly', () => {
|
||||
const content = `# Heading 1
|
||||
## Heading 2
|
||||
### Heading 3
|
||||
|
||||
Content`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
|
||||
expect(extracted.headings).toEqual(['Heading 1', 'Heading 2', 'Heading 3']);
|
||||
});
|
||||
|
||||
it('should extract embedded code blocks', () => {
|
||||
const content = `# Code Example
|
||||
|
||||
\`\`\`javascript
|
||||
console.log('hello world');
|
||||
\`\`\`
|
||||
|
||||
Some content`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
|
||||
expect(extracted.embeddedCodeBlocks).toHaveLength(1);
|
||||
expect(extracted.embeddedCodeBlocks[0]).toContain('console.log');
|
||||
});
|
||||
|
||||
it('should extract first paragraph', () => {
|
||||
const content = `First paragraph here.
|
||||
|
||||
Second paragraph here.
|
||||
|
||||
# Heading`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
|
||||
expect(extracted.firstParagraph).toBe('First paragraph here.');
|
||||
});
|
||||
|
||||
it('should extract raw text correctly', () => {
|
||||
const content = `---
|
||||
title: Test
|
||||
---
|
||||
|
||||
# Heading
|
||||
|
||||
Content with **bold** and [link](url).
|
||||
|
||||
\`\`\`javascript
|
||||
code
|
||||
\`\`\``;
|
||||
|
||||
const rawText = extractor.extractRawText(content);
|
||||
expect(rawText).not.toContain('---');
|
||||
expect(rawText).not.toContain('# Heading');
|
||||
expect(rawText).not.toContain('```javascript');
|
||||
expect(rawText).toContain('Content with bold and link');
|
||||
});
|
||||
});
|
||||
|
||||
describe('ContentNormalizer', () => {
|
||||
let normalizer: ContentNormalizer;
|
||||
let extractor: ContentExtractor;
|
||||
|
||||
beforeEach(() => {
|
||||
normalizer = new ContentNormalizer();
|
||||
extractor = new ContentExtractor();
|
||||
});
|
||||
|
||||
it('should normalize frontmatter dates to ISO format', () => {
|
||||
const content = `---
|
||||
title: Test
|
||||
date: 2023-01-01
|
||||
created: 2023-06-15
|
||||
updated: invalid-date
|
||||
tags: algorithm
|
||||
---
|
||||
|
||||
Content`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
const normalized = normalizer.normalize(extracted);
|
||||
|
||||
expect(normalized.frontmatter.date).toBe('2023-01-01T00:00:00.000Z');
|
||||
expect(normalized.frontmatter.created).toBe('2023-06-15T00:00:00.000Z');
|
||||
expect(normalized.frontmatter.updated).toBe('invalid-date'); // Should preserve invalid dates
|
||||
});
|
||||
|
||||
it('should convert tags to array format', () => {
|
||||
const content = `---
|
||||
title: Test
|
||||
tags: algorithm, programming, javascript
|
||||
---
|
||||
|
||||
Content`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
const normalized = normalizer.normalize(extracted);
|
||||
|
||||
expect(normalized.frontmatter.tags).toEqual(['algorithm', 'programming', 'javascript']);
|
||||
});
|
||||
|
||||
it('should calculate word count correctly', () => {
|
||||
const content = `# Title
|
||||
|
||||
This is a test document with several words to count.
|
||||
It has multiple sentences and words to make it longer.`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
const normalized = normalizer.normalize(extracted);
|
||||
|
||||
expect(normalized.wordCount).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should extract tokens correctly', () => {
|
||||
const content = `# Test Document
|
||||
|
||||
This is a test document with important keywords.`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
const normalized = normalizer.normalize(extracted);
|
||||
|
||||
expect(normalized.tokens).toContain('test');
|
||||
expect(normalized.tokens).toContain('document');
|
||||
expect(normalized.tokens).toContain('important');
|
||||
expect(normalized.tokens).toContain('keywords');
|
||||
});
|
||||
|
||||
it('should extract title correctly', () => {
|
||||
const content = `# Test Document
|
||||
|
||||
Content`;
|
||||
|
||||
const file: MockVaultFile = { basename: 'test.md', path: 'test.md' };
|
||||
const extracted = extractor.extractFromFile(file, content);
|
||||
const normalized = normalizer.normalize(extracted);
|
||||
|
||||
expect(normalized.title).toBe('test');
|
||||
});
|
||||
});
|
||||
|
||||
describe('ContentVectorizer', () => {
|
||||
let vectorizer: ContentVectorizer;
|
||||
|
||||
beforeEach(() => {
|
||||
vectorizer = new ContentVectorizer({
|
||||
model: 'nomic-embed-text',
|
||||
ollamaUrl: 'http://localhost:11434',
|
||||
});
|
||||
});
|
||||
|
||||
it('should create a proper prompt from content chunk', () => {
|
||||
const mockChunk = {
|
||||
id: 'test',
|
||||
path: 'test.md',
|
||||
title: 'Test',
|
||||
content: 'Test content',
|
||||
tokens: ['test', 'content'],
|
||||
headings: ['Heading'],
|
||||
frontmatter: { tags: ['test'] },
|
||||
firstParagraph: 'First paragraph',
|
||||
wordCount: 2,
|
||||
chunkIndex: 0,
|
||||
chunkSize: 100
|
||||
};
|
||||
|
||||
const prompt = (vectorizer as any).createPrompt(mockChunk);
|
||||
|
||||
expect(prompt).toContain('Test');
|
||||
expect(prompt).toContain('First paragraph');
|
||||
expect(prompt).toContain('Heading');
|
||||
expect(prompt).toContain('tags');
|
||||
});
|
||||
|
||||
// Note: Actual embedding tests would require mocking fetch or integration testing
|
||||
it('should handle vectorization errors gracefully', async () => {
|
||||
// This test would require mocking fetch to simulate error responses
|
||||
// For now, we're just ensuring the method exists and doesn't crash
|
||||
const mockChunk = {
|
||||
id: 'test',
|
||||
path: 'test.md',
|
||||
title: 'Test',
|
||||
content: 'Test content',
|
||||
tokens: ['test', 'content'],
|
||||
headings: ['Heading'],
|
||||
frontmatter: { tags: ['test'] },
|
||||
firstParagraph: 'First paragraph',
|
||||
wordCount: 2,
|
||||
chunkIndex: 0,
|
||||
chunkSize: 100
|
||||
};
|
||||
|
||||
// Mock fetch to simulate an error
|
||||
const originalFetch = global.fetch;
|
||||
(global.fetch as any) = jest.fn().mockRejectedValue(new Error('Network error'));
|
||||
|
||||
try {
|
||||
const result = await vectorizer.vectorize(mockChunk);
|
||||
expect(result).toEqual([]);
|
||||
} finally {
|
||||
global.fetch = originalFetch;
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('IndexingPipeline', () => {
|
||||
let pipeline: IndexingPipeline;
|
||||
|
||||
beforeEach(() => {
|
||||
pipeline = new IndexingPipeline({
|
||||
ollamaUrl: 'http://localhost:11434',
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
});
|
||||
});
|
||||
|
||||
it('should process files through the pipeline', async () => {
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
const content = `---
|
||||
title: Test Document
|
||||
tags: test, example
|
||||
---
|
||||
|
||||
# Introduction
|
||||
|
||||
This is a test document for pipeline processing.`;
|
||||
|
||||
const result = await pipeline.processFile(file, content);
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result?.title).toBe('test');
|
||||
expect(result?.path).toBe('test.md');
|
||||
expect(result?.content).toContain('This is a test document for pipeline processing');
|
||||
});
|
||||
|
||||
it('should handle processing errors gracefully', async () => {
|
||||
const file: MockVaultFile = { basename: 'test', path: 'test.md' };
|
||||
|
||||
const result = await pipeline.processFile(file, '');
|
||||
|
||||
// Should not crash, but might return null or incomplete result
|
||||
expect(result).toBeNull(); // Empty content should return null
|
||||
});
|
||||
|
||||
it('should process files in batches', async () => {
|
||||
const files: MockVaultFile[] = [
|
||||
{ basename: 'file1', path: 'file1.md' },
|
||||
{ basename: 'file2', path: 'file2.md' }
|
||||
];
|
||||
|
||||
const fileContents = {
|
||||
'file1.md': '# File 1\n\nContent 1',
|
||||
'file2.md': '# File 2\n\nContent 2'
|
||||
};
|
||||
|
||||
const results = await pipeline.processFilesInBatches(files, fileContents, 1);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0].title).toBe('file1');
|
||||
expect(results[1].title).toBe('file2');
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
+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