@zenera/rag 1.1.9 → 1.1.10

package/dist/docs/store.js

@@ -0,0 +1,171 @@
+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;
+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 writeChunks(dir, rows, vectors) {
+    if (rows.length === 0) {
+        throw new CliError('the documents hold nothing to index', EXIT.invalid, 'they are empty, or every one of them is blank');
+    }
+    const db = await connect(lancePath(dir));
+    // Every column is always populated — never null — so the Arrow schema is
+    // inferred from the first row without a declaration to keep in step.
+    const table = await db.createTable(TABLE, rows.map((row, i) => ({ ...row, vector: vectors[i] })), { mode: 'overwrite' });
+    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 = rows.length >= VECTOR_INDEX_MIN_ROWS;
+    if (vector) {
+        await table.createIndex('vector');
+    }
+    db.close();
+    return { rows: rows.length, fts: true, vector };
+}
+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

package/dist/docs/tools.js

@@ -0,0 +1,300 @@
+import { tool } from '@zenera/neo';
+import { PatternError } from "../common/match.js";
+import { assemble } from "./assemble.js";
+import { CHUNK_KINDS } from "./chunk.js";
+import { grepLines, listFiles, listSections, listTables, readRange, readSection, } from "./lookup.js";
+import { renderAssembly } from "./render.js";
+import { SEARCH_MODES } from "./search.js";
+// ---------------------------------------------------------------------------
+// The same index, given to an agent
+//
+// Four tools over one engine, and only one of them ranks anything. `search_docs`
+// is the way in when the question is vague; the other three are exact, because
+// a model told "no results" by a vector search has learned nothing — a ranking
+// returns the top of a list, so an empty answer and an absent thing look
+// identical.
+//
+// `search_docs` is shaped for the second call rather than the first. The first
+// is always a sentence and always returns too much of the wrong tree; the
+// second is the same sentence with `files: ["nsx_4.2*/api/**"]`, or
+// `section: "Rate limits"`, or `kind: ["table"]` because the answer is a table
+// and not the prose around it. Those are parameters and not separate tools, so
+// narrowing costs one call and not three.
+//
+// Every answer carries line numbers, and `read_docs` takes them. That is the
+// loop the whole thing exists for: find the passage, read around it, then edit
+// the file the passage came from — and a passage that cannot be pointed at is a
+// passage nothing can be done with.
+// ---------------------------------------------------------------------------
+const GROUP = 'docs';
+/** Kept small on purpose: a tool result is prompt, and the model asked for one thing. */
+const DEFAULT_LIMIT = 5;
+const MAX_LIMIT = 25;
+/** A listing is lines rather than passages, so it can afford more of them. */
+const DEFAULT_ROWS = 50;
+const MAX_ROWS = 200;
+/** A ceiling on what one answer may quote, so one long section cannot eat it. */
+const DEFAULT_MAX_LINES = 200;
+const MAX_LINES = 800;
+export function docsTools(index, options = {}) {
+    const limitOf = (asked) => clamp(asked ?? options.limit ?? DEFAULT_LIMIT, 1, MAX_LIMIT);
+    const linesOf = (asked) => clamp(asked ?? options.maxLines ?? DEFAULT_MAX_LINES, 20, MAX_LINES);
+    const searchDocs = tool({
+        name: 'search_docs',
+        group: GROUP,
+        description: 'Searches the documents and answers with the passages that matched, quoted ' +
+            'verbatim with their line numbers and with a marker wherever something between ' +
+            'two passages was left out. When the first answer is from the wrong part of the ' +
+            'corpus, ask again with the same query and a narrowing: `files` for a path ' +
+            'pattern, `section` for a heading, `kind` for tables or code only. Pass the ids ' +
+            'from an earlier answer in exclude_ids to be shown something new instead.',
+        parameters: {
+            type: 'object',
+            properties: {
+                query: {
+                    type: 'string',
+                    description: 'What is wanted, as a sentence rather than keywords.',
+                },
+                files: strings('Only documents whose name matches. A glob if it has * or ?, e.g. ' +
+                    '"guides/**" or "nsx_4.2*/api/**"; otherwise a substring.'),
+                exclude_files: strings('Documents to leave out, matched the same way.'),
+                section: strings('Only under these headings, and whatever nests inside them. A heading ' +
+                    'title, or a structure path from an earlier answer.'),
+                kind: {
+                    type: 'array',
+                    items: { type: 'string', enum: [...CHUNK_KINDS] },
+                    description: 'Only these kinds of block. Use ["table"] or ["table_row"] when the ' +
+                        'answer is tabular, ["code"] for examples.',
+                },
+                mode: {
+                    type: 'string',
+                    enum: [...SEARCH_MODES],
+                    description: 'hybrid blends meaning and wording; text is exact wording only, for ' +
+                        'an error string or an identifier.',
+                },
+                exclude_ids: strings('Passage ids already seen, as printed in an earlier answer.'),
+                limit: {
+                    type: 'integer',
+                    description: `Passages kept. Default ${DEFAULT_LIMIT}, at most ${MAX_LIMIT}.`,
+                },
+                before: { type: 'integer', description: 'Extra lines quoted before each passage.' },
+                after: { type: 'integer', description: 'Extra lines quoted after each passage.' },
+                max_lines: {
+                    type: 'integer',
+                    description: `A ceiling on the whole answer. Default ${DEFAULT_MAX_LINES}.`,
+                },
+            },
+            required: ['query'],
+            additionalProperties: false,
+        },
+        execute: async (args) => {
+            const query = {
+                query: args.query,
+                files: args.files,
+                exclude_files: args.exclude_files,
+                section: args.section,
+                kinds: args.kind,
+                mode: args.mode,
+                exclude_ids: args.exclude_ids,
+                limit: limitOf(args.limit),
+            };
+            const result = await guard(() => index.search(query));
+            if ('error' in result) {
+                return result;
+            }
+            if (result.files.length === 0) {
+                return {
+                    found: 0,
+                    hint: 'no document matched `files` — call list_docs to see their names',
+                };
+            }
+            if (result.matches.length === 0) {
+                return {
+                    found: 0,
+                    scope: { documents: result.files.length, sections: result.sections.length },
+                    hint: args.section?.length
+                        ? 'nothing under that section — drop `section` and search the whole document'
+                        : 'try fewer words, or mode "text" if it is an exact string',
+                };
+            }
+            const excerpt = await assemble(index, result.matches, {
+                before: args.before,
+                after: args.after,
+                maxLines: linesOf(args.max_lines),
+            });
+            return {
+                found: result.matches.length,
+                ids: result.matches.map((m) => m.id),
+                documents: excerpt.files.map((f) => f.path),
+                truncated: excerpt.truncated,
+                passages: renderAssembly(excerpt, { colour: false }),
+            };
+        },
+    });
+    const listDocs = tool({
+        name: 'list_docs',
+        group: GROUP,
+        description: 'Lists what is in the index — the documents, their headings, or their tables — ' +
+            'without searching or ranking anything. Call it first to learn the document ' +
+            "names that `files` patterns are matched against, or to see a document's " +
+            'structure before asking about one part of it.',
+        parameters: {
+            type: 'object',
+            properties: {
+                what: {
+                    type: 'string',
+                    enum: ['files', 'sections', 'tables'],
+                    description: 'Default files.',
+                },
+                files: strings('Only these documents, matched by glob or substring.'),
+                section: strings('Only under these headings.'),
+                depth: {
+                    type: 'integer',
+                    description: 'sections only: the deepest heading level to report.',
+                },
+                limit: { type: 'integer', description: `Rows kept. Default ${DEFAULT_ROWS}.` },
+            },
+            additionalProperties: false,
+        },
+        execute: async (args) => {
+            const options = {
+                files: args.files,
+                section: args.section,
+                depth: args.depth,
+                limit: clamp(args.limit ?? DEFAULT_ROWS, 1, MAX_ROWS),
+            };
+            const what = args.what ?? 'files';
+            const result = await guard(() => what === 'sections'
+                ? listSections(index, options)
+                : what === 'tables'
+                    ? listTables(index, options)
+                    : listFiles(index, options));
+            if ('error' in result) {
+                return result;
+            }
+            return {
+                found: result.found,
+                truncated: result.truncated,
+                [what]: result.rows,
+                ...(result.found === 0
+                    ? { hint: 'nothing matched — widen `files`, or call it with no arguments' }
+                    : {}),
+            };
+        },
+    });
+    const grepDocs = tool({
+        name: 'grep_docs',
+        group: GROUP,
+        description: 'Every line matching a pattern, with the document and line number, and the ' +
+            'section it sits in. Exact, not ranked, and it reports the true total even when ' +
+            'the rows are cut — so unlike search_docs it can answer "does this string appear ' +
+            'anywhere". Reach for it with an identifier, an error message, or a flag name.',
+        parameters: {
+            type: 'object',
+            properties: {
+                pattern: {
+                    type: 'string',
+                    description: 'A substring, a glob if it has * or ?, or a regex with regex.',
+                },
+                files: strings('Only these documents.'),
+                section: strings('Only under these headings.'),
+                regex: { type: 'boolean', description: 'Read the pattern as a regex.' },
+                case_sensitive: { type: 'boolean', description: 'Match the capitals too.' },
+                limit: { type: 'integer', description: `Lines kept. Default ${DEFAULT_ROWS}.` },
+            },
+            required: ['pattern'],
+            additionalProperties: false,
+        },
+        execute: async (args) => {
+            const result = await guard(() => grepLines(index, args.pattern, {
+                files: args.files,
+                section: args.section,
+                regex: args.regex,
+                caseSensitive: args.case_sensitive,
+                limit: clamp(args.limit ?? DEFAULT_ROWS, 1, MAX_ROWS),
+            }));
+            if ('error' in result) {
+                return result;
+            }
+            return {
+                found: result.found,
+                truncated: result.truncated,
+                lines: result.rows,
+                ...(result.found === 0 ? { hint: 'nothing matched anywhere in scope' } : {}),
+            };
+        },
+    });
+    const readDocs = tool({
+        name: 'read_docs',
+        group: GROUP,
+        description: 'Reads a document verbatim: a whole named section, or a line range as printed by ' +
+            'search_docs or grep_docs. Use it when a passage was found and the lines around ' +
+            'it are needed in full, with nothing omitted and nothing summarised.',
+        parameters: {
+            type: 'object',
+            properties: {
+                file: {
+                    type: 'string',
+                    description: 'A document name, exactly as list_docs prints it.',
+                },
+                section: {
+                    type: 'string',
+                    description: 'A heading title or structure path. Overrides from/to.',
+                },
+                from: { type: 'integer', description: 'First line, 1-based. Default 1.' },
+                to: { type: 'integer', description: 'Last line. Default the end of the document.' },
+            },
+            required: ['file'],
+            additionalProperties: false,
+        },
+        execute: async (args) => {
+            const file = index.resolveFiles([args.file])[0];
+            if (!file) {
+                return {
+                    error: `no document called ${args.file}`,
+                    hint: 'call list_docs for the names',
+                };
+            }
+            const result = await guard(async () => args.section
+                ? await readSection(index, file, args.section)
+                : await readRange(index, file, args.from ?? 1, args.to ?? Infinity));
+            if ('error' in result) {
+                return result;
+            }
+            // A whole large document is prompt spent on lines nobody asked for.
+            const kept = result.lines.slice(0, MAX_LINES);
+            return {
+                file: result.file,
+                start: result.start,
+                end: result.start + kept.length - 1,
+                total: result.total,
+                truncated: kept.length < result.lines.length,
+                text: kept.map((line, at) => `${result.start + at} | ${line}`).join('\n'),
+            };
+        },
+    });
+    return [searchDocs, listDocs, grepDocs, readDocs];
+}
+// ---------------------------------------------------------------------------
+const strings = (description) => ({
+    type: 'array',
+    items: { type: 'string' },
+    description,
+});
+const clamp = (value, low, high) => Math.min(Math.max(Math.trunc(value), low), high);
+/**
+ * A bad pattern is a bad argument, not a failure. A model handed a thrown
+ * exception retries the same call; one handed a sentence fixes it.
+ */
+async function guard(run) {
+    try {
+        return await run();
+    }
+    catch (err) {
+        if (err instanceof PatternError) {
+            return { error: err.message };
+        }
+        throw err;
+    }
+}
+//# sourceMappingURL=tools.js.map

package/dist/index.d.ts

@@ -1,6 +1,7 @@
 export * from './common/locate.ts';
 export * from './common/manifest.ts';
 export * from './common/match.ts';
+export * as docs from './docs/index.ts';
 export * from './schema/build.ts';
 export * from './schema/entities.ts';
 export * from './schema/files.ts';

package/dist/index.js

@@ -1,6 +1,9 @@
 export * from "./common/locate.js";
 export * from "./common/manifest.js";
 export * from "./common/match.js";
+// A namespace rather than a re-export: both subjects have a `Manifest`, a
+// `writeIndex` and a `search`, and flattening them would collide on every one.
+export * as docs from "./docs/index.js";
 export * from "./schema/build.js";
 export * from "./schema/entities.js";
 export * from "./schema/files.js";

package/dist/schema/command.js

@@ -99,6 +99,7 @@ export const command = {
             ['  --exclude-method <name>', dim('Drop an operation by name.')],
             ['  --exclude-type <name>', dim('Drop a schema by name.')],
             ['  --exclude-property <name>', dim('Drop a field by name.')],
+            ['  --source <name>', dim('Only this document, as `stats` names it. Repeatable.')],
             ['  --limit <n>', dim('Seeds kept per term. Default 5.')],
             ['  --max-hops <n>', dim('How far apart two hits may be. Default 3.')],
             ['  --max-nodes <n>', dim('Nodes per result. Default 200.')],
@@ -282,6 +283,7 @@ const SEARCH_OPTIONS = {
     'exclude-method': MANY,
     'exclude-type': MANY,
     'exclude-property': MANY,
+    source: MANY,
     limit: { type: 'string' },
     'max-hops': { type: 'string' },
     'max-nodes': { type: 'string' },
@@ -362,6 +364,7 @@ function fromFlags(values, positionals = []) {
     put('exclude_methods', values['exclude-method']);
     put('exclude_types', values['exclude-type']);
     put('exclude_properties', values['exclude-property']);
+    put('sources', values.source);
     put('direction', values.direction);
     put('method_type', values['method-type']);
     put('limit', values.limit && count(values.limit, '--limit'));

package/dist/schema/query.js

@@ -19,6 +19,7 @@ const LISTS = [
     'exclude_methods',
     'exclude_types',
     'exclude_properties',
+    'sources',
 ];
 const NUMBERS = ['limit', 'max_hops', 'max_nodes'];
 const DIRECTIONS = ['input', 'output', 'any'];

package/dist/schema/search.d.ts

@@ -24,6 +24,8 @@ export interface SchemaQuery {
     exclude_methods?: readonly string[];
     exclude_types?: readonly string[];
     exclude_properties?: readonly string[];
+    /** document names, as `stats` prints them; any one of them is enough */
+    sources?: readonly string[];
     /** seeds kept per query string */
     limit?: number;
     max_hops?: number;

package/dist/schema/search.js

@@ -1,3 +1,4 @@
+import { CliError, EXIT } from '@zenera/cli/lib';
 import { assertSameEmbedding } from "../common/manifest.js";
 import { openIndex } from "./files.js";
 import { EntityStore } from "./store.js";
@@ -21,7 +22,8 @@ export class SchemaIndex {
     static async open(dir, embedder) {
         const index = await openIndex(dir);
         assertSameEmbedding(index.manifest, embedder.id);
-        return new SchemaIndex(index, await EntityStore.open(dir), embedder);
+        const names = index.manifest.sources.map((s) => s.name);
+        return new SchemaIndex(index, await EntityStore.open(dir, names), embedder);
     }
     schemas() {
         return this.#index.schemas();
@@ -33,6 +35,7 @@ export class SchemaIndex {
         this.#store.close();
     }
     async search(query, signal) {
+        this.#assertSources(query.sources);
         const terms = termsOf(query);
         if (terms.length === 0) {
             return { seeds: [], subgraphs: [], empty: [] };
@@ -70,11 +73,24 @@ export class SchemaIndex {
         });
         return { seeds, subgraphs, empty };
     }
+    /** Settled before the embedder is called, so a typo costs no credential. */
+    #assertSources(wanted) {
+        if (!wanted || wanted.length === 0) {
+            return;
+        }
+        const known = this.manifest.sources.map((s) => s.name);
+        const missing = wanted.filter((name) => !known.includes(name));
+        if (missing.length > 0) {
+            throw new CliError(`this index holds no document called ${missing.join(', ')}`, EXIT.failed, `it has: ${known.join(', ')}`);
+        }
+    }
 }
 // ---------------------------------------------------------------------------
 function termsOf(query) {
     const method = methodTypes(query.method_type);
     const loose = query.direction ?? 'any';
+    // A document is a constraint on the whole question, not on one field of it.
+    const sources = query.sources?.length ? query.sources : undefined;
     return [
         ...group(query.all, 'all', { methodTypes: method.mixed }),
         ...group(query.methods, 'methods', { kinds: ['method'], methodTypes: method.only }),
@@ -96,7 +112,7 @@ function termsOf(query) {
             kinds: ['property'],
             directions: sides('output'),
         }),
-    ];
+    ].map((term) => (sources ? { ...term, filter: { ...term.filter, sources } } : term));
 }
 function group(texts, field, filter) {
     return (texts ?? [])

package/dist/schema/store.d.ts

@@ -4,6 +4,7 @@ export interface StoreFilter {
     kinds?: readonly string[];
     directions?: readonly string[];
     methodTypes?: readonly string[];
+    sources?: readonly string[];
 }
 export interface Hit {
     record: EntityRecord;
@@ -20,8 +21,9 @@ export interface WriteResult {
 export declare function writeStore(dir: string, rows: readonly EntityRecord[], vectors: readonly Float32Array[]): Promise<WriteResult>;
 export declare class EntityStore {
     #private;
-    constructor(db: Connection, table: Table);
-    static open(dir: string): Promise<EntityStore>;
+    constructor(db: Connection, table: Table, sources?: readonly string[]);
+    /** `sources` is the document vocabulary a `sources` filter is checked against. */
+    static open(dir: string, sources?: readonly string[]): Promise<EntityStore>;
     /**
      * One hybrid query: the same string goes to the full-text side and, as a
      * vector, to the nearest-neighbour side, and LanceDB fuses the two.