memind/scripts/embed-memory-v2-openai-compat.mjs

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;