@zenera/rag 1.1.9 → 1.1.11

package/dist/docs/search.js

@@ -0,0 +1,251 @@
+import { assertSameEmbedding } from "../common/manifest.js";
+import { loose, PatternError } from "../common/match.js";
+import { CHUNK_KINDS, parseLines } from "./chunk.js";
+import { openIndex, readLines, } from "./files.js";
+import { ChunkStore } from "./store.js";
+export const SEARCH_MODES = ['hybrid', 'vector', 'text'];
+export const DEFAULT_LIMIT = 8;
+/** Reciprocal rank fusion; the constant is the usual one and damps the top. */
+const RRF_K = 60;
+/** How much to ask each leg for, so fusion and the caps have room to work. */
+const OVERFETCH = 6;
+/**
+ * And an absolute ceiling on it. The multiple exists to survive the JavaScript
+ * filter and the diversifier, which is a fixed cost, not one that grows with
+ * the limit — at `--limit 32` the multiple alone asked for 192 rows, which on
+ * a small corpus is most of it. Every row fetched votes in the fusion at full
+ * strength, so an over-deep leg does not add recall, it adds an electorate.
+ */
+const MAX_FETCH = 96;
+/**
+ * Floors for the per-section and per-document caps, and the share of the limit
+ * they grow by once it is large enough to matter.
+ *
+ * The caps were constants, which made them a diversity nudge at the default
+ * limit of eight and a straitjacket above it: `--limit 20` for a question whose
+ * answer is one section of one file returned five chunks of the answer and
+ * fifteen of whatever else scored, because the sixth chunk of the right file
+ * was evicted in favour of the first chunk of a worse one. Asking for more
+ * results is asking to go deeper, and depth on a documented topic means more
+ * of the document that covers it.
+ */
+const MAX_PER_SECTION = 3;
+const MAX_PER_FILE = 5;
+const SECTION_SHARE = 0.25;
+const FILE_SHARE = 0.5;
+export class DocsIndex {
+    manifest;
+    outline;
+    #index;
+    #store;
+    #embedder;
+    constructor(index, store, embedder) {
+        this.manifest = index.manifest;
+        this.outline = index.outline;
+        this.#index = index;
+        this.#store = store;
+        this.#embedder = embedder;
+    }
+    /**
+     * An embedder is optional because half of what this index is for needs no
+     * model at all: listing, grepping and reading are exact, and asking for a
+     * credential to do them would be asking for nothing.
+     */
+    static async open(dir, embedder) {
+        const index = await openIndex(dir);
+        if (embedder) {
+            assertSameEmbedding(index.manifest, embedder.id);
+        }
+        return new DocsIndex(index, await ChunkStore.open(dir), embedder);
+    }
+    /** A document, split the way every line number in this index counts it. */
+    lines(name) {
+        return readLines(this.#index, name);
+    }
+    file(name) {
+        return this.outline.files.find((f) => f.name === name);
+    }
+    /** Document names matching the patterns, or all of them when there are none. */
+    resolveFiles(patterns = [], exclude = []) {
+        const names = this.manifest.sources.map((s) => s.name);
+        const keep = patterns.length === 0 ? names : names.filter(anyMatch(patterns));
+        return exclude.length === 0 ? keep : keep.filter((name) => !anyMatch(exclude)(name));
+    }
+    /**
+     * Headings the terms name, over the documents still in play. A term is tried
+     * as a structure id, then a structure path, then as a pattern over the
+     * title — so `--section overview` and `--section sec:3` are the same
+     * argument, and neither has to be explained.
+     */
+    resolveSections(terms, files) {
+        if (terms.length === 0) {
+            return [];
+        }
+        const within = new Set(files);
+        const out = [];
+        for (const file of this.outline.files) {
+            if (!within.has(file.name)) {
+                continue;
+            }
+            for (const heading of file.headings) {
+                const hit = terms.some((term) => term === heading.id ||
+                    term === heading.path ||
+                    anyMatch([term])(heading.title));
+                if (hit) {
+                    out.push(heading);
+                }
+            }
+        }
+        return out;
+    }
+    close() {
+        this.#store.close();
+    }
+    async search(query, signal) {
+        const mode = query.mode ?? 'hybrid';
+        const text = (query.query ?? '').trim();
+        const limit = Math.max(1, query.limit ?? DEFAULT_LIMIT);
+        const files = this.resolveFiles(query.files, query.exclude_files);
+        const headings = this.resolveSections(query.section ?? [], files);
+        const sections = headings.map((h) => h.path);
+        // Two ways to ask for nothing: a pattern no document answers to, and a
+        // section no document has. Both are told, not silently widened.
+        const impossible = files.length === 0 || ((query.section?.length ?? 0) > 0 && sections.length === 0);
+        if (impossible || !text) {
+            return { matches: [], mode, files, sections, considered: 0 };
+        }
+        const filter = this.#filter(query, files, sections);
+        const excluded = new Set(query.exclude_ids ?? []);
+        const fetch = Math.min(limit * OVERFETCH, MAX_FETCH) + excluded.size;
+        const vector = mode === 'text' ? [] : await this.#nearest(text, filter, fetch, signal);
+        const lexical = mode === 'vector' ? [] : await this.#store.matching(text, filter, fetch);
+        const fused = fuse(vector, lexical).filter((match) => !excluded.has(match.id) &&
+            files.includes(match.path) &&
+            under(match.structurePath, sections));
+        return {
+            matches: diversify(fused, limit),
+            mode,
+            files,
+            sections,
+            considered: fused.length,
+        };
+    }
+    async #nearest(text, filter, limit, signal) {
+        if (!this.#embedder) {
+            throw new Error('this index was opened without an embedder, so it cannot be searched');
+        }
+        const response = await this.#embedder.embed({
+            input: [text],
+            taskType: 'query',
+            signal,
+        });
+        return await this.#store.nearest(Float32Array.from(response.vectors[0]), filter, limit);
+    }
+    /**
+     * The SQL side of the narrowing. A short, safe file list becomes a
+     * predicate; a long or awkward one is left to the JavaScript filter that
+     * runs afterwards regardless, since that is what makes the answer correct.
+     */
+    #filter(query, files, sections) {
+        const all = this.manifest.sources.length;
+        return {
+            kinds: kindsOf(query.kinds),
+            paths: files.length < all && files.length <= 64 ? files : undefined,
+            prefixes: sections.length > 0 && sections.length <= 64 ? sections : undefined,
+        };
+    }
+}
+// ---------------------------------------------------------------------------
+const anyMatch = (patterns) => {
+    const matchers = patterns.map((pattern) => loose(pattern));
+    return (value) => matchers.some((match) => match(value));
+};
+/**
+ * Inside one of these sections, or anywhere when none were asked for. The
+ * boundary is a path separator and not a prefix, because `doc/sec:1` is a
+ * prefix of `doc/sec:10` and they are different sections.
+ */
+export const under = (path, sections) => sections.length === 0 ||
+    sections.some((section) => path === section || path.startsWith(`${section}/`));
+function kindsOf(kinds) {
+    if (!kinds || kinds.length === 0) {
+        return undefined;
+    }
+    for (const kind of kinds) {
+        if (!CHUNK_KINDS.includes(kind)) {
+            throw new PatternError(`${kind} is not a kind of block: it is one of ${CHUNK_KINDS.join(', ')}`);
+        }
+    }
+    return [...kinds];
+}
+/**
+ * Reciprocal rank fusion over the chunk id. Fusing on the id rather than on a
+ * physical row means the ordering survives a compaction of the store, and it is
+ * also the only key the two legs are guaranteed to agree on.
+ */
+function fuse(vector, lexical) {
+    const merged = new Map();
+    const add = (hits, leg) => {
+        for (const hit of hits) {
+            const existing = merged.get(hit.record.id) ?? match(hit);
+            existing.score += 1 / (RRF_K + hit.rank);
+            existing.ranks[leg] = hit.rank;
+            existing.relevance[leg] = hit.relevance;
+            merged.set(hit.record.id, existing);
+        }
+    };
+    add(vector, 'vector');
+    add(lexical, 'text');
+    return [...merged.values()].sort((a, b) => b.score - a.score || a.id.localeCompare(b.id));
+}
+function match(hit) {
+    const { record } = hit;
+    return {
+        id: record.id,
+        path: record.path,
+        kind: record.kind,
+        headings: record.headings,
+        structureId: record.structureId,
+        structurePath: record.structurePath,
+        lineNumbers: parseLines(record.lineSpec),
+        bodyStart: record.bodyStart,
+        bodyEnd: record.bodyEnd,
+        text: record.text,
+        score: 0,
+        ranks: {},
+        relevance: {},
+    };
+}
+/**
+ * Caps per section and per document, applied in score order. A question whose
+ * answer really is one section still gets it — the cap only stops a run of
+ * near-identical neighbours from crowding out the second place a thing is said.
+ */
+function diversify(matches, limit) {
+    const sectionCap = Math.max(MAX_PER_SECTION, Math.ceil(limit * SECTION_SHARE));
+    const fileCap = Math.max(MAX_PER_FILE, Math.ceil(limit * FILE_SHARE));
+    const perSection = new Map();
+    const perFile = new Map();
+    const kept = [];
+    const held = [];
+    for (const item of matches) {
+        const section = `${item.path}\u0000${parentOf(item.structurePath)}`;
+        const sections = perSection.get(section) ?? 0;
+        const files = perFile.get(item.path) ?? 0;
+        if (sections >= sectionCap || files >= fileCap) {
+            held.push(item);
+            continue;
+        }
+        perSection.set(section, sections + 1);
+        perFile.set(item.path, files + 1);
+        kept.push(item);
+        if (kept.length === limit) {
+            return kept;
+        }
+    }
+    // Better a crowded answer than a short one: what the caps held back fills
+    // the rest, still in score order.
+    return [...kept, ...held].slice(0, limit);
+}
+const parentOf = (path) => path.slice(0, path.lastIndexOf('/')) || path;
+//# sourceMappingURL=search.js.map

package/dist/docs/store.d.ts

@@ -0,0 +1,79 @@
+import { type Connection, type Table } from '@lancedb/lancedb';
+/**
+ * Rows per write.
+ *
+ * Arrow addresses a batch's buffers with 32-bit offsets, so one batch carrying
+ * more than 2 GiB does not error — it panics inside the reader, on a thread
+ * whose panic never reaches this one. At 3072 dimensions the vector column
+ * alone crosses that at about 175k rows, which a documentation tree reaches.
+ */
+export declare const WRITE_BATCH = 8192;
+export interface ChunkRecord {
+    /** `${path}#c${ordinal}` — the fusion key, and stable across compaction */
+    id: string;
+    /** the document's name in the index: its relative path */
+    path: string;
+    ordinal: number;
+    kind: string;
+    /** the full-text document — wider */
+    text: string;
+    /** what the vector was made from — tighter */
+    embedText: string;
+    /** the render set, run-length encoded: `1,5,12-40` */
+    lineSpec: string;
+    bodyStart: number;
+    bodyEnd: number;
+    structureId: string;
+    structurePath: string;
+    headings: string;
+    tokens: number;
+}
+export interface StoreFilter {
+    kinds?: readonly string[];
+    /** document names, exactly as the manifest spells them */
+    paths?: readonly string[];
+    /** structure path prefixes, as `structurePath LIKE 'x%'` */
+    prefixes?: readonly string[];
+}
+export interface Hit {
+    record: ChunkRecord;
+    /** 0-based position in this leg's own result list */
+    rank: number;
+    relevance: number;
+}
+export interface WriteResult {
+    rows: number;
+    fts: boolean;
+    vector: boolean;
+}
+/**
+ * A table built a window at a time.
+ *
+ * The corpus used to be embedded into one array and handed over in one call,
+ * which made peak memory a function of corpus size and, past 2 GiB of vectors,
+ * a panic. Rows go in as they are paid for instead, so what is resident is one
+ * window rather than all of it.
+ */
+export interface ChunkWriter {
+    add(rows: readonly ChunkRecord[], vectors: readonly Float32Array[]): Promise<void>;
+    /** Builds the indexes and closes. Throws if nothing was ever added. */
+    finish(): Promise<WriteResult>;
+    /** For a build that failed, so the connection does not outlive it. */
+    close(): void;
+}
+export declare function openChunks(dir: string): Promise<ChunkWriter>;
+export declare class ChunkStore {
+    #private;
+    constructor(db: Connection, table: Table);
+    static open(dir: string): Promise<ChunkStore>;
+    /** Nearest neighbours. Every row has a vector, so nothing has to be excluded. */
+    nearest(vector: Float32Array, filter: StoreFilter, limit: number): Promise<Hit[]>;
+    /**
+     * The lexical leg. A missing full-text index, or a query the tokenizer
+     * makes nothing of, answers with nothing rather than throwing: a hybrid
+     * search that loses one leg is a worse search, not a failed one.
+     */
+    matching(text: string, filter: StoreFilter, limit: number): Promise<Hit[]>;
+    close(): void;
+}
+//# sourceMappingURL=store.d.ts.map

package/dist/docs/store.js

@@ -0,0 +1,214 @@
+import { connect, Index } from '@lancedb/lancedb';
+import { CliError, EXIT } from '@zenera/cli/lib';
+import { CHUNK_KINDS } from "./chunk.js";
+import { lancePath } from "./files.js";
+// ---------------------------------------------------------------------------
+// The chunk table
+//
+// One row per chunk, carrying both retrieval texts: `text` is what the
+// full-text index reads and `embedText` is what the vector was made from. They
+// are different on purpose and are kept side by side so an index can be
+// re-embedded without re-reading the documents.
+//
+// The two legs are run separately rather than through the built-in hybrid
+// query, and fused in `search.ts`. Three reasons, none of them cosmetic: the
+// built-in fuses on a physical row id, which moves when the table is compacted,
+// where a chunk id does not; `mode` already has to offer a vector-only and a
+// text-only path, so making hybrid the same shape as those two is one mechanism
+// instead of three; and when there is no full-text index, or the query is
+// nonsense to it, degrading has to be a decision rather than an exception.
+//
+// Only closed vocabularies reach the SQL predicate. `kind` is one of seven
+// words. A document name is compared against a safe character set first and, if
+// it does not pass, is simply not put in the predicate at all — the JavaScript
+// filter that runs afterwards is what makes the answer correct, so the clause is
+// only ever an optimisation and there is nothing to escape.
+// ---------------------------------------------------------------------------
+const TABLE = 'chunks';
+/** Below this an IVF index has nothing to train on, and a flat scan is faster. */
+const VECTOR_INDEX_MIN_ROWS = 2000;
+/**
+ * Rows per write.
+ *
+ * Arrow addresses a batch's buffers with 32-bit offsets, so one batch carrying
+ * more than 2 GiB does not error — it panics inside the reader, on a thread
+ * whose panic never reaches this one. At 3072 dimensions the vector column
+ * alone crosses that at about 175k rows, which a documentation tree reaches.
+ */
+export const WRITE_BATCH = 8192;
+const KINDS = new Set(CHUNK_KINDS);
+/** What may go into a string literal in a predicate, and nothing else. */
+const SAFE = /^[\w.:/ -]+$/;
+export async function openChunks(dir) {
+    return new Writer(await connect(lancePath(dir)));
+}
+class Writer {
+    #db;
+    #table;
+    #rows = 0;
+    constructor(db) {
+        this.#db = db;
+    }
+    async add(rows, vectors) {
+        for (let from = 0; from < rows.length; from += WRITE_BATCH) {
+            const batch = rows.slice(from, from + WRITE_BATCH).map((row, i) => ({
+                ...row,
+                vector: vectors[from + i],
+            }));
+            if (this.#table) {
+                await this.#table.add(batch);
+            }
+            else {
+                // Every column is always populated — never null — so the Arrow
+                // schema is inferred from the first row without a declaration.
+                this.#table = await this.#db.createTable(TABLE, batch, { mode: 'overwrite' });
+            }
+            this.#rows += batch.length;
+        }
+    }
+    async finish() {
+        const table = this.#table;
+        if (!table) {
+            throw new CliError('the documents hold nothing to index', EXIT.invalid, 'they are empty, or every one of them is blank');
+        }
+        // A writer that panicked did not reject, so a short table is the only
+        // evidence left that the rows never landed.
+        const stored = await table.countRows();
+        if (stored !== this.#rows) {
+            throw new CliError(`the table holds ${stored} of ${this.#rows} chunks`, EXIT.failed, 'the write did not finish, and a partial index answers as if it were whole');
+        }
+        await table.createIndex('text', { config: Index.fts() });
+        await table.createIndex('kind', { config: Index.bitmap() });
+        for (const column of ['path', 'structurePath']) {
+            await table.createIndex(column, { config: Index.btree() });
+        }
+        const vector = this.#rows >= VECTOR_INDEX_MIN_ROWS;
+        if (vector) {
+            await table.createIndex('vector');
+        }
+        this.#db.close();
+        return { rows: this.#rows, fts: true, vector };
+    }
+    close() {
+        this.#db.close();
+    }
+}
+export class ChunkStore {
+    #db;
+    #table;
+    constructor(db, table) {
+        this.#db = db;
+        this.#table = table;
+    }
+    static async open(dir) {
+        const db = await connect(lancePath(dir));
+        try {
+            return new ChunkStore(db, await db.openTable(TABLE));
+        }
+        catch {
+            db.close();
+            throw new CliError(`${dir} holds no searchable table`, EXIT.invalid, 'rebuild it with `zen rag docs index`');
+        }
+    }
+    /** Nearest neighbours. Every row has a vector, so nothing has to be excluded. */
+    async nearest(vector, filter, limit) {
+        let query = this.#table.query().nearestTo(vector).limit(limit);
+        const predicate = where(filter);
+        if (predicate) {
+            query = query.where(predicate);
+        }
+        return hits(await query.toArray());
+    }
+    /**
+     * The lexical leg. A missing full-text index, or a query the tokenizer
+     * makes nothing of, answers with nothing rather than throwing: a hybrid
+     * search that loses one leg is a worse search, not a failed one.
+     */
+    async matching(text, filter, limit) {
+        try {
+            let query = this.#table
+                .query()
+                .fullTextSearch(text, { columns: ['text'] })
+                .limit(limit);
+            const predicate = where(filter);
+            if (predicate) {
+                query = query.where(predicate);
+            }
+            return hits(await query.toArray());
+        }
+        catch {
+            return [];
+        }
+    }
+    close() {
+        this.#db.close();
+    }
+}
+// ---------------------------------------------------------------------------
+function where(filter) {
+    return [clause('kind', filter.kinds), clause('path', filter.paths), prefixes(filter.prefixes)]
+        .filter(Boolean)
+        .join(' AND ');
+}
+function clause(column, values) {
+    if (!values || values.length === 0) {
+        return '';
+    }
+    if (column === 'kind') {
+        for (const value of values) {
+            if (!KINDS.has(value)) {
+                throw new Error(`kind cannot be ${JSON.stringify(value)}`);
+            }
+        }
+    }
+    else if (!values.every((value) => SAFE.test(value))) {
+        // Left to the JavaScript filter, which is what makes it correct anyway.
+        return '';
+    }
+    return `${column} IN (${values.map((v) => `'${v}'`).join(', ')})`;
+}
+/**
+ * One prefix covers a section and everything nested inside it. `LIKE` alone is
+ * too generous — `doc/sec:1` prefixes `doc/sec:10` as well — so this narrows
+ * the scan and the caller settles the boundary in JavaScript.
+ */
+function prefixes(values) {
+    if (!values || values.length === 0 || !values.every((value) => SAFE.test(value))) {
+        return '';
+    }
+    return `(${values.map((v) => `structurePath LIKE '${v}%'`).join(' OR ')})`;
+}
+function hits(rows) {
+    return rows.map((row, rank) => ({
+        record: strip(row),
+        rank,
+        relevance: score(row),
+    }));
+}
+/** A lexical query reports a score; a vector one reports a distance. */
+function score(row) {
+    const relevance = row._relevance_score ?? row._score;
+    if (typeof relevance === 'number') {
+        return relevance;
+    }
+    const distance = row._distance;
+    return typeof distance === 'number' ? 1 / (1 + distance) : 0;
+}
+function strip(row) {
+    return {
+        id: row.id,
+        path: row.path,
+        ordinal: row.ordinal,
+        kind: row.kind,
+        text: row.text,
+        embedText: row.embedText,
+        lineSpec: row.lineSpec,
+        bodyStart: row.bodyStart,
+        bodyEnd: row.bodyEnd,
+        structureId: row.structureId,
+        structurePath: row.structurePath,
+        headings: row.headings,
+        tokens: row.tokens,
+    };
+}
+//# sourceMappingURL=store.js.map

package/dist/docs/tools.d.ts

@@ -0,0 +1,10 @@
+import { type AnyTool } from '@zenera/neo';
+import { type DocsIndex } from './search.ts';
+export interface DocsToolOptions {
+    /** passages per search when the model does not say */
+    limit?: number;
+    /** lines quoted per answer when the model does not say */
+    maxLines?: number;
+}
+export declare function docsTools<TCtx = unknown>(index: DocsIndex, options?: DocsToolOptions): AnyTool<TCtx>[];
+//# sourceMappingURL=tools.d.ts.map