Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 13 additions & 12 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -280,21 +280,22 @@ the user profile to be injected when enabled.

Embeddings power similarity search for memories and the user profile. Configure them in the same file (`~/.config/opencode/opencode-mem.jsonc`). There is **no MLX backend** — local embeddings use `@huggingface/transformers` with ONNX, not Apple MLX.

**Local (default):** set only `embeddingModel`. On first use the model is downloaded from Hugging Face and cached under `{storagePath}/.cache` (default `~/.opencode-mem/data/.cache`).
**Local (default):** set only `embeddingModel`. On first use the model is downloaded from Hugging Face and cached under `{storagePath}/.cache` (default `~/.opencode-mem/data/.cache`). The weights stay resident for the life of the host process; set `embeddingIdleTimeoutMs` to release them after that many idle milliseconds, and the next memory operation re-warms them from the local cache (about a second, not a re-download).

**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 |
| `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`) |
| 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 })` |
| `embeddingIdleTimeoutMs` | Release the local model after this many idle milliseconds (e.g. `600000` for 10 minutes). `0` (default) keeps it loaded for the process lifetime |
| `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:

Expand Down
21 changes: 21 additions & 0 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,11 @@ interface OpenCodeMemConfig {
* When set, passed to `pipeline({ dtype })` and outranks auto-detection.
*/
embeddingDtype?: string;
/**
* Release the local ONNX pipeline after this many idle milliseconds.
* `0` (default) keeps it loaded for the process lifetime.
*/
embeddingIdleTimeoutMs?: number;
similarityThreshold?: number;
maxMemories?: number;
maxProfileItems?: number;
Expand Down Expand Up @@ -286,6 +291,7 @@ const DEFAULTS: Required<
embeddingDimensions: 768,
embeddingUseTaskPrefixes: false,
embeddingPooling: "mean",
embeddingIdleTimeoutMs: 0,
similarityThreshold: 0.6,
maxMemories: 10,
maxProfileItems: 5,
Expand Down Expand Up @@ -444,6 +450,10 @@ const CONFIG_TEMPLATE = `{

// Optional local ONNX dtype override (e.g. "q8", "fp32"). When unset, transformers.js defaults apply.
// "embeddingDtype": "q8",

// Release the local model after this many ms without embedding work.
// 0 (default) keeps it loaded for the whole process lifetime.
// "embeddingIdleTimeoutMs": 600000,

// Auto-detected dimensions (no need to set manually)
// "embeddingDimensions": 768,
Expand Down Expand Up @@ -941,6 +951,13 @@ export function normalizeEmbeddingDtype(value: unknown): LocalOnnxDtype | undefi
);
}

export function normalizeEmbeddingIdleTimeoutMs(value: number): number {
if (!Number.isInteger(value) || value < 0) {
throw new Error(`Invalid embeddingIdleTimeoutMs config: ${String(value)}`);
}
return value;
}

function getEmbeddingDimensions(model: string): number {
const dimensionMap: Record<string, number> = {
// Local Xenova models
Expand Down Expand Up @@ -1122,6 +1139,9 @@ function buildConfig(fileConfig: OpenCodeMemConfig) {
? fileConfig.embeddingDocumentPrefix
: embeddingPreset?.documentPrefix;
const embeddingDtype = normalizeEmbeddingDtype(fileConfig.embeddingDtype);
const embeddingIdleTimeoutMs = normalizeEmbeddingIdleTimeoutMs(
fileConfig.embeddingIdleTimeoutMs ?? DEFAULTS.embeddingIdleTimeoutMs
);
const autoCaptureMaxContextBytes = normalizeAutoCaptureMaxContextBytes(
fileConfig.autoCaptureMaxContextBytes ?? DEFAULTS.autoCaptureMaxContextBytes
);
Expand Down Expand Up @@ -1163,6 +1183,7 @@ function buildConfig(fileConfig: OpenCodeMemConfig) {
embeddingQueryPrefix,
embeddingDocumentPrefix,
embeddingDtype,
embeddingIdleTimeoutMs,
embeddingApiUrl: fileConfig.embeddingApiUrl,
embeddingApiKey: fileConfig.embeddingApiUrl
? resolveSecretValue(fileConfig.embeddingApiKey ?? process.env.OPENAI_API_KEY)
Expand Down
2 changes: 2 additions & 0 deletions src/services/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,8 @@ export class LocalMemoryClient {
const { releasePluginTurso } = await import("./turso/lifecycle.js");
const closed = await releasePluginTurso();
if (closed) {
// Process-wide like the pool: only the last instance may drop the model.
await embeddingService.dispose("client-close");
this.reset();
}
}
Expand Down
113 changes: 94 additions & 19 deletions src/services/embedding.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { CONFIG } from "../config.js";
import { log } from "./logger.js";
import { log, safeLog } from "./logger.js";
import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import {
Expand Down Expand Up @@ -354,6 +354,12 @@ export class EmbeddingService {
public initError: string | null = null;
private cache: Map<string, Float32Array> = new Map();
private cachedFingerprint: string | null = null;
private idleTimer: ReturnType<typeof setTimeout> | null = null;
private inFlight = 0;
/** Release requested while an embed held the pipeline; runs when it drains. */
private pendingDispose: string | null = null;
/** Bumped by dispose so a slow initialize cannot resurrect a released pipeline. */
private generation = 0;

static getInstance(): EmbeddingService {
if (!(globalThis as any)[GLOBAL_EMBEDDING_KEY]) {
Expand All @@ -362,6 +368,51 @@ export class EmbeddingService {
return (globalThis as any)[GLOBAL_EMBEDDING_KEY];
}

private clearIdleTimer(): void {
if (this.idleTimer !== null) {
clearTimeout(this.idleTimer);
this.idleTimer = null;
}
}

private scheduleIdleDispose(): void {
this.clearIdleTimer();
// Remote embeddings hold no model in-process, so there is nothing to release.
if ((CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) || !this.isWarmedUp) return;
const ms = CONFIG.embeddingIdleTimeoutMs;
if (!(ms > 0)) return;
this.idleTimer = setTimeout(() => {
this.idleTimer = null;
void this.dispose("idle-timeout");
}, ms);
(this.idleTimer as any).unref?.();
}

/** Drop the local pipeline so its weights can be reclaimed; the next embed() re-warms it. */
async dispose(reason: string = "dispose"): Promise<void> {
this.clearIdleTimer();
if (this.inFlight > 0) {
// A running embed still holds the pipeline; release once it drains.
this.pendingDispose = reason;
return;
}
this.generation++;
const pipe = this.pipe;
this.pipe = null;
this.isWarmedUp = false;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the existing cold-buffer path (buffered while embeddings are unavailable, drained on the next warm merge, cap 50); warming on demand here would reload the model at session idle, which is when profile learning runs, so it is out of scope for this PR.

this.initPromise = null;
this.initError = null;
if (pipe?.dispose) {
try {
await pipe.dispose();
} catch (error) {
safeLog("Failed to dispose embedding pipeline", { error: String(error) });
}
}
// Teardown path: a logger failure must not escape the idle timer or close().
safeLog("Embedding model released", { reason, model: CONFIG.embeddingModel });
}

async warmup(progressCallback?: (progress: any) => void): Promise<void> {
if (this.isWarmedUp) return;
if (this.initError) throw new Error(this.initError);
Expand All @@ -384,6 +435,7 @@ export class EmbeddingService {
}

// Local model path
const generation = this.generation;
const { pipeline } = await ensureTransformersLoaded();
// Explicit config wins; otherwise auto-resolve from model declaration / cache (#393).
const dtype = CONFIG.embeddingDtype ?? resolveLocalOnnxDtype(CONFIG.embeddingModel);
Expand All @@ -394,12 +446,19 @@ export class EmbeddingService {
source: CONFIG.embeddingDtype ? "config" : "auto",
});
}
this.pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, {
const pipe = await pipeline("feature-extraction", CONFIG.embeddingModel, {
...(dtype ? { dtype } : {}),
progress_callback: progressCallback,
});
if (generation !== this.generation) {
// dispose() ran while the model was loading; keep nothing resident.
await (pipe as any).dispose?.();
return;
}
this.pipe = pipe;
this.isWarmedUp = true;
this.initError = null;
this.scheduleIdleDispose();
log("Embedding model warmed up", {
model: CONFIG.embeddingModel,
pooling: CONFIG.embeddingPooling,
Expand Down Expand Up @@ -433,26 +492,42 @@ export class EmbeddingService {
await this.initPromise;
}

let result: Float32Array;
this.inFlight++;
try {
let result: Float32Array;

if (CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) {
const data: any = await fetchRemoteEmbeddingWithRetry(fetch, input);
result = new Float32Array(data.data[0].embedding);
} else {
const output = await this.pipe(input, {
pooling: CONFIG.embeddingPooling,
normalize: true,
});
result = new Float32Array(output.data);
}
if (CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) {
const data: any = await fetchRemoteEmbeddingWithRetry(fetch, input);
result = new Float32Array(data.data[0].embedding);
} else {
const pipe = this.pipe;
if (!pipe) {
throw new Error("Embedding model was released before the request ran");
}
const output = await pipe(input, {
pooling: CONFIG.embeddingPooling,
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);
if (this.cache.size >= MAX_CACHE_SIZE) {
const firstKey = this.cache.keys().next().value;
if (firstKey !== undefined) this.cache.delete(firstKey);
}
this.cache.set(input, result);

return result;
} finally {
this.inFlight--;
if (this.pendingDispose) {
const reason = this.pendingDispose;
this.pendingDispose = null;
void this.dispose(reason);
} else {
this.scheduleIdleDispose();
}
}
this.cache.set(input, result);

return result;
}

async embedWithTimeout(text: string, options?: EmbedOptions): Promise<Float32Array> {
Expand Down
14 changes: 14 additions & 0 deletions tests/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ const {
isPlaceholderApiKey,
normalizeAutoCleanupRetentionDays,
normalizeAutoCaptureMaxContextBytes,
normalizeEmbeddingIdleTimeoutMs,
} = await import("../src/config.js");

afterAll(() => {
Expand Down Expand Up @@ -51,6 +52,19 @@ describe("config", () => {
expect(CONFIG.embeddingDtype).toBeUndefined();
});

it("should keep the local embedding model loaded by default", async () => {
// Cache-busted: CONFIG otherwise carries the developer's real user config.
const fresh = await import(`../src/config.js?cachebust=${Math.random()}`);
expect(fresh.CONFIG.embeddingIdleTimeoutMs).toBe(0);
});

it("should normalize embeddingIdleTimeoutMs and allow 0 to disable", () => {
expect(normalizeEmbeddingIdleTimeoutMs(0)).toBe(0);
expect(normalizeEmbeddingIdleTimeoutMs(600000)).toBe(600000);
expect(() => normalizeEmbeddingIdleTimeoutMs(-1)).toThrow(/embeddingIdleTimeoutMs/);
expect(() => normalizeEmbeddingIdleTimeoutMs(1.5)).toThrow(/embeddingIdleTimeoutMs/);
});

it("should have numeric embeddingDimensions", () => {
expect(typeof CONFIG.embeddingDimensions).toBe("number");
expect(CONFIG.embeddingDimensions).toBeGreaterThan(0);
Expand Down
103 changes: 103 additions & 0 deletions tests/embedding-idle-dispose.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
import { afterAll, describe, expect, it } from "bun:test";
import { mkdtempSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";

const home = mkdtempSync(join(tmpdir(), "opencode-mem-idle-"));
const originalHome = process.env.HOME;
const originalUserProfile = process.env.USERPROFILE;
process.env.HOME = home;
process.env.USERPROFILE = home;

const { CONFIG } = await import("../src/config.js");
const { EmbeddingService } = await import("../src/services/embedding.js");

afterAll(() => {
process.env.HOME = originalHome;
process.env.USERPROFILE = originalUserProfile;
});

function makeService(): { service: any; disposed: string[] } {
const disposed: string[] = [];
const service = new EmbeddingService();
service.pipe = Object.assign(async () => ({ data: new Float32Array([1, 2, 3]) }), {
dispose: async () => {
disposed.push("dispose");
},
});
service.isWarmedUp = true;
return { service, disposed };
}

async function withIdleTimeout(ms: number, run: () => Promise<void>): Promise<void> {
const previous = CONFIG.embeddingIdleTimeoutMs;
CONFIG.embeddingIdleTimeoutMs = ms;
CONFIG.embeddingApiUrl = undefined;
CONFIG.embeddingApiKey = undefined;
try {
await run();
} finally {
CONFIG.embeddingIdleTimeoutMs = previous;
}
}

describe("EmbeddingService idle release", () => {
it("disposes the local pipeline after the idle timeout", async () => {
await withIdleTimeout(20, async () => {
const { service, disposed } = makeService();
await service.embed("hello");
expect(disposed).toEqual([]);

await new Promise((resolve) => setTimeout(resolve, 80));

expect(disposed).toEqual(["dispose"]);
expect(service.isWarmedUp).toBe(false);
expect(service.pipe).toBeNull();
});
});

it("keeps the model loaded when the timeout is 0", async () => {
await withIdleTimeout(0, async () => {
const { service, disposed } = makeService();
await service.embed("hello");

await new Promise((resolve) => setTimeout(resolve, 80));

expect(disposed).toEqual([]);
expect(service.isWarmedUp).toBe(true);
});
});

it("defers disposal while an embed is in flight and releases it afterwards", async () => {
await withIdleTimeout(0, async () => {
const { service, disposed } = makeService();
let release!: () => void;
const gate = new Promise<void>((resolve) => {
release = resolve;
});
service.pipe = Object.assign(
async () => {
await gate;
return { data: new Float32Array([1]) };
},
{ dispose: async () => disposed.push("dispose") }
);

const running = service.embed("in-flight");
await new Promise((resolve) => setTimeout(resolve, 10));
await service.dispose("client-close");

// The running embed still holds the pipeline, so nothing is released yet.
expect(disposed).toEqual([]);

release();
await running;
await new Promise((resolve) => setTimeout(resolve, 10));

// Release happens without a second dispose() call, even with the timer off.
expect(disposed).toEqual(["dispose"]);
expect(service.isWarmedUp).toBe(false);
expect(service.pipe).toBeNull();
});
});
});
Loading
Loading