Add semantic cache support using ChromaDB #3
+19
-7
@@ -72,6 +72,15 @@ export class ChatView extends ItemView {
|
|||||||
undefined,
|
undefined,
|
||||||
newSettings.cacheConfig
|
newSettings.cacheConfig
|
||||||
);
|
);
|
||||||
|
void this.ollamaClient.initializeCache().catch(() => {
|
||||||
|
new Notice(
|
||||||
|
'Semantic cache unavailable: could not connect to ChromaDB. Check the ChromaDB URL in settings.'
|
||||||
|
);
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
public async clearCache(): Promise<void> {
|
||||||
|
await this.ollamaClient.clearCache();
|
||||||
}
|
}
|
||||||
|
|
||||||
getViewType(): string {
|
getViewType(): string {
|
||||||
@@ -83,7 +92,13 @@ export class ChatView extends ItemView {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async onOpen(): Promise<void> {
|
async onOpen(): Promise<void> {
|
||||||
|
try {
|
||||||
await this.ollamaClient.initializeCache();
|
await this.ollamaClient.initializeCache();
|
||||||
|
} catch {
|
||||||
|
new Notice(
|
||||||
|
'Semantic cache unavailable: could not connect to ChromaDB. Check the ChromaDB URL in settings.'
|
||||||
|
);
|
||||||
|
}
|
||||||
this.render();
|
this.render();
|
||||||
this.removeEventListeners(); // Clean up any existing listeners before reattaching
|
this.removeEventListeners(); // Clean up any existing listeners before reattaching
|
||||||
this.setupEventListeners();
|
this.setupEventListeners();
|
||||||
@@ -93,7 +108,7 @@ export class ChatView extends ItemView {
|
|||||||
this.updateSettings(newSettings);
|
this.updateSettings(newSettings);
|
||||||
}
|
}
|
||||||
|
|
||||||
async onClose(): Promise<void> {
|
onClose(): Promise<void> {
|
||||||
this.ollamaClient.cancelStream();
|
this.ollamaClient.cancelStream();
|
||||||
this.removeEventListeners();
|
this.removeEventListeners();
|
||||||
this.cleanupStreamingResources();
|
this.cleanupStreamingResources();
|
||||||
@@ -268,7 +283,7 @@ export class ChatView extends ItemView {
|
|||||||
if (streamingMessage && !this.lastMessageEl) {
|
if (streamingMessage && !this.lastMessageEl) {
|
||||||
this.lastMessageEl = (this.chatContainer ?? this.contentEl).createEl('div', {
|
this.lastMessageEl = (this.chatContainer ?? this.contentEl).createEl('div', {
|
||||||
cls: `ollama-message assistant`,
|
cls: `ollama-message assistant`,
|
||||||
}) as HTMLElement;
|
});
|
||||||
this.lastMessageEl.setAttribute('data-msg-id', streamingMessage.id);
|
this.lastMessageEl.setAttribute('data-msg-id', streamingMessage.id);
|
||||||
}
|
}
|
||||||
if (this.lastMessageEl) {
|
if (this.lastMessageEl) {
|
||||||
@@ -314,14 +329,11 @@ export class ChatView extends ItemView {
|
|||||||
|
|
||||||
return [
|
return [
|
||||||
systemMessage,
|
systemMessage,
|
||||||
...this.messages.map(
|
...this.messages.map((m) => ({
|
||||||
(m) =>
|
|
||||||
({
|
|
||||||
role: m.role,
|
role: m.role,
|
||||||
content: m.content,
|
content: m.content,
|
||||||
tool_calls: m.tool_calls,
|
tool_calls: m.tool_calls,
|
||||||
}) as OllamaMessage
|
})),
|
||||||
),
|
|
||||||
userMessageWithContext,
|
userMessageWithContext,
|
||||||
];
|
];
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,5 +11,6 @@ export const DEFAULT_SETTINGS = {
|
|||||||
similarityThreshold: 0.85,
|
similarityThreshold: 0.85,
|
||||||
collectionName: 'ollama_semantic_cache',
|
collectionName: 'ollama_semantic_cache',
|
||||||
embeddingModel: 'nomic-embed-text',
|
embeddingModel: 'nomic-embed-text',
|
||||||
|
chromaUrl: 'http://localhost:8000',
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|||||||
+71
-1
@@ -49,7 +49,11 @@ export default class OllamaPlugin extends Plugin {
|
|||||||
const data = (await this.loadData()) as Partial<PluginSettings> | null;
|
const data = (await this.loadData()) as Partial<PluginSettings> | null;
|
||||||
if (data) {
|
if (data) {
|
||||||
Logger.debug('Loading saved settings', 'settings');
|
Logger.debug('Loading saved settings', 'settings');
|
||||||
this.settings = Object.assign({}, this.settings, data);
|
this.settings = {
|
||||||
|
...DEFAULT_SETTINGS,
|
||||||
|
...data,
|
||||||
|
cacheConfig: { ...DEFAULT_SETTINGS.cacheConfig, ...data.cacheConfig },
|
||||||
|
};
|
||||||
}
|
}
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
ErrorHandler.handleError(error, 'settings load');
|
ErrorHandler.handleError(error, 'settings load');
|
||||||
@@ -89,6 +93,17 @@ export default class OllamaPlugin extends Plugin {
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
public async clearSemanticCache(): Promise<void> {
|
||||||
|
const leaves = this.app.workspace.getLeavesOfType('ollama-chat-view');
|
||||||
|
for (const leaf of leaves) {
|
||||||
|
const view = leaf.view;
|
||||||
|
if (view instanceof ChatView) {
|
||||||
|
await view.clearCache();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
class OllamaSettingTab extends PluginSettingTab {
|
class OllamaSettingTab extends PluginSettingTab {
|
||||||
@@ -152,6 +167,61 @@ class OllamaSettingTab extends PluginSettingTab {
|
|||||||
this.plugin.notifyChatViews();
|
this.plugin.notifyChatViews();
|
||||||
})
|
})
|
||||||
);
|
);
|
||||||
|
|
||||||
|
new Setting(container)
|
||||||
|
.setName('ChromaDB URL')
|
||||||
|
.setDesc('URL of your ChromaDB instance (used for semantic cache)')
|
||||||
|
.addText((text) =>
|
||||||
|
text.setValue(this.plugin.settings.cacheConfig.chromaUrl).onChange(async (value) => {
|
||||||
|
this.plugin.settings.cacheConfig.chromaUrl = value;
|
||||||
|
await this.plugin.saveSettings();
|
||||||
|
this.plugin.notifyChatViews();
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
new Setting(container)
|
||||||
|
.setName('Cache Embedding Model')
|
||||||
|
.setDesc('Ollama model used to generate embeddings for the semantic cache')
|
||||||
|
.addText((text) =>
|
||||||
|
text.setValue(this.plugin.settings.cacheConfig.embeddingModel).onChange(async (value) => {
|
||||||
|
this.plugin.settings.cacheConfig.embeddingModel = value;
|
||||||
|
await this.plugin.saveSettings();
|
||||||
|
this.plugin.notifyChatViews();
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
new Setting(container)
|
||||||
|
.setName('Cache Similarity Threshold')
|
||||||
|
.setDesc(
|
||||||
|
'Minimum cosine similarity (0–1) for a cache hit. Higher values require closer matches.'
|
||||||
|
)
|
||||||
|
.addText((text) =>
|
||||||
|
text
|
||||||
|
.setValue(String(this.plugin.settings.cacheConfig.similarityThreshold))
|
||||||
|
.onChange(async (value) => {
|
||||||
|
const parsed = parseFloat(value);
|
||||||
|
if (!isNaN(parsed) && parsed >= 0 && parsed <= 1) {
|
||||||
|
this.plugin.settings.cacheConfig.similarityThreshold = parsed;
|
||||||
|
await this.plugin.saveSettings();
|
||||||
|
} else {
|
||||||
|
new Notice('Similarity threshold must be a number between 0 and 1.');
|
||||||
|
}
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
new Setting(container)
|
||||||
|
.setName('Clear Semantic Cache')
|
||||||
|
.setDesc('Delete all cached responses from ChromaDB')
|
||||||
|
.addButton((button) =>
|
||||||
|
button.setButtonText('Clear Cache').onClick(async () => {
|
||||||
|
try {
|
||||||
|
await this.plugin.clearSemanticCache();
|
||||||
|
new Notice('Semantic cache cleared.');
|
||||||
|
} catch {
|
||||||
|
new Notice('Failed to clear semantic cache. Is ChromaDB running?');
|
||||||
|
}
|
||||||
|
})
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
hide(): void {
|
hide(): void {
|
||||||
|
|||||||
+16
-8
@@ -1,9 +1,9 @@
|
|||||||
// src/ollama-client.ts
|
// src/ollama-client.ts
|
||||||
|
|
||||||
import type { OllamaMessage, OllamaTool } from './types';
|
import type { OllamaMessage, OllamaTool } from './types';
|
||||||
import { ApiError } from './types';
|
import { ApiError, CacheConfig } from './types';
|
||||||
import { Logger } from './utils';
|
import { Logger } from './utils';
|
||||||
import { SemanticCacheService, CacheConfig } from './semantic-cache';
|
import { SemanticCacheService } from './semantic-cache';
|
||||||
|
|
||||||
interface OllamaChatResponse {
|
interface OllamaChatResponse {
|
||||||
message?: Partial<OllamaMessage>;
|
message?: Partial<OllamaMessage>;
|
||||||
@@ -33,6 +33,12 @@ export class OllamaClient {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async clearCache(): Promise<void> {
|
||||||
|
if (this.cacheService) {
|
||||||
|
await this.cacheService.clearCache();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
cancelStream(): void {
|
cancelStream(): void {
|
||||||
if (this.currentStreamController) {
|
if (this.currentStreamController) {
|
||||||
this.currentStreamController.abort();
|
this.currentStreamController.abort();
|
||||||
@@ -50,7 +56,7 @@ export class OllamaClient {
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
|
const lastUserMsg = messages.findLast((m) => m.role === 'user');
|
||||||
if (lastUserMsg && this.cacheService) {
|
if (lastUserMsg && this.cacheService) {
|
||||||
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
||||||
if (cached) {
|
if (cached) {
|
||||||
@@ -59,16 +65,16 @@ export class OllamaClient {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (this.cacheService && lastUserMsg) {
|
||||||
const chunks: OllamaMessage[] = [];
|
const chunks: OllamaMessage[] = [];
|
||||||
for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) {
|
for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) {
|
||||||
chunks.push(chunk);
|
chunks.push(chunk);
|
||||||
yield chunk;
|
yield chunk;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Populate cache in background after successful stream
|
|
||||||
if (this.cacheService && lastUserMsg) {
|
|
||||||
const fullContent = chunks.map((c) => c.content).join('');
|
const fullContent = chunks.map((c) => c.content).join('');
|
||||||
void this.cacheService.setCache(lastUserMsg.content, fullContent);
|
void this.cacheService.setCache(lastUserMsg.content, fullContent);
|
||||||
|
} else {
|
||||||
|
yield* this.streamChatWithRetry(messages, tools, 0);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -89,7 +95,7 @@ export class OllamaClient {
|
|||||||
return this.chatWithRetry(messages, tools, 0);
|
return this.chatWithRetry(messages, tools, 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user');
|
const lastUserMsg = messages.findLast((m) => m.role === 'user');
|
||||||
if (lastUserMsg && this.cacheService) {
|
if (lastUserMsg && this.cacheService) {
|
||||||
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
const cached = await this.cacheService.getCache(lastUserMsg.content);
|
||||||
if (cached) {
|
if (cached) {
|
||||||
@@ -334,7 +340,9 @@ export class OllamaClient {
|
|||||||
|
|
||||||
private throwIfOllamaError(parsed: Record<string, unknown>): void {
|
private throwIfOllamaError(parsed: Record<string, unknown>): void {
|
||||||
if (parsed.error) {
|
if (parsed.error) {
|
||||||
throw new Error(`Ollama error: ${String(parsed.error)}`);
|
const errorMsg =
|
||||||
|
typeof parsed.error === 'string' ? parsed.error : JSON.stringify(parsed.error);
|
||||||
|
throw new Error(`Ollama error: ${errorMsg}`);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+36
-9
@@ -1,22 +1,24 @@
|
|||||||
// src/semantic-cache.ts
|
// src/semantic-cache.ts
|
||||||
|
|
||||||
import { ChromaClient } from 'chromadb';
|
import { ChromaClient, Collection, IncludeEnum } from 'chromadb';
|
||||||
import { Logger } from './utils';
|
import { Logger } from './utils';
|
||||||
import { CacheConfig } from './types';
|
import { CacheConfig } from './types';
|
||||||
|
|
||||||
|
export { CacheConfig } from './types';
|
||||||
|
|
||||||
export class SemanticCacheService {
|
export class SemanticCacheService {
|
||||||
private client: ChromaClient;
|
private client: ChromaClient;
|
||||||
private collection: ReturnType<ChromaClient['getOrCreateCollection']> | null = null;
|
private collection: Collection | null = null;
|
||||||
private config: CacheConfig;
|
private config: CacheConfig;
|
||||||
private ollamaURL: string;
|
private ollamaURL: string;
|
||||||
|
|
||||||
constructor(ollamaURL: string, config: CacheConfig) {
|
constructor(ollamaURL: string, config: CacheConfig) {
|
||||||
this.ollamaURL = ollamaURL.replace(/\/+$/, '');
|
this.ollamaURL = ollamaURL.replace(/\/+$/, '');
|
||||||
this.config = config;
|
this.config = config;
|
||||||
this.client = new ChromaClient({ path: 'http://localhost:8000' });
|
this.client = new ChromaClient({ path: config.chromaUrl });
|
||||||
}
|
}
|
||||||
|
|
||||||
async initialize() {
|
async initialize(): Promise<void> {
|
||||||
if (!this.config.enabled) return;
|
if (!this.config.enabled) return;
|
||||||
|
|
||||||
try {
|
try {
|
||||||
@@ -27,9 +29,32 @@ export class SemanticCacheService {
|
|||||||
Logger.info(`Semantic cache initialized: ${this.config.collectionName}`, 'semantic-cache');
|
Logger.info(`Semantic cache initialized: ${this.config.collectionName}`, 'semantic-cache');
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
Logger.error(`Failed to initialize semantic cache: ${String(error)}`, 'semantic-cache');
|
Logger.error(`Failed to initialize semantic cache: ${String(error)}`, 'semantic-cache');
|
||||||
|
throw error;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async clearCache(): Promise<void> {
|
||||||
|
await this.client.deleteCollection({ name: this.config.collectionName });
|
||||||
|
this.collection = null;
|
||||||
|
Logger.info('Semantic cache cleared', 'semantic-cache');
|
||||||
|
|
||||||
|
try {
|
||||||
|
await this.initialize();
|
||||||
|
} catch {
|
||||||
|
// best-effort re-init — swallow errors
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FNV-1a 32-bit hash — deterministic and collision-resistant enough for cache keys
|
||||||
|
private computeId(text: string): string {
|
||||||
|
let hash = 0x811c9dc5;
|
||||||
|
for (let i = 0; i < text.length; i++) {
|
||||||
|
hash ^= text.charCodeAt(i);
|
||||||
|
hash = Math.imul(hash, 0x01000193) >>> 0;
|
||||||
|
}
|
||||||
|
return hash.toString(16).padStart(8, '0');
|
||||||
|
}
|
||||||
|
|
||||||
private async getEmbedding(text: string): Promise<number[]> {
|
private async getEmbedding(text: string): Promise<number[]> {
|
||||||
try {
|
try {
|
||||||
const response = await fetch(`${this.ollamaURL}/api/embeddings`, {
|
const response = await fetch(`${this.ollamaURL}/api/embeddings`, {
|
||||||
@@ -45,7 +70,7 @@ export class SemanticCacheService {
|
|||||||
throw new Error(`Embedding failed with status ${response.status}`);
|
throw new Error(`Embedding failed with status ${response.status}`);
|
||||||
}
|
}
|
||||||
|
|
||||||
const data = await response.json();
|
const data = (await response.json()) as { embedding: number[] };
|
||||||
return data.embedding;
|
return data.embedding;
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
Logger.warn(`Failed to generate embedding: ${String(error)}`, 'semantic-cache');
|
Logger.warn(`Failed to generate embedding: ${String(error)}`, 'semantic-cache');
|
||||||
@@ -65,7 +90,7 @@ export class SemanticCacheService {
|
|||||||
const results = await this.collection.query({
|
const results = await this.collection.query({
|
||||||
queryEmbeddings: [embedding],
|
queryEmbeddings: [embedding],
|
||||||
nResults: 1,
|
nResults: 1,
|
||||||
include: ['metadatas', 'distances'],
|
include: [IncludeEnum.Metadatas, IncludeEnum.Distances],
|
||||||
});
|
});
|
||||||
|
|
||||||
// Cosine distance = 1 - cosine_similarity
|
// Cosine distance = 1 - cosine_similarity
|
||||||
@@ -76,7 +101,8 @@ export class SemanticCacheService {
|
|||||||
results.distances[0][0] < 1 - this.config.similarityThreshold
|
results.distances[0][0] < 1 - this.config.similarityThreshold
|
||||||
) {
|
) {
|
||||||
Logger.debug('Semantic cache hit', 'semantic-cache');
|
Logger.debug('Semantic cache hit', 'semantic-cache');
|
||||||
return results.metadatas?.[0]?.[0]?.fullResponse ?? null;
|
const fullResponse = results.metadatas?.[0]?.[0]?.fullResponse;
|
||||||
|
return typeof fullResponse === 'string' ? fullResponse : null;
|
||||||
}
|
}
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
Logger.warn(`Cache lookup failed: ${String(error)}`, 'semantic-cache');
|
Logger.warn(`Cache lookup failed: ${String(error)}`, 'semantic-cache');
|
||||||
@@ -94,8 +120,9 @@ export class SemanticCacheService {
|
|||||||
const embedding = await this.getEmbedding(prompt);
|
const embedding = await this.getEmbedding(prompt);
|
||||||
if (!embedding.length) return;
|
if (!embedding.length) return;
|
||||||
|
|
||||||
await this.collection.add({
|
const id = this.computeId(prompt);
|
||||||
ids: [crypto.randomUUID()],
|
await this.collection.upsert({
|
||||||
|
ids: [id],
|
||||||
embeddings: [embedding],
|
embeddings: [embedding],
|
||||||
metadatas: [{ fullResponse: response }],
|
metadatas: [{ fullResponse: response }],
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -95,6 +95,7 @@ export interface CacheConfig {
|
|||||||
similarityThreshold: number;
|
similarityThreshold: number;
|
||||||
collectionName: string;
|
collectionName: string;
|
||||||
embeddingModel: string;
|
embeddingModel: string;
|
||||||
|
chromaUrl: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface PluginSettings {
|
export interface PluginSettings {
|
||||||
|
|||||||
@@ -32,6 +32,13 @@ const mockSettings: PluginSettings = {
|
|||||||
vaultSearchLimit: 3,
|
vaultSearchLimit: 3,
|
||||||
maxMessageHistory: 50,
|
maxMessageHistory: 50,
|
||||||
lastIndexTime: 0,
|
lastIndexTime: 0,
|
||||||
|
cacheConfig: {
|
||||||
|
enabled: false,
|
||||||
|
similarityThreshold: 0.9,
|
||||||
|
collectionName: 'test-cache',
|
||||||
|
embeddingModel: 'nomic-embed-text',
|
||||||
|
chromaUrl: 'http://localhost:8000',
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
describe('ChatView', () => {
|
describe('ChatView', () => {
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import { OllamaMessage, OllamaTool, CacheConfig } from '../src/types';
|
|||||||
const mockInitialize = jest.fn().mockResolvedValue(undefined);
|
const mockInitialize = jest.fn().mockResolvedValue(undefined);
|
||||||
const mockGetCache = jest.fn().mockResolvedValue(null);
|
const mockGetCache = jest.fn().mockResolvedValue(null);
|
||||||
const mockSetCache = jest.fn().mockResolvedValue(undefined);
|
const mockSetCache = jest.fn().mockResolvedValue(undefined);
|
||||||
|
const mockClearCache = jest.fn().mockResolvedValue(undefined);
|
||||||
|
|
||||||
// Mock the semantic cache service BEFORE importing OllamaClient
|
// Mock the semantic cache service BEFORE importing OllamaClient
|
||||||
jest.mock('../src/semantic-cache', () => ({
|
jest.mock('../src/semantic-cache', () => ({
|
||||||
@@ -13,6 +14,7 @@ jest.mock('../src/semantic-cache', () => ({
|
|||||||
initialize: mockInitialize,
|
initialize: mockInitialize,
|
||||||
getCache: mockGetCache,
|
getCache: mockGetCache,
|
||||||
setCache: mockSetCache,
|
setCache: mockSetCache,
|
||||||
|
clearCache: mockClearCache,
|
||||||
})),
|
})),
|
||||||
}));
|
}));
|
||||||
|
|
||||||
@@ -68,6 +70,7 @@ describe('OllamaClient with Semantic Cache', () => {
|
|||||||
similarityThreshold: 0.85,
|
similarityThreshold: 0.85,
|
||||||
collectionName: 'test_cache',
|
collectionName: 'test_cache',
|
||||||
embeddingModel: 'nomic-embed-text',
|
embeddingModel: 'nomic-embed-text',
|
||||||
|
chromaUrl: 'http://localhost:8000',
|
||||||
};
|
};
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
@@ -101,6 +104,7 @@ describe('OllamaClient with Semantic Cache', () => {
|
|||||||
});
|
});
|
||||||
|
|
||||||
it('should not create cache service when no config provided', () => {
|
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);
|
new OllamaClient('http://localhost:11434', 'llama3', mockFetch);
|
||||||
|
|
||||||
expect(SemanticCacheService).not.toHaveBeenCalled();
|
expect(SemanticCacheService).not.toHaveBeenCalled();
|
||||||
@@ -350,19 +354,8 @@ describe('OllamaClient with Semantic Cache', () => {
|
|||||||
// Mock cache service to throw an error
|
// Mock cache service to throw an error
|
||||||
mockGetCache.mockRejectedValueOnce(new Error('Cache error'));
|
mockGetCache.mockRejectedValueOnce(new Error('Cache error'));
|
||||||
|
|
||||||
const mockResponse = {
|
// Note: fetch is never reached because the cache throws first.
|
||||||
ok: true,
|
// The chat method does not handle cache errors, so it should propagate.
|
||||||
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');
|
await expect(client.chat(mockMessages)).rejects.toThrow('Cache error');
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -414,6 +407,18 @@ describe('OllamaClient with Semantic Cache', () => {
|
|||||||
expect(mockGetCache).toHaveBeenCalledWith('Second question');
|
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 () => {
|
it('should skip cache when no user message found', async () => {
|
||||||
const onlyAssistantMessages: OllamaMessage[] = [
|
const onlyAssistantMessages: OllamaMessage[] = [
|
||||||
{ role: 'system', content: 'You are helpful.' },
|
{ role: 'system', content: 'You are helpful.' },
|
||||||
|
|||||||
@@ -7,12 +7,20 @@ import { CacheConfig } from '../src/types';
|
|||||||
jest.mock('chromadb', () => ({
|
jest.mock('chromadb', () => ({
|
||||||
ChromaClient: jest.fn().mockImplementation(() => {
|
ChromaClient: jest.fn().mockImplementation(() => {
|
||||||
return {
|
return {
|
||||||
getOrCreateCollection: jest.fn().mockResolvedValue({
|
getOrCreateCollection: jest.fn().mockReturnValue({
|
||||||
query: jest.fn(),
|
query: jest.fn(),
|
||||||
add: jest.fn(),
|
add: jest.fn(),
|
||||||
|
upsert: jest.fn(),
|
||||||
}),
|
}),
|
||||||
|
deleteCollection: jest.fn(),
|
||||||
};
|
};
|
||||||
}),
|
}),
|
||||||
|
IncludeEnum: {
|
||||||
|
Documents: 'documents',
|
||||||
|
Embeddings: 'embeddings',
|
||||||
|
Metadatas: 'metadatas',
|
||||||
|
Distances: 'distances',
|
||||||
|
},
|
||||||
}));
|
}));
|
||||||
|
|
||||||
// Now import SemanticCacheService after mocking
|
// Now import SemanticCacheService after mocking
|
||||||
@@ -21,10 +29,12 @@ import { SemanticCacheService } from '../src/semantic-cache';
|
|||||||
jest.spyOn(global, 'fetch').mockImplementation(jest.fn());
|
jest.spyOn(global, 'fetch').mockImplementation(jest.fn());
|
||||||
|
|
||||||
const mockChromaClient = {
|
const mockChromaClient = {
|
||||||
getOrCreateCollection: jest.fn().mockResolvedValue({
|
getOrCreateCollection: jest.fn().mockReturnValue({
|
||||||
query: jest.fn(),
|
query: jest.fn(),
|
||||||
add: jest.fn(),
|
add: jest.fn(),
|
||||||
|
upsert: jest.fn(),
|
||||||
}),
|
}),
|
||||||
|
deleteCollection: jest.fn(),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Set up mock instance
|
// Set up mock instance
|
||||||
@@ -44,6 +54,7 @@ describe('SemanticCacheService', () => {
|
|||||||
similarityThreshold: 0.85,
|
similarityThreshold: 0.85,
|
||||||
collectionName: 'test_cache',
|
collectionName: 'test_cache',
|
||||||
embeddingModel: 'nomic-embed-text',
|
embeddingModel: 'nomic-embed-text',
|
||||||
|
chromaUrl: 'http://localhost:8000',
|
||||||
};
|
};
|
||||||
|
|
||||||
service = new SemanticCacheService('http://localhost:11434', config);
|
service = new SemanticCacheService('http://localhost:11434', config);
|
||||||
@@ -211,32 +222,30 @@ describe('SemanticCacheService', () => {
|
|||||||
json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
|
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');
|
await service.setCache('test prompt', 'test response');
|
||||||
|
|
||||||
const mockCollection = mockChromaClient.getOrCreateCollection();
|
const mockCollection = mockChromaClient.getOrCreateCollection();
|
||||||
expect(mockCollection.add).toHaveBeenCalledWith({
|
expect(mockCollection.upsert).toHaveBeenCalledWith(
|
||||||
ids: [mockUuid],
|
expect.objectContaining({
|
||||||
|
ids: [expect.any(String)],
|
||||||
embeddings: [[0.1, 0.2, 0.3]],
|
embeddings: [[0.1, 0.2, 0.3]],
|
||||||
metadatas: [{ fullResponse: 'test response' }],
|
metadatas: [{ fullResponse: 'test response' }],
|
||||||
});
|
})
|
||||||
|
);
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should not add entry when prompt is empty', async () => {
|
it('should not add entry when prompt is empty', async () => {
|
||||||
await service.setCache(' ', 'test response');
|
await service.setCache(' ', 'test response');
|
||||||
|
|
||||||
expect(mockFetch).not.toHaveBeenCalled();
|
expect(mockFetch).not.toHaveBeenCalled();
|
||||||
expect(mockChromaClient.getOrCreateCollection().add).not.toHaveBeenCalled();
|
expect(mockChromaClient.getOrCreateCollection().upsert).not.toHaveBeenCalled();
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should not add entry when response is empty', async () => {
|
it('should not add entry when response is empty', async () => {
|
||||||
await service.setCache('test prompt', ' ');
|
await service.setCache('test prompt', ' ');
|
||||||
|
|
||||||
expect(mockFetch).not.toHaveBeenCalled();
|
expect(mockFetch).not.toHaveBeenCalled();
|
||||||
expect(mockChromaClient.getOrCreateCollection().add).not.toHaveBeenCalled();
|
expect(mockChromaClient.getOrCreateCollection().upsert).not.toHaveBeenCalled();
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should not add entry when cache is disabled', async () => {
|
it('should not add entry when cache is disabled', async () => {
|
||||||
@@ -255,13 +264,33 @@ describe('SemanticCacheService', () => {
|
|||||||
json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
|
json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
|
||||||
});
|
});
|
||||||
|
|
||||||
const mockUuid = 'mock-uuid-456' as any;
|
mockChromaClient
|
||||||
jest.spyOn(crypto, 'randomUUID').mockReturnValue(mockUuid);
|
.getOrCreateCollection()
|
||||||
|
.upsert.mockRejectedValueOnce(new Error('Add failed'));
|
||||||
mockChromaClient.getOrCreateCollection().add.mockRejectedValueOnce(new Error('Add failed'));
|
|
||||||
|
|
||||||
// Should not throw
|
// Should not throw
|
||||||
await expect(service.setCache('test prompt', 'test response')).resolves.toBeUndefined();
|
await expect(service.setCache('test prompt', 'test response')).resolves.toBeUndefined();
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('clearCache', () => {
|
||||||
|
beforeEach(async () => {
|
||||||
|
await service.initialize();
|
||||||
|
mockChromaClient.deleteCollection.mockResolvedValue(undefined);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should delete the collection and re-initialize', async () => {
|
||||||
|
await service.clearCache();
|
||||||
|
|
||||||
|
expect(mockChromaClient.deleteCollection).toHaveBeenCalledWith({ name: 'test_cache' });
|
||||||
|
// getOrCreateCollection called once in beforeEach initialize, once in clearCache re-init
|
||||||
|
expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledTimes(2);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should propagate errors from deleteCollection', async () => {
|
||||||
|
mockChromaClient.deleteCollection.mockRejectedValueOnce(new Error('Delete failed'));
|
||||||
|
|
||||||
|
await expect(service.clearCache()).rejects.toThrow('Delete failed');
|
||||||
|
});
|
||||||
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
Reference in New Issue
Block a user