socraticode/dist/services/embeddings.js

SocratiCode — MCP server for local codebase indexing, semantic search, and code dependency graphs. All private, all local via Docker.

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Giancarlo Erra - Altaire Limited
import { documentIncludesPath, documentPrefix, getEmbeddingConfig, queryPrefix, } from "./embedding-config.js";
import { getEmbeddingProvider } from "./embedding-provider.js";
import { logger } from "./logger.js";
// Number of texts to embed per provider request.
// Configurable via env var EMBEDDING_BATCH_SIZE (positive integer); defaults to 32.
const BATCH_SIZE = (() => {
    const raw = process.env.EMBEDDING_BATCH_SIZE;
    if (raw === undefined)
        return 32;
    const num = Number(raw);
    if (!Number.isInteger(num) || num <= 0) {
        throw new Error(`Invalid EMBEDDING_BATCH_SIZE: "${raw}". Must be a positive integer.`);
    }
    return num;
})();
const MAX_RETRIES = 3;
const BASE_DELAY_MS = 500;
/**
 * Inter-batch delay (ms) to respect provider rate limits.
 * Google free tier: 5 RPM → need ~12s between requests.
 * We use a conservative 15s. OpenAI/Ollama have generous limits, so 0.
 */
const PROVIDER_BATCH_DELAY = {
    ollama: 0,
    openai: 0,
    google: 15_000, // 15s — stay safely under the free-tier 5 RPM
};
/** Retry an async operation with exponential backoff */
async function withRetry(operation, label, maxRetries = MAX_RETRIES, baseDelay = BASE_DELAY_MS) {
    let lastError;
    for (let attempt = 1; attempt <= maxRetries; attempt++) {
        try {
            return await operation();
        }
        catch (err) {
            lastError = err;
            if (attempt < maxRetries) {
                // Use longer backoff for rate-limit errors (429 / RESOURCE_EXHAUSTED)
                const errMsg = err instanceof Error ? err.message : String(err);
                const isRateLimit = errMsg.includes("429") || errMsg.includes("RESOURCE_EXHAUSTED") || errMsg.includes("quota");
                const delay = isRateLimit
                    ? Math.max(baseDelay * 2 ** (attempt - 1), 15_000) // at least 15s for rate limits
                    : baseDelay * 2 ** (attempt - 1);
                logger.warn(`${label} failed (attempt ${attempt}/${maxRetries}), retrying in ${delay}ms`, {
                    error: errMsg,
                });
                await new Promise((r) => setTimeout(r, delay));
            }
        }
    }
    throw lastError;
}
/**
 * Generate embeddings for a batch of texts, handling batching automatically.
 * Texts are pre-truncated to the model's context window inside the provider.
 *
 * Takes the texts as given: task prefixes are the caller's business, added by
 * {@link prepareDocumentText} before the texts get here.
 */
export async function generateEmbeddings(texts, onBatchComplete) {
    if (texts.length === 0)
        return [];
    const provider = await getEmbeddingProvider();
    const config = getEmbeddingConfig();
    const batchDelay = PROVIDER_BATCH_DELAY[config.embeddingProvider] ?? 0;
    const results = [];
    for (let i = 0; i < texts.length; i += BATCH_SIZE) {
        // Inter-batch delay to respect provider rate limits (skip first batch)
        if (batchDelay > 0 && i > 0) {
            logger.info(`Rate-limit pause: waiting ${batchDelay / 1000}s before next batch`);
            await new Promise((r) => setTimeout(r, batchDelay));
        }
        const batch = texts.slice(i, i + BATCH_SIZE);
        const batchLabel = `Embedding batch ${Math.floor(i / BATCH_SIZE) + 1}`;
        const embeddings = await withRetry(() => provider.embed(batch), batchLabel);
        results.push(...embeddings);
        onBatchComplete?.(Math.min(i + batch.length, texts.length), texts.length);
    }
    return results;
}
/**
 * Generate a single query embedding, prefixed with the configured query task
 * prefix (`EMBEDDING_QUERY_PREFIX`; see `queryPrefix` in embedding-config.ts).
 */
export async function generateQueryEmbedding(query, effectiveQueryPrefix = queryPrefix()) {
    const provider = await getEmbeddingProvider();
    return withRetry(() => provider.embedSingle(`${effectiveQueryPrefix}${query}`), "Query embedding");
}
export function prepareDocumentText(content, filePath, profile = {
    documentPrefix: documentPrefix(),
    documentIncludesPath: documentIncludesPath(),
}) {
    const head = profile.documentIncludesPath ? `${filePath}\n` : "";
    return `${profile.documentPrefix}${head}${content}`;
}
//# sourceMappingURL=embeddings.js.map