@adonis-agora/agent 0.36.0 → 0.37.0

package/dist/src/rag/tree/pg-tree-store.js

@@ -0,0 +1,205 @@
+import { buildMetadataWhere, normalizeRows, stripNulBytes } from '../pg-vector-store.js';
+import { headerOf } from './store.js';
+const IDENTIFIER = /^[A-Za-z_][A-Za-z0-9_]*$/;
+function assertIdentifier(name) {
+    if (!IDENTIFIER.test(name)) {
+        throw new Error(`PgDocumentTreeStore: invalid table name ${JSON.stringify(name)}`);
+    }
+    return name;
+}
+/** A `jsonb` column comes back parsed on `pg`, as a string on some drivers — accept both. */
+function json(value) {
+    if (value === null || value === undefined) {
+        return undefined;
+    }
+    return (typeof value === 'string' ? JSON.parse(value) : value);
+}
+/** `WHERE a AND b` → `['a', 'b']`-style: the conditions of {@link buildMetadataWhere}, without the keyword. */
+function metadataConditions(filter) {
+    const where = buildMetadataWhere(filter, 'metadata');
+    return {
+        sql: where.sql === '' ? undefined : where.sql.replace(/^WHERE\s+/, ''),
+        bindings: where.bindings,
+    };
+}
+/**
+ * A Postgres {@link DocumentTreeStore} over Lucid's raw runner (`db` or `db.connection(name)` — the
+ * same structural {@link LucidRawRunner} `PgVectorStore` takes, so `@adonisjs/lucid` stays an
+ * optional peer): one row per tree (header columns, `metadata` as `jsonb` for filtering, the nodes as
+ * `jsonb`) and one row per unit of text. Filters compile through the same builder as `PgVectorStore`'s
+ * — scalar `@>` containment, array match-any, empty array denies — so the filter you pass the vector
+ * store works here unchanged.
+ *
+ * Provision it with the published `create_agent_rag_trees` migration, or {@link ensureSchema} /
+ * {@link schemaStatements}. Needs no extension. Writes run in one transaction when the runner has one.
+ */
+export class PgDocumentTreeStore {
+    db;
+    table;
+    unitsTable;
+    insertBatchSize;
+    constructor(db, options = {}) {
+        this.db = db;
+        this.table = assertIdentifier(options.table ?? 'agent_rag_trees');
+        this.unitsTable = assertIdentifier(options.unitsTable ?? `${this.table}_units`);
+        this.insertBatchSize = Math.max(1, Math.floor(options.insertBatchSize ?? 200));
+    }
+    /** The idempotent DDL {@link ensureSchema} runs, in order. */
+    schemaStatements() {
+        return [
+            `CREATE TABLE IF NOT EXISTS ${this.table} (
+        document_id TEXT PRIMARY KEY,
+        title TEXT,
+        description TEXT,
+        structure TEXT NOT NULL,
+        fingerprint TEXT NOT NULL,
+        unit_count INTEGER NOT NULL,
+        source TEXT,
+        metadata JSONB,
+        tree JSONB NOT NULL,
+        built_at TIMESTAMPTZ NOT NULL
+      )`,
+            `CREATE INDEX IF NOT EXISTS ${this.table}_metadata_idx ON ${this.table} USING gin (metadata)`,
+            `CREATE TABLE IF NOT EXISTS ${this.unitsTable} (
+        document_id TEXT NOT NULL,
+        idx INTEGER NOT NULL,
+        page INTEGER,
+        text TEXT NOT NULL,
+        PRIMARY KEY (document_id, idx)
+      )`,
+        ];
+    }
+    async ensureSchema() {
+        for (const statement of this.schemaStatements()) {
+            await this.db.rawQuery(statement);
+        }
+    }
+    async write(work) {
+        const transaction = this.db.transaction;
+        if (typeof transaction !== 'function') {
+            await work(this.db);
+            return;
+        }
+        await transaction.call(this.db, async (trx) => {
+            const runner = trx;
+            await work(typeof runner.rawQuery === 'function' ? runner : this.db);
+        });
+    }
+    /** Replace the tree and its units (one transaction when the runner has one). */
+    async put(tree, units) {
+        const clean = stripNulBytes(tree);
+        const { metadata, ...rest } = clean;
+        await this.write(async (runner) => {
+            await runner.rawQuery(`DELETE FROM ${this.unitsTable} WHERE document_id = ?`, [
+                clean.documentId,
+            ]);
+            for (let start = 0; start < units.length; start += this.insertBatchSize) {
+                const batch = units.slice(start, start + this.insertBatchSize);
+                const bindings = [];
+                for (const unit of batch) {
+                    bindings.push(clean.documentId, unit.index, unit.page ?? null, stripNulBytes(unit.text));
+                }
+                await runner.rawQuery(`INSERT INTO ${this.unitsTable} (document_id, idx, page, text) VALUES ${batch
+                    .map(() => '(?, ?, ?, ?)')
+                    .join(', ')}`, bindings);
+            }
+            await runner.rawQuery(`INSERT INTO ${this.table}
+           (document_id, title, description, structure, fingerprint, unit_count, source, metadata, tree, built_at)
+         VALUES (?, ?, ?, ?, ?, ?, ?, ?::jsonb, ?::jsonb, ?)
+         ON CONFLICT (document_id) DO UPDATE SET
+           title = EXCLUDED.title, description = EXCLUDED.description, structure = EXCLUDED.structure,
+           fingerprint = EXCLUDED.fingerprint, unit_count = EXCLUDED.unit_count, source = EXCLUDED.source,
+           metadata = EXCLUDED.metadata, tree = EXCLUDED.tree, built_at = EXCLUDED.built_at`, [
+                clean.documentId,
+                clean.title ?? null,
+                clean.description ?? null,
+                clean.structure,
+                clean.fingerprint,
+                clean.unitCount,
+                clean.source ?? null,
+                metadata === undefined ? null : JSON.stringify(metadata),
+                JSON.stringify(rest),
+                clean.builtAt,
+            ]);
+        });
+    }
+    async get(documentId, filter) {
+        return (await this.getMany([documentId], filter))[0];
+    }
+    async getMany(documentIds, filter) {
+        if (documentIds.length === 0) {
+            return [];
+        }
+        const where = metadataConditions(filter);
+        const rows = normalizeRows(await this.db.rawQuery(`SELECT document_id, metadata, tree FROM ${this.table}
+          WHERE document_id = ANY(?::text[])${where.sql !== undefined ? ` AND ${where.sql}` : ''}`, [[...new Set(documentIds)], ...where.bindings]));
+        return rows.map((row) => {
+            const tree = json(row.tree);
+            const metadata = json(row.metadata);
+            return { ...tree, ...(metadata !== undefined ? { metadata } : {}) };
+        });
+    }
+    async list(options = {}) {
+        const where = metadataConditions(options.filter);
+        const conditions = [];
+        const bindings = [];
+        if (where.sql !== undefined) {
+            conditions.push(where.sql);
+            bindings.push(...where.bindings);
+        }
+        if (options.documentIds !== undefined) {
+            conditions.push('document_id = ANY(?::text[])');
+            bindings.push([...options.documentIds]);
+        }
+        let limit = '';
+        if (options.limit !== undefined) {
+            limit = 'LIMIT ?';
+            bindings.push(Math.max(0, Math.floor(options.limit)));
+        }
+        const rows = normalizeRows(await this.db.rawQuery(`SELECT document_id, title, description, structure, fingerprint, unit_count, source, metadata, built_at
+           FROM ${this.table}
+           ${conditions.length > 0 ? `WHERE ${conditions.join(' AND ')}` : ''}
+           ORDER BY document_id ${limit}`, bindings));
+        return rows.map((row) => {
+            const metadata = json(row.metadata);
+            const builtAt = row.built_at;
+            return headerOf({
+                documentId: String(row.document_id),
+                version: 1,
+                ...(typeof row.title === 'string' ? { title: row.title } : {}),
+                ...(typeof row.description === 'string' ? { description: row.description } : {}),
+                structure: row.structure,
+                fingerprint: String(row.fingerprint),
+                unitCount: Number(row.unit_count),
+                nodes: [],
+                stats: {
+                    llmCalls: 0,
+                    inputTokens: 0,
+                    outputTokens: 0,
+                    durationMs: 0,
+                    budgetExhausted: false,
+                    reusedSummaries: 0,
+                },
+                ...(metadata !== undefined ? { metadata } : {}),
+                ...(typeof row.source === 'string' ? { source: row.source } : {}),
+                builtAt: builtAt instanceof Date ? builtAt.toISOString() : new Date(String(builtAt)).toISOString(),
+            });
+        });
+    }
+    async readUnits(documentId, start, end) {
+        const rows = normalizeRows(await this.db.rawQuery(`SELECT idx, page, text FROM ${this.unitsTable}
+          WHERE document_id = ? AND idx BETWEEN ? AND ? ORDER BY idx`, [documentId, Math.max(0, start), end]));
+        return rows.map((row) => ({
+            index: Number(row.idx),
+            text: String(row.text),
+            ...(row.page !== null && row.page !== undefined ? { page: Number(row.page) } : {}),
+        }));
+    }
+    async remove(documentId) {
+        await this.write(async (runner) => {
+            await runner.rawQuery(`DELETE FROM ${this.unitsTable} WHERE document_id = ?`, [documentId]);
+            await runner.rawQuery(`DELETE FROM ${this.table} WHERE document_id = ?`, [documentId]);
+        });
+    }
+}
+//# sourceMappingURL=pg-tree-store.js.map

package/dist/src/rag/tree/pg-tree-store.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"pg-tree-store.js","sourceRoot":"","sources":["../../../../src/rag/tree/pg-tree-store.ts"],"names":[],"mappings":"AACA,OAAO,EAAE,kBAAkB,EAAE,aAAa,EAAE,aAAa,EAAE,MAAM,uBAAuB,CAAC;AACzF,OAAO,EAAmD,QAAQ,EAAE,MAAM,YAAY,CAAC;AAYvF,MAAM,UAAU,GAAG,0BAA0B,CAAC;AAE9C,SAAS,gBAAgB,CAAC,IAAY;IACpC,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,EAAE,CAAC;QAC3B,MAAM,IAAI,KAAK,CAAC,2CAA2C,IAAI,CAAC,SAAS,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;IACrF,CAAC;IACD,OAAO,IAAI,CAAC;AACd,CAAC;AAOD,6FAA6F;AAC7F,SAAS,IAAI,CAAI,KAAc;IAC7B,IAAI,KAAK,KAAK,IAAI,IAAI,KAAK,KAAK,SAAS,EAAE,CAAC;QAC1C,OAAO,SAAS,CAAC;IACnB,CAAC;IACD,OAAO,CAAC,OAAO,KAAK,KAAK,QAAQ,CAAC,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,KAAK,CAAM,CAAC;AACtE,CAAC;AAED,+GAA+G;AAC/G,SAAS,kBAAkB,CAAC,MAA2C;IAIrE,MAAM,KAAK,GAAG,kBAAkB,CAAC,MAAM,EAAE,UAAU,CAAC,CAAC;IACrD,OAAO;QACL,GAAG,EAAE,KAAK,CAAC,GAAG,KAAK,EAAE,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,WAAW,EAAE,EAAE,CAAC;QACtE,QAAQ,EAAE,KAAK,CAAC,QAAQ;KACzB,CAAC;AACJ,CAAC;AAED;;;;;;;;;;GAUG;AACH,MAAM,OAAO,mBAAmB;IAMT,EAAE;IALJ,KAAK,CAAS;IACd,UAAU,CAAS;IACrB,eAAe,CAAS;IAEzC,YACqB,EAAkB,EACrC,OAAO,GAA+B,EAAE;kBADrB,EAAE;QAGrB,IAAI,CAAC,KAAK,GAAG,gBAAgB,CAAC,OAAO,CAAC,KAAK,IAAI,iBAAiB,CAAC,CAAC;QAClE,IAAI,CAAC,UAAU,GAAG,gBAAgB,CAAC,OAAO,CAAC,UAAU,IAAI,GAAG,IAAI,CAAC,KAAK,QAAQ,CAAC,CAAC;QAChF,IAAI,CAAC,eAAe,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,IAAI,CAAC,KAAK,CAAC,OAAO,CAAC,eAAe,IAAI,GAAG,CAAC,CAAC,CAAC;IACjF,CAAC;IAED,8DAA8D;IAC9D,gBAAgB;QACd,OAAO;YACL,8BAA8B,IAAI,CAAC,KAAK;;;;;;;;;;;QAWtC;YACF,8BAA8B,IAAI,CAAC,KAAK,oBAAoB,IAAI,CAAC,KAAK,uBAAuB;YAC7F,8BAA8B,IAAI,CAAC,UAAU;;;;;;QAM3C;SACH,CAAC;IACJ,CAAC;IAED,KAAK,CAAC,YAAY;QAChB,KAAK,MAAM,SAAS,IAAI,IAAI,CAAC,gBAAgB,EAAE,EAAE,CAAC;YAChD,MAAM,IAAI,CAAC,EAAE,CAAC,QAAQ,CAAC,SAAS,CAAC,CAAC;QACpC,CAAC;IACH,CAAC;IAEO,KAAK,CAAC,KAAK,CAAC,IAA+C;QACjE,MAAM,WAAW,GAAI,IAAI,CAAC,EAAiC,CAAC,WAAW,CAAC;QACxE,IAAI,OAAO,WAAW,KAAK,UAAU,EAAE,CAAC;YACtC,MAAM,IAAI,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;YACpB,OAAO;QACT,CAAC;QACD,MAAM,WAAW,CAAC,IAAI,CAAC,IAAI,CAAC,EAAE,EAAE,KAAK,EAAE,GAAY,EAAE,EAAE;YACrD,MAAM,MAAM,GAAG,GAA8B,CAAC;YAC9C,MAAM,IAAI,CAAC,OAAO,MAAM,CAAC,QAAQ,KAAK,UAAU,CAAC,CAAC,CAAE,MAAyB,CAAC,CAAC,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;QAC3F,CAAC,CAAC,CAAC;IACL,CAAC;IAED,gFAAgF;IAChF,KAAK,CAAC,GAAG,CAAC,IAAkB,EAAE,KAAiB;QAC7C,MAAM,KAAK,GAAG,aAAa,CAAC,IAAI,CAAC,CAAC;QAClC,MAAM,EAAE,QAAQ,EAAE,GAAG,IAAI,EAAE,GAAG,KAAK,CAAC;QACpC,MAAM,IAAI,CAAC,KAAK,CAAC,KAAK,EAAE,MAAM,EAAE,EAAE;YAChC,MAAM,MAAM,CAAC,QAAQ,CAAC,eAAe,IAAI,CAAC,UAAU,wBAAwB,EAAE;gBAC5E,KAAK,CAAC,UAAU;aACjB,CAAC,CAAC;YACH,KAAK,IAAI,KAAK,GAAG,CAAC,EAAE,KAAK,GAAG,KAAK,CAAC,MAAM,EAAE,KAAK,IAAI,IAAI,CAAC,eAAe,EAAE,CAAC;gBACxE,MAAM,KAAK,GAAG,KAAK,CAAC,KAAK,CAAC,KAAK,EAAE,KAAK,GAAG,IAAI,CAAC,eAAe,CAAC,CAAC;gBAC/D,MAAM,QAAQ,GAAc,EAAE,CAAC;gBAC/B,KAAK,MAAM,IAAI,IAAI,KAAK,EAAE,CAAC;oBACzB,QAAQ,CAAC,IAAI,CAAC,KAAK,CAAC,UAAU,EAAE,IAAI,CAAC,KAAK,EAAE,IAAI,CAAC,IAAI,IAAI,IAAI,EAAE,aAAa,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC;gBAC3F,CAAC;gBACD,MAAM,MAAM,CAAC,QAAQ,CACnB,eAAe,IAAI,CAAC,UAAU,0CAA0C,KAAK;qBAC1E,GAAG,CAAC,GAAG,EAAE,CAAC,cAAc,CAAC;qBACzB,IAAI,CAAC,IAAI,CAAC,EAAE,EACf,QAAQ,CACT,CAAC;YACJ,CAAC;YACD,MAAM,MAAM,CAAC,QAAQ,CACnB,eAAe,IAAI,CAAC,KAAK;;;;;;4FAM2D,EACpF;gBACE,KAAK,CAAC,UAAU;gBAChB,KAAK,CAAC,KAAK,IAAI,IAAI;gBACnB,KAAK,CAAC,WAAW,IAAI,IAAI;gBACzB,KAAK,CAAC,SAAS;gBACf,KAAK,CAAC,WAAW;gBACjB,KAAK,CAAC,SAAS;gBACf,KAAK,CAAC,MAAM,IAAI,IAAI;gBACpB,QAAQ,KAAK,SAAS,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAC;gBACxD,IAAI,CAAC,SAAS,CAAC,IAAI,CAAC;gBACpB,KAAK,CAAC,OAAO;aACd,CACF,CAAC;QACJ,CAAC,CAAC,CAAC;IACL,CAAC;IAED,KAAK,CAAC,GAAG,CACP,UAAkB,EAClB,MAAgC;QAEhC,OAAO,CAAC,MAAM,IAAI,CAAC,OAAO,CAAC,CAAC,UAAU,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;IACvD,CAAC;IAED,KAAK,CAAC,OAAO,CACX,WAA8B,EAC9B,MAAgC;QAEhC,IAAI,WAAW,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YAC7B,OAAO,EAAE,CAAC;QACZ,CAAC;QACD,MAAM,KAAK,GAAG,kBAAkB,CAAC,MAAM,CAAC,CAAC;QACzC,MAAM,IAAI,GAAG,aAAa,CACxB,MAAM,IAAI,CAAC,EAAE,CAAC,QAAQ,CACpB,2CAA2C,IAAI,CAAC,KAAK;8CACf,KAAK,CAAC,GAAG,KAAK,SAAS,CAAC,CAAC,CAAC,QAAQ,KAAK,CAAC,GAAG,EAAE,CAAC,CAAC,CAAC,EAAE,EAAE,EAC1F,CAAC,CAAC,GAAG,IAAI,GAAG,CAAC,WAAW,CAAC,CAAC,EAAE,GAAG,KAAK,CAAC,QAAQ,CAAC,CAC/C,CACF,CAAC;QACF,OAAO,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,EAAE;YACtB,MAAM,IAAI,GAAG,IAAI,CAAe,GAAG,CAAC,IAAI,CAAiB,CAAC;YAC1D,MAAM,QAAQ,GAAG,IAAI,CAA0B,GAAG,CAAC,QAAQ,CAAC,CAAC;YAC7D,OAAO,EAAE,GAAG,IAAI,EAAE,GAAG,CAAC,QAAQ,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,QAAQ,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC;QACtE,CAAC,CAAC,CAAC;IACL,CAAC;IAED,KAAK,CAAC,IAAI,CACR,OAAO,GAIH,EAAE;QAEN,MAAM,KAAK,GAAG,kBAAkB,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC;QACjD,MAAM,UAAU,GAAa,EAAE,CAAC;QAChC,MAAM,QAAQ,GAAc,EAAE,CAAC;QAC/B,IAAI,KAAK,CAAC,GAAG,KAAK,SAAS,EAAE,CAAC;YAC5B,UAAU,CAAC,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;YAC3B,QAAQ,CAAC,IAAI,CAAC,GAAG,KAAK,CAAC,QAAQ,CAAC,CAAC;QACnC,CAAC;QACD,IAAI,OAAO,CAAC,WAAW,KAAK,SAAS,EAAE,CAAC;YACtC,UAAU,CAAC,IAAI,CAAC,8BAA8B,CAAC,CAAC;YAChD,QAAQ,CAAC,IAAI,CAAC,CAAC,GAAG,OAAO,CAAC,WAAW,CAAC,CAAC,CAAC;QAC1C,CAAC;QACD,IAAI,KAAK,GAAG,EAAE,CAAC;QACf,IAAI,OAAO,CAAC,KAAK,KAAK,SAAS,EAAE,CAAC;YAChC,KAAK,GAAG,SAAS,CAAC;YAClB,QAAQ,CAAC,IAAI,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,IAAI,CAAC,KAAK,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;QACxD,CAAC;QACD,MAAM,IAAI,GAAG,aAAa,CACxB,MAAM,IAAI,CAAC,EAAE,CAAC,QAAQ,CACpB;kBACU,IAAI,CAAC,KAAK;aACf,UAAU,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,CAAC,SAAS,UAAU,CAAC,IAAI,CAAC,OAAO,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE;kCAC3C,KAAK,EAAE,EACjC,QAAQ,CACT,CACF,CAAC;QACF,OAAO,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,EAAE;YACtB,MAAM,QAAQ,GAAG,IAAI,CAA0B,GAAG,CAAC,QAAQ,CAAC,CAAC;YAC7D,MAAM,OAAO,GAAG,GAAG,CAAC,QAAQ,CAAC;YAC7B,OAAO,QAAQ,CAAC;gBACd,UAAU,EAAE,MAAM,CAAC,GAAG,CAAC,WAAW,CAAC;gBACnC,OAAO,EAAE,CAAC;gBACV,GAAG,CAAC,OAAO,GAAG,CAAC,KAAK,KAAK,QAAQ,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,GAAG,CAAC,KAAK,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;gBAC9D,GAAG,CAAC,OAAO,GAAG,CAAC,WAAW,KAAK,QAAQ,CAAC,CAAC,CAAC,EAAE,WAAW,EAAE,GAAG,CAAC,WAAW,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;gBAChF,SAAS,EAAE,GAAG,CAAC,SAAsC;gBACrD,WAAW,EAAE,MAAM,CAAC,GAAG,CAAC,WAAW,CAAC;gBACpC,SAAS,EAAE,MAAM,CAAC,GAAG,CAAC,UAAU,CAAC;gBACjC,KAAK,EAAE,EAAE;gBACT,KAAK,EAAE;oBACL,QAAQ,EAAE,CAAC;oBACX,WAAW,EAAE,CAAC;oBACd,YAAY,EAAE,CAAC;oBACf,UAAU,EAAE,CAAC;oBACb,eAAe,EAAE,KAAK;oBACtB,eAAe,EAAE,CAAC;iBACnB;gBACD,GAAG,CAAC,QAAQ,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,QAAQ,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;gBAC/C,GAAG,CAAC,OAAO,GAAG,CAAC,MAAM,KAAK,QAAQ,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,GAAG,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;gBACjE,OAAO,EACL,OAAO,YAAY,IAAI,CAAC,CAAC,CAAC,OAAO,CAAC,WAAW,EAAE,CAAC,CAAC,CAAC,IAAI,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,CAAC,CAAC,WAAW,EAAE;aAC5F,CAAC,CAAC;QACL,CAAC,CAAC,CAAC;IACL,CAAC;IAED,KAAK,CAAC,SAAS,CAAC,UAAkB,EAAE,KAAa,EAAE,GAAW;QAC5D,MAAM,IAAI,GAAG,aAAa,CACxB,MAAM,IAAI,CAAC,EAAE,CAAC,QAAQ,CACpB,+BAA+B,IAAI,CAAC,UAAU;qEACe,EAC7D,CAAC,UAAU,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,KAAK,CAAC,EAAE,GAAG,CAAC,CACtC,CACF,CAAC;QACF,OAAO,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,EAAE,CAAC,CAAC;YACxB,KAAK,EAAE,MAAM,CAAC,GAAG,CAAC,GAAG,CAAC;YACtB,IAAI,EAAE,MAAM,CAAC,GAAG,CAAC,IAAI,CAAC;YACtB,GAAG,CAAC,GAAG,CAAC,IAAI,KAAK,IAAI,IAAI,GAAG,CAAC,IAAI,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,IAAI,EAAE,MAAM,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;SACnF,CAAC,CAAC,CAAC;IACN,CAAC;IAED,KAAK,CAAC,MAAM,CAAC,UAAkB;QAC7B,MAAM,IAAI,CAAC,KAAK,CAAC,KAAK,EAAE,MAAM,EAAE,EAAE;YAChC,MAAM,MAAM,CAAC,QAAQ,CAAC,eAAe,IAAI,CAAC,UAAU,wBAAwB,EAAE,CAAC,UAAU,CAAC,CAAC,CAAC;YAC5F,MAAM,MAAM,CAAC,QAAQ,CAAC,eAAe,IAAI,CAAC,KAAK,wBAAwB,EAAE,CAAC,UAAU,CAAC,CAAC,CAAC;QACzF,CAAC,CAAC,CAAC;IACL,CAAC;CACF"}

package/dist/src/rag/tree/store.d.ts

@@ -0,0 +1,57 @@
+import type { DocumentTree, TreeUnit } from './types.js';
+/** A stored tree without its nodes — what listing and document selection need. */
+export interface DocumentTreeHeader {
+    documentId: string;
+    title?: string;
+    description?: string;
+    structure: DocumentTree['structure'];
+    fingerprint: string;
+    unitCount: number;
+    metadata?: Record<string, unknown>;
+    source?: string;
+    builtAt: string;
+}
+/**
+ * Where navigation trees live: one tree per document plus the units (pages/sections) its node
+ * ranges point at, so navigation can read a node's text without going back to the source file.
+ *
+ * `filter` has **exactly** the vector stores' semantics (see `VectorStore`): every key must match
+ * the tree's `metadata`, an array value is match-any, and an empty array denies. It is an access
+ * boundary — a navigator given a tenant filter must not read another tenant's tree — so every read
+ * that takes one applies it. `readUnits` takes none: call it only for a tree a filtered `get` just
+ * returned (which is what the navigator does).
+ */
+export interface DocumentTreeStore {
+    /** Store (replace) a document's tree and units. */
+    put(tree: DocumentTree, units: TreeUnit[]): Promise<void>;
+    /** A document's tree, when it exists **and** passes `filter`. */
+    get(documentId: string, filter?: Record<string, unknown>): Promise<DocumentTree | undefined>;
+    /** The trees of those of `documentIds` that exist and pass `filter` (any order). */
+    getMany(documentIds: readonly string[], filter?: Record<string, unknown>): Promise<DocumentTree[]>;
+    /** Headers (no nodes) of every tree passing `filter`, or of the listed `documentIds` only. */
+    list(options?: {
+        filter?: Record<string, unknown>;
+        documentIds?: readonly string[];
+        limit?: number;
+    }): Promise<DocumentTreeHeader[]>;
+    /** Units `start..end` (inclusive) of a document, in order. */
+    readUnits(documentId: string, start: number, end: number): Promise<TreeUnit[]>;
+    /** Drop a document's tree and units. Unknown ids are a no-op. */
+    remove(documentId: string): Promise<void>;
+}
+export declare function headerOf(tree: DocumentTree): DocumentTreeHeader;
+/** An in-process {@link DocumentTreeStore} — tests, and single-process apps with small corpora. */
+export declare class MemoryDocumentTreeStore implements DocumentTreeStore {
+    private readonly trees;
+    put(tree: DocumentTree, units: TreeUnit[]): Promise<void>;
+    get(documentId: string, filter?: Record<string, unknown>): Promise<DocumentTree | undefined>;
+    getMany(documentIds: readonly string[], filter?: Record<string, unknown>): Promise<DocumentTree[]>;
+    list(options?: {
+        filter?: Record<string, unknown>;
+        documentIds?: readonly string[];
+        limit?: number;
+    }): Promise<DocumentTreeHeader[]>;
+    readUnits(documentId: string, start: number, end: number): Promise<TreeUnit[]>;
+    remove(documentId: string): Promise<void>;
+}
+//# sourceMappingURL=store.d.ts.map

package/dist/src/rag/tree/store.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"store.d.ts","sourceRoot":"","sources":["../../../../src/rag/tree/store.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EAAE,YAAY,EAAE,QAAQ,EAAE,MAAM,YAAY,CAAC;AAEzD,kFAAkF;AAClF,MAAM,WAAW,kBAAkB;IACjC,UAAU,EAAE,MAAM,CAAC;IACnB,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,SAAS,EAAE,YAAY,CAAC,WAAW,CAAC,CAAC;IACrC,WAAW,EAAE,MAAM,CAAC;IACpB,SAAS,EAAE,MAAM,CAAC;IAClB,QAAQ,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;IACnC,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,OAAO,EAAE,MAAM,CAAC;CACjB;AAED;;;;;;;;;GASG;AACH,MAAM,WAAW,iBAAiB;IAChC,mDAAmD;IACnD,GAAG,CAAC,IAAI,EAAE,YAAY,EAAE,KAAK,EAAE,QAAQ,EAAE,GAAG,OAAO,CAAC,IAAI,CAAC,CAAC;IAC1D,iEAAiE;IACjE,GAAG,CAAC,UAAU,EAAE,MAAM,EAAE,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAAG,OAAO,CAAC,YAAY,GAAG,SAAS,CAAC,CAAC;IAC7F,oFAAoF;IACpF,OAAO,CACL,WAAW,EAAE,SAAS,MAAM,EAAE,EAC9B,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAC/B,OAAO,CAAC,YAAY,EAAE,CAAC,CAAC;IAC3B,8FAA8F;IAC9F,IAAI,CAAC,OAAO,CAAC,EAAE;QACb,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;QACjC,WAAW,CAAC,EAAE,SAAS,MAAM,EAAE,CAAC;QAChC,KAAK,CAAC,EAAE,MAAM,CAAC;KAChB,GAAG,OAAO,CAAC,kBAAkB,EAAE,CAAC,CAAC;IAClC,8DAA8D;IAC9D,SAAS,CAAC,UAAU,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC,QAAQ,EAAE,CAAC,CAAC;IAC/E,iEAAiE;IACjE,MAAM,CAAC,UAAU,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC,CAAC;CAC3C;AAED,wBAAgB,QAAQ,CAAC,IAAI,EAAE,YAAY,GAAG,kBAAkB,CAY/D;AAYD,mGAAmG;AACnG,qBAAa,uBAAwB,YAAW,iBAAiB;IAC/D,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAgE;IAEhF,GAAG,CAAC,IAAI,EAAE,YAAY,EAAE,KAAK,EAAE,QAAQ,EAAE,GAAG,OAAO,CAAC,IAAI,CAAC,CAK9D;IAEK,GAAG,CACP,UAAU,EAAE,MAAM,EAClB,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAC/B,OAAO,CAAC,YAAY,GAAG,SAAS,CAAC,CAKnC;IAEK,OAAO,CACX,WAAW,EAAE,SAAS,MAAM,EAAE,EAC9B,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GAC/B,OAAO,CAAC,YAAY,EAAE,CAAC,CASzB;IAEK,IAAI,CACR,OAAO,GAAE;QACP,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;QACjC,WAAW,CAAC,EAAE,SAAS,MAAM,EAAE,CAAC;QAChC,KAAK,CAAC,EAAE,MAAM,CAAC;KACX,GACL,OAAO,CAAC,kBAAkB,EAAE,CAAC,CAU/B;IAEK,SAAS,CAAC,UAAU,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC,QAAQ,EAAE,CAAC,CAGnF;IAEK,MAAM,CAAC,UAAU,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC,CAE9C;CACF"}

package/dist/src/rag/tree/store.js

@@ -0,0 +1,66 @@
+import { matchesFilter } from '../filter.js';
+import { filterDeniesAll } from '../vector-store.js';
+export function headerOf(tree) {
+    return {
+        documentId: tree.documentId,
+        ...(tree.title !== undefined ? { title: tree.title } : {}),
+        ...(tree.description !== undefined ? { description: tree.description } : {}),
+        structure: tree.structure,
+        fingerprint: tree.fingerprint,
+        unitCount: tree.unitCount,
+        ...(tree.metadata !== undefined ? { metadata: tree.metadata } : {}),
+        ...(tree.source !== undefined ? { source: tree.source } : {}),
+        builtAt: tree.builtAt,
+    };
+}
+function passes(tree, filter) {
+    if (filter === undefined || Object.keys(filter).length === 0) {
+        return true;
+    }
+    if (filterDeniesAll(filter)) {
+        return false;
+    }
+    return matchesFilter(tree.metadata, filter);
+}
+/** An in-process {@link DocumentTreeStore} — tests, and single-process apps with small corpora. */
+export class MemoryDocumentTreeStore {
+    trees = new Map();
+    async put(tree, units) {
+        this.trees.set(tree.documentId, {
+            tree: structuredClone(tree),
+            units: units.map((unit) => ({ ...unit })),
+        });
+    }
+    async get(documentId, filter) {
+        const entry = this.trees.get(documentId);
+        return entry !== undefined && passes(entry.tree, filter)
+            ? structuredClone(entry.tree)
+            : undefined;
+    }
+    async getMany(documentIds, filter) {
+        const trees = [];
+        for (const id of new Set(documentIds)) {
+            const tree = await this.get(id, filter);
+            if (tree !== undefined) {
+                trees.push(tree);
+            }
+        }
+        return trees;
+    }
+    async list(options = {}) {
+        const ids = options.documentIds !== undefined ? new Set(options.documentIds) : undefined;
+        const headers = [...this.trees.values()]
+            .filter(({ tree }) => (ids === undefined || ids.has(tree.documentId)) && passes(tree, options.filter))
+            .map(({ tree }) => headerOf(tree))
+            .sort((a, b) => a.documentId.localeCompare(b.documentId));
+        return options.limit !== undefined ? headers.slice(0, options.limit) : headers;
+    }
+    async readUnits(documentId, start, end) {
+        const units = this.trees.get(documentId)?.units ?? [];
+        return units.slice(Math.max(0, start), Math.max(0, end + 1)).map((unit) => ({ ...unit }));
+    }
+    async remove(documentId) {
+        this.trees.delete(documentId);
+    }
+}
+//# sourceMappingURL=store.js.map

package/dist/src/rag/tree/store.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"store.js","sourceRoot":"","sources":["../../../../src/rag/tree/store.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,aAAa,EAAE,MAAM,cAAc,CAAC;AAC7C,OAAO,EAAE,eAAe,EAAE,MAAM,oBAAoB,CAAC;AAgDrD,MAAM,UAAU,QAAQ,CAAC,IAAkB;IACzC,OAAO;QACL,UAAU,EAAE,IAAI,CAAC,UAAU;QAC3B,GAAG,CAAC,IAAI,CAAC,KAAK,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,IAAI,CAAC,KAAK,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;QAC1D,GAAG,CAAC,IAAI,CAAC,WAAW,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,WAAW,EAAE,IAAI,CAAC,WAAW,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;QAC5E,SAAS,EAAE,IAAI,CAAC,SAAS;QACzB,WAAW,EAAE,IAAI,CAAC,WAAW;QAC7B,SAAS,EAAE,IAAI,CAAC,SAAS;QACzB,GAAG,CAAC,IAAI,CAAC,QAAQ,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,QAAQ,EAAE,IAAI,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;QACnE,GAAG,CAAC,IAAI,CAAC,MAAM,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;QAC7D,OAAO,EAAE,IAAI,CAAC,OAAO;KACtB,CAAC;AACJ,CAAC;AAED,SAAS,MAAM,CAAC,IAAkB,EAAE,MAA2C;IAC7E,IAAI,MAAM,KAAK,SAAS,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;QAC7D,OAAO,IAAI,CAAC;IACd,CAAC;IACD,IAAI,eAAe,CAAC,MAAM,CAAC,EAAE,CAAC;QAC5B,OAAO,KAAK,CAAC;IACf,CAAC;IACD,OAAO,aAAa,CAAC,IAAI,CAAC,QAAQ,EAAE,MAAM,CAAC,CAAC;AAC9C,CAAC;AAED,mGAAmG;AACnG,MAAM,OAAO,uBAAuB;IACjB,KAAK,GAAG,IAAI,GAAG,EAAqD,CAAC;IAEtF,KAAK,CAAC,GAAG,CAAC,IAAkB,EAAE,KAAiB;QAC7C,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,UAAU,EAAE;YAC9B,IAAI,EAAE,eAAe,CAAC,IAAI,CAAC;YAC3B,KAAK,EAAE,KAAK,CAAC,GAAG,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,GAAG,IAAI,EAAE,CAAC,CAAC;SAC1C,CAAC,CAAC;IACL,CAAC;IAED,KAAK,CAAC,GAAG,CACP,UAAkB,EAClB,MAAgC;QAEhC,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,UAAU,CAAC,CAAC;QACzC,OAAO,KAAK,KAAK,SAAS,IAAI,MAAM,CAAC,KAAK,CAAC,IAAI,EAAE,MAAM,CAAC;YACtD,CAAC,CAAC,eAAe,CAAC,KAAK,CAAC,IAAI,CAAC;YAC7B,CAAC,CAAC,SAAS,CAAC;IAChB,CAAC;IAED,KAAK,CAAC,OAAO,CACX,WAA8B,EAC9B,MAAgC;QAEhC,MAAM,KAAK,GAAmB,EAAE,CAAC;QACjC,KAAK,MAAM,EAAE,IAAI,IAAI,GAAG,CAAC,WAAW,CAAC,EAAE,CAAC;YACtC,MAAM,IAAI,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,EAAE,EAAE,MAAM,CAAC,CAAC;YACxC,IAAI,IAAI,KAAK,SAAS,EAAE,CAAC;gBACvB,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;YACnB,CAAC;QACH,CAAC;QACD,OAAO,KAAK,CAAC;IACf,CAAC;IAED,KAAK,CAAC,IAAI,CACR,OAAO,GAIH,EAAE;QAEN,MAAM,GAAG,GAAG,OAAO,CAAC,WAAW,KAAK,SAAS,CAAC,CAAC,CAAC,IAAI,GAAG,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC;QACzF,MAAM,OAAO,GAAG,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,MAAM,EAAE,CAAC;aACrC,MAAM,CACL,CAAC,EAAE,IAAI,EAAE,EAAE,EAAE,CACX,CAAC,GAAG,KAAK,SAAS,IAAI,GAAG,CAAC,GAAG,CAAC,IAAI,CAAC,UAAU,CAAC,CAAC,IAAI,MAAM,CAAC,IAAI,EAAE,OAAO,CAAC,MAAM,CAAC,CAClF;aACA,GAAG,CAAC,CAAC,EAAE,IAAI,EAAE,EAAE,EAAE,CAAC,QAAQ,CAAC,IAAI,CAAC,CAAC;aACjC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,UAAU,CAAC,aAAa,CAAC,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC;QAC5D,OAAO,OAAO,CAAC,KAAK,KAAK,SAAS,CAAC,CAAC,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,EAAE,OAAO,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC;IACjF,CAAC;IAED,KAAK,CAAC,SAAS,CAAC,UAAkB,EAAE,KAAa,EAAE,GAAW;QAC5D,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,UAAU,CAAC,EAAE,KAAK,IAAI,EAAE,CAAC;QACtD,OAAO,KAAK,CAAC,KAAK,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,KAAK,CAAC,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,GAAG,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,IAAI,EAAE,EAAE,CAAC,CAAC,EAAE,GAAG,IAAI,EAAE,CAAC,CAAC,CAAC;IAC5F,CAAC;IAED,KAAK,CAAC,MAAM,CAAC,UAAkB;QAC7B,IAAI,CAAC,KAAK,CAAC,MAAM,CAAC,UAAU,CAAC,CAAC;IAChC,CAAC;CACF"}

package/dist/src/rag/tree/two-stage.d.ts

@@ -0,0 +1,54 @@
+import type { Passage, RetrieveOptions, Retriever } from '../../spi/retriever.js';
+import type { TreeBudget } from './llm.js';
+import type { DocumentNavigation, TreeNavigationRetriever } from './navigate.js';
+export interface TwoStageRetrieverOptions {
+    /** Passages pulled from the first stage. Default 20. */
+    fetchTopK?: number;
+    /**
+     * Only the first this-many distinct documents of the first stage are considered for navigation.
+     * Default 3 — past the top few, a document is more likely a large irrelevant tree (the costliest
+     * failure measured: navigating FAR Part 52 because it ranked fourth) than the answer.
+     */
+    considerDocuments?: number;
+    /** Of those, navigate at most this many (the ones that have a tree). Default 2. */
+    navigateDocuments?: number;
+    /** Per-document navigation budget (overrides the navigator's). */
+    budget?: TreeBudget;
+    /** Keep the first stage's chunks of a navigated document after its tree passages. Default `false`. */
+    keepFirstStageChunks?: boolean;
+    /**
+     * The document a first-stage passage belongs to. Default: `metadata.documentId` when it is a
+     * string, else the chunk id without its `#<n>` suffix (the id scheme `chunkDocuments` produces).
+     */
+    documentIdOf?: (passage: Passage) => string;
+    /** Called with each navigation's audit trail (log it, trace it, show it). */
+    onNavigation?: (navigation: DocumentNavigation) => void;
+}
+/**
+ * Hybrid first stage, tree navigation second — the design the benchmark recommends: keep fast
+ * chunk search (hybrid/lexical/dense, 0.3 s, no LLM) for picking **documents**, and navigate the
+ * table of contents only inside the top **long** documents, where chunk ranking is what fails.
+ *
+ * 1. The first stage retrieves `fetchTopK` passages.
+ * 2. Among its first `considerDocuments` distinct documents, those with a tree in the navigator's
+ *    store (build trees only for long documents, see `indexDocumentTree`'s `minUnits`) are long;
+ *    the first `navigateDocuments` of them are navigated, in parallel, each on its own budget.
+ * 3. The output keeps the first stage's order: a navigated document's section passages take the
+ *    place of its chunks (at its best chunk's rank); short documents keep their chunks. A navigation
+ *    that fails or finds nothing leaves that document's chunks in place — never worse than stage one.
+ *
+ * `filter` is passed to both stages, so the access boundary holds on either path.
+ */
+export declare class TwoStageRetriever implements Retriever {
+    private readonly firstStage;
+    private readonly navigator;
+    private readonly options;
+    constructor(firstStage: Retriever, navigator: TreeNavigationRetriever, options?: TwoStageRetrieverOptions);
+    retrieve(query: string, options?: RetrieveOptions): Promise<Passage[]>;
+    /** {@link retrieve} with the navigations that were run. */
+    search(query: string, options?: RetrieveOptions): Promise<{
+        passages: Passage[];
+        navigations: DocumentNavigation[];
+    }>;
+}
+//# sourceMappingURL=two-stage.d.ts.map

package/dist/src/rag/tree/two-stage.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"two-stage.d.ts","sourceRoot":"","sources":["../../../../src/rag/tree/two-stage.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,OAAO,EAAE,eAAe,EAAE,SAAS,EAAE,MAAM,wBAAwB,CAAC;AAElF,OAAO,KAAK,EAAE,UAAU,EAAE,MAAM,UAAU,CAAC;AAC3C,OAAO,KAAK,EAAE,kBAAkB,EAAE,uBAAuB,EAAE,MAAM,eAAe,CAAC;AAEjF,MAAM,WAAW,wBAAwB;IACvC,wDAAwD;IACxD,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB;;;;OAIG;IACH,iBAAiB,CAAC,EAAE,MAAM,CAAC;IAC3B,mFAAmF;IACnF,iBAAiB,CAAC,EAAE,MAAM,CAAC;IAC3B,kEAAkE;IAClE,MAAM,CAAC,EAAE,UAAU,CAAC;IACpB,sGAAsG;IACtG,oBAAoB,CAAC,EAAE,OAAO,CAAC;IAC/B;;;OAGG;IACH,YAAY,CAAC,EAAE,CAAC,OAAO,EAAE,OAAO,KAAK,MAAM,CAAC;IAC5C,6EAA6E;IAC7E,YAAY,CAAC,EAAE,CAAC,UAAU,EAAE,kBAAkB,KAAK,IAAI,CAAC;CACzD;AAOD;;;;;;;;;;;;;;GAcG;AACH,qBAAa,iBAAkB,YAAW,SAAS;IAE/C,OAAO,CAAC,QAAQ,CAAC,UAAU;IAC3B,OAAO,CAAC,QAAQ,CAAC,SAAS;IAC1B,OAAO,CAAC,QAAQ,CAAC,OAAO;IAH1B,YACmB,UAAU,EAAE,SAAS,EACrB,SAAS,EAAE,uBAAuB,EAClC,OAAO,GAAE,wBAA6B,EACrD;IAEE,QAAQ,CAAC,KAAK,EAAE,MAAM,EAAE,OAAO,GAAE,eAAoB,GAAG,OAAO,CAAC,OAAO,EAAE,CAAC,CAE/E;IAED,2DAA2D;IACrD,MAAM,CACV,KAAK,EAAE,MAAM,EACb,OAAO,GAAE,eAAoB,GAC5B,OAAO,CAAC;QAAE,QAAQ,EAAE,OAAO,EAAE,CAAC;QAAC,WAAW,EAAE,kBAAkB,EAAE,CAAA;KAAE,CAAC,CA8DrE;CACF"}

package/dist/src/rag/tree/two-stage.js

@@ -0,0 +1,92 @@
+import { documentIdOf } from '../vector-store.js';
+function defaultDocumentIdOf(passage) {
+    const fromMetadata = passage.metadata?.documentId;
+    return typeof fromMetadata === 'string' ? fromMetadata : documentIdOf(passage.id);
+}
+/**
+ * Hybrid first stage, tree navigation second — the design the benchmark recommends: keep fast
+ * chunk search (hybrid/lexical/dense, 0.3 s, no LLM) for picking **documents**, and navigate the
+ * table of contents only inside the top **long** documents, where chunk ranking is what fails.
+ *
+ * 1. The first stage retrieves `fetchTopK` passages.
+ * 2. Among its first `considerDocuments` distinct documents, those with a tree in the navigator's
+ *    store (build trees only for long documents, see `indexDocumentTree`'s `minUnits`) are long;
+ *    the first `navigateDocuments` of them are navigated, in parallel, each on its own budget.
+ * 3. The output keeps the first stage's order: a navigated document's section passages take the
+ *    place of its chunks (at its best chunk's rank); short documents keep their chunks. A navigation
+ *    that fails or finds nothing leaves that document's chunks in place — never worse than stage one.
+ *
+ * `filter` is passed to both stages, so the access boundary holds on either path.
+ */
+export class TwoStageRetriever {
+    firstStage;
+    navigator;
+    options;
+    constructor(firstStage, navigator, options = {}) {
+        this.firstStage = firstStage;
+        this.navigator = navigator;
+        this.options = options;
+    }
+    async retrieve(query, options = {}) {
+        return (await this.search(query, options)).passages;
+    }
+    /** {@link retrieve} with the navigations that were run. */
+    async search(query, options = {}) {
+        const topK = options.topK ?? 5;
+        const docOf = this.options.documentIdOf ?? defaultDocumentIdOf;
+        const first = await this.firstStage.retrieve(query, {
+            topK: Math.max(topK, this.options.fetchTopK ?? 20),
+            ...(options.filter !== undefined ? { filter: options.filter } : {}),
+        });
+        const ranked = [];
+        for (const passage of first) {
+            const id = docOf(passage);
+            if (!ranked.includes(id)) {
+                ranked.push(id);
+            }
+        }
+        const considered = ranked.slice(0, Math.max(1, this.options.considerDocuments ?? 3));
+        const withTrees = new Set((await this.navigator.store.list({
+            documentIds: considered,
+            ...(options.filter !== undefined ? { filter: options.filter } : {}),
+        })).map((header) => header.documentId));
+        const toNavigate = considered
+            .filter((id) => withTrees.has(id))
+            .slice(0, Math.max(0, this.options.navigateDocuments ?? 2));
+        const navigated = new Map();
+        let navigations = [];
+        if (toNavigate.length > 0) {
+            const result = await this.navigator.navigate(query, toNavigate, {
+                ...(options.filter !== undefined ? { filter: options.filter } : {}),
+                ...(this.options.budget !== undefined ? { budget: this.options.budget } : {}),
+            });
+            navigations = result.documents;
+            for (const passage of result.passages) {
+                const id = String(passage.metadata?.documentId ?? '');
+                navigated.set(id, [...(navigated.get(id) ?? []), passage]);
+            }
+            for (const navigation of navigations) {
+                this.options.onNavigation?.(navigation);
+            }
+        }
+        const merged = [];
+        const emitted = new Set();
+        for (const passage of first) {
+            const id = docOf(passage);
+            const replacement = navigated.get(id);
+            if (replacement === undefined || replacement.length === 0) {
+                merged.push(passage);
+                continue;
+            }
+            if (!emitted.has(id)) {
+                emitted.add(id);
+                merged.push(...replacement);
+            }
+            if (this.options.keepFirstStageChunks) {
+                merged.push(passage);
+            }
+        }
+        return { passages: merged.slice(0, topK), navigations };
+    }
+}
+//# sourceMappingURL=two-stage.js.map

package/dist/src/rag/tree/two-stage.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"two-stage.js","sourceRoot":"","sources":["../../../../src/rag/tree/two-stage.ts"],"names":[],"mappings":"AACA,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AA4BlD,SAAS,mBAAmB,CAAC,OAAgB;IAC3C,MAAM,YAAY,GAAG,OAAO,CAAC,QAAQ,EAAE,UAAU,CAAC;IAClD,OAAO,OAAO,YAAY,KAAK,QAAQ,CAAC,CAAC,CAAC,YAAY,CAAC,CAAC,CAAC,YAAY,CAAC,OAAO,CAAC,EAAE,CAAC,CAAC;AACpF,CAAC;AAED;;;;;;;;;;;;;;GAcG;AACH,MAAM,OAAO,iBAAiB;IAET,UAAU;IACV,SAAS;IACT,OAAO;IAH1B,YACmB,UAAqB,EACrB,SAAkC,EAClC,OAAO,GAA6B,EAAE;0BAFtC,UAAU;yBACV,SAAS;uBACT,OAAO;IACvB,CAAC;IAEJ,KAAK,CAAC,QAAQ,CAAC,KAAa,EAAE,OAAO,GAAoB,EAAE;QACzD,OAAO,CAAC,MAAM,IAAI,CAAC,MAAM,CAAC,KAAK,EAAE,OAAO,CAAC,CAAC,CAAC,QAAQ,CAAC;IACtD,CAAC;IAED,2DAA2D;IAC3D,KAAK,CAAC,MAAM,CACV,KAAa,EACb,OAAO,GAAoB,EAAE;QAE7B,MAAM,IAAI,GAAG,OAAO,CAAC,IAAI,IAAI,CAAC,CAAC;QAC/B,MAAM,KAAK,GAAG,IAAI,CAAC,OAAO,CAAC,YAAY,IAAI,mBAAmB,CAAC;QAC/D,MAAM,KAAK,GAAG,MAAM,IAAI,CAAC,UAAU,CAAC,QAAQ,CAAC,KAAK,EAAE;YAClD,IAAI,EAAE,IAAI,CAAC,GAAG,CAAC,IAAI,EAAE,IAAI,CAAC,OAAO,CAAC,SAAS,IAAI,EAAE,CAAC;YAClD,GAAG,CAAC,OAAO,CAAC,MAAM,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;SACpE,CAAC,CAAC;QACH,MAAM,MAAM,GAAa,EAAE,CAAC;QAC5B,KAAK,MAAM,OAAO,IAAI,KAAK,EAAE,CAAC;YAC5B,MAAM,EAAE,GAAG,KAAK,CAAC,OAAO,CAAC,CAAC;YAC1B,IAAI,CAAC,MAAM,CAAC,QAAQ,CAAC,EAAE,CAAC,EAAE,CAAC;gBACzB,MAAM,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;YAClB,CAAC;QACH,CAAC;QACD,MAAM,UAAU,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,IAAI,CAAC,OAAO,CAAC,iBAAiB,IAAI,CAAC,CAAC,CAAC,CAAC;QACrF,MAAM,SAAS,GAAG,IAAI,GAAG,CACvB,CACE,MAAM,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,IAAI,CAAC;YAC9B,WAAW,EAAE,UAAU;YACvB,GAAG,CAAC,OAAO,CAAC,MAAM,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;SACpE,CAAC,CACH,CAAC,GAAG,CAAC,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,CAAC,UAAU,CAAC,CACrC,CAAC;QACF,MAAM,UAAU,GAAG,UAAU;aAC1B,MAAM,CAAC,CAAC,EAAE,EAAE,EAAE,CAAC,SAAS,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;aACjC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,IAAI,CAAC,OAAO,CAAC,iBAAiB,IAAI,CAAC,CAAC,CAAC,CAAC;QAE9D,MAAM,SAAS,GAAG,IAAI,GAAG,EAAqB,CAAC;QAC/C,IAAI,WAAW,GAAyB,EAAE,CAAC;QAC3C,IAAI,UAAU,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YAC1B,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAC,KAAK,EAAE,UAAU,EAAE;gBAC9D,GAAG,CAAC,OAAO,CAAC,MAAM,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;gBACnE,GAAG,CAAC,IAAI,CAAC,OAAO,CAAC,MAAM,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,IAAI,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;aAC9E,CAAC,CAAC;YACH,WAAW,GAAG,MAAM,CAAC,SAAS,CAAC;YAC/B,KAAK,MAAM,OAAO,IAAI,MAAM,CAAC,QAAQ,EAAE,CAAC;gBACtC,MAAM,EAAE,GAAG,MAAM,CAAC,OAAO,CAAC,QAAQ,EAAE,UAAU,IAAI,EAAE,CAAC,CAAC;gBACtD,SAAS,CAAC,GAAG,CAAC,EAAE,EAAE,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,EAAE,CAAC,EAAE,OAAO,CAAC,CAAC,CAAC;YAC7D,CAAC;YACD,KAAK,MAAM,UAAU,IAAI,WAAW,EAAE,CAAC;gBACrC,IAAI,CAAC,OAAO,CAAC,YAAY,EAAE,CAAC,UAAU,CAAC,CAAC;YAC1C,CAAC;QACH,CAAC;QAED,MAAM,MAAM,GAAc,EAAE,CAAC;QAC7B,MAAM,OAAO,GAAG,IAAI,GAAG,EAAU,CAAC;QAClC,KAAK,MAAM,OAAO,IAAI,KAAK,EAAE,CAAC;YAC5B,MAAM,EAAE,GAAG,KAAK,CAAC,OAAO,CAAC,CAAC;YAC1B,MAAM,WAAW,GAAG,SAAS,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;YACtC,IAAI,WAAW,KAAK,SAAS,IAAI,WAAW,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;gBAC1D,MAAM,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;gBACrB,SAAS;YACX,CAAC;YACD,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC;gBACrB,OAAO,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;gBAChB,MAAM,CAAC,IAAI,CAAC,GAAG,WAAW,CAAC,CAAC;YAC9B,CAAC;YACD,IAAI,IAAI,CAAC,OAAO,CAAC,oBAAoB,EAAE,CAAC;gBACtC,MAAM,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;YACvB,CAAC;QACH,CAAC;QACD,OAAO,EAAE,QAAQ,EAAE,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,EAAE,WAAW,EAAE,CAAC;IAC1D,CAAC;CACF"}

package/dist/src/rag/tree/types.d.ts

@@ -0,0 +1,135 @@
+/**
+ * The data model of tree navigation retrieval: a hierarchical table of contents per long document
+ * (PageIndex-style: sections with page ranges and summaries), built once at ingestion and read by an
+ * LLM at query time to walk to the relevant sections instead of ranking hundreds of look-alike chunks.
+ */
+/**
+ * One addressable piece of a document's text, in reading order — a PDF page, or a section of a
+ * markdown/HTML document. Tree nodes point at ranges of units, and navigation reads the text of the
+ * units a chosen node spans. `index` is the unit's zero-based position; `page` its 1-based page
+ * number when the document has pages.
+ */
+export interface TreeUnit {
+    index: number;
+    text: string;
+    /** 1-based page number (PDF). Defaults to `index + 1` for page input; absent for section input. */
+    page?: number;
+}
+/**
+ * A heading the document declares about itself: a PDF outline (bookmark) entry, a markdown header,
+ * a section boundary your extractor already knows. `unit` is the zero-based unit it starts in;
+ * `page` (1-based) is accepted instead for page input. `level` 1 is top level.
+ */
+export interface TreeHeading {
+    title: string;
+    level: number;
+    unit?: number;
+    page?: number;
+}
+/** A section as produced by an extractor that already splits on headings. */
+export interface TreeSection {
+    text: string;
+    title?: string;
+    /** Heading depth, 1 = top level. Sections without a level are treated as level 1. */
+    level?: number;
+    /** 1-based page the section starts on, when known. */
+    page?: number;
+}
+/**
+ * What {@link import('./build.js').buildDocumentTree} builds from. Give exactly one of `pages`
+ * (one string per page, in order) or `sections` (heading-split text). `outline`, when present, is the
+ * cheapest and most reliable structure there is (a PDF's bookmarks) and is used before any heuristic.
+ */
+export interface DocumentTreeInput {
+    pages?: string[];
+    sections?: TreeSection[];
+    /** Headings the document declares (PDF outline/bookmarks), in reading order. */
+    outline?: TreeHeading[];
+    /** Document title, shown to the navigating LLM and used as the citation source. */
+    title?: string;
+}
+/** How the structure of a tree was obtained, cheapest first. */
+export type TreeStructureSource = 
+/** From {@link DocumentTreeInput.outline}. */
+'outline'
+/** From section titles/levels or headings detected in the text — no LLM. */
+ | 'headings'
+/** An LLM proposed the sections (the document had no usable outline). */
+ | 'llm'
+/** Fallback: fixed-size page groups. Always succeeds; used when everything above failed or was skipped. */
+ | 'flat';
+/** Where a node's summary came from. */
+export type TreeSummarySource = 'llm' | 'text' | 'reused';
+export interface DocumentTreeNode {
+    /** Stable within the tree: `0001`, `0002`… in depth-first order (PageIndex's scheme). */
+    id: string;
+    title: string;
+    /** What the section covers — what the navigating LLM reads to decide whether to go in. */
+    summary?: string;
+    summarySource?: TreeSummarySource;
+    /** First and last unit (inclusive, zero-based) the node spans, children included. */
+    unitStart: number;
+    unitEnd: number;
+    /** First and last page (1-based, inclusive) when the document has pages. */
+    pageStart?: number;
+    pageEnd?: number;
+    /**
+     * Titles of descendants removed to respect `maxDepth` — kept because they are routing information:
+     * the pages are still reachable by reading this node, the titles tell the navigator they are here.
+     */
+    keyItems?: string[];
+    /** Hash of the node's title and the text it spans; an incremental rebuild reuses a summary when it matches. */
+    contentHash: string;
+    children: DocumentTreeNode[];
+}
+/** What a build cost and how it ended — persisted with the tree, so an operator can see it. */
+export interface TreeBuildStats {
+    llmCalls: number;
+    inputTokens: number;
+    outputTokens: number;
+    durationMs: number;
+    /** Set when the structure step fell back (to `flat`) or summaries were cut short, and why. */
+    fallbackReason?: string;
+    /** True when the per-document budget or timeout ran out at some point during the build. */
+    budgetExhausted: boolean;
+    /** Summaries carried over unchanged from the previous tree (incremental rebuild). */
+    reusedSummaries: number;
+}
+export interface DocumentTree {
+    documentId: string;
+    /** Format version of this structure. */
+    version: 1;
+    title?: string;
+    /** One-paragraph description of the whole document (LLM, when budget allows). */
+    description?: string;
+    structure: TreeStructureSource;
+    /**
+     * Hash of the input text and every build option that shapes the tree. Equal fingerprint ⇒ the
+     * build would produce the same tree, so it is skipped (see `previous` in the build options).
+     */
+    fingerprint: string;
+    unitCount: number;
+    /** Top-level nodes. */
+    nodes: DocumentTreeNode[];
+    stats: TreeBuildStats;
+    /**
+     * Document metadata used for access filtering, exactly like chunk metadata in a vector store
+     * (tenant, audience, collection…). Stores apply `filter` to it with the vector stores' semantics.
+     */
+    metadata?: Record<string, unknown>;
+    /** Citation-facing origin (title, URL), copied onto returned passages. */
+    source?: string;
+    /** ISO timestamp of the build. */
+    builtAt: string;
+}
+/** Walk a tree depth-first (pre-order). */
+export declare function walkTree(nodes: readonly DocumentTreeNode[], path?: readonly DocumentTreeNode[]): Generator<{
+    node: DocumentTreeNode;
+    path: readonly DocumentTreeNode[];
+}>;
+/** Index a tree's nodes by id, each with its ancestors. */
+export declare function indexTree(tree: Pick<DocumentTree, 'nodes'>): Map<string, {
+    node: DocumentTreeNode;
+    path: readonly DocumentTreeNode[];
+}>;
+//# sourceMappingURL=types.d.ts.map