arceus-s 1.6.4

package/dist/core/embeddings/embedder.js

@@ -0,0 +1,349 @@
+/**
+ * Embedder Module
+ *
+ * Singleton factory for transformers.js embedding pipeline.
+ * Handles model loading, caching, and both single and batch embedding operations.
+ *
+ * Uses snowflake-arctic-embed-xs by default (22M params, 384 dims, ~90MB)
+ */
+// Suppress ONNX Runtime native warnings (e.g. VerifyEachNodeIsAssignedToAnEp)
+// Must be set BEFORE onnxruntime-node is imported by transformers.js
+// Level 3 = Error only (skips Warning/Info)
+if (!process.env.ORT_LOG_LEVEL) {
+    process.env.ORT_LOG_LEVEL = '3';
+}
+import { pipeline, env, } from '@huggingface/transformers';
+import { existsSync } from 'fs';
+import { execFileSync } from 'child_process';
+import { join, dirname } from 'path';
+import { createRequire } from 'module';
+import { DEFAULT_EMBEDDING_CONFIG } from './types.js';
+import { isHttpMode, getHttpDimensions, httpEmbed } from './http-client.js';
+import { resolveEmbeddingConfig } from './config.js';
+import { applyHfEnvOverrides, isHfDownloadFailure, withHfDownloadRetry } from './hf-env.js';
+import { logger } from '../logger.js';
+/**
+ * Check whether the onnxruntime-node package that @huggingface/transformers
+ * will actually load at runtime ships the CUDA execution provider.
+ *
+ * Critical: we resolve from transformers' own module scope, NOT from ours.
+ * npm may install two copies — a top-level 1.24.x (our dep) and a nested
+ * 1.21.0 (transformers' pinned dep). The guard must inspect whichever copy
+ * transformers.js will dlopen, otherwise the check is meaningless.
+ */
+function hasOrtCudaProvider() {
+    try {
+        const require = createRequire(import.meta.url);
+        // Resolve from @huggingface/transformers' scope so we find the same
+        // onnxruntime-node binary that transformers.js will use at runtime
+        const transformersDir = dirname(require.resolve('@huggingface/transformers/package.json'));
+        const ortRequire = createRequire(join(transformersDir, 'package.json'));
+        const ortPath = dirname(ortRequire.resolve('onnxruntime-node/package.json'));
+        // ORT 1.24.x only ships CUDA binaries for linux/x64 (downloaded from NuGet
+        // at postinstall). arm64 will correctly return false here until ORT adds support.
+        const arch = process.arch;
+        return existsSync(join(ortPath, 'bin', 'napi-v6', 'linux', arch, 'libonnxruntime_providers_cuda.so'));
+    }
+    catch {
+        return false;
+    }
+}
+/**
+ * Check whether CUDA libraries are actually available on this system.
+ * ONNX Runtime's native layer crashes (uncatchable) if we attempt CUDA
+ * without the required shared libraries, so we probe first.
+ *
+ * Checks both:
+ * 1. That system CUDA libraries (libcublasLt) are present
+ * 2. That onnxruntime-node ships the CUDA execution provider binary
+ *
+ * Both conditions must be true — system CUDA libs alone are not enough
+ * if onnxruntime-node is a CPU-only build (versions < 1.24.0).
+ */
+function isCudaAvailable() {
+    // First, verify onnxruntime-node has the CUDA provider binary.
+    // Without this, requesting CUDA causes an uncatchable native crash.
+    if (!hasOrtCudaProvider())
+        return false;
+    // Primary: query the dynamic linker cache — covers all architectures,
+    // distro layouts, and custom install paths registered with ldconfig
+    try {
+        const out = execFileSync('ldconfig', ['-p'], { timeout: 3000, encoding: 'utf-8' });
+        if (out.includes('libcublasLt.so.12'))
+            return true;
+    }
+    catch {
+        // ldconfig not available (e.g. non-standard container)
+    }
+    // Fallback: check CUDA_PATH and LD_LIBRARY_PATH for environments where
+    // ldconfig doesn't know about the CUDA install (conda, manual /opt/cuda, etc.)
+    for (const envVar of ['CUDA_PATH', 'LD_LIBRARY_PATH']) {
+        const val = process.env[envVar];
+        if (!val)
+            continue;
+        for (const dir of val.split(':').filter(Boolean)) {
+            if (existsSync(join(dir, 'lib64', 'libcublasLt.so.12')) ||
+                existsSync(join(dir, 'lib', 'libcublasLt.so.12')) ||
+                existsSync(join(dir, 'libcublasLt.so.12')))
+                return true;
+        }
+    }
+    return false;
+}
+// Module-level state for singleton pattern
+let embedderInstance = null;
+let isInitializing = false;
+let initPromise = null;
+let currentDevice = null;
+/**
+ * Get the current device being used for inference
+ */
+export const getCurrentDevice = () => currentDevice;
+/**
+ * Initialize the embedding model
+ * Uses singleton pattern - only loads once, subsequent calls return cached instance
+ *
+ * @param onProgress - Optional callback for model download progress
+ * @param config - Optional configuration override
+ * @param forceDevice - Force a specific device
+ * @returns Promise resolving to the embedder pipeline
+ */
+export const initEmbedder = async (onProgress, config = {}, forceDevice) => {
+    if (isHttpMode()) {
+        throw new Error('initEmbedder() should not be called in HTTP mode. ' +
+            'Use embedText()/embedBatch() which handle HTTP transparently.');
+    }
+    // Return existing instance if available
+    if (embedderInstance) {
+        return embedderInstance;
+    }
+    // If already initializing, wait for that promise
+    if (isInitializing && initPromise) {
+        return initPromise;
+    }
+    isInitializing = true;
+    const finalConfig = resolveEmbeddingConfig(config);
+    // CUDA is probe-gated because ONNX Runtime can crash in native code when
+    // provider libraries are missing. DirectML stays opt-in for the same reason.
+    // Probe for CUDA first — ONNX Runtime crashes (uncatchable native error)
+    // if we attempt CUDA without the required shared libraries
+    const gpuDevice = isCudaAvailable() ? 'cuda' : 'cpu';
+    const requestedDevice = forceDevice || (finalConfig.device === 'auto' ? gpuDevice : finalConfig.device);
+    initPromise = (async () => {
+        try {
+            // Configure transformers.js environment
+            env.allowLocalModels = false;
+            // Bridge user-controlled env vars to transformers.js: HF_HOME →
+            // env.cacheDir, HF_ENDPOINT → env.remoteHost (#1205). Centralised in
+            // applyHfEnvOverrides so the MCP embedder entry point behaves
+            // identically.
+            applyHfEnvOverrides(env);
+            const isDev = process.env.NODE_ENV === 'development';
+            if (isDev) {
+                logger.info(`🧠 Loading embedding model: ${finalConfig.modelId}`);
+            }
+            const progressCallback = onProgress
+                ? (data) => {
+                    const progress = {
+                        // Map the `progress_total` aggregate event (not in ModelProgress.status)
+                        // back to 'progress' so callers don't need to handle it separately.
+                        status: data.status === 'progress_total'
+                            ? 'progress'
+                            : (data.status ?? 'progress'),
+                        file: 'file' in data ? data.file : undefined,
+                        progress: 'progress' in data ? data.progress : undefined,
+                        loaded: 'loaded' in data ? data.loaded : undefined,
+                        total: 'total' in data ? data.total : undefined,
+                    };
+                    onProgress(progress);
+                }
+                : undefined;
+            // Try GPU first if auto, fall back to CPU
+            // Windows: dml (DirectML/DirectX12), Linux: cuda
+            const devicesToTry = requestedDevice === 'dml' || requestedDevice === 'cuda'
+                ? [requestedDevice, 'cpu']
+                : [requestedDevice];
+            for (const device of devicesToTry) {
+                try {
+                    if (isDev && device === 'dml') {
+                        logger.info('🔧 Trying DirectML (DirectX12) GPU backend...');
+                    }
+                    else if (isDev && device === 'cuda') {
+                        logger.info('🔧 Trying CUDA GPU backend...');
+                    }
+                    else if (isDev && device === 'cpu') {
+                        logger.info('🔧 Using CPU backend...');
+                    }
+                    else if (isDev && device === 'wasm') {
+                        logger.info('🔧 Using WASM backend (slower)...');
+                    }
+                    embedderInstance = await withHfDownloadRetry(() => pipeline('feature-extraction', finalConfig.modelId, {
+                        device: device,
+                        dtype: 'fp32',
+                        progress_callback: progressCallback,
+                        session_options: {
+                            logSeverityLevel: 3,
+                            intraOpNumThreads: finalConfig.threads,
+                            interOpNumThreads: 1,
+                            executionMode: 'sequential',
+                        },
+                    }), {
+                        onRetry: isDev
+                            ? (attempt, max, err) => logger.warn({ attempt, max, err: err.message }, `⚠️  Model download network error (attempt ${attempt}/${max}), retrying…`)
+                            : undefined,
+                    });
+                    currentDevice = device;
+                    if (isDev) {
+                        const label = device === 'dml'
+                            ? 'GPU (DirectML/DirectX12)'
+                            : device === 'cuda'
+                                ? 'GPU (CUDA)'
+                                : device.toUpperCase();
+                        logger.info(`✅ Using ${label} backend`);
+                        logger.info('✅ Embedding model loaded successfully');
+                    }
+                    return embedderInstance;
+                }
+                catch (deviceError) {
+                    // Network errors and circuit-open errors are not device-specific —
+                    // they will fail the same way on every device. Rethrow immediately
+                    // with actionable HF_ENDPOINT guidance rather than silently falling
+                    // back to the next device.
+                    const errMsg = deviceError instanceof Error ? deviceError.message : String(deviceError);
+                    if (isHfDownloadFailure(errMsg)) {
+                        const endpointHint = process.env.HF_ENDPOINT
+                            ? `The configured endpoint (${process.env.HF_ENDPOINT}) may be unreachable.`
+                            : `huggingface.co may be unreachable from your network.\n` +
+                                `  Set HF_ENDPOINT to a mirror and retry:\n` +
+                                `    HF_ENDPOINT=https://hf-mirror.com npx arc analyze --embeddings\n` +
+                                `    (Windows: set HF_ENDPOINT=https://hf-mirror.com && npx arc analyze --embeddings)`;
+                        throw new Error(`Failed to download embedding model: ${errMsg}\n  ${endpointHint}`);
+                    }
+                    if (isDev && (device === 'cuda' || device === 'dml')) {
+                        const gpuType = device === 'dml' ? 'DirectML' : 'CUDA';
+                        logger.info(`⚠️  ${gpuType} not available, falling back to CPU...`);
+                    }
+                    // Continue to next device in list
+                    if (device === devicesToTry[devicesToTry.length - 1]) {
+                        throw deviceError; // Last device failed, propagate error
+                    }
+                }
+            }
+            throw new Error('No suitable device found for embedding model');
+        }
+        catch (error) {
+            isInitializing = false;
+            initPromise = null;
+            embedderInstance = null;
+            throw error;
+        }
+        finally {
+            isInitializing = false;
+        }
+    })();
+    return initPromise;
+};
+/**
+ * Check if the embedder is initialized and ready
+ */
+export const isEmbedderReady = () => {
+    return isHttpMode() || embedderInstance !== null;
+};
+/**
+ * Get the effective embedding dimensions.
+ * In HTTP mode, uses ARC_EMBEDDING_DIMS if set, otherwise the default.
+ */
+export const getEmbeddingDimensions = () => {
+    if (isHttpMode()) {
+        return getHttpDimensions() ?? DEFAULT_EMBEDDING_CONFIG.dimensions;
+    }
+    return DEFAULT_EMBEDDING_CONFIG.dimensions;
+};
+/**
+ * Get the embedder instance (throws if not initialized)
+ */
+export const getEmbedder = () => {
+    if (isHttpMode()) {
+        throw new Error('getEmbedder() is not available in HTTP embedding mode. Use embedText()/embedBatch() instead.');
+    }
+    if (!embedderInstance) {
+        throw new Error('Embedder not initialized. Call initEmbedder() first.');
+    }
+    return embedderInstance;
+};
+/**
+ * Embed a single text string
+ *
+ * @param text - Text to embed
+ * @returns Float32Array of embedding vector
+ */
+export const embedText = async (text) => {
+    if (isHttpMode()) {
+        const [vec] = await httpEmbed([text]);
+        return vec;
+    }
+    const embedder = getEmbedder();
+    const result = await embedder(text, {
+        pooling: 'mean',
+        normalize: true,
+    });
+    // Result is a Tensor, convert to Float32Array
+    return new Float32Array(result.data);
+};
+/**
+ * Embed multiple texts in a single batch
+ * More efficient than calling embedText multiple times
+ *
+ * @param texts - Array of texts to embed
+ * @returns Array of Float32Array embedding vectors
+ */
+export const embedBatch = async (texts) => {
+    if (texts.length === 0) {
+        return [];
+    }
+    if (isHttpMode()) {
+        return httpEmbed(texts);
+    }
+    const embedder = getEmbedder();
+    // Process batch
+    const result = await embedder(texts, {
+        pooling: 'mean',
+        normalize: true,
+    });
+    // Result shape is [batch_size, dimensions]
+    // Need to split into individual vectors
+    const data = result.data;
+    const dimensions = DEFAULT_EMBEDDING_CONFIG.dimensions;
+    const embeddings = [];
+    for (let i = 0; i < texts.length; i++) {
+        const start = i * dimensions;
+        const end = start + dimensions;
+        embeddings.push(new Float32Array(Array.prototype.slice.call(data, start, end)));
+    }
+    return embeddings;
+};
+/**
+ * Convert Float32Array to regular number array (for LadybugDB storage)
+ */
+export const embeddingToArray = (embedding) => {
+    return Array.from(embedding);
+};
+/**
+ * Cleanup the embedder (free memory)
+ * Call this when done with embeddings
+ */
+export const disposeEmbedder = async () => {
+    if (embedderInstance) {
+        // transformers.js pipelines may have a dispose method
+        try {
+            if ('dispose' in embedderInstance && typeof embedderInstance.dispose === 'function') {
+                await embedderInstance.dispose();
+            }
+        }
+        catch {
+            // Ignore disposal errors
+        }
+        embedderInstance = null;
+        initPromise = null;
+    }
+};

package/dist/core/embeddings/embedding-pipeline.d.ts

@@ -0,0 +1,68 @@
+/**
+ * Embedding Pipeline Module
+ *
+ * Orchestrates the background embedding process:
+ * 1. Query embeddable nodes from LadybugDB
+ * 2. Generate text representations with enriched metadata
+ * 3. Chunk long nodes, batch embed
+ * 4. Update LadybugDB with chunk-aware embeddings
+ * 5. Create vector index for semantic search
+ */
+import { type EmbeddingProgress, type EmbeddingConfig, type EmbeddableNode, type SemanticSearchResult, type EmbeddingContext } from './types.js';
+/**
+ * Bump this when the embedding text template changes in a way that should
+ * invalidate existing vectors, such as metadata/header shape changes,
+ * structural container context changes, or preceding-context formatting rules.
+ */
+export declare const EMBEDDING_TEXT_VERSION = "v2";
+/**
+ * Compute a stable content fingerprint for an embeddable node.
+ * Used to detect when the underlying text has changed so stale vectors
+ * can be replaced (DELETE-then-INSERT, the Kuzu-sanctioned pattern for
+ * vector-indexed rows).
+ */
+export declare const contentHashForNode: (node: EmbeddableNode, config?: Partial<EmbeddingConfig>) => string;
+/**
+ * Progress callback type
+ */
+export type EmbeddingProgressCallback = (progress: EmbeddingProgress) => void;
+/**
+ * Batch INSERT chunk-aware embeddings into CodeEmbedding table
+ */
+export declare const batchInsertEmbeddings: (executeWithReusedStatement: (cypher: string, paramsList: Array<Record<string, any>>) => Promise<void>, updates: Array<{
+    nodeId: string;
+    chunkIndex: number;
+    startLine: number;
+    endLine: number;
+    embedding: number[];
+    contentHash?: string;
+}>) => Promise<void>;
+export interface EmbeddingPipelineResult {
+    nodesProcessed: number;
+    chunksProcessed: number;
+    vectorIndexReady: boolean;
+    semanticMode: 'vector-index' | 'exact-scan';
+}
+/**
+ * Run the embedding pipeline
+ *
+ * @param executeQuery - Function to execute Cypher queries against LadybugDB
+ * @param executeWithReusedStatement - Function to execute with reused prepared statement
+ * @param onProgress - Callback for progress updates
+ * @param config - Optional configuration override
+ * @param skipNodeIds - Optional set of node IDs that already have embeddings (incremental mode)
+ * @param context - Optional repo/server context for metadata enrichment
+ * @param existingEmbeddings - Optional map of nodeId → contentHash for incremental mode.
+ *        Nodes whose hash matches are skipped; nodes with a changed hash are DELETE'd
+ *        and re-embedded; nodes not in the map are embedded fresh.
+
+ */
+export declare const runEmbeddingPipeline: (executeQuery: (cypher: string) => Promise<any[]>, executeWithReusedStatement: (cypher: string, paramsList: Array<Record<string, any>>) => Promise<void>, onProgress: EmbeddingProgressCallback, config?: Partial<EmbeddingConfig>, skipNodeIds?: Set<string>, context?: EmbeddingContext, existingEmbeddings?: Map<string, string>) => Promise<EmbeddingPipelineResult>;
+/**
+ * Perform semantic search using the vector index with chunk deduplication
+ */
+export declare const semanticSearch: (executeQuery: (cypher: string) => Promise<any[]>, query: string, k?: number, maxDistance?: number) => Promise<SemanticSearchResult[]>;
+/**
+ * Semantic search with graph expansion (flattened results)
+ */
+export declare const semanticSearchWithContext: (executeQuery: (cypher: string) => Promise<any[]>, query: string, k?: number, _hops?: number) => Promise<any[]>;