feat(embedding): add opt-in Nomic task prefixes

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",

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)

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);
 

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);

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<EmbeddingTask, string> = {
+  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<Float32Array> {
+  async embed(text: string, options?: EmbedOptions): Promise<Float32Array> {
     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,21 +174,21 @@ 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);
     }
 
     if (this.cache.size >= MAX_CACHE_SIZE) {
       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<Float32Array> {
-    return withTimeout(this.embed(text), TIMEOUT_MS);
+  async embedWithTimeout(text: string, options?: EmbedOptions): Promise<Float32Array> {
+    return withTimeout(this.embed(text, options), TIMEOUT_MS);
   }
 
   clearCache(): void {

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({

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);
 

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);

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"
+    );
+  });
+});