From e0b4f653636e907b58524e85c91678feb62d2b94 Mon Sep 17 00:00:00 2001 From: EyJunge1 <149941075+EyJunge1@users.noreply.github.com> Date: Sat, 10 Oct 2026 12:39:56 +0200 Subject: [PATCH] feat(embedding): configurable pooling and task prefixes Allow mean/cls/last_token pooling, custom query/document prefixes, optional dtype override, and known-model presets so bge-m3/e5/Qwen3 work without patching dist. Closes #392. --- README.md | 45 +++++--- opencode-mem.example.jsonc | 14 +++ src/config.ts | 156 +++++++++++++++++++++++++- src/services/embedding.ts | 83 ++++++++++++-- tests/config.test.ts | 8 ++ tests/embedding-model-presets.test.ts | 75 +++++++++++++ tests/embedding-task-prefixes.test.ts | 78 ++++++++++++- 7 files changed, 431 insertions(+), 28 deletions(-) create mode 100644 tests/embedding-model-presets.test.ts diff --git a/README.md b/README.md index f83390b2..269db980 100644 --- a/README.md +++ b/README.md @@ -231,22 +231,29 @@ Embeddings power similarity search for memories and the user profile. Configure **Remote (OpenAI-compatible):** set both `embeddingApiUrl` and `embeddingApiKey`. The plugin then calls `{embeddingApiUrl}/embeddings` with a Bearer token. `embeddingApiKey` accepts the same secret formats as `memoryApiKey` (`literal`, `env://…`, `file://…`). -| Key | Role | -| --------------------- | ----------------------------------------------------------------------------------------- | -| `embeddingModel` | Hugging Face id (local) or API model name (remote). Default: `Xenova/nomic-embed-text-v1` | -| `embeddingDimensions` | Optional override; usually omit — dimensions are looked up from a built-in map | -| `embeddingApiUrl` | Base URL for an OpenAI-compatible embeddings API (no trailing path beyond `/v1`) | -| `embeddingApiKey` | API key for that endpoint (required together with `embeddingApiUrl`) | +| Key | Role | +| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------- | +| `embeddingModel` | Hugging Face id (local) or API model name (remote). Default: `Xenova/nomic-embed-text-v1` | +| `embeddingDimensions` | Optional override; usually omit — dimensions are looked up from a built-in map | +| `embeddingPooling` | Local pooling: `"mean"` (default), `"cls"`, or `"last_token"`. Unset → small known-model preset or `"mean"` | +| `embeddingQueryPrefix` | Prefix for query-task embeddings. Unset → model preset or Nomic when `embeddingUseTaskPrefixes` is true. Explicit `""` disables | +| `embeddingDocumentPrefix` | Prefix for document-task embeddings (same resolution rules as `embeddingQueryPrefix`) | +| `embeddingUseTaskPrefixes` | Opt-in Nomic `search_query:` / `search_document:` prefixes when custom/preset prefixes are unset. Default `false` | +| `embeddingDtype` | Optional local ONNX dtype override (e.g. `"q8"`, `"fp32"`). When set, passed to transformers.js `pipeline({ dtype })` | +| `embeddingApiUrl` | Base URL for an OpenAI-compatible embeddings API (no trailing path beyond `/v1`) | +| `embeddingApiKey` | API key for that endpoint (required together with `embeddingApiUrl`) | Recommended local models: -| Model | Dims | Notes | -| ------------------------------------ | ---- | ----------------------------------- | -| `Xenova/nomic-embed-text-v1` | 768 | Default; multilingual, 8192 context | -| `Xenova/jina-embeddings-v2-base-en` | 768 | English-only, 8192 context | -| `Xenova/jina-embeddings-v2-small-en` | 512 | Faster, 8192 context | -| `Xenova/all-MiniLM-L6-v2` | 384 | Very fast, 512 context | -| `Xenova/all-mpnet-base-v2` | 768 | Good quality, 512 context | +| Model | Dims | Notes | +| ------------------------------------ | ---- | ------------------------------------------------- | +| `Xenova/nomic-embed-text-v1` | 768 | Default; multilingual, 8192 context; mean pooling | +| `Xenova/jina-embeddings-v2-base-en` | 768 | English-only, 8192 context | +| `Xenova/jina-embeddings-v2-small-en` | 512 | Faster, 8192 context | +| `Xenova/all-MiniLM-L6-v2` | 384 | Very fast, 512 context | +| `Xenova/all-mpnet-base-v2` | 768 | Good quality, 512 context | +| `Xenova/bge-m3` | 1024 | Multilingual; auto CLS pooling | +| `intfloat/multilingual-e5-large` | 1024 | Auto `query:` / `passage:` prefixes | Example — remote OpenAI embeddings: @@ -258,7 +265,17 @@ Example — remote OpenAI embeddings: } ``` -Changing `embeddingModel` (or dimensions) can trigger re-embedding of stored memories on next startup. Prefer picking a model once and sticking with it for a given data directory. +Example — local bge-m3 (pooling defaults to CLS via preset; override if needed): + +```jsonc +{ + "embeddingModel": "Xenova/bge-m3", + // "embeddingPooling": "cls", + // "embeddingDtype": "q8", +} +``` + +Changing `embeddingModel`, dimensions, pooling, or task prefixes can require re-embedding stored memories so store and query vectors stay aligned. Prefer picking a model (and pooling/prefix settings) once and sticking with them for a given data directory. **Unsupported — Intel Mac (`darwin/x64`):** Local persistence requires `@tursodatabase/database`, which does not publish an Intel Mac native binding. Fixed `onnxruntime-node` releases (`1.24.1+`, including the pinned `1.30.0`) also lack darwin/x64. Use an Apple Silicon Mac, Linux, or Windows, or a remote endpoint via `embeddingApiUrl` + `embeddingApiKey` (example above). On supported platforms, `opencode-mem` pins `onnxruntime-node@1.30.0` (Ort::Env teardown fix from `1.24.1` / #225) and loads transformers through a CJS resolve shim so OpenCode nested installs keep that binding. Transformers is resolved to an absolute path before that shim is installed so OpenCode's Bun `--compile` host does not fail with `Cannot find module '@huggingface/transformers' from ''`. After upgrading, clear OpenCode's nested plugin cache (`~/.cache/opencode/packages/opencode-mem@*`) and reinstall. diff --git a/opencode-mem.example.jsonc b/opencode-mem.example.jsonc index 7967207b..50f6fb9c 100644 --- a/opencode-mem.example.jsonc +++ b/opencode-mem.example.jsonc @@ -26,6 +26,18 @@ // re-index existing memories so store and query vectors stay aligned. "embeddingUseTaskPrefixes": false, + // Local pooling: "mean" (default) | "cls" | "last_token". + // Unset → model preset (bge-m3=cls, Qwen3-Embedding=last_token) or "mean". + // "embeddingPooling": "mean", + + // Custom task prefixes (override presets / Nomic). Explicit "" disables a side. + // Changing pooling or prefixes requires re-embedding stored memories. + // "embeddingQueryPrefix": "query: ", + // "embeddingDocumentPrefix": "passage: ", + + // Optional local ONNX dtype override (e.g. "q8", "fp32"). When unset, transformers.js defaults apply. + // "embeddingDtype": "q8", + // Auto-detected dimensions (no need to set manually) // "embeddingDimensions": 768, @@ -34,6 +46,8 @@ // "embeddingModel": "Xenova/jina-embeddings-v2-small-en", // 512 dims, faster, 8192 context // "embeddingModel": "Xenova/all-MiniLM-L6-v2", // 384 dims, very fast, 512 context // "embeddingModel": "Xenova/all-mpnet-base-v2", // 768 dims, good quality, 512 context + // "embeddingModel": "Xenova/bge-m3", // 1024 dims, CLS pooling (auto) + // "embeddingModel": "intfloat/multilingual-e5-large", // 1024 dims, query:/passage: prefixes (auto) // Optional: OpenAI-compatible API for embeddings (both URL and key required) // "embeddingApiUrl": "https://api.openai.com/v1", diff --git a/src/config.ts b/src/config.ts index c672b9ac..ac9ac210 100644 --- a/src/config.ts +++ b/src/config.ts @@ -34,6 +34,26 @@ interface OpenCodeMemConfig { embeddingApiKey?: string; /** Opt-in Nomic task prefixes (`search_document:` / `search_query:`). Default false. */ embeddingUseTaskPrefixes?: boolean; + /** + * Local transformers.js pooling strategy. Default `"mean"`. + * Model presets (bge-m3 → cls, Qwen3-Embedding → last_token) apply when unset. + */ + embeddingPooling?: "mean" | "cls" | "last_token"; + /** + * Prefix applied to query-task embeddings. When unset, model presets or + * `embeddingUseTaskPrefixes` (Nomic) may supply one. Explicit `""` disables. + */ + embeddingQueryPrefix?: string; + /** + * Prefix applied to document-task embeddings. Same resolution rules as + * `embeddingQueryPrefix`. + */ + embeddingDocumentPrefix?: string; + /** + * Optional dtype for the local ONNX pipeline (e.g. `"q8"`, `"fp32"`). + * When set, passed to `pipeline({ dtype })` and outranks auto-detection. + */ + embeddingDtype?: string; similarityThreshold?: number; maxMemories?: number; maxProfileItems?: number; @@ -158,6 +178,9 @@ const DEFAULTS: Required< OpenCodeMemConfig, | "embeddingApiUrl" | "embeddingApiKey" + | "embeddingQueryPrefix" + | "embeddingDocumentPrefix" + | "embeddingDtype" | "memoryModel" | "memoryApiUrl" | "memoryApiKey" @@ -180,6 +203,9 @@ const DEFAULTS: Required< > & { embeddingApiUrl?: string; embeddingApiKey?: string; + embeddingQueryPrefix?: string; + embeddingDocumentPrefix?: string; + embeddingDtype?: string; memoryModel?: string; memoryApiUrl?: string; memoryApiKey?: string; @@ -207,6 +233,7 @@ const DEFAULTS: Required< embeddingModel: "Xenova/nomic-embed-text-v1", embeddingDimensions: 768, embeddingUseTaskPrefixes: false, + embeddingPooling: "mean", similarityThreshold: 0.6, maxMemories: 10, maxProfileItems: 5, @@ -341,6 +368,18 @@ const CONFIG_TEMPLATE = `{ // Opt-in Nomic task prefixes (search_document: / search_query:). After enabling, // re-index existing memories so store and query vectors stay aligned. // "embeddingUseTaskPrefixes": true, + + // Local pooling: "mean" (default) | "cls" | "last_token". + // Unset → model preset (bge-m3=cls, Qwen3-Embedding=last_token) or "mean". + // "embeddingPooling": "mean", + + // Custom task prefixes (override presets / Nomic). Explicit "" disables a side. + // Changing pooling or prefixes requires re-embedding stored memories. + // "embeddingQueryPrefix": "query: ", + // "embeddingDocumentPrefix": "passage: ", + + // Optional local ONNX dtype override (e.g. "q8", "fp32"). When unset, transformers.js defaults apply. + // "embeddingDtype": "q8", // Auto-detected dimensions (no need to set manually) // "embeddingDimensions": 768, @@ -350,6 +389,8 @@ const CONFIG_TEMPLATE = `{ // "embeddingModel": "Xenova/jina-embeddings-v2-small-en", // 512 dims, faster, 8192 context // "embeddingModel": "Xenova/all-MiniLM-L6-v2", // 384 dims, very fast, 512 context // "embeddingModel": "Xenova/all-mpnet-base-v2", // 768 dims, good quality, 512 context + // "embeddingModel": "Xenova/bge-m3", // 1024 dims, CLS pooling (auto) + // "embeddingModel": "intfloat/multilingual-e5-large", // 1024 dims, query:/passage: prefixes (auto) // Optional: OpenAI-compatible API for embeddings (both URL and key required) // "embeddingApiUrl": "https://api.openai.com/v1", @@ -696,6 +737,91 @@ function ensureConfigExists(): void { ensureConfigExists(); +export type EmbeddingPooling = "mean" | "cls" | "last_token"; + +export type EmbeddingModelPreset = { + pooling: EmbeddingPooling; + /** + * When present (including `""`), used as the default query/document prefix + * when the corresponding config key is unset. Omitting both keeps Nomic + * opt-in via `embeddingUseTaskPrefixes` (nomic models). + */ + queryPrefix?: string; + documentPrefix?: string; +}; + +const QWEN3_EMBEDDING_QUERY_PREFIX = + "Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: "; + +/** + * Known local embedding model presets (substring match on `embeddingModel`). + * Explicit config keys always win over these defaults. + */ +export function getEmbeddingModelPreset(model: string): EmbeddingModelPreset | undefined { + const id = model.toLowerCase(); + if (id.includes("bge-m3")) { + return { pooling: "cls", queryPrefix: "", documentPrefix: "" }; + } + if (id.includes("multilingual-e5") || id.includes("e5-")) { + return { pooling: "mean", queryPrefix: "query: ", documentPrefix: "passage: " }; + } + if (id.includes("qwen3-embedding")) { + return { + pooling: "last_token", + queryPrefix: QWEN3_EMBEDDING_QUERY_PREFIX, + documentPrefix: "", + }; + } + if (id.includes("jina-embeddings")) { + return { pooling: "mean", queryPrefix: "", documentPrefix: "" }; + } + if (id.includes("nomic-embed")) { + return { pooling: "mean" }; + } + return undefined; +} + +export function normalizeEmbeddingPooling(value: unknown): EmbeddingPooling { + if (value === "mean" || value === "cls" || value === "last_token") return value; + throw new Error( + `Invalid embeddingPooling config: ${String(value)} (expected "mean" | "cls" | "last_token")` + ); +} + +/** Dtypes accepted by transformers.js `pipeline({ dtype })`. */ +export const LOCAL_ONNX_DTYPES = [ + "auto", + "fp32", + "fp16", + "q8", + "int8", + "uint8", + "q4", + "bnb4", + "q4f16", + "q2", + "q2f16", + "q1", + "q1f16", +] as const; + +export type LocalOnnxDtype = (typeof LOCAL_ONNX_DTYPES)[number]; + +export function normalizeEmbeddingDtype(value: unknown): LocalOnnxDtype | undefined { + if (value === undefined || value === null) return undefined; + if (typeof value !== "string") { + throw new Error(`Invalid embeddingDtype config: ${String(value)}`); + } + const trimmed = value.trim(); + if (!trimmed) return undefined; + if ((LOCAL_ONNX_DTYPES as readonly string[]).includes(trimmed)) { + return trimmed as LocalOnnxDtype; + } + throw new Error( + `Invalid embeddingDtype config: ${trimmed} (expected one of ${LOCAL_ONNX_DTYPES.join(", ")})` + ); +} + function getEmbeddingDimensions(model: string): number { const dimensionMap: Record = { // Local Xenova models @@ -711,10 +837,17 @@ function getEmbeddingDimensions(model: string): number { "Xenova/all-mpnet-base-v2": 768, "Xenova/bge-base-en-v1.5": 768, "Xenova/bge-small-en-v1.5": 384, + "Xenova/bge-m3": 1024, "Xenova/gte-small": 384, "Xenova/GIST-small-Embedding-v0": 384, "Xenova/text-embedding-ada-002": 1536, + // Other local / ONNX community models + "intfloat/multilingual-e5-large": 1024, + "intfloat/multilingual-e5-base": 768, + "intfloat/e5-large-v2": 1024, + "onnx-community/Qwen3-Embedding-0.6B-ONNX": 1024, + // OpenAI API models "text-embedding-3-small": 1536, "text-embedding-3-large": 3072, @@ -831,9 +964,22 @@ function buildConfig(fileConfig: OpenCodeMemConfig) { const memoryApiKey = resolveSecretValue(fileConfig.memoryApiKey) ?? (memoryProvider === "atlas-cloud" ? process.env.ATLASCLOUD_API_KEY : undefined); + const embeddingModel = fileConfig.embeddingModel ?? DEFAULTS.embeddingModel; + const embeddingPreset = getEmbeddingModelPreset(embeddingModel); const embeddingDimensions = - fileConfig.embeddingDimensions ?? - getEmbeddingDimensions(fileConfig.embeddingModel ?? DEFAULTS.embeddingModel); + fileConfig.embeddingDimensions ?? getEmbeddingDimensions(embeddingModel); + const embeddingPooling = normalizeEmbeddingPooling( + fileConfig.embeddingPooling ?? embeddingPreset?.pooling ?? DEFAULTS.embeddingPooling + ); + const embeddingQueryPrefix = + fileConfig.embeddingQueryPrefix !== undefined + ? fileConfig.embeddingQueryPrefix + : embeddingPreset?.queryPrefix; + const embeddingDocumentPrefix = + fileConfig.embeddingDocumentPrefix !== undefined + ? fileConfig.embeddingDocumentPrefix + : embeddingPreset?.documentPrefix; + const embeddingDtype = normalizeEmbeddingDtype(fileConfig.embeddingDtype); const autoCaptureMaxContextBytes = normalizeAutoCaptureMaxContextBytes( fileConfig.autoCaptureMaxContextBytes ?? DEFAULTS.autoCaptureMaxContextBytes ); @@ -857,10 +1003,14 @@ function buildConfig(fileConfig: OpenCodeMemConfig) { storagePath: expandPath(fileConfig.storagePath ?? DEFAULTS.storagePath), userEmailOverride: fileConfig.userEmailOverride, userNameOverride: fileConfig.userNameOverride, - embeddingModel: fileConfig.embeddingModel ?? DEFAULTS.embeddingModel, + embeddingModel, embeddingDimensions, embeddingUseTaskPrefixes: fileConfig.embeddingUseTaskPrefixes ?? DEFAULTS.embeddingUseTaskPrefixes, + embeddingPooling, + embeddingQueryPrefix, + embeddingDocumentPrefix, + embeddingDtype, embeddingApiUrl: fileConfig.embeddingApiUrl, embeddingApiKey: fileConfig.embeddingApiUrl ? resolveSecretValue(fileConfig.embeddingApiKey ?? process.env.OPENAI_API_KEY) diff --git a/src/services/embedding.ts b/src/services/embedding.ts index 789594cd..efbb9ffb 100644 --- a/src/services/embedding.ts +++ b/src/services/embedding.ts @@ -19,26 +19,79 @@ export type EmbedOptions = { task?: EmbeddingTask; }; -const TASK_PREFIXES: Record = { +const NOMIC_TASK_PREFIXES: Record = { document: "search_document: ", query: "search_query: ", }; +export type EmbeddingTaskPrefixOptions = { + useTaskPrefixes?: boolean; + queryPrefix?: string; + documentPrefix?: string; +}; + +/** + * Resolve the prefix for a task: explicit/custom → Nomic (when enabled) → none. + * Empty string means "no prefix" (disables Nomic fallback for that task). + */ +export function resolveEmbeddingTaskPrefix( + task: EmbeddingTask, + options: EmbeddingTaskPrefixOptions = {} +): string { + const custom = task === "query" ? options.queryPrefix : options.documentPrefix; + if (custom !== undefined) return custom; + if (options.useTaskPrefixes) return NOMIC_TASK_PREFIXES[task]; + return ""; +} + /** - * Apply Nomic-style task prefixes when enabled. + * Apply task prefixes when configured. * Cache keys and API/local inputs should use the returned string. + * + * Third argument accepts the legacy boolean (`useTaskPrefixes`) or an options + * object. Defaults read from CONFIG (custom/preset prefixes + Nomic flag). */ export function applyEmbeddingTaskPrefix( text: string, options?: EmbedOptions, - useTaskPrefixes: boolean = CONFIG.embeddingUseTaskPrefixes + useTaskPrefixesOrOpts: boolean | EmbeddingTaskPrefixOptions = { + useTaskPrefixes: CONFIG.embeddingUseTaskPrefixes, + queryPrefix: CONFIG.embeddingQueryPrefix, + documentPrefix: CONFIG.embeddingDocumentPrefix, + } ): string { - if (!useTaskPrefixes || !options?.task) return text; - const prefix = TASK_PREFIXES[options.task]; + if (!options?.task) return text; + + const prefixOpts: EmbeddingTaskPrefixOptions = + typeof useTaskPrefixesOrOpts === "boolean" + ? { + useTaskPrefixes: useTaskPrefixesOrOpts, + queryPrefix: CONFIG.embeddingQueryPrefix, + documentPrefix: CONFIG.embeddingDocumentPrefix, + } + : { + useTaskPrefixes: useTaskPrefixesOrOpts.useTaskPrefixes ?? CONFIG.embeddingUseTaskPrefixes, + queryPrefix: useTaskPrefixesOrOpts.queryPrefix ?? CONFIG.embeddingQueryPrefix, + documentPrefix: useTaskPrefixesOrOpts.documentPrefix ?? CONFIG.embeddingDocumentPrefix, + }; + + const prefix = resolveEmbeddingTaskPrefix(options.task, prefixOpts); + if (!prefix) return text; if (text.startsWith(prefix)) return text; return `${prefix}${text}`; } +function embeddingCacheFingerprint(): string { + return [ + CONFIG.embeddingModel, + CONFIG.embeddingPooling, + CONFIG.embeddingQueryPrefix ?? "", + CONFIG.embeddingDocumentPrefix ?? "", + CONFIG.embeddingUseTaskPrefixes ? "1" : "0", + CONFIG.embeddingDtype ?? "", + ].join("\0"); +} + type HfTransformers = typeof import("@huggingface/transformers"); let _transformers: { @@ -114,7 +167,7 @@ export class EmbeddingService { /** Set when warmup fails permanently; prevents "initializing forever" (#184). */ public initError: string | null = null; private cache: Map = new Map(); - private cachedModelName: string | null = null; + private cachedFingerprint: string | null = null; static getInstance(): EmbeddingService { if (!(globalThis as any)[GLOBAL_EMBEDDING_KEY]) { @@ -164,12 +217,18 @@ export class EmbeddingService { // Local model path const { pipeline } = await ensureTransformersLoaded(); + const dtype = CONFIG.embeddingDtype; this.pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, { + ...(dtype ? { dtype } : {}), progress_callback: progressCallback, }); this.isWarmedUp = true; this.initError = null; - log("Embedding model warmed up", { model: CONFIG.embeddingModel }); + log("Embedding model warmed up", { + model: CONFIG.embeddingModel, + pooling: CONFIG.embeddingPooling, + dtype: dtype ?? null, + }); } catch (error) { const rewritten = formatOnnxruntimeInitError(error); this.initPromise = null; @@ -180,9 +239,10 @@ export class EmbeddingService { } async embed(text: string, options?: EmbedOptions): Promise { - if (this.cachedModelName !== CONFIG.embeddingModel) { + const fingerprint = embeddingCacheFingerprint(); + if (this.cachedFingerprint !== fingerprint) { this.clearCache(); - this.cachedModelName = CONFIG.embeddingModel; + this.cachedFingerprint = fingerprint; } const input = applyEmbeddingTaskPrefix(text, options); @@ -219,7 +279,10 @@ export class EmbeddingService { const data: any = await response.json(); result = new Float32Array(data.data[0].embedding); } else { - const output = await this.pipe(input, { pooling: "mean", normalize: true }); + const output = await this.pipe(input, { + pooling: CONFIG.embeddingPooling, + normalize: true, + }); result = new Float32Array(output.data); } diff --git a/tests/config.test.ts b/tests/config.test.ts index 94c408c4..d1b569c3 100644 --- a/tests/config.test.ts +++ b/tests/config.test.ts @@ -39,6 +39,14 @@ describe("config", () => { expect(CONFIG.embeddingUseTaskPrefixes).toBe(false); }); + it("should default embeddingPooling to mean", () => { + expect(CONFIG.embeddingPooling).toBe("mean"); + }); + + it("should leave embeddingDtype unset by default", () => { + expect(CONFIG.embeddingDtype).toBeUndefined(); + }); + it("should have numeric embeddingDimensions", () => { expect(typeof CONFIG.embeddingDimensions).toBe("number"); expect(CONFIG.embeddingDimensions).toBeGreaterThan(0); diff --git a/tests/embedding-model-presets.test.ts b/tests/embedding-model-presets.test.ts new file mode 100644 index 00000000..9911c579 --- /dev/null +++ b/tests/embedding-model-presets.test.ts @@ -0,0 +1,75 @@ +import { describe, expect, it } from "bun:test"; +import { + getEmbeddingModelPreset, + normalizeEmbeddingDtype, + normalizeEmbeddingPooling, +} from "../src/config.js"; + +describe("getEmbeddingModelPreset", () => { + it("returns mean pooling for nomic without forcing task prefixes", () => { + expect(getEmbeddingModelPreset("Xenova/nomic-embed-text-v1")).toEqual({ pooling: "mean" }); + }); + + it("returns cls pooling and empty prefixes for bge-m3", () => { + expect(getEmbeddingModelPreset("Xenova/bge-m3")).toEqual({ + pooling: "cls", + queryPrefix: "", + documentPrefix: "", + }); + }); + + it("returns e5 query/passage prefixes", () => { + expect(getEmbeddingModelPreset("intfloat/multilingual-e5-large")).toEqual({ + pooling: "mean", + queryPrefix: "query: ", + documentPrefix: "passage: ", + }); + expect(getEmbeddingModelPreset("intfloat/e5-large-v2")?.queryPrefix).toBe("query: "); + }); + + it("returns last_token pooling and instruct query prefix for Qwen3-Embedding", () => { + const preset = getEmbeddingModelPreset("onnx-community/Qwen3-Embedding-0.6B-ONNX"); + expect(preset?.pooling).toBe("last_token"); + expect(preset?.documentPrefix).toBe(""); + expect(preset?.queryPrefix).toContain("Instruct:"); + expect(preset?.queryPrefix).toContain("Query: "); + }); + + it("returns mean pooling and empty prefixes for jina", () => { + expect(getEmbeddingModelPreset("Xenova/jina-embeddings-v2-base-en")).toEqual({ + pooling: "mean", + queryPrefix: "", + documentPrefix: "", + }); + }); + + it("returns undefined for unknown models", () => { + expect(getEmbeddingModelPreset("Xenova/all-MiniLM-L6-v2")).toBeUndefined(); + }); +}); + +describe("normalizeEmbeddingPooling", () => { + it("accepts mean, cls, and last_token", () => { + expect(normalizeEmbeddingPooling("mean")).toBe("mean"); + expect(normalizeEmbeddingPooling("cls")).toBe("cls"); + expect(normalizeEmbeddingPooling("last_token")).toBe("last_token"); + }); + + it("rejects invalid values", () => { + expect(() => normalizeEmbeddingPooling("max")).toThrow(/embeddingPooling/); + expect(() => normalizeEmbeddingPooling(1)).toThrow(/embeddingPooling/); + }); +}); + +describe("normalizeEmbeddingDtype", () => { + it("accepts known transformers.js dtypes and trims whitespace", () => { + expect(normalizeEmbeddingDtype("q8")).toBe("q8"); + expect(normalizeEmbeddingDtype(" fp32 ")).toBe("fp32"); + expect(normalizeEmbeddingDtype("")).toBeUndefined(); + expect(normalizeEmbeddingDtype(undefined)).toBeUndefined(); + }); + + it("rejects unknown dtypes", () => { + expect(() => normalizeEmbeddingDtype("bfloat16")).toThrow(/embeddingDtype/); + }); +}); diff --git a/tests/embedding-task-prefixes.test.ts b/tests/embedding-task-prefixes.test.ts index a48509cc..da5bdad3 100644 --- a/tests/embedding-task-prefixes.test.ts +++ b/tests/embedding-task-prefixes.test.ts @@ -1,5 +1,43 @@ import { describe, it, expect } from "bun:test"; -import { applyEmbeddingTaskPrefix } from "../src/services/embedding.js"; +import { applyEmbeddingTaskPrefix, resolveEmbeddingTaskPrefix } from "../src/services/embedding.js"; + +describe("resolveEmbeddingTaskPrefix", () => { + it("returns empty when prefixes are disabled and no custom prefix is set", () => { + expect(resolveEmbeddingTaskPrefix("document", { useTaskPrefixes: false })).toBe(""); + expect(resolveEmbeddingTaskPrefix("query", { useTaskPrefixes: false })).toBe(""); + }); + + it("returns Nomic prefixes when useTaskPrefixes is true", () => { + expect(resolveEmbeddingTaskPrefix("document", { useTaskPrefixes: true })).toBe( + "search_document: " + ); + expect(resolveEmbeddingTaskPrefix("query", { useTaskPrefixes: true })).toBe("search_query: "); + }); + + it("prefers custom prefixes over Nomic", () => { + expect( + resolveEmbeddingTaskPrefix("query", { + useTaskPrefixes: true, + queryPrefix: "query: ", + }) + ).toBe("query: "); + expect( + resolveEmbeddingTaskPrefix("document", { + useTaskPrefixes: true, + documentPrefix: "passage: ", + }) + ).toBe("passage: "); + }); + + it("treats an explicit empty custom prefix as no prefix", () => { + expect( + resolveEmbeddingTaskPrefix("query", { + useTaskPrefixes: true, + queryPrefix: "", + }) + ).toBe(""); + }); +}); describe("applyEmbeddingTaskPrefix", () => { it("returns text unchanged when prefixes are disabled", () => { @@ -38,4 +76,42 @@ describe("applyEmbeddingTaskPrefix", () => { "search_query: hello" ); }); + + it("applies custom prefixes from options object", () => { + expect( + applyEmbeddingTaskPrefix( + "hello", + { task: "query" }, + { + useTaskPrefixes: false, + queryPrefix: "query: ", + documentPrefix: "passage: ", + } + ) + ).toBe("query: hello"); + expect( + applyEmbeddingTaskPrefix( + "hello", + { task: "document" }, + { + useTaskPrefixes: false, + queryPrefix: "query: ", + documentPrefix: "passage: ", + } + ) + ).toBe("passage: hello"); + }); + + it("skips prefixing when custom prefix is empty even if Nomic is enabled", () => { + expect( + applyEmbeddingTaskPrefix( + "hello", + { task: "query" }, + { + useTaskPrefixes: true, + queryPrefix: "", + } + ) + ).toBe("hello"); + }); });