@knowall-ai/reverie 0.4.0

package/build/embeddings.js

@@ -0,0 +1,303 @@
+export const EMBEDDING_FIELDS = ['embedding', 'name_embedding', 'embedding_model', 'embedded_at'];
+const DEFAULT_MODELS = {
+    local: 'Xenova/all-MiniLM-L6-v2',
+    openai: 'text-embedding-3-small',
+    azure: 'text-embedding-3-small',
+    ollama: 'nomic-embed-text',
+    voyage: 'voyage-3-lite'
+};
+function getProvider(env) {
+    const raw = env.REVERIE_EMBEDDINGS?.trim().toLowerCase() ?? 'local';
+    if (raw === 'local' || raw === 'openai' || raw === 'azure' || raw === 'ollama' || raw === 'voyage') {
+        return raw;
+    }
+    if (raw === 'none') {
+        return 'none';
+    }
+    console.error(`Embeddings disabled: unknown provider "${env.REVERIE_EMBEDDINGS}"`);
+    return 'none';
+}
+function getModel(env, provider) {
+    return env.REVERIE_EMBEDDING_MODEL?.trim() || DEFAULT_MODELS[provider];
+}
+const DEFAULT_TIMEOUT_MS = 30_000;
+const MAX_TIMEOUT_MS = 120_000;
+/** REVERIE_EMBED_TIMEOUT_MS, whole milliseconds, clamped to 1s..120s. */
+function configuredTimeoutMs(env) {
+    const raw = (env.REVERIE_EMBED_TIMEOUT_MS ?? '').trim();
+    const value = /^\d+$/.test(raw) ? Number(raw) : DEFAULT_TIMEOUT_MS;
+    return Math.min(Math.max(value, 1_000), MAX_TIMEOUT_MS);
+}
+/** Credentials only travel over TLS; plain http is allowed for loopback hosts (local proxies). */
+export function secureEndpoint(raw, name) {
+    let url;
+    try {
+        url = new URL(raw);
+    }
+    catch {
+        throw new Error(`${name} is not a valid URL`);
+    }
+    const loopback = ['localhost', '127.0.0.1', '::1', '[::1]'].includes(url.hostname);
+    if (url.protocol !== 'https:' && !(url.protocol === 'http:' && loopback)) {
+        throw new Error(`${name} must use https (http is only accepted for localhost)`);
+    }
+    return url.toString().replace(/\/$/, '');
+}
+function withTimeout(promise, ms, message) {
+    let timer;
+    const deadline = new Promise((_, reject) => {
+        timer = setTimeout(() => reject(new Error(message)), ms);
+    });
+    return Promise.race([promise, deadline]).finally(() => clearTimeout(timer));
+}
+/** Credentials never follow a redirect (a cross-origin or https→http hop would leak them); every request is bounded by timeoutMs. */
+async function fetchJson(url, init, timeoutMs) {
+    const controller = new AbortController();
+    const timer = setTimeout(() => controller.abort(), timeoutMs);
+    try {
+        const response = await fetch(url, { ...init, signal: controller.signal, redirect: 'error' });
+        if (!response.ok) {
+            const body = (await response.text()).slice(0, 200);
+            throw new Error(`Embedding request failed ${response.status}: ${body}`);
+        }
+        return await response.json();
+    }
+    finally {
+        clearTimeout(timer);
+    }
+}
+async function embedRemote(texts, doRequest) {
+    const vectors = [];
+    for (let index = 0; index < texts.length; index += 64) {
+        const batch = texts.slice(index, index + 64);
+        const batchVectors = await doRequest(batch);
+        // Vectors are paired back to inputs by position, so a short batch would shift every later pair.
+        if (batchVectors.length !== batch.length) {
+            throw new Error(`Embedding provider returned ${batchVectors.length} vectors for ${batch.length} inputs`);
+        }
+        vectors.push(...batchVectors);
+    }
+    return vectors;
+}
+/** OpenAI-style `{ data: [{ index, embedding }] }`: order by the declared index, not array position. */
+function vectorsFromData(data) {
+    if (!Array.isArray(data)) {
+        return [];
+    }
+    const items = data;
+    const ordered = items.every((item) => typeof item.index === 'number')
+        ? [...items].sort((a, b) => a.index - b.index)
+        : items;
+    return ordered.map((item) => normalizeVector(item.embedding));
+}
+function normalizeVector(value) {
+    if (!Array.isArray(value)) {
+        return [];
+    }
+    return value.map((item) => Number(item));
+}
+export function createEmbedder(env = process.env) {
+    const provider = getProvider(env);
+    if (provider === 'none') {
+        if ((env.REVERIE_EMBEDDINGS?.trim().toLowerCase() ?? '') === 'none') {
+            console.error('Embeddings disabled: provider set to none');
+        }
+        return null;
+    }
+    const model = getModel(env, provider);
+    const timeoutMs = configuredTimeoutMs(env);
+    if (provider === 'local') {
+        let extractorPromise = null;
+        return {
+            id: `${provider}/${model}`,
+            async embed(texts) {
+                if (texts.length === 0) {
+                    return [];
+                }
+                if (!extractorPromise) {
+                    // Loading may download the model; bound it, and forget a failed load so a later
+                    // call can retry instead of falling back to keyword search until the process restarts.
+                    const loading = (async () => {
+                        const transformers = await import('@huggingface/transformers');
+                        if (env.REVERIE_MODEL_CACHE?.trim()) {
+                            transformers.env.cacheDir = env.REVERIE_MODEL_CACHE.trim();
+                        }
+                        return transformers.pipeline('feature-extraction', model, { dtype: 'q8' });
+                    })();
+                    extractorPromise = withTimeout(loading, timeoutMs, `local embedding model ${model} did not load within ${timeoutMs}ms`)
+                        .catch((error) => {
+                        extractorPromise = null;
+                        throw error;
+                    });
+                }
+                const extractor = await extractorPromise;
+                const output = await extractor(texts, { pooling: 'mean', normalize: true });
+                const vectors = Array.isArray(output) ? output : output.tolist?.();
+                if (!Array.isArray(vectors)) {
+                    throw new Error('Local embedding pipeline returned an unexpected result');
+                }
+                return vectors.map((vector) => normalizeVector(vector));
+            }
+        };
+    }
+    if (provider === 'openai') {
+        const apiKey = env.OPENAI_API_KEY?.trim();
+        if (!apiKey) {
+            throw new Error('OPENAI_API_KEY is required for OpenAI embeddings');
+        }
+        const baseUrl = secureEndpoint(env.OPENAI_BASE_URL?.trim() || 'https://api.openai.com/v1', 'OPENAI_BASE_URL');
+        return {
+            id: `${provider}/${model}`,
+            async embed(texts) {
+                return embedRemote(texts, async (batch) => {
+                    const data = await fetchJson(`${baseUrl}/embeddings`, {
+                        method: 'POST',
+                        headers: {
+                            'Authorization': `Bearer ${apiKey}`,
+                            'Content-Type': 'application/json'
+                        },
+                        body: JSON.stringify({ model, input: batch })
+                    }, timeoutMs);
+                    return vectorsFromData(data.data);
+                });
+            }
+        };
+    }
+    if (provider === 'azure') {
+        const rawEndpoint = env.AZURE_OPENAI_ENDPOINT?.trim();
+        const apiKey = env.AZURE_OPENAI_API_KEY?.trim();
+        const deployment = env.AZURE_OPENAI_EMBEDDING_DEPLOYMENT?.trim();
+        const apiVersion = env.AZURE_OPENAI_API_VERSION?.trim() || '2024-10-21';
+        if (!rawEndpoint) {
+            throw new Error('AZURE_OPENAI_ENDPOINT is required for Azure embeddings');
+        }
+        const endpoint = secureEndpoint(rawEndpoint, 'AZURE_OPENAI_ENDPOINT');
+        if (!apiKey) {
+            throw new Error('AZURE_OPENAI_API_KEY is required for Azure embeddings');
+        }
+        if (!deployment) {
+            throw new Error('AZURE_OPENAI_EMBEDDING_DEPLOYMENT is required for Azure embeddings');
+        }
+        return {
+            id: `${provider}/${deployment}`,
+            async embed(texts) {
+                return embedRemote(texts, async (batch) => {
+                    const data = await fetchJson(`${endpoint}/openai/deployments/${deployment}/embeddings?api-version=${apiVersion}`, {
+                        method: 'POST',
+                        headers: {
+                            'api-key': apiKey,
+                            'Content-Type': 'application/json'
+                        },
+                        body: JSON.stringify({ model, input: batch })
+                    }, timeoutMs);
+                    return vectorsFromData(data.data);
+                });
+            }
+        };
+    }
+    if (provider === 'ollama') {
+        const host = secureEndpoint(env.OLLAMA_HOST?.trim() || 'http://127.0.0.1:11434', 'OLLAMA_HOST');
+        return {
+            id: `${provider}/${model}`,
+            async embed(texts) {
+                return embedRemote(texts, async (batch) => {
+                    const data = await fetchJson(`${host}/api/embed`, {
+                        method: 'POST',
+                        headers: {
+                            'Content-Type': 'application/json'
+                        },
+                        body: JSON.stringify({ model, input: batch })
+                    }, timeoutMs);
+                    return Array.isArray(data.embeddings)
+                        ? data.embeddings.map((item) => normalizeVector(item))
+                        : [];
+                });
+            }
+        };
+    }
+    const apiKey = env.VOYAGE_API_KEY?.trim();
+    if (!apiKey) {
+        throw new Error('VOYAGE_API_KEY is required for Voyage embeddings');
+    }
+    return {
+        id: `${provider}/${model}`,
+        async embed(texts) {
+            return embedRemote(texts, async (batch) => {
+                const data = await fetchJson('https://api.voyageai.com/v1/embeddings', {
+                    method: 'POST',
+                    headers: {
+                        'Authorization': `Bearer ${apiKey}`,
+                        'Content-Type': 'application/json'
+                    },
+                    body: JSON.stringify({ model, input: batch })
+                }, timeoutMs);
+                return vectorsFromData(data.data);
+            });
+        }
+    };
+}
+export function embeddingText(label, props) {
+    const skipKeys = new Set([...EMBEDDING_FIELDS, 'created_at', 'updated_at']);
+    const lines = [`${label}: ${String(props.name ?? '')}`.trim()];
+    for (const [key, value] of Object.entries(props)) {
+        if (key === 'name' || skipKeys.has(key) || key.endsWith('_id') || key.endsWith('Id')) {
+            continue;
+        }
+        if (typeof value === 'string') {
+            lines.push(`${key}: ${value}`);
+            continue;
+        }
+        if (Array.isArray(value) && value.every((item) => typeof item === 'string' || typeof item === 'number')) {
+            lines.push(`${key}: ${value.join(', ')}`);
+        }
+    }
+    return lines.join('\n').slice(0, 1000);
+}
+/** Label, name and aliases only. Kept separate because descriptive text dilutes name similarity badly. */
+export function nameText(label, props) {
+    const lines = [`${label}: ${String(props.name ?? '')}`.trim()];
+    const aliases = props.aliases;
+    if (Array.isArray(aliases) && aliases.length > 0) {
+        lines.push(`aliases: ${aliases.map((alias) => String(alias)).join(', ')}`);
+    }
+    return lines.join('\n');
+}
+/** Embed a batch of nodes, two texts each, and pair the vectors back up per node. */
+export async function embedNodes(embedder, nodes) {
+    const texts = nodes.flatMap((node) => [embeddingText(node.label, node.props), nameText(node.label, node.props)]);
+    const vectors = await embedder.embed(texts);
+    return nodes.map((_, index) => ({
+        embedding: vectors[index * 2] ?? [],
+        name_embedding: vectors[index * 2 + 1] ?? []
+    }));
+}
+export function cosine(a, b) {
+    if (a.length === 0 || a.length !== b.length) {
+        return 0;
+    }
+    let dot = 0;
+    let aNorm = 0;
+    let bNorm = 0;
+    for (let index = 0; index < a.length; index += 1) {
+        dot += a[index] * b[index];
+        aNorm += a[index] * a[index];
+        bNorm += b[index] * b[index];
+    }
+    if (aNorm === 0 || bNorm === 0) {
+        return 0;
+    }
+    return dot / (Math.sqrt(aNorm) * Math.sqrt(bNorm));
+}
+export function scrub(value) {
+    const blocked = new Set(EMBEDDING_FIELDS);
+    if (Array.isArray(value)) {
+        return value.map((item) => scrub(item));
+    }
+    if (value && typeof value === 'object') {
+        const entries = Object.entries(value)
+            .filter(([key]) => !blocked.has(key))
+            .map(([key, item]) => [key, scrub(item)]);
+        return Object.fromEntries(entries);
+    }
+    return value;
+}