Add semantic cache support using ChromaDB
This commit is contained in:
Generated
+423
-243
File diff suppressed because it is too large
Load Diff
+3
-2
@@ -22,8 +22,8 @@
|
||||
"devDependencies": {
|
||||
"@types/jest": "^29.5.14",
|
||||
"@types/node": "^20.11.19",
|
||||
"@typescript-eslint/eslint-plugin": "^6.19.1",
|
||||
"@typescript-eslint/parser": "^6.19.1",
|
||||
"@typescript-eslint/eslint-plugin": "^8.59.2",
|
||||
"@typescript-eslint/parser": "^8.59.2",
|
||||
"eslint": "^8.56.0",
|
||||
"jest": "^29.7.0",
|
||||
"jest-environment-jsdom": "^30.3.0",
|
||||
@@ -32,6 +32,7 @@
|
||||
"typescript": "^5.3.3"
|
||||
},
|
||||
"dependencies": {
|
||||
"chromadb": "^1.5.3",
|
||||
"node-fetch": "^3.3.2",
|
||||
"obsidian": "^1.4.11"
|
||||
}
|
||||
|
||||
+15
-5
@@ -54,14 +54,24 @@ export class ChatView extends ItemView {
|
||||
constructor(leaf: WorkspaceLeaf, settings: PluginSettings) {
|
||||
super(leaf);
|
||||
this.settings = settings;
|
||||
this.ollamaClient = new OllamaClient(settings.ollamaUrl, settings.model);
|
||||
this.ollamaClient = new OllamaClient(
|
||||
settings.ollamaUrl,
|
||||
settings.model,
|
||||
undefined,
|
||||
settings.cacheConfig
|
||||
);
|
||||
this.vaultIndexer = new VaultIndexer(this.app.vault);
|
||||
this.toolExecutor = new ToolExecutor(this.app.vault, this.app);
|
||||
}
|
||||
|
||||
public updateSettings(newSettings: PluginSettings): void {
|
||||
this.settings = newSettings;
|
||||
this.ollamaClient = new OllamaClient(newSettings.ollamaUrl, newSettings.model);
|
||||
this.ollamaClient = new OllamaClient(
|
||||
newSettings.ollamaUrl,
|
||||
newSettings.model,
|
||||
undefined,
|
||||
newSettings.cacheConfig
|
||||
);
|
||||
}
|
||||
|
||||
getViewType(): string {
|
||||
@@ -72,18 +82,18 @@ export class ChatView extends ItemView {
|
||||
return 'Ollama Chat';
|
||||
}
|
||||
|
||||
onOpen(): Promise<void> {
|
||||
async onOpen(): Promise<void> {
|
||||
await this.ollamaClient.initializeCache();
|
||||
this.render();
|
||||
this.removeEventListeners(); // Clean up any existing listeners before reattaching
|
||||
this.setupEventListeners();
|
||||
return Promise.resolve();
|
||||
}
|
||||
|
||||
public onSettingsChange(newSettings: PluginSettings): void {
|
||||
this.updateSettings(newSettings);
|
||||
}
|
||||
|
||||
onClose(): Promise<void> {
|
||||
async onClose(): Promise<void> {
|
||||
this.ollamaClient.cancelStream();
|
||||
this.removeEventListeners();
|
||||
this.cleanupStreamingResources();
|
||||
|
||||
@@ -6,4 +6,10 @@ export const DEFAULT_SETTINGS = {
|
||||
vaultSearchLimit: 3,
|
||||
maxMessageHistory: 50,
|
||||
lastIndexTime: 0,
|
||||
cacheConfig: {
|
||||
enabled: false,
|
||||
similarityThreshold: 0.85,
|
||||
collectionName: 'ollama_semantic_cache',
|
||||
embeddingModel: 'nomic-embed-text',
|
||||
},
|
||||
};
|
||||
|
||||
+11
@@ -141,6 +141,17 @@ class OllamaSettingTab extends PluginSettingTab {
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
new Setting(container)
|
||||
.setName('Enable Semantic Cache')
|
||||
.setDesc('Cache responses semantically to speed up repeated queries')
|
||||
.addToggle((toggle) =>
|
||||
toggle.setValue(this.plugin.settings.cacheConfig.enabled).onChange(async (value) => {
|
||||
this.plugin.settings.cacheConfig.enabled = value;
|
||||
await this.plugin.saveSettings();
|
||||
this.plugin.notifyChatViews();
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
hide(): void {
|
||||
|
||||
+62
-6
@@ -3,6 +3,7 @@
|
||||
import type { OllamaMessage, OllamaTool } from './types';
|
||||
import { ApiError } from './types';
|
||||
import { Logger } from './utils';
|
||||
import { SemanticCacheService, CacheConfig } from './semantic-cache';
|
||||
|
||||
interface OllamaChatResponse {
|
||||
message?: Partial<OllamaMessage>;
|
||||
@@ -14,11 +15,22 @@ export class OllamaClient {
|
||||
private fetchFn: typeof fetch;
|
||||
private readonly maxRetries: number = 3;
|
||||
private currentStreamController: AbortController | null = null;
|
||||
private cacheService?: SemanticCacheService;
|
||||
|
||||
constructor(baseURL: string, model: string, fetchFn?: typeof fetch) {
|
||||
constructor(baseURL: string, model: string, fetchFn?: typeof fetch, cacheConfig?: CacheConfig) {
|
||||
this.baseURL = baseURL;
|
||||
this.model = model;
|
||||
this.fetchFn = fetchFn ?? fetch;
|
||||
|
||||
if (cacheConfig?.enabled) {
|
||||
this.cacheService = new SemanticCacheService(baseURL, cacheConfig);
|
||||
}
|
||||
}
|
||||
|
||||
async initializeCache(): Promise<void> {
|
||||
if (this.cacheService) {
|
||||
await this.cacheService.initialize();
|
||||
}
|
||||
}
|
||||
|
||||
cancelStream(): void {
|
||||
@@ -32,7 +44,32 @@ export class OllamaClient {
|
||||
messages: OllamaMessage[],
|
||||
tools: OllamaTool[] = []
|
||||
): AsyncGenerator<OllamaMessage, void, unknown> {
|
||||
yield* this.streamChatWithRetry(messages, tools, 0);
|
||||
// Bypass cache if tools are involved to prevent state corruption
|
||||
if (tools.length > 0) {
|
||||
yield* this.streamChatWithRetry(messages, tools, 0);
|
||||
return;
|
||||
}
|
||||
|
||||
const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
|
||||
if (lastUserMsg && this.cacheService) {
|
||||
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
||||
if (cached) {
|
||||
yield { role: 'assistant', content: cached, tool_calls: [] };
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
const chunks: OllamaMessage[] = [];
|
||||
for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) {
|
||||
chunks.push(chunk);
|
||||
yield chunk;
|
||||
}
|
||||
|
||||
// Populate cache in background after successful stream
|
||||
if (this.cacheService && lastUserMsg) {
|
||||
const fullContent = chunks.map((c) => c.content).join('');
|
||||
void this.cacheService.setCache(lastUserMsg.content, fullContent);
|
||||
}
|
||||
}
|
||||
|
||||
async streamChatAsPromise(
|
||||
@@ -46,6 +83,29 @@ export class OllamaClient {
|
||||
return chunks;
|
||||
}
|
||||
|
||||
async chat(messages: OllamaMessage[], tools: OllamaTool[] = []): Promise<OllamaMessage> {
|
||||
// Bypass cache if tools are involved
|
||||
if (tools.length > 0) {
|
||||
return this.chatWithRetry(messages, tools, 0);
|
||||
}
|
||||
|
||||
const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
|
||||
if (lastUserMsg && this.cacheService) {
|
||||
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
||||
if (cached) {
|
||||
return { role: 'assistant', content: cached, tool_calls: [] };
|
||||
}
|
||||
}
|
||||
|
||||
const response = await this.chatWithRetry(messages, tools, 0);
|
||||
|
||||
if (this.cacheService && lastUserMsg) {
|
||||
void this.cacheService.setCache(lastUserMsg.content, response.content);
|
||||
}
|
||||
|
||||
return response;
|
||||
}
|
||||
|
||||
private async *streamChatWithRetry(
|
||||
messages: OllamaMessage[],
|
||||
tools: OllamaTool[] = [],
|
||||
@@ -200,10 +260,6 @@ export class OllamaClient {
|
||||
}
|
||||
}
|
||||
|
||||
async chat(messages: OllamaMessage[], tools: OllamaTool[] = []): Promise<OllamaMessage> {
|
||||
return this.chatWithRetry(messages, tools, 0);
|
||||
}
|
||||
|
||||
private async chatWithRetry(
|
||||
messages: OllamaMessage[],
|
||||
tools: OllamaTool[] = [],
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
// src/semantic-cache.ts
|
||||
|
||||
import { ChromaClient } from 'chromadb';
|
||||
import { Logger } from './utils';
|
||||
import { CacheConfig } from './types';
|
||||
|
||||
export class SemanticCacheService {
|
||||
private client: ChromaClient;
|
||||
private collection: ReturnType<ChromaClient['getOrCreateCollection']> | null = null;
|
||||
private config: CacheConfig;
|
||||
private ollamaURL: string;
|
||||
|
||||
constructor(ollamaURL: string, config: CacheConfig) {
|
||||
this.ollamaURL = ollamaURL.replace(/\/+$/, '');
|
||||
this.config = config;
|
||||
this.client = new ChromaClient({ path: 'http://localhost:8000' });
|
||||
}
|
||||
|
||||
async initialize() {
|
||||
if (!this.config.enabled) return;
|
||||
|
||||
try {
|
||||
this.collection = await this.client.getOrCreateCollection({
|
||||
name: this.config.collectionName,
|
||||
metadata: { 'hnsw:space': 'cosine' },
|
||||
});
|
||||
Logger.info(`Semantic cache initialized: ${this.config.collectionName}`, 'semantic-cache');
|
||||
} catch (error) {
|
||||
Logger.error(`Failed to initialize semantic cache: ${String(error)}`, 'semantic-cache');
|
||||
}
|
||||
}
|
||||
|
||||
private async getEmbedding(text: string): Promise<number[]> {
|
||||
try {
|
||||
const response = await fetch(`${this.ollamaURL}/api/embeddings`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
model: this.config.embeddingModel,
|
||||
prompt: text,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`Embedding failed with status ${response.status}`);
|
||||
}
|
||||
|
||||
const data = await response.json();
|
||||
return data.embedding;
|
||||
} catch (error) {
|
||||
Logger.warn(`Failed to generate embedding: ${String(error)}`, 'semantic-cache');
|
||||
return [];
|
||||
}
|
||||
}
|
||||
|
||||
async getCache(prompt: string): Promise<string | null> {
|
||||
if (!this.collection || !this.config.enabled || !prompt.trim()) {
|
||||
return null;
|
||||
}
|
||||
|
||||
try {
|
||||
const embedding = await this.getEmbedding(prompt);
|
||||
if (!embedding.length) return null;
|
||||
|
||||
const results = await this.collection.query({
|
||||
queryEmbeddings: [embedding],
|
||||
nResults: 1,
|
||||
include: ['metadatas', 'distances'],
|
||||
});
|
||||
|
||||
// Cosine distance = 1 - cosine_similarity
|
||||
// We want distance < (1 - threshold)
|
||||
if (
|
||||
results.distances &&
|
||||
results.distances[0] &&
|
||||
results.distances[0][0] < 1 - this.config.similarityThreshold
|
||||
) {
|
||||
Logger.debug('Semantic cache hit', 'semantic-cache');
|
||||
return results.metadatas?.[0]?.[0]?.fullResponse ?? null;
|
||||
}
|
||||
} catch (error) {
|
||||
Logger.warn(`Cache lookup failed: ${String(error)}`, 'semantic-cache');
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
async setCache(prompt: string, response: string): Promise<void> {
|
||||
if (!this.collection || !this.config.enabled || !prompt.trim() || !response.trim()) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const embedding = await this.getEmbedding(prompt);
|
||||
if (!embedding.length) return;
|
||||
|
||||
await this.collection.add({
|
||||
ids: [crypto.randomUUID()],
|
||||
embeddings: [embedding],
|
||||
metadatas: [{ fullResponse: response }],
|
||||
});
|
||||
Logger.debug('Cached new response', 'semantic-cache');
|
||||
} catch (error) {
|
||||
Logger.warn(`Cache write failed: ${String(error)}`, 'semantic-cache');
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -90,12 +90,20 @@ export class PathValidationError extends OllamaError {
|
||||
// Plugin Configuration
|
||||
// ============================================================
|
||||
|
||||
export interface CacheConfig {
|
||||
enabled: boolean;
|
||||
similarityThreshold: number;
|
||||
collectionName: string;
|
||||
embeddingModel: string;
|
||||
}
|
||||
|
||||
export interface PluginSettings {
|
||||
ollamaUrl: string;
|
||||
model: string;
|
||||
vaultSearchLimit: number;
|
||||
maxMessageHistory: number;
|
||||
lastIndexTime: number;
|
||||
cacheConfig: CacheConfig;
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
|
||||
@@ -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