doc2vec 2.16.0 → 2.17.1

package/dist/eval.js

@@ -0,0 +1,375 @@
+"use strict";
+// Retrieval evaluation: runs a set of questions with known relevant URLs
+// against a sqlite-vec database and reports recall@k and MRR. It measures
+// the same thing the MCP server serves (nearest chunks by distance), so a
+// change to the embedding model, chunking or ranking can be judged by
+// numbers instead of by eyeballing a few queries.
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
+    Object.defineProperty(o, "default", { enumerable: true, value: v });
+}) : function(o, v) {
+    o["default"] = v;
+});
+var __importStar = (this && this.__importStar) || (function () {
+    var ownKeys = function(o) {
+        ownKeys = Object.getOwnPropertyNames || function (o) {
+            var ar = [];
+            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
+            return ar;
+        };
+        return ownKeys(o);
+    };
+    return function (mod) {
+        if (mod && mod.__esModule) return mod;
+        var result = {};
+        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
+        __setModuleDefault(result, mod);
+        return result;
+    };
+})();
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.DEFAULT_K_VALUES = exports.DEFAULT_RERANK_CANDIDATES = exports.DEFAULT_RERANK_MODEL = exports.RRF_K = void 0;
+exports.createDeepInfraReranker = createDeepInfraReranker;
+exports.normalizeUrl = normalizeUrl;
+exports.firstRelevantRank = firstRelevantRank;
+exports.computeMetrics = computeMetrics;
+exports.loadEvalCases = loadEvalCases;
+exports.runEval = runEval;
+exports.onePerPage = onePerPage;
+exports.fuseRrf = fuseRrf;
+exports.toFtsQuery = toFtsQuery;
+exports.createSqliteRetriever = createSqliteRetriever;
+exports.createEmbedder = createEmbedder;
+exports.formatReport = formatReport;
+exports.runEvalCommand = runEvalCommand;
+const fs = __importStar(require("fs"));
+const path = __importStar(require("path"));
+const yaml = __importStar(require("js-yaml"));
+const openai_1 = require("openai");
+exports.RRF_K = 60;
+exports.DEFAULT_RERANK_MODEL = 'nvidia/Nemotron-3.5-Rerank-8B-BF16';
+exports.DEFAULT_RERANK_CANDIDATES = 40;
+// Same request the MCP server's reranker sends (mcp/src/server.ts)
+function createDeepInfraReranker(apiKey, model = exports.DEFAULT_RERANK_MODEL) {
+    return async (query, documents) => {
+        const response = await fetch(`https://api.deepinfra.com/v1/inference/${model}`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${apiKey}` },
+            body: JSON.stringify({ queries: [query], documents }),
+        });
+        if (!response.ok) {
+            throw new Error(`reranker returned HTTP ${response.status}: ${(await response.text()).slice(0, 200)}`);
+        }
+        const body = await response.json();
+        if (!Array.isArray(body.scores) || body.scores.length !== documents.length) {
+            throw new Error('reranker did not return one score per document');
+        }
+        return body.scores.map(Number);
+    };
+}
+exports.DEFAULT_K_VALUES = [1, 4, 10];
+// Treat https://x/a/, https://x/a and https://x/a#section as the same page:
+// eval sets are written by hand and the stored URLs vary on these points.
+function normalizeUrl(url) {
+    let normalized = url.trim();
+    const hashIndex = normalized.indexOf('#');
+    if (hashIndex !== -1)
+        normalized = normalized.substring(0, hashIndex);
+    while (normalized.length > 1 && normalized.endsWith('/'))
+        normalized = normalized.slice(0, -1);
+    return normalized.toLowerCase();
+}
+function firstRelevantRank(retrievedUrls, relevantUrls) {
+    const relevant = new Set(relevantUrls.map(normalizeUrl));
+    for (let i = 0; i < retrievedUrls.length; i++) {
+        const url = retrievedUrls[i];
+        if (url && relevant.has(normalizeUrl(url)))
+            return i + 1;
+    }
+    return null;
+}
+function computeMetrics(results, kValues = exports.DEFAULT_K_VALUES) {
+    const recall = {};
+    for (const k of kValues) {
+        const hits = results.filter(r => r.rank !== null && r.rank <= k).length;
+        recall[k] = results.length ? hits / results.length : 0;
+    }
+    const reciprocalSum = results.reduce((sum, r) => sum + (r.rank ? 1 / r.rank : 0), 0);
+    return { cases: results.length, recall, mrr: results.length ? reciprocalSum / results.length : 0 };
+}
+// Accepts .jsonl (one case per line), .json (an array) or .yaml/.yml (a list).
+function loadEvalCases(filePath) {
+    const raw = fs.readFileSync(filePath, 'utf8');
+    const ext = path.extname(filePath).toLowerCase();
+    let parsed;
+    if (ext === '.jsonl') {
+        parsed = raw.split('\n').map(l => l.trim()).filter(l => l && !l.startsWith('//')).map(l => JSON.parse(l));
+    }
+    else if (ext === '.json') {
+        parsed = JSON.parse(raw);
+    }
+    else {
+        parsed = yaml.load(raw);
+    }
+    if (!Array.isArray(parsed)) {
+        throw new Error(`Eval file ${filePath} must contain a list of cases`);
+    }
+    return parsed.map((c, i) => {
+        if (!c || typeof c.query !== 'string' || !c.query.trim()) {
+            throw new Error(`Eval case #${i + 1} has no query`);
+        }
+        if (!Array.isArray(c.relevant_urls) || c.relevant_urls.length === 0) {
+            throw new Error(`Eval case #${i + 1} ("${c.query}") has no relevant_urls`);
+        }
+        return { ...c, id: c.id ?? String(i + 1) };
+    });
+}
+async function runEval(cases, deps, options) {
+    const results = [];
+    for (const evalCase of cases) {
+        const embedding = await deps.embed(evalCase.query);
+        const chunks = await deps.retrieve(embedding, evalCase, options.topK);
+        const retrievedUrls = chunks.map(c => c.url ?? '');
+        results.push({
+            id: evalCase.id ?? '',
+            query: evalCase.query,
+            rank: firstRelevantRank(retrievedUrls, evalCase.relevant_urls),
+            retrieved_urls: retrievedUrls,
+        });
+    }
+    const kValues = (options.kValues ?? exports.DEFAULT_K_VALUES).filter(k => k <= options.topK);
+    return {
+        model: options.model,
+        top_k: options.topK,
+        created_at: new Date().toISOString(),
+        ...computeMetrics(results, kValues),
+        results,
+    };
+}
+// Keeps the first (best) chunk of each page, in order.
+function onePerPage(chunks) {
+    const seen = new Set();
+    return chunks.filter(chunk => {
+        const key = normalizeUrl(chunk.url ?? '');
+        if (seen.has(key))
+            return false;
+        seen.add(key);
+        return true;
+    });
+}
+// Reciprocal Rank Fusion: each list contributes 1 / (RRF_K + rank) per chunk.
+// Rank-based, so dense distances and bm25 scores need no calibration.
+function fuseRrf(lists) {
+    const scores = new Map();
+    for (const list of lists) {
+        list.forEach((chunk, index) => {
+            if (chunk.rowid === undefined)
+                return;
+            const entry = scores.get(chunk.rowid) ?? { chunk, score: 0 };
+            entry.score += 1 / (exports.RRF_K + index + 1);
+            scores.set(chunk.rowid, entry);
+        });
+    }
+    return [...scores.values()].sort((a, b) => b.score - a.score).map(e => e.chunk);
+}
+// Turns free text into an FTS5 OR-query of quoted terms, so punctuation in a
+// question (colons, quotes, parentheses) cannot break the MATCH syntax.
+function toFtsQuery(text) {
+    const terms = text.toLowerCase().match(/[\p{L}\p{N}_]+/gu) ?? [];
+    return [...new Set(terms)].map(t => `"${t}"`).join(' OR ');
+}
+// Queries the database the MCP server's sqlite provider reads
+// (mcp/src/server.ts); dense mode runs exactly its query, so the numbers
+// reflect what an agent receives today.
+function createSqliteRetriever(defaultDbPath, options = {}) {
+    const Database = require('better-sqlite3');
+    const sqliteVec = require('sqlite-vec');
+    const mode = options.mode ?? 'dense';
+    const connections = new Map();
+    const open = (dbPath) => {
+        let conn = connections.get(dbPath);
+        if (!conn) {
+            if (!fs.existsSync(dbPath))
+                throw new Error(`Database file not found at ${dbPath}`);
+            const db = new Database(dbPath, { readonly: true });
+            sqliteVec.load(db);
+            conn = { db };
+            connections.set(dbPath, conn);
+        }
+        return conn;
+    };
+    // The keyword index is built in memory from the stored chunk text, so an
+    // existing database can be evaluated without re-ingesting it.
+    const ftsFor = (conn) => {
+        if (!conn.fts) {
+            const fts = new Database(':memory:');
+            fts.exec(`CREATE VIRTUAL TABLE chunks USING fts5(content, url UNINDEXED, product_name UNINDEXED, version UNINDEXED)`);
+            const insert = fts.prepare('INSERT INTO chunks (rowid, content, url, product_name, version) VALUES (?, ?, ?, ?, ?)');
+            const rows = conn.db.prepare('SELECT rowid, content, url, product_name, version FROM vec_items').all();
+            fts.transaction(() => {
+                for (const r of rows)
+                    insert.run(BigInt(r.rowid), r.content ?? '', r.url, r.product_name, r.version);
+            })();
+            conn.fts = fts;
+        }
+        return conn.fts;
+    };
+    const dense = (conn, queryEmbedding, evalCase, limit) => {
+        let query = `
+            SELECT rowid, url, distance, content
+            FROM vec_items
+            WHERE embedding MATCH @query_embedding`;
+        if (evalCase.product_name)
+            query += ` AND product_name = @product_name`;
+        if (evalCase.version)
+            query += ` AND version = @version`;
+        query += `
+            ORDER BY distance
+            LIMIT @top_k;`;
+        return conn.db.prepare(query).all({
+            query_embedding: new Float32Array(queryEmbedding),
+            product_name: evalCase.product_name,
+            version: evalCase.version,
+            top_k: limit,
+        });
+    };
+    const bm25 = (conn, evalCase, limit) => {
+        const match = toFtsQuery(evalCase.query);
+        if (!match)
+            return [];
+        let query = `SELECT rowid, url, content FROM chunks WHERE chunks MATCH @match`;
+        if (evalCase.product_name)
+            query += ` AND product_name = @product_name`;
+        if (evalCase.version)
+            query += ` AND version = @version`;
+        query += ` ORDER BY bm25(chunks) LIMIT @limit`;
+        return ftsFor(conn).prepare(query).all({
+            match,
+            product_name: evalCase.product_name,
+            version: evalCase.version,
+            limit,
+        });
+    };
+    return async (queryEmbedding, evalCase, topK) => {
+        const dbPath = evalCase.db ?? defaultDbPath;
+        if (!dbPath)
+            throw new Error(`No database for eval case ${evalCase.id}: pass --db or set "db" on the case`);
+        const conn = open(dbPath);
+        const rerankCandidates = options.rerankCandidates ?? exports.DEFAULT_RERANK_CANDIDATES;
+        const depth = Math.max(topK, options.candidates ?? (mode === 'dense' && !options.perPage ? topK : 50), options.rerank ? rerankCandidates : 0);
+        let chunks;
+        if (mode === 'bm25') {
+            chunks = bm25(conn, evalCase, depth);
+        }
+        else if (mode === 'hybrid') {
+            chunks = fuseRrf([dense(conn, queryEmbedding, evalCase, depth), bm25(conn, evalCase, depth)]);
+        }
+        else {
+            chunks = dense(conn, queryEmbedding, evalCase, depth);
+        }
+        if (options.rerank && chunks.length > 1) {
+            const pool = chunks.slice(0, rerankCandidates);
+            const scores = await options.rerank(evalCase.query, pool.map(c => c.content ?? ''));
+            const reranked = pool.map((c, i) => ({ c, s: scores[i] })).sort((a, b) => b.s - a.s).map(x => x.c);
+            chunks = [...reranked, ...chunks.slice(rerankCandidates)];
+        }
+        if (options.perPage)
+            chunks = onePerPage(chunks);
+        return chunks.slice(0, topK);
+    };
+}
+// Mirrors the provider selection in Doc2Vec's constructor and the MCP server's
+// env vars, so the query is embedded with the model the database was built with.
+function createEmbedder(embeddingConfig) {
+    const provider = embeddingConfig?.provider || process.env.EMBEDDING_PROVIDER || 'openai';
+    let client;
+    let model;
+    if (provider === 'azure') {
+        const apiKey = embeddingConfig?.azure?.api_key || process.env.AZURE_OPENAI_KEY;
+        const endpoint = embeddingConfig?.azure?.endpoint || process.env.AZURE_OPENAI_ENDPOINT;
+        model = embeddingConfig?.azure?.deployment_name || process.env.AZURE_OPENAI_DEPLOYMENT_NAME || 'text-embedding-3-large';
+        if (!apiKey || !endpoint)
+            throw new Error('Azure OpenAI requires api_key and endpoint to be configured');
+        client = new openai_1.AzureOpenAI({
+            apiKey,
+            endpoint,
+            deployment: model,
+            apiVersion: embeddingConfig?.azure?.api_version || process.env.AZURE_OPENAI_API_VERSION || '2024-10-21',
+        });
+    }
+    else {
+        const apiKey = embeddingConfig?.openai?.api_key || process.env.OPENAI_API_KEY;
+        const baseURL = embeddingConfig?.openai?.base_url || process.env.OPENAI_BASE_URL;
+        model = embeddingConfig?.openai?.model || process.env.OPENAI_MODEL || 'text-embedding-3-large';
+        if (!apiKey)
+            throw new Error('OpenAI requires api_key to be configured');
+        client = new openai_1.OpenAI({ apiKey, ...(baseURL && { baseURL }) });
+    }
+    const embed = async (text) => {
+        const response = await client.embeddings.create({ model, input: text });
+        return response.data[0].embedding;
+    };
+    return { embed, model };
+}
+function formatReport(report) {
+    const lines = [];
+    lines.push(`Model: ${report.model}   Cases: ${report.cases}   top_k: ${report.top_k}`);
+    const recallParts = Object.entries(report.recall).map(([k, v]) => `recall@${k} ${(v * 100).toFixed(1)}%`);
+    lines.push(`${recallParts.join('   ')}   MRR ${report.mrr.toFixed(3)}`);
+    const misses = report.results.filter(r => r.rank === null);
+    if (misses.length) {
+        lines.push('');
+        lines.push(`Misses (${misses.length}):`);
+        for (const m of misses)
+            lines.push(`  [${m.id}] ${m.query}`);
+    }
+    return lines.join('\n');
+}
+async function runEvalCommand(options) {
+    const mode = options.mode ?? 'dense';
+    if (!['dense', 'bm25', 'hybrid'].includes(mode)) {
+        throw new Error(`Unknown mode '${mode}': use dense, bm25 or hybrid`);
+    }
+    let rerank;
+    if (options.rerank) {
+        if (options.rerank !== 'deepinfra')
+            throw new Error(`Unknown reranker '${options.rerank}': use deepinfra`);
+        const apiKey = process.env.DEEPINFRA_API_KEY;
+        if (!apiKey)
+            throw new Error('--rerank deepinfra requires the DEEPINFRA_API_KEY environment variable');
+        rerank = createDeepInfraReranker(apiKey, options.rerankModel);
+    }
+    let embeddingConfig;
+    if (options.configPath) {
+        const config = yaml.load(fs.readFileSync(options.configPath, 'utf8'));
+        embeddingConfig = config?.embedding;
+    }
+    // Keyword-only retrieval never looks at the query vector, so skip the API call
+    const { embed, model } = mode === 'bm25'
+        ? { embed: async () => [], model: 'none (bm25)' }
+        : createEmbedder(embeddingConfig);
+    const cases = loadEvalCases(options.casesPath);
+    const retrieve = createSqliteRetriever(options.dbPath, {
+        mode, perPage: options.perPage, rerank, rerankCandidates: options.rerankCandidates,
+    });
+    const label = `${model}, ${mode}` +
+        `${rerank ? `, reranked by ${options.rerankModel ?? exports.DEFAULT_RERANK_MODEL} over ${options.rerankCandidates ?? exports.DEFAULT_RERANK_CANDIDATES}` : ''}` +
+        `${options.perPage ? ', one chunk per page' : ''}`;
+    const report = await runEval(cases, { embed, retrieve }, { topK: options.topK, model: label });
+    console.log(formatReport(report));
+    if (options.outPath) {
+        fs.writeFileSync(options.outPath, JSON.stringify(report, null, 2));
+        console.log(`\nReport written to ${options.outPath}`);
+    }
+    return report;
+}