import crypto from 'node:crypto'; import fs from 'node:fs/promises'; import path from 'node:path'; import { loadMemindEnvFiles } from './memind-runtime-profile.mjs'; import { resolveLlmEmbeddingCredentials } from './resolve-llm-embedding-credentials.mjs'; loadMemindEnvFiles(process.cwd()); const memoryCache = new Map(); let diskEntries = null; let diskCachePath = null; let diskDirty = false; let llmCredentialCache = null; function envFlag(value) { return ['1', 'true', 'yes', 'on'].includes(String(value ?? '').trim().toLowerCase()); } async function resolveConfig(env = process.env) { const provider = String(env.MEMIND_EMBEDDING_PROVIDER ?? 'openai').trim().toLowerCase(); let apiKey = String( env.MEMIND_EMBEDDING_API_KEY ?? env.DASHSCOPE_API_KEY ?? env.OPENAI_API_KEY ?? env.OPENROUTER_API_KEY ?? '', ).trim(); let baseUrl = String( env.MEMIND_EMBEDDING_BASE_URL ?? (provider === 'ollama' ? 'http://127.0.0.1:11434' : provider === 'dashscope' ? 'https://dashscope.aliyuncs.com/compatible-mode/v1' : env.OPENAI_API_BASE_URL ?? 'https://api.openai.com/v1'), ).trim().replace(/\/$/, ''); let model = String( env.MEMIND_EMBEDDING_MODEL ?? (provider === 'ollama' ? 'nomic-embed-text' : provider === 'dashscope' ? 'text-embedding-v3' : 'text-embedding-3-small'), ).trim(); const cachePath = String( env.MEMIND_EMBEDDING_CACHE_PATH ?? '.release-gate/memfuse-embedding-cache.json', ).trim(); const batchSize = Math.max(1, Math.min(256, Number(env.MEMIND_EMBEDDING_BATCH_SIZE ?? 64) || 64)); const dimensions = Number(env.MEMIND_EMBEDDING_DIMENSIONS ?? 0) || null; const useLlmKeys = envFlag(env.MEMIND_EMBEDDING_FROM_LLM_KEYS) || (provider === 'dashscope' && !apiKey); if (useLlmKeys && !apiKey) { if (!llmCredentialCache) { llmCredentialCache = await resolveLlmEmbeddingCredentials({ env }); } if (llmCredentialCache.available) { apiKey = llmCredentialCache.apiKey; baseUrl = llmCredentialCache.baseUrl ?? baseUrl; if (!env.MEMIND_EMBEDDING_MODEL) model = llmCredentialCache.model ?? model; } } return { provider, apiKey, baseUrl, model, cachePath, batchSize, dimensions, llmKeyName: llmCredentialCache?.available ? llmCredentialCache.keyName : null, }; } function hashCacheKey(text, model) { return crypto.createHash('sha256').update(`${model}\0${text}`).digest('hex'); } async function ensureDiskCache(config) { if (diskEntries && diskCachePath === config.cachePath) return diskEntries; diskCachePath = config.cachePath; diskEntries = new Map(); diskDirty = false; if (!config.cachePath) return diskEntries; try { const raw = await fs.readFile(config.cachePath, 'utf8'); const parsed = JSON.parse(raw); for (const [key, value] of Object.entries(parsed?.entries ?? {})) { if (Array.isArray(value) && value.every((item) => Number.isFinite(Number(item)))) { diskEntries.set(key, value.map(Number)); } } } catch (err) { if (err && typeof err === 'object' && err.code === 'ENOENT') return diskEntries; if (err instanceof SyntaxError) { const corruptPath = `${path.resolve(config.cachePath)}.corrupt`; try { await fs.rename(path.resolve(config.cachePath), corruptPath); } catch { // ignore rename failure; start with empty cache } return diskEntries; } throw err; } return diskEntries; } async function persistDiskCache(config) { if (!diskDirty || !config.cachePath || !diskEntries) return; const target = path.resolve(config.cachePath); await fs.mkdir(path.dirname(target), { recursive: true }); const payload = `${JSON.stringify({ model: config.model, entries: Object.fromEntries(diskEntries) }, null, 2)}\n`; const tempPath = `${target}.tmp`; await fs.writeFile(tempPath, payload, 'utf8'); await fs.rename(tempPath, target); diskDirty = false; } function normalizeVector(value) { if (!Array.isArray(value)) return null; const numbers = value.map((item) => Number(item)); if (!numbers.length || numbers.some((item) => !Number.isFinite(item))) return null; return numbers; } function sleep(ms) { return new Promise((resolve) => setTimeout(resolve, ms)); } async function fetchWithRetry(label, fn, { retries = 6, baseDelayMs = 500, env = process.env, } = {}) { const maxRetries = Math.max(0, Number(env.MEMIND_EMBEDDING_MAX_RETRIES ?? retries) || retries); const delayMs = Math.max(100, Number(env.MEMIND_EMBEDDING_RETRY_BASE_MS ?? baseDelayMs) || baseDelayMs); let lastError = null; for (let attempt = 0; attempt <= maxRetries; attempt += 1) { try { return await fn(); } catch (err) { lastError = err; const message = err instanceof Error ? err.message : String(err); const retriable = /(?:429|503|502|504|rate limit|limit_requests|timeout|fetch failed)/i.test(message); if (!retriable || attempt === maxRetries) throw err; await sleep(delayMs * (2 ** attempt)); } } throw lastError; } async function fetchOpenAiEmbeddingsBatch(texts, config, fetchImpl = fetch) { if (!config.apiKey) { throw new Error( 'MEMIND_EMBEDDING_API_KEY / OPENAI_API_KEY is required for OpenAI-compatible embeddings', ); } const requestDelayMs = Math.max( 0, Number(process.env.MEMIND_EMBEDDING_REQUEST_DELAY_MS ?? 300) || 0, ); const payload = await fetchWithRetry('openai-embedding-batch', async () => { const response = await fetchImpl(`${config.baseUrl}/embeddings`, { method: 'POST', headers: { Authorization: `Bearer ${config.apiKey}`, 'Content-Type': 'application/json', }, body: JSON.stringify({ model: config.model, input: texts, encoding_format: 'float', ...(config.dimensions ? { dimensions: config.dimensions } : {}), }), }); if (!response.ok) { const detail = await response.text().catch(() => ''); throw new Error( `Embedding request failed (${response.status}): ${detail.slice(0, 240)}`, ); } return response.json(); }, { env: process.env }); const rows = Array.isArray(payload?.data) ? payload.data : []; rows.sort((left, right) => Number(left?.index ?? 0) - Number(right?.index ?? 0)); const vectors = rows.map((row) => normalizeVector(row?.embedding)); if (vectors.length !== texts.length || vectors.some((vector) => !vector)) { throw new Error('Embedding batch response size mismatch'); } if (requestDelayMs > 0) await sleep(requestDelayMs); return vectors; } async function fetchOpenAiEmbedding(text, config, fetchImpl = fetch) { const [vector] = await fetchOpenAiEmbeddingsBatch([text], config, fetchImpl); return vector; } async function fetchOllamaEmbedding(text, config, fetchImpl = fetch) { const response = await fetchImpl(`${config.baseUrl}/api/embed`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ model: config.model, input: text, }), }); if (!response.ok) { const detail = await response.text().catch(() => ''); throw new Error( `Ollama embedding request failed (${response.status}): ${detail.slice(0, 240)}`, ); } const payload = await response.json(); const vector = normalizeVector(payload?.embeddings?.[0]); if (!vector) throw new Error('Ollama embedding response missing embeddings[0]'); return vector; } async function fetchRemoteEmbedding(text, config, fetchImpl = fetch) { if (config.provider === 'ollama') { return fetchOllamaEmbedding(text, config, fetchImpl); } return fetchOpenAiEmbedding(text, config, fetchImpl); } export async function embedText(text, options = {}) { const config = await resolveConfig(options.env ?? process.env); const normalized = String(text ?? '').trim(); if (!normalized) { throw new Error('embedText requires non-empty text'); } const key = hashCacheKey(normalized, config.model); if (memoryCache.has(key)) return memoryCache.get(key); const disk = await ensureDiskCache(config); if (disk.has(key)) { const cached = disk.get(key); memoryCache.set(key, cached); return cached; } const vector = await fetchRemoteEmbedding( normalized, config, options.fetchImpl ?? fetch, ); memoryCache.set(key, vector); disk.set(key, vector); diskDirty = true; if (options.persist !== false && !envFlag(process.env.MEMIND_EMBEDDING_DEFER_PERSIST)) { await persistDiskCache(config); } return vector; } function resolveBatchSize(config, env = process.env) { const configured = Number(env.MEMIND_EMBEDDING_BATCH_SIZE ?? config.batchSize ?? 10); const maxBatch = config.provider === 'dashscope' ? 10 : 64; return Math.max(1, Math.min(maxBatch, configured || 10)); } export async function prefetchEmbedTexts(texts, options = {}) { const env = options.env ?? process.env; const config = await resolveConfig(env); const disk = await ensureDiskCache(config); const normalized = [...new Set( (Array.isArray(texts) ? texts : []) .map((text) => String(text ?? '').trim()) .filter(Boolean), )]; const missing = normalized.filter((text) => { const key = hashCacheKey(text, config.model); return !memoryCache.has(key) && !disk.has(key); }); if (missing.length === 0) return { requested: normalized.length, fetched: 0 }; const batchSize = resolveBatchSize(config, env); let fetched = 0; for (let index = 0; index < missing.length; index += batchSize) { const chunk = missing.slice(index, index + batchSize); const vectors = config.provider === 'ollama' ? await Promise.all(chunk.map((text) => fetchOllamaEmbedding(text, config, options.fetchImpl ?? fetch))) : await fetchOpenAiEmbeddingsBatch(chunk, config, options.fetchImpl ?? fetch); for (let offset = 0; offset < chunk.length; offset += 1) { const text = chunk[offset]; const key = hashCacheKey(text, config.model); memoryCache.set(key, vectors[offset]); disk.set(key, vectors[offset]); fetched += 1; } diskDirty = true; } if (options.persist !== false) { await persistDiskCache(config); } return { requested: normalized.length, fetched }; } export async function embedQuery(query, input = {}, options = {}) { return embedText(query, { ...options, env: input.env ?? options.env, }); } export async function flushEmbeddingCache(options = {}) { const config = await resolveConfig(options.env ?? process.env); await persistDiskCache(config); } export async function probeEmbeddingAvailability(options = {}) { const config = await resolveConfig(options.env ?? process.env); if (config.provider !== 'ollama' && !config.apiKey) { return { available: false, reason: config.provider === 'dashscope' ? 'dashscope_api_key_missing_or_llm_keys_unavailable' : 'embedding_api_key_missing', provider: config.provider, model: config.model, baseUrl: config.baseUrl, }; } try { const vector = await embedText('memfuse embedding probe', { env: options.env ?? process.env, fetchImpl: options.fetchImpl ?? fetch, persist: false, }); await flushEmbeddingCache({ env: options.env ?? process.env }); return { available: true, provider: config.provider, model: config.model, baseUrl: config.baseUrl, llmKeyName: config.llmKeyName, dimensions: vector.length, }; } catch (err) { return { available: false, reason: err instanceof Error ? err.message : String(err), provider: config.provider, model: config.model, baseUrl: config.baseUrl, llmKeyName: config.llmKeyName, }; } } export function __resetEmbeddingCacheForTests() { memoryCache.clear(); diskEntries = null; diskCachePath = null; diskDirty = false; llmCredentialCache = null; } export default embedQuery;