diff --git a/README.md b/README.md index e363b9f..cf80a08 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,9 @@ The plugin creates a full commented template at this path on first startup. This "userEmailOverride": "user@example.com", "userNameOverride": "John Doe", "embeddingModel": "Xenova/nomic-embed-text-v1", + // Optional Nomic task prefixes (search_document: / search_query:). After enabling, + // re-index existing memories so store and query vectors stay aligned. + // "embeddingUseTaskPrefixes": true, // Optional OpenAI-compatible embedding endpoint: // "embeddingApiUrl": "https://api.openai.com/v1", // "embeddingApiKey": "env://OPENAI_API_KEY", diff --git a/src/config.ts b/src/config.ts index 920cd86..ed8591a 100644 --- a/src/config.ts +++ b/src/config.ts @@ -31,6 +31,8 @@ interface OpenCodeMemConfig { embeddingDimensions?: number; embeddingApiUrl?: string; embeddingApiKey?: string; + /** Opt-in Nomic task prefixes (`search_document:` / `search_query:`). Default false. */ + embeddingUseTaskPrefixes?: boolean; similarityThreshold?: number; maxMemories?: number; maxProfileItems?: number; @@ -140,6 +142,7 @@ const DEFAULTS: Required< storagePath: join(DATA_DIR, "data"), embeddingModel: "Xenova/nomic-embed-text-v1", embeddingDimensions: 768, + embeddingUseTaskPrefixes: false, similarityThreshold: 0.6, maxMemories: 10, maxProfileItems: 5, @@ -238,6 +241,10 @@ const CONFIG_TEMPLATE = `{ // Default: Nomic Embed v1 (768 dimensions, 8192 context, multilingual) "embeddingModel": "Xenova/nomic-embed-text-v1", + + // Opt-in Nomic task prefixes (search_document: / search_query:). After enabling, + // re-index existing memories so store and query vectors stay aligned. + // "embeddingUseTaskPrefixes": true, // Auto-detected dimensions (no need to set manually) // "embeddingDimensions": 768, @@ -572,6 +579,8 @@ function buildConfig(fileConfig: OpenCodeMemConfig) { userNameOverride: fileConfig.userNameOverride, embeddingModel: fileConfig.embeddingModel ?? DEFAULTS.embeddingModel, embeddingDimensions, + embeddingUseTaskPrefixes: + fileConfig.embeddingUseTaskPrefixes ?? DEFAULTS.embeddingUseTaskPrefixes, embeddingApiUrl: fileConfig.embeddingApiUrl, embeddingApiKey: fileConfig.embeddingApiUrl ? resolveSecretValue(fileConfig.embeddingApiKey ?? process.env.OPENAI_API_KEY) diff --git a/src/services/api-handlers.ts b/src/services/api-handlers.ts index 0964684..db9e231 100644 --- a/src/services/api-handlers.ts +++ b/src/services/api-handlers.ts @@ -326,10 +326,12 @@ export async function handleAddMemory(data: { const embeddingInput = tags.length > 0 ? `${data.content}\nTags: ${tags.join(", ")}` : data.content; - const vector = await embeddingService.embedWithTimeout(embeddingInput); + const vector = await embeddingService.embedWithTimeout(embeddingInput, { task: "document" }); let tagsVector: Float32Array | undefined = undefined; if (tags.length > 0) { - tagsVector = await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags)); + tagsVector = await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags), { + task: "document", + }); } const { scope, hash } = extractScopeFromTag(data.containerTag); @@ -446,10 +448,12 @@ export async function handleUpdateMemory( const existingTags = existingMemory.tags ? String(existingMemory.tags) : ""; const tags = data.tags || (existingTags ? existingTags.split(",").map((t: string) => t.trim()) : []); - const vector = await embeddingService.embedWithTimeout(newContent); + const vector = await embeddingService.embedWithTimeout(newContent, { task: "document" }); let tagsVector: Float32Array | undefined = undefined; if (tags.length > 0) { - tagsVector = await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags)); + tagsVector = await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags), { + task: "document", + }); } const db = await tursoConnectionManager.getConnection(foundShard.dbPath); @@ -536,7 +540,7 @@ export async function handleSearch( if (!query) return { success: false, error: "query is required" }; await ensureTursoReady(); await embeddingService.warmup(); - const queryVector = await embeddingService.embedWithTimeout(query); + const queryVector = await embeddingService.embedWithTimeout(query, { task: "query" }); let memoryResults: any[] = []; let promptResults: any[] = []; if (tag) { @@ -1434,9 +1438,11 @@ export async function handleRunTagMigrationBatch( } } - const vector = await embeddingService.embedWithTimeout(m.content); + const vector = await embeddingService.embedWithTimeout(m.content, { task: "document" }); const tagsVector = currentTags.length - ? await embeddingService.embedWithTimeout(formatTagsForEmbedding(currentTags)) + ? await embeddingService.embedWithTimeout(formatTagsForEmbedding(currentTags), { + task: "document", + }) : undefined; await tursoVectorSearch.updateVector(db, m.id, vector, tagsVector); diff --git a/src/services/client.ts b/src/services/client.ts index e439285..8ecafee 100644 --- a/src/services/client.ts +++ b/src/services/client.ts @@ -107,7 +107,7 @@ export class LocalMemoryClient { try { await this.initialize(); - const queryVector = await embeddingService.embedWithTimeout(query); + const queryVector = await embeddingService.embedWithTimeout(query, { task: "query" }); const resolved = resolveScopeValue(scope, containerTag); const shards = await tursoShardManager.getAllShards(resolved.scope, resolved.hash); @@ -162,11 +162,13 @@ export class LocalMemoryClient { await this.initialize(); const tags = metadata?.tags || []; - const vector = await embeddingService.embedWithTimeout(content); + const vector = await embeddingService.embedWithTimeout(content, { task: "document" }); let tagsVector: Float32Array | undefined = undefined; if (tags.length > 0) { - tagsVector = await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags)); + tagsVector = await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags), { + task: "document", + }); } const { scope, hash } = extractScopeFromContainerTag(containerTag); diff --git a/src/services/embedding.ts b/src/services/embedding.ts index f0d132e..21712aa 100644 --- a/src/services/embedding.ts +++ b/src/services/embedding.ts @@ -6,6 +6,32 @@ const TIMEOUT_MS = 30000; const GLOBAL_EMBEDDING_KEY = Symbol.for("opencode-mem.embedding.instance"); const MAX_CACHE_SIZE = 100; +export type EmbeddingTask = "document" | "query"; + +export type EmbedOptions = { + task?: EmbeddingTask; +}; + +const TASK_PREFIXES: Record = { + document: "search_document: ", + query: "search_query: ", +}; + +/** + * Apply Nomic-style task prefixes when enabled. + * Cache keys and API/local inputs should use the returned string. + */ +export function applyEmbeddingTaskPrefix( + text: string, + options?: EmbedOptions, + useTaskPrefixes: boolean = CONFIG.embeddingUseTaskPrefixes +): string { + if (!useTaskPrefixes || !options?.task) return text; + const prefix = TASK_PREFIXES[options.task]; + if (text.startsWith(prefix)) return text; + return `${prefix}${text}`; +} + type HfTransformers = typeof import("@huggingface/transformers"); let _transformers: { @@ -108,13 +134,15 @@ export class EmbeddingService { } } - async embed(text: string): Promise { + async embed(text: string, options?: EmbedOptions): Promise { if (this.cachedModelName !== CONFIG.embeddingModel) { this.clearCache(); this.cachedModelName = CONFIG.embeddingModel; } - const cached = this.cache.get(text); + const input = applyEmbeddingTaskPrefix(text, options); + + const cached = this.cache.get(input); if (cached) return cached; if (!this.isWarmedUp && !this.initPromise) { @@ -134,7 +162,7 @@ export class EmbeddingService { Authorization: `Bearer ${CONFIG.embeddingApiKey}`, }, body: JSON.stringify({ - input: text, + input, model: CONFIG.embeddingModel, }), }); @@ -146,7 +174,7 @@ export class EmbeddingService { const data: any = await response.json(); result = new Float32Array(data.data[0].embedding); } else { - const output = await this.pipe(text, { pooling: "mean", normalize: true }); + const output = await this.pipe(input, { pooling: "mean", normalize: true }); result = new Float32Array(output.data); } @@ -154,13 +182,13 @@ export class EmbeddingService { const firstKey = this.cache.keys().next().value; if (firstKey !== undefined) this.cache.delete(firstKey); } - this.cache.set(text, result); + this.cache.set(input, result); return result; } - async embedWithTimeout(text: string): Promise { - return withTimeout(this.embed(text), TIMEOUT_MS); + async embedWithTimeout(text: string, options?: EmbedOptions): Promise { + return withTimeout(this.embed(text, options), TIMEOUT_MS); } clearCache(): void { diff --git a/src/services/migration-service.ts b/src/services/migration-service.ts index 6484e13..95179f1 100644 --- a/src/services/migration-service.ts +++ b/src/services/migration-service.ts @@ -292,10 +292,14 @@ export class MigrationService { .filter(Boolean) : []; const embeddingInput = tags.length > 0 ? `${content}\nTags: ${tags.join(", ")}` : content; - const vector = await embeddingService.embedWithTimeout(embeddingInput); + const vector = await embeddingService.embedWithTimeout(embeddingInput, { + task: "document", + }); const tagsVector = tags.length > 0 - ? await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags)) + ? await embeddingService.embedWithTimeout(formatTagsForEmbedding(tags), { + task: "document", + }) : undefined; stagedBatch.push({ diff --git a/src/services/user-profile/user-profile-manager.ts b/src/services/user-profile/user-profile-manager.ts index 5473302..6e41e96 100644 --- a/src/services/user-profile/user-profile-manager.ts +++ b/src/services/user-profile/user-profile-manager.ts @@ -569,7 +569,9 @@ export class UserProfileManager { if (useEmbedding && item.description.length >= minDescLen && !(item as any).centroid) { try { - const emb = await embed.embed(normalizeDescription(item.description)); + const emb = await embed.embed(normalizeDescription(item.description), { + task: "document", + }); const arr = Array.from(emb); (item as any).centroid = arr; (item as any).anchor = arr; @@ -630,7 +632,9 @@ export class UserProfileManager { let top2SameCat = false; let top2Band: "strong" | "weak" | null = null; - const newEmb = await embed.embed(normalizeDescription(newItem.description)); + const newEmb = await embed.embed(normalizeDescription(newItem.description), { + task: "document", + }); for (let i = 0; i < existing.length; i++) { const existingItem = existing[i]; @@ -1060,7 +1064,9 @@ export class UserProfileManager { let initAnchor: number[] | undefined; if (useEmbedding && newItem.description.length >= minDescLen) { try { - const emb = await embed.embed(normalizeDescription(newItem.description)); + const emb = await embed.embed(normalizeDescription(newItem.description), { + task: "document", + }); initCentroid = Array.from(emb); initAnchor = initCentroid; } catch {} @@ -1730,7 +1736,9 @@ Generate a concise, abstract description of the user's general behavioral tenden if (evidence.length >= 3) { const oldDesc = item.description; const evEmbs = await Promise.all( - evidence.map((e: string) => embed.embed(normalizeDescription(e))) + evidence.map((e: string) => + embed.embed(normalizeDescription(e), { task: "document" }) + ) ); const sumVec = evEmbs.reduce( (acc: number[], e: any) => { @@ -1742,8 +1750,10 @@ Generate a concise, abstract description of the user's general behavioral tenden const evCentroid = l2Normalize(sumVec); let adopted = false; - const oldEmb = await embed.embed(normalizeDescription(item.description)); - const newEmb = await embed.embed(normalizeDescription(evolved)); + const oldEmb = await embed.embed(normalizeDescription(item.description), { + task: "document", + }); + const newEmb = await embed.embed(normalizeDescription(evolved), { task: "document" }); const cosOld = cosineSimilarityNumbers(Array.from(oldEmb) as number[], evCentroid); const cosNew = cosineSimilarityNumbers(Array.from(newEmb) as number[], evCentroid); diff --git a/tests/config.test.ts b/tests/config.test.ts index c0c4c02..4757a5b 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -33,6 +33,10 @@ describe("config", () => { expect(typeof CONFIG.embeddingModel).toBe("string"); }); + it("should default embeddingUseTaskPrefixes to false", () => { + expect(CONFIG.embeddingUseTaskPrefixes).toBe(false); + }); + it("should have numeric embeddingDimensions", () => { expect(typeof CONFIG.embeddingDimensions).toBe("number"); expect(CONFIG.embeddingDimensions).toBeGreaterThan(0); diff --git a/tests/embedding-task-prefixes.test.ts b/tests/embedding-task-prefixes.test.ts new file mode 100644 index 0000000..a48509c --- /dev/null +++ b/tests/embedding-task-prefixes.test.ts @@ -0,0 +1,41 @@ +import { describe, it, expect } from "bun:test"; +import { applyEmbeddingTaskPrefix } from "../src/services/embedding.js"; + +describe("applyEmbeddingTaskPrefix", () => { + it("returns text unchanged when prefixes are disabled", () => { + expect(applyEmbeddingTaskPrefix("hello", { task: "document" }, false)).toBe("hello"); + expect(applyEmbeddingTaskPrefix("hello", { task: "query" }, false)).toBe("hello"); + }); + + it("returns text unchanged when no task is provided", () => { + expect(applyEmbeddingTaskPrefix("hello", undefined, true)).toBe("hello"); + expect(applyEmbeddingTaskPrefix("hello", {}, true)).toBe("hello"); + }); + + it("applies search_document prefix for document task", () => { + expect(applyEmbeddingTaskPrefix("hello", { task: "document" }, true)).toBe( + "search_document: hello" + ); + }); + + it("applies search_query prefix for query task", () => { + expect(applyEmbeddingTaskPrefix("hello", { task: "query" }, true)).toBe("search_query: hello"); + }); + + it("keeps document and query cache keys distinct for the same text", () => { + const documentKey = applyEmbeddingTaskPrefix("same", { task: "document" }, true); + const queryKey = applyEmbeddingTaskPrefix("same", { task: "query" }, true); + expect(documentKey).not.toBe(queryKey); + expect(documentKey).toBe("search_document: same"); + expect(queryKey).toBe("search_query: same"); + }); + + it("does not double-prefix an already prefixed string", () => { + expect(applyEmbeddingTaskPrefix("search_document: hello", { task: "document" }, true)).toBe( + "search_document: hello" + ); + expect(applyEmbeddingTaskPrefix("search_query: hello", { task: "query" }, true)).toBe( + "search_query: hello" + ); + }); +});