From 056f52b054c4b6f69791543ea4c28c7a090442f2 Mon Sep 17 00:00:00 2001 From: capyBearista Date: Sun, 3 May 2026 03:46:48 -0400 Subject: [PATCH] feat(embedding): improve robustness with pre-flight checks and URL normalization --- src/config.ts | 24 ++++++++++- src/index.ts | 29 +++++++++++++ src/services/embedding.ts | 90 ++++++++++++++++++++++++++++++++++++--- 3 files changed, 134 insertions(+), 9 deletions(-) diff --git a/src/config.ts b/src/config.ts index 8fad5da..e9c4898 100644 --- a/src/config.ts +++ b/src/config.ts @@ -212,6 +212,8 @@ const CONFIG_TEMPLATE = `{ // "embeddingModel": "Xenova/all-mpnet-base-v2", // 768 dims, good quality, 512 context // Optional: Use OpenAI-compatible API for embeddings + // The /v1 prefix is auto-appended if missing, and trailing /embeddings is stripped. + // You can use a plain base URL: "embeddingApiUrl": "https://api.openai.com" // "embeddingApiUrl": "https://api.openai.com/v1", // "embeddingApiKey": "sk-...", // "embeddingModel": "text-embedding-3-small", // 1536 dims, auto-detected @@ -480,6 +482,22 @@ function getEmbeddingDimensions(model: string): number { return dimensionMap[model] || 768; } +function normalizeEmbeddingUrl(url: string): string { + let normalized = url.replace(/\/+$/, ""); + if (normalized.endsWith("/embeddings")) { + normalized = normalized.slice(0, -"/embeddings".length); + } + normalized = normalized.replace(/\/+$/, ""); + if (!normalized.endsWith("/v1")) { + normalized = `${normalized}/v1`; + } + return normalized; +} + +function stripTrailingSlash(url: string): string { + return url.replace(/\/+$/, ""); +} + function buildConfig(fileConfig: OpenCodeMemConfig) { return { storagePath: expandPath(fileConfig.storagePath ?? DEFAULTS.storagePath), @@ -489,7 +507,9 @@ function buildConfig(fileConfig: OpenCodeMemConfig) { embeddingDimensions: fileConfig.embeddingDimensions ?? getEmbeddingDimensions(fileConfig.embeddingModel ?? DEFAULTS.embeddingModel), - embeddingApiUrl: fileConfig.embeddingApiUrl, + embeddingApiUrl: fileConfig.embeddingApiUrl + ? normalizeEmbeddingUrl(fileConfig.embeddingApiUrl) + : undefined, embeddingApiKey: fileConfig.embeddingApiUrl ? resolveSecretValue(fileConfig.embeddingApiKey ?? process.env.OPENAI_API_KEY) : undefined, @@ -509,7 +529,7 @@ function buildConfig(fileConfig: OpenCodeMemConfig) { | "openai-responses" | "anthropic", memoryModel: fileConfig.memoryModel, - memoryApiUrl: fileConfig.memoryApiUrl, + memoryApiUrl: fileConfig.memoryApiUrl ? stripTrailingSlash(fileConfig.memoryApiUrl) : undefined, memoryApiKey: resolveSecretValue(fileConfig.memoryApiKey), memoryTemperature: fileConfig.memoryTemperature, memoryExtraParams: fileConfig.memoryExtraParams, diff --git a/src/index.ts b/src/index.ts index 141c89c..c57a79a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,6 +3,7 @@ import type { Part } from "@opencode-ai/sdk"; import { tool } from "@opencode-ai/plugin"; import { memoryClient } from "./services/client.js"; +import { embeddingService } from "./services/embedding.js"; import { formatContextForPrompt } from "./services/context.js"; import { getTags } from "./services/tags.js"; import { stripPrivateContent, isFullyPrivate } from "./services/privacy.js"; @@ -33,8 +34,36 @@ export const OpenCodeMemPlugin: Plugin = async (ctx: PluginInput) => { try { await memoryClient.warmup(); (globalThis as any)[GLOBAL_PLUGIN_WARMUP_KEY] = true; + + if (!embeddingService.localEmbeddingsAvailable && !CONFIG.embeddingApiUrl) { + if (ctx.client?.tui) { + ctx.client.tui + .showToast({ + body: { + title: "Memory System", + message: + "Local embedding model unavailable. Configure embeddingApiUrl and embeddingApiKey in opencode-mem.jsonc to use API-based embeddings.", + variant: "warning", + duration: 10000, + }, + }) + .catch(() => {}); + } + } } catch (error) { log("Plugin warmup failed", { error: String(error) }); + if (ctx.client?.tui) { + ctx.client.tui + .showToast({ + body: { + title: "Memory System Error", + message: `Failed to initialize: ${error instanceof Error ? error.message : String(error)}`, + variant: "error", + duration: 10000, + }, + }) + .catch(() => {}); + } } } diff --git a/src/services/embedding.ts b/src/services/embedding.ts index 8473c09..514e6b3 100644 --- a/src/services/embedding.ts +++ b/src/services/embedding.ts @@ -1,11 +1,14 @@ import { CONFIG } from "../config.js"; import { log } from "./logger.js"; import { join } from "node:path"; +import { spawn } from "node:child_process"; const TIMEOUT_MS = 30000; const GLOBAL_EMBEDDING_KEY = Symbol.for("opencode-mem.embedding.instance"); const MAX_CACHE_SIZE = 100; +let localEmbeddingsFailed = false; + let _transformers: { pipeline: (typeof import("@huggingface/transformers"))["pipeline"]; env: (typeof import("@huggingface/transformers"))["env"]; @@ -13,12 +16,21 @@ let _transformers: { async function ensureTransformersLoaded(): Promise> { if (_transformers !== null) return _transformers; - const mod = await import("@huggingface/transformers"); - mod.env.allowLocalModels = true; - mod.env.allowRemoteModels = true; - mod.env.cacheDir = join(CONFIG.storagePath, ".cache"); - _transformers = mod; - return _transformers!; + if (localEmbeddingsFailed) { + throw new Error("Local embedding model previously failed to load (native dependency issue)"); + } + try { + const mod = await import("@huggingface/transformers"); + mod.env.allowLocalModels = true; + mod.env.allowRemoteModels = true; + mod.env.cacheDir = join(CONFIG.storagePath, ".cache"); + _transformers = mod; + return _transformers!; + } catch (error) { + localEmbeddingsFailed = true; + log("Failed to load @huggingface/transformers", { error: String(error) }); + throw error; + } } function withTimeout(promise: Promise, ms: number): Promise { @@ -28,6 +40,44 @@ function withTimeout(promise: Promise, ms: number): Promise { ]); } +function runTransformersPreflight(): Promise { + return new Promise((resolve) => { + const child = spawn( + process.execPath, + ["-e", "await import('@huggingface/transformers');process.stdout.write('OK')"], + { stdio: ["ignore", "pipe", "pipe"], timeout: 10000 } + ); + + let stdout = ""; + let stderr = ""; + + child.stdout.on("data", (d: Buffer) => { + stdout += d.toString(); + }); + child.stderr.on("data", (d: Buffer) => { + stderr += d.toString(); + }); + + child.on("close", (code: number | null, signal: NodeJS.Signals | null) => { + if (code === 0 && stdout.includes("OK")) { + resolve(true); + } else { + log("Transformers pre-flight check failed", { + code, + signal, + stderr: stderr.slice(0, 500), + }); + resolve(false); + } + }); + + child.on("error", (err: Error) => { + log("Transformers pre-flight spawn error", { error: String(err) }); + resolve(false); + }); + }); +} + export class EmbeddingService { private pipe: any = null; private initPromise: Promise | null = null; @@ -35,6 +85,10 @@ export class EmbeddingService { private cache: Map = new Map(); private cachedModelName: string | null = null; + get localEmbeddingsAvailable(): boolean { + return !localEmbeddingsFailed; + } + static getInstance(): EmbeddingService { if (!(globalThis as any)[GLOBAL_EMBEDDING_KEY]) { (globalThis as any)[GLOBAL_EMBEDDING_KEY] = new EmbeddingService(); @@ -45,7 +99,24 @@ export class EmbeddingService { async warmup(progressCallback?: (progress: any) => void): Promise { if (this.isWarmedUp) return; if (this.initPromise) return this.initPromise; - this.initPromise = this.initializeModel(progressCallback); + + this.initPromise = (async () => { + if (CONFIG.embeddingApiUrl && CONFIG.embeddingApiKey) { + this.isWarmedUp = true; + return; + } + + const ok = await runTransformersPreflight(); + if (!ok) { + localEmbeddingsFailed = true; + log("Local embedding model unavailable (native dependency check failed)", {}); + this.isWarmedUp = true; + return; + } + + await this.initializeModel(progressCallback); + })(); + return this.initPromise; } @@ -105,6 +176,11 @@ export class EmbeddingService { const data: any = await response.json(); result = new Float32Array(data.data[0].embedding); } else { + if (localEmbeddingsFailed) { + throw new Error( + "Local embedding model unavailable. Configure embeddingApiUrl and embeddingApiKey in opencode-mem.jsonc to use an API-based embedding service." + ); + } const output = await this.pipe(text, { pooling: "mean", normalize: true }); result = new Float32Array(output.data); }