@dudousxd/nestjs-agent-rag 0.12.0 → 0.13.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/NOTICE +31 -0
- package/README.md +102 -1
- package/dist/index.cjs +2411 -66
- package/dist/index.cjs.map +1 -1
- package/dist/index.d.cts +807 -5
- package/dist/index.d.ts +807 -5
- package/dist/index.js +2390 -65
- package/dist/index.js.map +1 -1
- package/package.json +4 -3
package/dist/index.js
CHANGED
|
@@ -4,7 +4,7 @@ var __name = (target, value) => __defProp(target, "name", { value, configurable:
|
|
|
4
4
|
// src/chunk.ts
|
|
5
5
|
function chunkText(text, options = {}) {
|
|
6
6
|
const chunkSize = options.chunkSize ?? 800;
|
|
7
|
-
const
|
|
7
|
+
const overlap2 = options.overlap ?? 100;
|
|
8
8
|
const normalized = text.trim();
|
|
9
9
|
if (normalized.length === 0) {
|
|
10
10
|
return [];
|
|
@@ -36,7 +36,7 @@ function chunkText(text, options = {}) {
|
|
|
36
36
|
if (breakAt >= normalized.length) {
|
|
37
37
|
break;
|
|
38
38
|
}
|
|
39
|
-
start = Math.max(breakAt -
|
|
39
|
+
start = Math.max(breakAt - overlap2, start + 1);
|
|
40
40
|
}
|
|
41
41
|
return chunks;
|
|
42
42
|
}
|
|
@@ -63,12 +63,12 @@ function chunkByRecord(text, chunkSize, separator) {
|
|
|
63
63
|
}));
|
|
64
64
|
continue;
|
|
65
65
|
}
|
|
66
|
-
const
|
|
67
|
-
if (
|
|
66
|
+
const candidate2 = current.length > 0 ? `${current}${separator}${trimmed}` : trimmed;
|
|
67
|
+
if (candidate2.length > chunkSize) {
|
|
68
68
|
flush();
|
|
69
69
|
current = trimmed;
|
|
70
70
|
} else {
|
|
71
|
-
current =
|
|
71
|
+
current = candidate2;
|
|
72
72
|
}
|
|
73
73
|
}
|
|
74
74
|
flush();
|
|
@@ -418,16 +418,16 @@ function keywordTerms(query, max = 24, stopWords = DEFAULT_STOP_WORDS) {
|
|
|
418
418
|
const best = Math.max(0, ...counts.map(([, n]) => n));
|
|
419
419
|
if (best === 0) return tokens.slice(0, max);
|
|
420
420
|
const stop = counts.filter(([, n]) => n === best).map(([s]) => s);
|
|
421
|
-
const
|
|
422
|
-
return (
|
|
421
|
+
const terms2 = tokens.filter((t) => !stop.some((s) => s.has(t)));
|
|
422
|
+
return (terms2.length > 0 ? terms2 : tokens).slice(0, max);
|
|
423
423
|
}
|
|
424
424
|
__name(keywordTerms, "keywordTerms");
|
|
425
425
|
function hasSearchSyntax(query) {
|
|
426
426
|
return /"[^"]+"/.test(query) || /(^|\s)-[\p{L}\p{N}]/u.test(query);
|
|
427
427
|
}
|
|
428
428
|
__name(hasSearchSyntax, "hasSearchSyntax");
|
|
429
|
-
function anyTermTsquery(
|
|
430
|
-
return
|
|
429
|
+
function anyTermTsquery(terms2) {
|
|
430
|
+
return terms2.map((t) => `'${t.replace(/'/g, "")}'`).join(" | ");
|
|
431
431
|
}
|
|
432
432
|
__name(anyTermTsquery, "anyTermTsquery");
|
|
433
433
|
|
|
@@ -924,8 +924,8 @@ var PgLexicalVectorStore = class extends PgVectorStore {
|
|
|
924
924
|
return strict;
|
|
925
925
|
}
|
|
926
926
|
}
|
|
927
|
-
const
|
|
928
|
-
return
|
|
927
|
+
const terms2 = keywordTerms(query, 24, this.stopWords);
|
|
928
|
+
return terms2.length === 0 ? [] : this.runRankedTermSearch(terms2, options);
|
|
929
929
|
}
|
|
930
930
|
/** `websearch_to_tsquery` on the query as written: every word, quoted phrases, `-exclusions`. */
|
|
931
931
|
async runTextSearch(query, options) {
|
|
@@ -951,12 +951,12 @@ var PgLexicalVectorStore = class extends PgVectorStore {
|
|
|
951
951
|
* once, and `ts_rank_cd` only for the rows ranked within `topK` (ties included): a question full
|
|
952
952
|
* of common words can match most of the table.
|
|
953
953
|
*/
|
|
954
|
-
async runRankedTermSearch(
|
|
954
|
+
async runRankedTermSearch(terms2, options) {
|
|
955
955
|
const any = `to_tsquery('${this.tsConfig}', $1)`;
|
|
956
956
|
const params = [
|
|
957
|
-
anyTermTsquery(
|
|
957
|
+
anyTermTsquery(terms2),
|
|
958
958
|
options.topK,
|
|
959
|
-
|
|
959
|
+
terms2
|
|
960
960
|
];
|
|
961
961
|
const conditions = [
|
|
962
962
|
`${this.tsvector()} @@ ${any}`,
|
|
@@ -1660,12 +1660,12 @@ var RedisVectorStore = class {
|
|
|
1660
1660
|
if (filterMatchesNothing(options.filter)) {
|
|
1661
1661
|
return [];
|
|
1662
1662
|
}
|
|
1663
|
-
const
|
|
1664
|
-
if (
|
|
1663
|
+
const terms2 = tokenizeQuery(query);
|
|
1664
|
+
if (terms2.length === 0) {
|
|
1665
1665
|
return [];
|
|
1666
1666
|
}
|
|
1667
1667
|
const filterQuery = buildFilter(options.filter);
|
|
1668
|
-
const textQuery = `@text:(${
|
|
1668
|
+
const textQuery = `@text:(${terms2.join("|")})`;
|
|
1669
1669
|
const reply = await this.client.sendCommand([
|
|
1670
1670
|
"FT.SEARCH",
|
|
1671
1671
|
this.index,
|
|
@@ -1919,8 +1919,8 @@ function describesRetrieval(value) {
|
|
|
1919
1919
|
return typeof value === "object" && value !== null && typeof Reflect.get(value, "describeRetrieval") === "function";
|
|
1920
1920
|
}
|
|
1921
1921
|
__name(describesRetrieval, "describesRetrieval");
|
|
1922
|
-
function describeRetrieval(
|
|
1923
|
-
return describesRetrieval(
|
|
1922
|
+
function describeRetrieval(candidate2) {
|
|
1923
|
+
return describesRetrieval(candidate2) ? candidate2.describeRetrieval() : {};
|
|
1924
1924
|
}
|
|
1925
1925
|
__name(describeRetrieval, "describeRetrieval");
|
|
1926
1926
|
function describeSource(base) {
|
|
@@ -2453,53 +2453,94 @@ var HttpModelError = class extends Error {
|
|
|
2453
2453
|
this.name = "HttpModelError";
|
|
2454
2454
|
}
|
|
2455
2455
|
};
|
|
2456
|
+
function isBatchTooLarge(error) {
|
|
2457
|
+
if (!(error instanceof HttpModelError)) {
|
|
2458
|
+
return false;
|
|
2459
|
+
}
|
|
2460
|
+
if (error.status === 413) {
|
|
2461
|
+
return true;
|
|
2462
|
+
}
|
|
2463
|
+
return (error.status === 400 || error.status === 422) && /batch.?size|too many inputs|maximum allowed batch|payload too large|max.*(inputs|batch)/i.test(error.message);
|
|
2464
|
+
}
|
|
2465
|
+
__name(isBatchTooLarge, "isBatchTooLarge");
|
|
2466
|
+
var learnedBatchLimits = /* @__PURE__ */ new Map();
|
|
2456
2467
|
function openAiEmbeddings(options) {
|
|
2457
2468
|
const url = `${(options.baseUrl ?? "https://api.openai.com/v1").replace(/\/+$/, "")}/embeddings`;
|
|
2458
|
-
const
|
|
2469
|
+
const limitKey = `${url} ${options.model}`;
|
|
2470
|
+
const configured = Math.max(1, Math.floor(options.batchSize ?? 64));
|
|
2471
|
+
let batchSize = Math.min(configured, learnedBatchLimits.get(limitKey) ?? configured);
|
|
2459
2472
|
const maxInputChars = options.maxInputChars ?? 24e3;
|
|
2460
2473
|
const timeoutMs = options.timeoutMs ?? 12e4;
|
|
2461
2474
|
const doFetch = options.fetch ?? fetch;
|
|
2475
|
+
const request = /* @__PURE__ */ __name(async (input) => {
|
|
2476
|
+
const response = await doFetch(url, {
|
|
2477
|
+
method: "POST",
|
|
2478
|
+
headers: {
|
|
2479
|
+
"content-type": "application/json",
|
|
2480
|
+
...options.apiKey ? {
|
|
2481
|
+
authorization: `Bearer ${options.apiKey}`
|
|
2482
|
+
} : {},
|
|
2483
|
+
...options.headers
|
|
2484
|
+
},
|
|
2485
|
+
body: JSON.stringify({
|
|
2486
|
+
model: options.model,
|
|
2487
|
+
input,
|
|
2488
|
+
...options.dimensions !== void 0 ? {
|
|
2489
|
+
dimensions: options.dimensions
|
|
2490
|
+
} : {}
|
|
2491
|
+
}),
|
|
2492
|
+
signal: AbortSignal.timeout(timeoutMs)
|
|
2493
|
+
});
|
|
2494
|
+
const body = await response.text();
|
|
2495
|
+
if (!response.ok) {
|
|
2496
|
+
throw new HttpModelError(response.status, `Embeddings request failed (${response.status}): ${errorMessageOf(body)}`);
|
|
2497
|
+
}
|
|
2498
|
+
const json2 = parseJson(body);
|
|
2499
|
+
const items = [
|
|
2500
|
+
...json2?.data ?? []
|
|
2501
|
+
].sort((a, b) => (a.index ?? 0) - (b.index ?? 0));
|
|
2502
|
+
if (items.length !== input.length || !items.every((item) => Array.isArray(item.embedding))) {
|
|
2503
|
+
throw new HttpModelError(502, `Embeddings returned ${items.length} vectors for ${input.length} inputs`);
|
|
2504
|
+
}
|
|
2505
|
+
const tokens = json2?.usage?.total_tokens ?? json2?.usage?.prompt_tokens;
|
|
2506
|
+
if (typeof tokens === "number" && Number.isFinite(tokens)) {
|
|
2507
|
+
options.onUsage?.(tokens);
|
|
2508
|
+
}
|
|
2509
|
+
return items.map((item) => item.embedding);
|
|
2510
|
+
}, "request");
|
|
2511
|
+
const requestSplitting = /* @__PURE__ */ __name(async (input) => {
|
|
2512
|
+
if (input.length > batchSize) {
|
|
2513
|
+
const vectors = [];
|
|
2514
|
+
for (let start = 0; start < input.length; start += batchSize) {
|
|
2515
|
+
vectors.push(...await requestSplitting(input.slice(start, start + batchSize)));
|
|
2516
|
+
}
|
|
2517
|
+
return vectors;
|
|
2518
|
+
}
|
|
2519
|
+
try {
|
|
2520
|
+
return await request(input);
|
|
2521
|
+
} catch (error) {
|
|
2522
|
+
if (!isBatchTooLarge(error) || input.length <= 1) {
|
|
2523
|
+
throw error;
|
|
2524
|
+
}
|
|
2525
|
+
const half = Math.ceil(input.length / 2);
|
|
2526
|
+
if (half < batchSize) {
|
|
2527
|
+
batchSize = half;
|
|
2528
|
+
learnedBatchLimits.set(limitKey, half);
|
|
2529
|
+
options.onWarn?.(`embeddings server ${url} refused a batch of ${input.length} inputs (${error.message}); retrying in batches of ${half}. Raise the server's batch limit (TEI: --max-client-batch-size) or lower batchSize.`);
|
|
2530
|
+
}
|
|
2531
|
+
return [
|
|
2532
|
+
...await requestSplitting(input.slice(0, half)),
|
|
2533
|
+
...await requestSplitting(input.slice(half))
|
|
2534
|
+
];
|
|
2535
|
+
}
|
|
2536
|
+
}, "requestSplitting");
|
|
2462
2537
|
return {
|
|
2463
2538
|
async embed(texts) {
|
|
2464
2539
|
const vectors = [];
|
|
2465
|
-
for (let start = 0; start < texts.length;
|
|
2540
|
+
for (let start = 0; start < texts.length; ) {
|
|
2466
2541
|
const input = texts.slice(start, start + batchSize).map((text) => (text.length > maxInputChars ? text.slice(0, maxInputChars) : text) || " ");
|
|
2467
|
-
|
|
2468
|
-
|
|
2469
|
-
headers: {
|
|
2470
|
-
"content-type": "application/json",
|
|
2471
|
-
...options.apiKey ? {
|
|
2472
|
-
authorization: `Bearer ${options.apiKey}`
|
|
2473
|
-
} : {},
|
|
2474
|
-
...options.headers
|
|
2475
|
-
},
|
|
2476
|
-
body: JSON.stringify({
|
|
2477
|
-
model: options.model,
|
|
2478
|
-
input,
|
|
2479
|
-
...options.dimensions !== void 0 ? {
|
|
2480
|
-
dimensions: options.dimensions
|
|
2481
|
-
} : {}
|
|
2482
|
-
}),
|
|
2483
|
-
signal: AbortSignal.timeout(timeoutMs)
|
|
2484
|
-
});
|
|
2485
|
-
const body = await response.text();
|
|
2486
|
-
if (!response.ok) {
|
|
2487
|
-
throw new HttpModelError(response.status, `Embeddings request failed (${response.status}): ${errorMessageOf(body)}`);
|
|
2488
|
-
}
|
|
2489
|
-
const json = parseJson(body);
|
|
2490
|
-
const items = [
|
|
2491
|
-
...json?.data ?? []
|
|
2492
|
-
].sort((a, b) => (a.index ?? 0) - (b.index ?? 0));
|
|
2493
|
-
if (items.length !== input.length || !items.every((item) => Array.isArray(item.embedding))) {
|
|
2494
|
-
throw new HttpModelError(502, `Embeddings returned ${items.length} vectors for ${input.length} inputs`);
|
|
2495
|
-
}
|
|
2496
|
-
for (const item of items) {
|
|
2497
|
-
vectors.push(item.embedding);
|
|
2498
|
-
}
|
|
2499
|
-
const tokens = json?.usage?.total_tokens ?? json?.usage?.prompt_tokens;
|
|
2500
|
-
if (typeof tokens === "number" && Number.isFinite(tokens)) {
|
|
2501
|
-
options.onUsage?.(tokens);
|
|
2502
|
-
}
|
|
2542
|
+
vectors.push(...await requestSplitting(input));
|
|
2543
|
+
start += input.length;
|
|
2503
2544
|
}
|
|
2504
2545
|
return vectors;
|
|
2505
2546
|
}
|
|
@@ -2515,8 +2556,8 @@ function parseJson(body) {
|
|
|
2515
2556
|
}
|
|
2516
2557
|
__name(parseJson, "parseJson");
|
|
2517
2558
|
function errorMessageOf(body) {
|
|
2518
|
-
const
|
|
2519
|
-
const message = typeof
|
|
2559
|
+
const json2 = parseJson(body);
|
|
2560
|
+
const message = typeof json2?.error === "string" ? json2.error : json2?.error?.message ?? json2?.message ?? void 0;
|
|
2520
2561
|
return typeof message === "string" ? message : body.slice(0, 300);
|
|
2521
2562
|
}
|
|
2522
2563
|
__name(errorMessageOf, "errorMessageOf");
|
|
@@ -2591,20 +2632,2284 @@ var HttpReranker = class {
|
|
|
2591
2632
|
}
|
|
2592
2633
|
};
|
|
2593
2634
|
function itemsOf(body) {
|
|
2594
|
-
let
|
|
2635
|
+
let json2;
|
|
2595
2636
|
try {
|
|
2596
|
-
|
|
2637
|
+
json2 = JSON.parse(body);
|
|
2597
2638
|
} catch {
|
|
2598
2639
|
return void 0;
|
|
2599
2640
|
}
|
|
2600
|
-
if (Array.isArray(
|
|
2601
|
-
return
|
|
2641
|
+
if (Array.isArray(json2)) {
|
|
2642
|
+
return json2;
|
|
2602
2643
|
}
|
|
2603
|
-
const wrapped =
|
|
2644
|
+
const wrapped = json2;
|
|
2604
2645
|
const items = wrapped?.results ?? wrapped?.data;
|
|
2605
2646
|
return Array.isArray(items) ? items : void 0;
|
|
2606
2647
|
}
|
|
2607
2648
|
__name(itemsOf, "itemsOf");
|
|
2649
|
+
|
|
2650
|
+
// src/tree/types.ts
|
|
2651
|
+
function* walkTree(nodes, path = []) {
|
|
2652
|
+
for (const node of nodes) {
|
|
2653
|
+
yield {
|
|
2654
|
+
node,
|
|
2655
|
+
path
|
|
2656
|
+
};
|
|
2657
|
+
yield* walkTree(node.children, [
|
|
2658
|
+
...path,
|
|
2659
|
+
node
|
|
2660
|
+
]);
|
|
2661
|
+
}
|
|
2662
|
+
}
|
|
2663
|
+
__name(walkTree, "walkTree");
|
|
2664
|
+
function indexTree(tree) {
|
|
2665
|
+
const index = /* @__PURE__ */ new Map();
|
|
2666
|
+
for (const entry of walkTree(tree.nodes)) {
|
|
2667
|
+
index.set(entry.node.id, entry);
|
|
2668
|
+
}
|
|
2669
|
+
return index;
|
|
2670
|
+
}
|
|
2671
|
+
__name(indexTree, "indexTree");
|
|
2672
|
+
|
|
2673
|
+
// src/tree/headings.ts
|
|
2674
|
+
var MARKDOWN = /^(#{1,6})\s+(.+?)\s*#*\s*$/;
|
|
2675
|
+
var KEYWORD_LEVEL = {
|
|
2676
|
+
part: 1,
|
|
2677
|
+
chapter: 1,
|
|
2678
|
+
title: 1,
|
|
2679
|
+
appendix: 1,
|
|
2680
|
+
annex: 1,
|
|
2681
|
+
exhibit: 1,
|
|
2682
|
+
schedule: 1,
|
|
2683
|
+
subpart: 2,
|
|
2684
|
+
article: 2,
|
|
2685
|
+
item: 2,
|
|
2686
|
+
section: 3
|
|
2687
|
+
};
|
|
2688
|
+
var KEYWORD = /^(PART|Part|CHAPTER|Chapter|TITLE|APPENDIX|Appendix|ANNEX|Annex|EXHIBIT|Exhibit|SCHEDULE|Schedule|SUBPART|Subpart|ARTICLE|Article|ITEM|Item|SECTION|Section)\s+([0-9]{1,4}[A-Z]?(?:[.\-–][0-9A-Z]{1,4}){0,3}|[IVXLC]{1,6}|[A-Z])\b\.?(?:\s*[—–:\-.]\s*|\s+)?(.{0,110})$/;
|
|
2689
|
+
var NUMBERED = /^(\d{1,3}(?:\.\d{1,4}){1,4})([-–]\d{1,4})?\.?\s+([A-Z[][^\n]{1,100})$/;
|
|
2690
|
+
function clean(title) {
|
|
2691
|
+
return title.replace(/\s+/g, " ").replace(/[\s.:—–-]+$/, "").trim();
|
|
2692
|
+
}
|
|
2693
|
+
__name(clean, "clean");
|
|
2694
|
+
function looksLikeBody(rest) {
|
|
2695
|
+
return /(\d\)?|[%$])\s*$/.test(rest) || rest.split(" ").length > 16;
|
|
2696
|
+
}
|
|
2697
|
+
__name(looksLikeBody, "looksLikeBody");
|
|
2698
|
+
var DANGLING = new Set("a an the of to in for and or as by with on at from that which is are be than under".split(" "));
|
|
2699
|
+
function dangling(title) {
|
|
2700
|
+
if (/[,;(]$/.test(title)) {
|
|
2701
|
+
return true;
|
|
2702
|
+
}
|
|
2703
|
+
const last = title.split(/\s+/).pop()?.toLowerCase() ?? "";
|
|
2704
|
+
return DANGLING.has(last);
|
|
2705
|
+
}
|
|
2706
|
+
__name(dangling, "dangling");
|
|
2707
|
+
function detectHeadings(texts, markdownOnly = false) {
|
|
2708
|
+
const headings = [];
|
|
2709
|
+
let inFence = false;
|
|
2710
|
+
let run = [];
|
|
2711
|
+
const dropped = /* @__PURE__ */ new Set();
|
|
2712
|
+
const closeRun = /* @__PURE__ */ __name(() => {
|
|
2713
|
+
if (run.length >= 5) {
|
|
2714
|
+
run.forEach((entry, position) => {
|
|
2715
|
+
const heading = headings[entry.index];
|
|
2716
|
+
const keep = position === run.length - 1 || position === 0 && entry.keyword && heading !== void 0 && heading.level === 1;
|
|
2717
|
+
if (!keep) {
|
|
2718
|
+
dropped.add(entry.index);
|
|
2719
|
+
}
|
|
2720
|
+
});
|
|
2721
|
+
}
|
|
2722
|
+
run = [];
|
|
2723
|
+
}, "closeRun");
|
|
2724
|
+
texts.forEach((text, unit) => {
|
|
2725
|
+
let seenText = false;
|
|
2726
|
+
const lines = text.split("\n").map((raw) => raw.trim());
|
|
2727
|
+
for (let at = 0; at < lines.length; at++) {
|
|
2728
|
+
let line = lines[at] ?? "";
|
|
2729
|
+
if (line.startsWith("```")) {
|
|
2730
|
+
inFence = !inFence;
|
|
2731
|
+
seenText = true;
|
|
2732
|
+
closeRun();
|
|
2733
|
+
continue;
|
|
2734
|
+
}
|
|
2735
|
+
if (line === "") {
|
|
2736
|
+
continue;
|
|
2737
|
+
}
|
|
2738
|
+
if (inFence || line.length > 140) {
|
|
2739
|
+
seenText = true;
|
|
2740
|
+
closeRun();
|
|
2741
|
+
continue;
|
|
2742
|
+
}
|
|
2743
|
+
const atUnitStart = !seenText;
|
|
2744
|
+
seenText = true;
|
|
2745
|
+
const markdown = MARKDOWN.exec(line);
|
|
2746
|
+
if (markdown?.[1] !== void 0 && markdown[2] !== void 0) {
|
|
2747
|
+
closeRun();
|
|
2748
|
+
const title = clean(markdown[2].replace(/\*\*/g, ""));
|
|
2749
|
+
if (title !== "") {
|
|
2750
|
+
headings.push({
|
|
2751
|
+
title,
|
|
2752
|
+
level: markdown[1].length,
|
|
2753
|
+
unit,
|
|
2754
|
+
atUnitStart
|
|
2755
|
+
});
|
|
2756
|
+
}
|
|
2757
|
+
continue;
|
|
2758
|
+
}
|
|
2759
|
+
if (markdownOnly) {
|
|
2760
|
+
continue;
|
|
2761
|
+
}
|
|
2762
|
+
const candidate2 = KEYWORD.test(line) || NUMBERED.test(line);
|
|
2763
|
+
if (candidate2) {
|
|
2764
|
+
for (let joined = 0; joined < 2; joined++) {
|
|
2765
|
+
const next = lines[at + 1];
|
|
2766
|
+
if (!/[a-z]-$/.test(line) || next === void 0 || next === "" || next.length > 80) {
|
|
2767
|
+
break;
|
|
2768
|
+
}
|
|
2769
|
+
line = `${line.slice(0, -1)}${next}`;
|
|
2770
|
+
at += 1;
|
|
2771
|
+
}
|
|
2772
|
+
}
|
|
2773
|
+
const keyword = KEYWORD.exec(line);
|
|
2774
|
+
if (keyword?.[1] !== void 0 && keyword[2] !== void 0) {
|
|
2775
|
+
const rest = keyword[3] ?? "";
|
|
2776
|
+
const allCaps = keyword[1] === keyword[1].toUpperCase();
|
|
2777
|
+
if (!looksLikeBody(rest) && !/^[a-z,;)]/.test(rest) && (rest !== "" || allCaps) && !dangling(clean(line))) {
|
|
2778
|
+
const level = KEYWORD_LEVEL[keyword[1].toLowerCase()] ?? 2;
|
|
2779
|
+
run.push({
|
|
2780
|
+
index: headings.length,
|
|
2781
|
+
keyword: true
|
|
2782
|
+
});
|
|
2783
|
+
headings.push({
|
|
2784
|
+
title: clean(line),
|
|
2785
|
+
level,
|
|
2786
|
+
unit,
|
|
2787
|
+
atUnitStart
|
|
2788
|
+
});
|
|
2789
|
+
continue;
|
|
2790
|
+
}
|
|
2791
|
+
}
|
|
2792
|
+
const numbered = NUMBERED.exec(line);
|
|
2793
|
+
if (numbered?.[1] !== void 0 && numbered[3] !== void 0 && !looksLikeBody(numbered[3]) && !dangling(clean(numbered[3]))) {
|
|
2794
|
+
const level = Math.min(6, numbered[1].split(".").length + 1 + (numbered[2] ? 1 : 0));
|
|
2795
|
+
run.push({
|
|
2796
|
+
index: headings.length,
|
|
2797
|
+
keyword: false
|
|
2798
|
+
});
|
|
2799
|
+
headings.push({
|
|
2800
|
+
title: clean(line),
|
|
2801
|
+
level,
|
|
2802
|
+
unit,
|
|
2803
|
+
atUnitStart
|
|
2804
|
+
});
|
|
2805
|
+
continue;
|
|
2806
|
+
}
|
|
2807
|
+
if (!candidate2 && /[A-Za-z]{3,}/.test(line) && !isPageFurniture(line) && !isTitleCase(line)) {
|
|
2808
|
+
closeRun();
|
|
2809
|
+
}
|
|
2810
|
+
}
|
|
2811
|
+
});
|
|
2812
|
+
closeRun();
|
|
2813
|
+
return dropContentsDuplicates(dropRunningHeaders(headings.filter((_, index) => !dropped.has(index)), texts.length));
|
|
2814
|
+
}
|
|
2815
|
+
__name(detectHeadings, "detectHeadings");
|
|
2816
|
+
function dropContentsDuplicates(headings) {
|
|
2817
|
+
const last = /* @__PURE__ */ new Map();
|
|
2818
|
+
headings.forEach((heading, index) => {
|
|
2819
|
+
if (/\d/.test(heading.title)) {
|
|
2820
|
+
last.set(heading.title.toLowerCase(), index);
|
|
2821
|
+
}
|
|
2822
|
+
});
|
|
2823
|
+
return headings.filter((heading, index) => {
|
|
2824
|
+
const key = heading.title.toLowerCase();
|
|
2825
|
+
return !last.has(key) || last.get(key) === index;
|
|
2826
|
+
});
|
|
2827
|
+
}
|
|
2828
|
+
__name(dropContentsDuplicates, "dropContentsDuplicates");
|
|
2829
|
+
function isTitleCase(line) {
|
|
2830
|
+
if (line.length > 70 || /\.\s*$/.test(line)) {
|
|
2831
|
+
return false;
|
|
2832
|
+
}
|
|
2833
|
+
const words2 = line.split(/\s+/).filter((word) => /^[A-Za-z]{4,}/.test(word));
|
|
2834
|
+
return words2.length > 0 && words2.filter((word) => /^[A-Z]/.test(word)).length / words2.length >= 0.6;
|
|
2835
|
+
}
|
|
2836
|
+
__name(isTitleCase, "isTitleCase");
|
|
2837
|
+
function isPageFurniture(line) {
|
|
2838
|
+
return line.length < 50 && /^[\d\s\-–—.()]*(\d+\s*CFR|Federal Acquisition Regulation|Page \d+)/i.test(line);
|
|
2839
|
+
}
|
|
2840
|
+
__name(isPageFurniture, "isPageFurniture");
|
|
2841
|
+
function runningLines(texts) {
|
|
2842
|
+
const counts = /* @__PURE__ */ new Map();
|
|
2843
|
+
for (const text of texts) {
|
|
2844
|
+
const seen = /* @__PURE__ */ new Set();
|
|
2845
|
+
for (const raw of text.split("\n")) {
|
|
2846
|
+
const key = raw.trim().replace(/\d+/g, "#").toLowerCase();
|
|
2847
|
+
if (key && !seen.has(key)) {
|
|
2848
|
+
seen.add(key);
|
|
2849
|
+
counts.set(key, (counts.get(key) ?? 0) + 1);
|
|
2850
|
+
}
|
|
2851
|
+
}
|
|
2852
|
+
}
|
|
2853
|
+
const threshold = Math.max(3, Math.ceil(texts.length * 0.2));
|
|
2854
|
+
return new Set([
|
|
2855
|
+
...counts
|
|
2856
|
+
].filter(([, count]) => count >= threshold).map(([key]) => key));
|
|
2857
|
+
}
|
|
2858
|
+
__name(runningLines, "runningLines");
|
|
2859
|
+
function dropRunningHeaders(headings, unitCount) {
|
|
2860
|
+
const counts = /* @__PURE__ */ new Map();
|
|
2861
|
+
for (const heading of headings) {
|
|
2862
|
+
const key = heading.title.toLowerCase();
|
|
2863
|
+
counts.set(key, (counts.get(key) ?? 0) + 1);
|
|
2864
|
+
}
|
|
2865
|
+
const seen = /* @__PURE__ */ new Set();
|
|
2866
|
+
const threshold = Math.max(3, Math.ceil(unitCount * 0.2));
|
|
2867
|
+
return headings.filter((heading) => {
|
|
2868
|
+
const key = heading.title.toLowerCase();
|
|
2869
|
+
if ((counts.get(key) ?? 0) >= threshold) {
|
|
2870
|
+
if (seen.has(key)) {
|
|
2871
|
+
return false;
|
|
2872
|
+
}
|
|
2873
|
+
}
|
|
2874
|
+
seen.add(key);
|
|
2875
|
+
return true;
|
|
2876
|
+
});
|
|
2877
|
+
}
|
|
2878
|
+
__name(dropRunningHeaders, "dropRunningHeaders");
|
|
2879
|
+
|
|
2880
|
+
// src/tree/llm.ts
|
|
2881
|
+
var TreeBudgetExceededError = class extends Error {
|
|
2882
|
+
static {
|
|
2883
|
+
__name(this, "TreeBudgetExceededError");
|
|
2884
|
+
}
|
|
2885
|
+
limit;
|
|
2886
|
+
constructor(limit) {
|
|
2887
|
+
super(`tree LLM budget exhausted (${limit})`), this.limit = limit;
|
|
2888
|
+
this.name = "TreeBudgetExceededError";
|
|
2889
|
+
}
|
|
2890
|
+
};
|
|
2891
|
+
function estimateTokens(text) {
|
|
2892
|
+
return Math.ceil(text.length / 4);
|
|
2893
|
+
}
|
|
2894
|
+
__name(estimateTokens, "estimateTokens");
|
|
2895
|
+
var BudgetedLlm = class {
|
|
2896
|
+
static {
|
|
2897
|
+
__name(this, "BudgetedLlm");
|
|
2898
|
+
}
|
|
2899
|
+
llm;
|
|
2900
|
+
budget;
|
|
2901
|
+
outerSignal;
|
|
2902
|
+
calls = 0;
|
|
2903
|
+
inputTokens = 0;
|
|
2904
|
+
outputTokens = 0;
|
|
2905
|
+
exhausted = false;
|
|
2906
|
+
deadline;
|
|
2907
|
+
constructor(llm, budget, outerSignal) {
|
|
2908
|
+
this.llm = llm;
|
|
2909
|
+
this.budget = budget;
|
|
2910
|
+
this.outerSignal = outerSignal;
|
|
2911
|
+
this.deadline = budget.timeoutMs !== void 0 ? Date.now() + budget.timeoutMs : Number.POSITIVE_INFINITY;
|
|
2912
|
+
}
|
|
2913
|
+
/** Time left before the deadline, ms (`Infinity` without one). */
|
|
2914
|
+
remainingMs() {
|
|
2915
|
+
return this.deadline - Date.now();
|
|
2916
|
+
}
|
|
2917
|
+
/** Would a call with this prompt fit? Does not reserve anything; a `false` marks the budget exhausted. */
|
|
2918
|
+
fits(prompt, maxOutputTokens) {
|
|
2919
|
+
const fits = this.limitHit(prompt, maxOutputTokens) === void 0;
|
|
2920
|
+
if (!fits) {
|
|
2921
|
+
this.exhausted = true;
|
|
2922
|
+
}
|
|
2923
|
+
return fits;
|
|
2924
|
+
}
|
|
2925
|
+
limitHit(prompt, maxOutputTokens) {
|
|
2926
|
+
if (this.calls >= this.budget.maxCalls) {
|
|
2927
|
+
return "calls";
|
|
2928
|
+
}
|
|
2929
|
+
if (this.inputTokens + estimateTokens(prompt) > this.budget.maxInputTokens) {
|
|
2930
|
+
return "inputTokens";
|
|
2931
|
+
}
|
|
2932
|
+
if (this.outputTokens + Math.min(maxOutputTokens, 16) > this.budget.maxOutputTokens) {
|
|
2933
|
+
return "outputTokens";
|
|
2934
|
+
}
|
|
2935
|
+
if (this.remainingMs() <= 0 || this.outerSignal?.aborted) {
|
|
2936
|
+
return "timeout";
|
|
2937
|
+
}
|
|
2938
|
+
return void 0;
|
|
2939
|
+
}
|
|
2940
|
+
/**
|
|
2941
|
+
* Make one call, or throw {@link TreeBudgetExceededError} without calling. Provider errors
|
|
2942
|
+
* propagate (after being counted). Never retries.
|
|
2943
|
+
*/
|
|
2944
|
+
async call(task, system, prompt, maxOutputTokens) {
|
|
2945
|
+
const limit = this.limitHit(system + prompt, maxOutputTokens);
|
|
2946
|
+
if (limit !== void 0) {
|
|
2947
|
+
this.exhausted = true;
|
|
2948
|
+
throw new TreeBudgetExceededError(limit);
|
|
2949
|
+
}
|
|
2950
|
+
const outputCap = Math.max(16, Math.min(maxOutputTokens, this.budget.maxOutputTokens - this.outputTokens));
|
|
2951
|
+
const timeoutMs = Math.min(this.budget.callTimeoutMs ?? Number.POSITIVE_INFINITY, this.remainingMs());
|
|
2952
|
+
const reservedInput = estimateTokens(system + prompt);
|
|
2953
|
+
this.inputTokens += reservedInput;
|
|
2954
|
+
this.outputTokens += outputCap;
|
|
2955
|
+
const controller = new AbortController();
|
|
2956
|
+
const onOuterAbort = /* @__PURE__ */ __name(() => controller.abort(this.outerSignal?.reason), "onOuterAbort");
|
|
2957
|
+
this.outerSignal?.addEventListener("abort", onOuterAbort, {
|
|
2958
|
+
once: true
|
|
2959
|
+
});
|
|
2960
|
+
let timer;
|
|
2961
|
+
this.calls += 1;
|
|
2962
|
+
try {
|
|
2963
|
+
const response = await Promise.race([
|
|
2964
|
+
this.llm({
|
|
2965
|
+
task,
|
|
2966
|
+
system,
|
|
2967
|
+
prompt,
|
|
2968
|
+
maxOutputTokens: outputCap,
|
|
2969
|
+
signal: controller.signal
|
|
2970
|
+
}),
|
|
2971
|
+
new Promise((_, reject) => {
|
|
2972
|
+
if (Number.isFinite(timeoutMs)) {
|
|
2973
|
+
timer = setTimeout(() => {
|
|
2974
|
+
controller.abort(new TreeBudgetExceededError("timeout"));
|
|
2975
|
+
reject(new TreeBudgetExceededError("timeout"));
|
|
2976
|
+
}, Math.max(0, timeoutMs));
|
|
2977
|
+
}
|
|
2978
|
+
})
|
|
2979
|
+
]);
|
|
2980
|
+
this.inputTokens += (response.usage?.inputTokens ?? reservedInput) - reservedInput;
|
|
2981
|
+
this.outputTokens += (response.usage?.outputTokens ?? estimateTokens(response.text)) - outputCap;
|
|
2982
|
+
return response.text;
|
|
2983
|
+
} catch (error) {
|
|
2984
|
+
this.outputTokens -= outputCap;
|
|
2985
|
+
if (error instanceof TreeBudgetExceededError) {
|
|
2986
|
+
this.exhausted = true;
|
|
2987
|
+
}
|
|
2988
|
+
throw error;
|
|
2989
|
+
} finally {
|
|
2990
|
+
if (timer !== void 0) {
|
|
2991
|
+
clearTimeout(timer);
|
|
2992
|
+
}
|
|
2993
|
+
this.outerSignal?.removeEventListener("abort", onOuterAbort);
|
|
2994
|
+
}
|
|
2995
|
+
}
|
|
2996
|
+
};
|
|
2997
|
+
function parseJsonReply(text) {
|
|
2998
|
+
const tryParse = /* @__PURE__ */ __name((candidate2) => {
|
|
2999
|
+
try {
|
|
3000
|
+
const value = JSON.parse(candidate2);
|
|
3001
|
+
return value !== null && typeof value === "object" && !Array.isArray(value) ? value : void 0;
|
|
3002
|
+
} catch {
|
|
3003
|
+
return void 0;
|
|
3004
|
+
}
|
|
3005
|
+
}, "tryParse");
|
|
3006
|
+
const trimmed = text.trim();
|
|
3007
|
+
const direct = tryParse(trimmed);
|
|
3008
|
+
if (direct !== void 0) {
|
|
3009
|
+
return direct;
|
|
3010
|
+
}
|
|
3011
|
+
const fenced = /```(?:json)?\s*([\s\S]*?)```/i.exec(trimmed);
|
|
3012
|
+
if (fenced?.[1] !== void 0) {
|
|
3013
|
+
const inner = tryParse(fenced[1].trim());
|
|
3014
|
+
if (inner !== void 0) {
|
|
3015
|
+
return inner;
|
|
3016
|
+
}
|
|
3017
|
+
}
|
|
3018
|
+
const start = trimmed.indexOf("{");
|
|
3019
|
+
const end = trimmed.lastIndexOf("}");
|
|
3020
|
+
if (start >= 0 && end > start) {
|
|
3021
|
+
return tryParse(trimmed.slice(start, end + 1));
|
|
3022
|
+
}
|
|
3023
|
+
return void 0;
|
|
3024
|
+
}
|
|
3025
|
+
__name(parseJsonReply, "parseJsonReply");
|
|
3026
|
+
function stringList(value) {
|
|
3027
|
+
if (!Array.isArray(value)) {
|
|
3028
|
+
return [];
|
|
3029
|
+
}
|
|
3030
|
+
return value.map((item) => typeof item === "number" ? String(item) : item).filter((item) => typeof item === "string");
|
|
3031
|
+
}
|
|
3032
|
+
__name(stringList, "stringList");
|
|
3033
|
+
var NOOP_SINK = {
|
|
3034
|
+
write() {
|
|
3035
|
+
},
|
|
3036
|
+
end() {
|
|
3037
|
+
},
|
|
3038
|
+
fail() {
|
|
3039
|
+
}
|
|
3040
|
+
};
|
|
3041
|
+
function treeLlmFromModelProvider(provider) {
|
|
3042
|
+
return async (request) => {
|
|
3043
|
+
const result = await provider.runTurn({
|
|
3044
|
+
system: request.system,
|
|
3045
|
+
messages: [
|
|
3046
|
+
{
|
|
3047
|
+
role: "user",
|
|
3048
|
+
content: request.prompt
|
|
3049
|
+
}
|
|
3050
|
+
],
|
|
3051
|
+
tools: [],
|
|
3052
|
+
sink: NOOP_SINK,
|
|
3053
|
+
abortSignal: request.signal
|
|
3054
|
+
});
|
|
3055
|
+
return {
|
|
3056
|
+
text: result.text,
|
|
3057
|
+
usage: {
|
|
3058
|
+
inputTokens: result.usage.inputTokens,
|
|
3059
|
+
outputTokens: result.usage.outputTokens
|
|
3060
|
+
}
|
|
3061
|
+
};
|
|
3062
|
+
};
|
|
3063
|
+
}
|
|
3064
|
+
__name(treeLlmFromModelProvider, "treeLlmFromModelProvider");
|
|
3065
|
+
function openAiChatTreeLlm(options) {
|
|
3066
|
+
const url = `${(options.baseUrl ?? "https://api.openai.com/v1").replace(/\/+$/, "")}/chat/completions`;
|
|
3067
|
+
const doFetch = options.fetch ?? fetch;
|
|
3068
|
+
return async (request) => {
|
|
3069
|
+
const json2 = request.task.kind !== "summarize";
|
|
3070
|
+
const response = await doFetch(url, {
|
|
3071
|
+
method: "POST",
|
|
3072
|
+
headers: {
|
|
3073
|
+
"content-type": "application/json",
|
|
3074
|
+
...options.apiKey ? {
|
|
3075
|
+
authorization: `Bearer ${options.apiKey}`
|
|
3076
|
+
} : {},
|
|
3077
|
+
...options.headers
|
|
3078
|
+
},
|
|
3079
|
+
body: JSON.stringify({
|
|
3080
|
+
model: options.model,
|
|
3081
|
+
messages: [
|
|
3082
|
+
{
|
|
3083
|
+
role: "system",
|
|
3084
|
+
content: request.system
|
|
3085
|
+
},
|
|
3086
|
+
{
|
|
3087
|
+
role: "user",
|
|
3088
|
+
content: request.prompt
|
|
3089
|
+
}
|
|
3090
|
+
],
|
|
3091
|
+
temperature: options.temperature ?? 0,
|
|
3092
|
+
max_tokens: request.maxOutputTokens,
|
|
3093
|
+
...options.jsonMode && json2 ? {
|
|
3094
|
+
response_format: {
|
|
3095
|
+
type: "json_object"
|
|
3096
|
+
}
|
|
3097
|
+
} : {},
|
|
3098
|
+
...options.body
|
|
3099
|
+
}),
|
|
3100
|
+
signal: request.signal
|
|
3101
|
+
});
|
|
3102
|
+
const body = await response.text();
|
|
3103
|
+
if (!response.ok) {
|
|
3104
|
+
throw new HttpModelError(response.status, `Chat completion failed (${response.status}): ${errorMessageOf(body)}`);
|
|
3105
|
+
}
|
|
3106
|
+
const parsed = parseJsonReply(body);
|
|
3107
|
+
const content = parsed?.choices?.[0]?.message?.content;
|
|
3108
|
+
if (typeof content !== "string") {
|
|
3109
|
+
throw new HttpModelError(502, "Chat completion returned no message content");
|
|
3110
|
+
}
|
|
3111
|
+
const usage = {};
|
|
3112
|
+
if (typeof parsed?.usage?.prompt_tokens === "number") {
|
|
3113
|
+
usage.inputTokens = parsed.usage.prompt_tokens;
|
|
3114
|
+
}
|
|
3115
|
+
if (typeof parsed?.usage?.completion_tokens === "number") {
|
|
3116
|
+
usage.outputTokens = parsed.usage.completion_tokens;
|
|
3117
|
+
}
|
|
3118
|
+
return {
|
|
3119
|
+
text: content,
|
|
3120
|
+
usage
|
|
3121
|
+
};
|
|
3122
|
+
};
|
|
3123
|
+
}
|
|
3124
|
+
__name(openAiChatTreeLlm, "openAiChatTreeLlm");
|
|
3125
|
+
function cachedTreeLlm(llm, cache = /* @__PURE__ */ new Map()) {
|
|
3126
|
+
return async (request) => {
|
|
3127
|
+
const key = stableHash(`${request.task.kind}\0${request.maxOutputTokens}\0${request.system}\0${request.prompt}`);
|
|
3128
|
+
const hit = await cache.get(key);
|
|
3129
|
+
if (hit !== void 0) {
|
|
3130
|
+
return {
|
|
3131
|
+
text: hit.text,
|
|
3132
|
+
usage: {
|
|
3133
|
+
inputTokens: 0,
|
|
3134
|
+
outputTokens: 0
|
|
3135
|
+
}
|
|
3136
|
+
};
|
|
3137
|
+
}
|
|
3138
|
+
const response = await llm(request);
|
|
3139
|
+
await cache.set(key, response);
|
|
3140
|
+
return response;
|
|
3141
|
+
};
|
|
3142
|
+
}
|
|
3143
|
+
__name(cachedTreeLlm, "cachedTreeLlm");
|
|
3144
|
+
function stableHash(text) {
|
|
3145
|
+
let h1 = 2166136261;
|
|
3146
|
+
let h2 = 16777619 ^ text.length;
|
|
3147
|
+
for (let i = 0; i < text.length; i++) {
|
|
3148
|
+
const c = text.charCodeAt(i);
|
|
3149
|
+
h1 = Math.imul(h1 ^ c, 16777619);
|
|
3150
|
+
h2 = Math.imul(h2 ^ c, 1540483477);
|
|
3151
|
+
h2 ^= h2 >>> 13;
|
|
3152
|
+
}
|
|
3153
|
+
h1 = Math.imul(h1 ^ h1 >>> 16, 2246822507);
|
|
3154
|
+
h2 = Math.imul(h2 ^ h2 >>> 15, 3266489909);
|
|
3155
|
+
return (h1 >>> 0).toString(16).padStart(8, "0") + (h2 >>> 0).toString(16).padStart(8, "0");
|
|
3156
|
+
}
|
|
3157
|
+
__name(stableHash, "stableHash");
|
|
3158
|
+
|
|
3159
|
+
// src/tree/build.ts
|
|
3160
|
+
var DEFAULT_BUDGET = {
|
|
3161
|
+
maxCalls: 60,
|
|
3162
|
+
maxInputTokens: 4e5,
|
|
3163
|
+
maxOutputTokens: 4e4,
|
|
3164
|
+
timeoutMs: 3e5,
|
|
3165
|
+
callTimeoutMs: 6e4
|
|
3166
|
+
};
|
|
3167
|
+
var STRUCTURE_SYSTEM = "You extract the table of contents of a document. You answer with JSON only, no prose.";
|
|
3168
|
+
var SUMMARY_SYSTEM = "You summarize sections of a document so a reader can decide whether a section answers their question. Be specific: name the topics, clauses, figures and defined terms it covers.";
|
|
3169
|
+
function toUnits(input) {
|
|
3170
|
+
if (input.pages !== void 0 && input.sections !== void 0) {
|
|
3171
|
+
throw new Error("buildDocumentTree: give `pages` or `sections`, not both");
|
|
3172
|
+
}
|
|
3173
|
+
if (input.sections !== void 0) {
|
|
3174
|
+
const units = input.sections.map((section, index) => ({
|
|
3175
|
+
index,
|
|
3176
|
+
text: section.text,
|
|
3177
|
+
...section.page !== void 0 ? {
|
|
3178
|
+
page: section.page
|
|
3179
|
+
} : {}
|
|
3180
|
+
}));
|
|
3181
|
+
const titled = input.sections.some((section) => section.title?.trim());
|
|
3182
|
+
const declared = titled ? input.sections.flatMap((section, unit) => section.title?.trim() ? [
|
|
3183
|
+
{
|
|
3184
|
+
title: section.title.trim(),
|
|
3185
|
+
level: Math.max(1, Math.min(6, section.level ?? 1)),
|
|
3186
|
+
unit,
|
|
3187
|
+
atUnitStart: true
|
|
3188
|
+
}
|
|
3189
|
+
] : []) : void 0;
|
|
3190
|
+
return {
|
|
3191
|
+
units,
|
|
3192
|
+
declared,
|
|
3193
|
+
kind: "sections"
|
|
3194
|
+
};
|
|
3195
|
+
}
|
|
3196
|
+
const pages2 = input.pages ?? [];
|
|
3197
|
+
return {
|
|
3198
|
+
units: pages2.map((text, index) => ({
|
|
3199
|
+
index,
|
|
3200
|
+
text,
|
|
3201
|
+
page: index + 1
|
|
3202
|
+
})),
|
|
3203
|
+
declared: void 0,
|
|
3204
|
+
kind: "pages"
|
|
3205
|
+
};
|
|
3206
|
+
}
|
|
3207
|
+
__name(toUnits, "toUnits");
|
|
3208
|
+
function outlineHeadings(input, units) {
|
|
3209
|
+
const headings = [];
|
|
3210
|
+
for (const entry of input.outline ?? []) {
|
|
3211
|
+
let unit = entry.unit;
|
|
3212
|
+
if (unit === void 0 && entry.page !== void 0) {
|
|
3213
|
+
unit = units.findIndex((candidate2) => candidate2.page === entry.page);
|
|
3214
|
+
if (unit < 0) {
|
|
3215
|
+
unit = entry.page - 1;
|
|
3216
|
+
}
|
|
3217
|
+
}
|
|
3218
|
+
if (unit === void 0 || unit < 0 || unit >= units.length || !entry.title.trim()) {
|
|
3219
|
+
continue;
|
|
3220
|
+
}
|
|
3221
|
+
headings.push({
|
|
3222
|
+
title: entry.title.trim(),
|
|
3223
|
+
level: Math.max(1, Math.min(6, Math.floor(entry.level))),
|
|
3224
|
+
unit,
|
|
3225
|
+
atUnitStart: false
|
|
3226
|
+
});
|
|
3227
|
+
}
|
|
3228
|
+
return headings.sort((a, b) => a.unit - b.unit);
|
|
3229
|
+
}
|
|
3230
|
+
__name(outlineHeadings, "outlineHeadings");
|
|
3231
|
+
function nestHeadings(headings, unitCount) {
|
|
3232
|
+
const roots = [];
|
|
3233
|
+
const stack = [];
|
|
3234
|
+
headings.forEach((heading, index) => {
|
|
3235
|
+
let end = unitCount - 1;
|
|
3236
|
+
for (let next = index + 1; next < headings.length; next++) {
|
|
3237
|
+
const candidate2 = headings[next];
|
|
3238
|
+
if (candidate2 !== void 0 && candidate2.level <= heading.level) {
|
|
3239
|
+
end = candidate2.atUnitStart ? candidate2.unit - 1 : candidate2.unit;
|
|
3240
|
+
break;
|
|
3241
|
+
}
|
|
3242
|
+
}
|
|
3243
|
+
const node = {
|
|
3244
|
+
title: heading.title,
|
|
3245
|
+
unitStart: heading.unit,
|
|
3246
|
+
unitEnd: Math.max(heading.unit, end),
|
|
3247
|
+
children: []
|
|
3248
|
+
};
|
|
3249
|
+
while (stack.length > 0 && (stack[stack.length - 1]?.level ?? 0) >= heading.level) {
|
|
3250
|
+
stack.pop();
|
|
3251
|
+
}
|
|
3252
|
+
const parent = stack[stack.length - 1];
|
|
3253
|
+
if (parent === void 0) {
|
|
3254
|
+
roots.push(node);
|
|
3255
|
+
} else {
|
|
3256
|
+
parent.node.children.push(node);
|
|
3257
|
+
parent.node.unitEnd = Math.max(parent.node.unitEnd, node.unitEnd);
|
|
3258
|
+
}
|
|
3259
|
+
stack.push({
|
|
3260
|
+
node,
|
|
3261
|
+
level: heading.level
|
|
3262
|
+
});
|
|
3263
|
+
});
|
|
3264
|
+
const first = roots[0];
|
|
3265
|
+
if (first !== void 0 && first.unitStart > 0) {
|
|
3266
|
+
roots.unshift({
|
|
3267
|
+
title: "Front matter",
|
|
3268
|
+
unitStart: 0,
|
|
3269
|
+
unitEnd: first.unitStart - 1,
|
|
3270
|
+
children: []
|
|
3271
|
+
});
|
|
3272
|
+
}
|
|
3273
|
+
return roots;
|
|
3274
|
+
}
|
|
3275
|
+
__name(nestHeadings, "nestHeadings");
|
|
3276
|
+
function firstLine(text, running) {
|
|
3277
|
+
const line = text.split("\n").map((candidate2) => candidate2.trim()).find((candidate2) => /[A-Za-z]{3,}/.test(candidate2) && !running.has(candidate2.replace(/\d+/g, "#").toLowerCase()));
|
|
3278
|
+
return line === void 0 ? "" : line.length > 80 ? `${line.slice(0, 77)}\u2026` : line;
|
|
3279
|
+
}
|
|
3280
|
+
__name(firstLine, "firstLine");
|
|
3281
|
+
function rangeLabel(units, start, end) {
|
|
3282
|
+
const first = units[start];
|
|
3283
|
+
const last = units[end];
|
|
3284
|
+
if (first?.page !== void 0 && last?.page !== void 0) {
|
|
3285
|
+
return first.page === last.page ? `Page ${first.page}` : `Pages ${first.page}\u2013${last.page}`;
|
|
3286
|
+
}
|
|
3287
|
+
return start === end ? `Section ${start + 1}` : `Sections ${start + 1}\u2013${end + 1}`;
|
|
3288
|
+
}
|
|
3289
|
+
__name(rangeLabel, "rangeLabel");
|
|
3290
|
+
function groupNodes(units, start, end, groupSize, maxChildren, levelsLeft, running) {
|
|
3291
|
+
let level = [];
|
|
3292
|
+
for (let from = start; from <= end; from += groupSize) {
|
|
3293
|
+
const to = Math.min(end, from + groupSize - 1);
|
|
3294
|
+
const lead = firstLine(units[from]?.text ?? "", running);
|
|
3295
|
+
level.push({
|
|
3296
|
+
title: `${rangeLabel(units, from, to)}${lead ? `: ${lead}` : ""}`,
|
|
3297
|
+
unitStart: from,
|
|
3298
|
+
unitEnd: to,
|
|
3299
|
+
children: []
|
|
3300
|
+
});
|
|
3301
|
+
}
|
|
3302
|
+
let depth = 1;
|
|
3303
|
+
while (level.length > maxChildren && depth < levelsLeft) {
|
|
3304
|
+
const parents = [];
|
|
3305
|
+
const per = Math.max(maxChildren, Math.ceil(level.length / maxChildren));
|
|
3306
|
+
for (let index = 0; index < level.length; index += per) {
|
|
3307
|
+
const children = level.slice(index, index + per);
|
|
3308
|
+
const from = children[0]?.unitStart ?? start;
|
|
3309
|
+
const to = children[children.length - 1]?.unitEnd ?? end;
|
|
3310
|
+
parents.push({
|
|
3311
|
+
title: rangeLabel(units, from, to),
|
|
3312
|
+
unitStart: from,
|
|
3313
|
+
unitEnd: to,
|
|
3314
|
+
children
|
|
3315
|
+
});
|
|
3316
|
+
}
|
|
3317
|
+
level = parents;
|
|
3318
|
+
depth += 1;
|
|
3319
|
+
}
|
|
3320
|
+
return level;
|
|
3321
|
+
}
|
|
3322
|
+
__name(groupNodes, "groupNodes");
|
|
3323
|
+
function descendantTitles(nodes, into = []) {
|
|
3324
|
+
for (const { node } of walkDraft(nodes)) {
|
|
3325
|
+
if (into.length < 30) {
|
|
3326
|
+
into.push(node.title);
|
|
3327
|
+
}
|
|
3328
|
+
}
|
|
3329
|
+
return into;
|
|
3330
|
+
}
|
|
3331
|
+
__name(descendantTitles, "descendantTitles");
|
|
3332
|
+
function shape(nodes, units, depth, options) {
|
|
3333
|
+
const shaped = nodes.map((node) => {
|
|
3334
|
+
if (depth >= options.maxDepth || node.unitStart === node.unitEnd && node.children.length > 0) {
|
|
3335
|
+
const keyItems = descendantTitles(node.children, [
|
|
3336
|
+
...node.keyItems ?? []
|
|
3337
|
+
]);
|
|
3338
|
+
return {
|
|
3339
|
+
...node,
|
|
3340
|
+
...keyItems.length > 0 ? {
|
|
3341
|
+
keyItems
|
|
3342
|
+
} : {},
|
|
3343
|
+
children: []
|
|
3344
|
+
};
|
|
3345
|
+
}
|
|
3346
|
+
let children = node.children;
|
|
3347
|
+
const only = children[0];
|
|
3348
|
+
if (children.length === 1 && only !== void 0 && only.unitStart === node.unitStart && only.unitEnd === node.unitEnd) {
|
|
3349
|
+
children = only.children;
|
|
3350
|
+
}
|
|
3351
|
+
if (children.length === 0 && node.unitEnd - node.unitStart + 1 > options.maxLeafUnits) {
|
|
3352
|
+
children = groupNodes(units, node.unitStart, node.unitEnd, options.groupSize, options.maxChildren, options.maxDepth - depth, options.running);
|
|
3353
|
+
}
|
|
3354
|
+
return {
|
|
3355
|
+
...node,
|
|
3356
|
+
children: shape(children, units, depth + 1, options)
|
|
3357
|
+
};
|
|
3358
|
+
});
|
|
3359
|
+
const merged = [];
|
|
3360
|
+
for (const node of shaped) {
|
|
3361
|
+
const previous = merged[merged.length - 1];
|
|
3362
|
+
if (previous !== void 0 && previous.children.length === 0 && node.children.length === 0 && previous.unitStart === node.unitStart && previous.unitEnd === node.unitEnd) {
|
|
3363
|
+
previous.keyItems = [
|
|
3364
|
+
...previous.keyItems ?? [],
|
|
3365
|
+
node.title,
|
|
3366
|
+
...node.keyItems ?? []
|
|
3367
|
+
].slice(0, 30);
|
|
3368
|
+
continue;
|
|
3369
|
+
}
|
|
3370
|
+
merged.push(node);
|
|
3371
|
+
}
|
|
3372
|
+
return merged;
|
|
3373
|
+
}
|
|
3374
|
+
__name(shape, "shape");
|
|
3375
|
+
function* walkDraft(nodes) {
|
|
3376
|
+
for (const node of nodes) {
|
|
3377
|
+
yield {
|
|
3378
|
+
node
|
|
3379
|
+
};
|
|
3380
|
+
yield* walkDraft(node.children);
|
|
3381
|
+
}
|
|
3382
|
+
}
|
|
3383
|
+
__name(walkDraft, "walkDraft");
|
|
3384
|
+
function spanText(units, start, end) {
|
|
3385
|
+
const parts = [];
|
|
3386
|
+
for (let index = Math.max(0, start); index <= end && index < units.length; index++) {
|
|
3387
|
+
parts.push(units[index]?.text ?? "");
|
|
3388
|
+
}
|
|
3389
|
+
return parts.join("\n\n");
|
|
3390
|
+
}
|
|
3391
|
+
__name(spanText, "spanText");
|
|
3392
|
+
function ownTexts(units, node) {
|
|
3393
|
+
const first = node.children[0];
|
|
3394
|
+
const end = first === void 0 ? node.unitEnd : Math.max(node.unitStart, first.unitStart - 1);
|
|
3395
|
+
const texts = [];
|
|
3396
|
+
for (let index = node.unitStart; index <= end && index < units.length; index++) {
|
|
3397
|
+
texts.push(units[index]?.text ?? "");
|
|
3398
|
+
}
|
|
3399
|
+
const head = texts[0];
|
|
3400
|
+
if (head !== void 0) {
|
|
3401
|
+
const probe = node.title.slice(0, 30);
|
|
3402
|
+
const at = probe.length >= 4 ? head.indexOf(probe) : -1;
|
|
3403
|
+
if (at > 0) {
|
|
3404
|
+
texts[0] = head.slice(at);
|
|
3405
|
+
}
|
|
3406
|
+
}
|
|
3407
|
+
return texts;
|
|
3408
|
+
}
|
|
3409
|
+
__name(ownTexts, "ownTexts");
|
|
3410
|
+
function excerptOf(texts, max) {
|
|
3411
|
+
const total = texts.reduce((sum, text) => sum + text.length, 0);
|
|
3412
|
+
if (total <= max) {
|
|
3413
|
+
return texts.join("\n\n");
|
|
3414
|
+
}
|
|
3415
|
+
const per = Math.max(400, Math.floor(max / texts.length));
|
|
3416
|
+
const parts = [];
|
|
3417
|
+
let used = 0;
|
|
3418
|
+
for (const text of texts) {
|
|
3419
|
+
if (used >= max) {
|
|
3420
|
+
break;
|
|
3421
|
+
}
|
|
3422
|
+
const part = text.length > per ? `${text.slice(0, per)}\u2026` : text;
|
|
3423
|
+
parts.push(part);
|
|
3424
|
+
used += part.length;
|
|
3425
|
+
}
|
|
3426
|
+
return parts.join("\n\n");
|
|
3427
|
+
}
|
|
3428
|
+
__name(excerptOf, "excerptOf");
|
|
3429
|
+
function collapse(text, max) {
|
|
3430
|
+
const flat = text.replace(/\s+/g, " ").trim();
|
|
3431
|
+
return flat.length > max ? `${flat.slice(0, max - 1)}\u2026` : flat;
|
|
3432
|
+
}
|
|
3433
|
+
__name(collapse, "collapse");
|
|
3434
|
+
async function llmHeadings(llm, units, options) {
|
|
3435
|
+
const perUnit = Math.max(500, Math.floor(options.windowChars / 4));
|
|
3436
|
+
const windows = [];
|
|
3437
|
+
let current = [];
|
|
3438
|
+
let size = 0;
|
|
3439
|
+
for (const unit of units) {
|
|
3440
|
+
const length = Math.min(unit.text.length, perUnit);
|
|
3441
|
+
if (current.length > 0 && size + length > options.windowChars) {
|
|
3442
|
+
windows.push(current);
|
|
3443
|
+
current = [];
|
|
3444
|
+
size = 0;
|
|
3445
|
+
}
|
|
3446
|
+
current.push(unit);
|
|
3447
|
+
size += length;
|
|
3448
|
+
}
|
|
3449
|
+
if (current.length > 0) {
|
|
3450
|
+
windows.push(current);
|
|
3451
|
+
}
|
|
3452
|
+
if (windows.length > options.maxCalls) {
|
|
3453
|
+
throw new Error(`document needs ${windows.length} structuring calls, more than maxStructureCalls (${options.maxCalls})`);
|
|
3454
|
+
}
|
|
3455
|
+
const headings = [];
|
|
3456
|
+
for (const window of windows) {
|
|
3457
|
+
const shown = window.map((unit) => ({
|
|
3458
|
+
unit: unit.index,
|
|
3459
|
+
...unit.page !== void 0 ? {
|
|
3460
|
+
page: unit.page
|
|
3461
|
+
} : {},
|
|
3462
|
+
text: unit.text.length > perUnit ? `${unit.text.slice(0, perUnit)}\u2026` : unit.text
|
|
3463
|
+
}));
|
|
3464
|
+
const body = shown.map((unit) => `<unit ${unit.unit}>
|
|
3465
|
+
${unit.text}
|
|
3466
|
+
</unit ${unit.unit}>`).join("\n");
|
|
3467
|
+
const prompt = `Below is part of a document, split into numbered units (pages or sections).
|
|
3468
|
+
List the headings of the sections that START in these units, in reading order, with their nesting level (1 = top-level chapter/part, 2 = its subsections, ...) and the unit number they start in.
|
|
3469
|
+
Only list real headings (titles of parts, chapters, sections, clauses, items), not running page headers, table rows or sentences. If no section starts here, return an empty list.
|
|
3470
|
+
|
|
3471
|
+
${body}
|
|
3472
|
+
|
|
3473
|
+
Reply with JSON only: {"sections": [{"title": "...", "level": 1, "unit": 0}]}`;
|
|
3474
|
+
const reply = parseJsonReply(await llm.call({
|
|
3475
|
+
kind: "structure",
|
|
3476
|
+
units: shown
|
|
3477
|
+
}, STRUCTURE_SYSTEM, prompt, 2e3));
|
|
3478
|
+
const sections = reply?.sections;
|
|
3479
|
+
if (!Array.isArray(sections)) {
|
|
3480
|
+
throw new Error("structuring reply had no `sections` list");
|
|
3481
|
+
}
|
|
3482
|
+
const first = window[0]?.index ?? 0;
|
|
3483
|
+
const last = window[window.length - 1]?.index ?? first;
|
|
3484
|
+
for (const entry of sections) {
|
|
3485
|
+
if (entry === null || typeof entry !== "object") {
|
|
3486
|
+
continue;
|
|
3487
|
+
}
|
|
3488
|
+
const { title, level, unit } = entry;
|
|
3489
|
+
const at = typeof unit === "number" ? unit : Number(unit);
|
|
3490
|
+
if (typeof title !== "string" || !title.trim() || !Number.isInteger(at) || at < first || at > last) {
|
|
3491
|
+
continue;
|
|
3492
|
+
}
|
|
3493
|
+
const previous = headings[headings.length - 1];
|
|
3494
|
+
if (previous !== void 0 && at < previous.unit) {
|
|
3495
|
+
continue;
|
|
3496
|
+
}
|
|
3497
|
+
headings.push({
|
|
3498
|
+
title: collapse(title, 160),
|
|
3499
|
+
level: Math.max(1, Math.min(6, Math.floor(typeof level === "number" ? level : 1))),
|
|
3500
|
+
unit: at,
|
|
3501
|
+
atUnitStart: false
|
|
3502
|
+
});
|
|
3503
|
+
}
|
|
3504
|
+
}
|
|
3505
|
+
return headings;
|
|
3506
|
+
}
|
|
3507
|
+
__name(llmHeadings, "llmHeadings");
|
|
3508
|
+
async function mapLimit(items, limit, work) {
|
|
3509
|
+
const results = new Array(items.length);
|
|
3510
|
+
let next = 0;
|
|
3511
|
+
const workers = Array.from({
|
|
3512
|
+
length: Math.max(1, Math.min(limit, items.length))
|
|
3513
|
+
}, async () => {
|
|
3514
|
+
while (next < items.length) {
|
|
3515
|
+
const index = next++;
|
|
3516
|
+
results[index] = await work(items[index], index);
|
|
3517
|
+
}
|
|
3518
|
+
});
|
|
3519
|
+
await Promise.all(workers);
|
|
3520
|
+
return results;
|
|
3521
|
+
}
|
|
3522
|
+
__name(mapLimit, "mapLimit");
|
|
3523
|
+
async function buildDocumentTree(input, options) {
|
|
3524
|
+
const started = Date.now();
|
|
3525
|
+
const { units, declared, kind } = toUnits(input);
|
|
3526
|
+
const maxDepth = Math.max(1, options.maxDepth ?? 4);
|
|
3527
|
+
const groupSize = Math.max(1, options.groupSize ?? 5);
|
|
3528
|
+
const maxLeafUnits = Math.max(groupSize, options.maxLeafUnits ?? 10);
|
|
3529
|
+
const maxChildren = Math.max(2, options.maxChildren ?? 12);
|
|
3530
|
+
const minHeadings = Math.max(1, options.minHeadings ?? 3);
|
|
3531
|
+
const llmEnabled = options.llm !== void 0;
|
|
3532
|
+
const summaries = options.summaries ?? llmEnabled;
|
|
3533
|
+
const describe = options.describe ?? llmEnabled;
|
|
3534
|
+
const llmStructure = options.llmStructure ?? llmEnabled;
|
|
3535
|
+
const detect = options.detectHeadings ?? true;
|
|
3536
|
+
const summaryMinChars = options.summaryMinChars ?? 600;
|
|
3537
|
+
const summaryInputChars = options.summaryInputChars ?? 8e3;
|
|
3538
|
+
const unitHashes = units.map((unit) => stableHash(unit.text));
|
|
3539
|
+
const fingerprint = stableHash(JSON.stringify({
|
|
3540
|
+
v: 1,
|
|
3541
|
+
kind,
|
|
3542
|
+
title: input.title ?? null,
|
|
3543
|
+
outline: input.outline ?? null,
|
|
3544
|
+
sections: input.sections?.map((section) => [
|
|
3545
|
+
section.title ?? null,
|
|
3546
|
+
section.level ?? null,
|
|
3547
|
+
section.page ?? null
|
|
3548
|
+
]) ?? null,
|
|
3549
|
+
units: unitHashes,
|
|
3550
|
+
maxDepth,
|
|
3551
|
+
groupSize,
|
|
3552
|
+
maxLeafUnits,
|
|
3553
|
+
maxChildren,
|
|
3554
|
+
minHeadings,
|
|
3555
|
+
detect,
|
|
3556
|
+
summaries,
|
|
3557
|
+
describe,
|
|
3558
|
+
llmStructure: llmStructure && llmEnabled,
|
|
3559
|
+
summaryMinChars,
|
|
3560
|
+
summaryInputChars
|
|
3561
|
+
}));
|
|
3562
|
+
if (options.previous?.fingerprint === fingerprint) {
|
|
3563
|
+
const { metadata: _metadata, source: _source, ...previous } = options.previous;
|
|
3564
|
+
const metadata = options.metadata ?? options.previous.metadata;
|
|
3565
|
+
const source = options.source ?? options.previous.source;
|
|
3566
|
+
return {
|
|
3567
|
+
tree: {
|
|
3568
|
+
...previous,
|
|
3569
|
+
...metadata !== void 0 ? {
|
|
3570
|
+
metadata
|
|
3571
|
+
} : {},
|
|
3572
|
+
...source !== void 0 ? {
|
|
3573
|
+
source
|
|
3574
|
+
} : {}
|
|
3575
|
+
},
|
|
3576
|
+
units,
|
|
3577
|
+
unchanged: true
|
|
3578
|
+
};
|
|
3579
|
+
}
|
|
3580
|
+
const budget = new BudgetedLlm(options.llm ?? (() => Promise.reject(new Error("no llm"))), {
|
|
3581
|
+
...DEFAULT_BUDGET,
|
|
3582
|
+
...options.budget
|
|
3583
|
+
}, options.signal);
|
|
3584
|
+
let fallbackReason;
|
|
3585
|
+
let structure = "flat";
|
|
3586
|
+
let headings = [];
|
|
3587
|
+
const running = runningLines(units.map((unit) => unit.text));
|
|
3588
|
+
const outline = outlineHeadings(input, units);
|
|
3589
|
+
if (outline.length >= Math.min(minHeadings, 2)) {
|
|
3590
|
+
structure = "outline";
|
|
3591
|
+
headings = outline;
|
|
3592
|
+
} else if (declared !== void 0 && declared.length >= 1) {
|
|
3593
|
+
structure = "headings";
|
|
3594
|
+
headings = declared;
|
|
3595
|
+
} else if (detect !== false && units.length > 0) {
|
|
3596
|
+
const found = detectHeadings(units.map((unit) => unit.text), detect === "markdown");
|
|
3597
|
+
if (found.length >= minHeadings) {
|
|
3598
|
+
structure = "headings";
|
|
3599
|
+
headings = found;
|
|
3600
|
+
}
|
|
3601
|
+
}
|
|
3602
|
+
if (structure === "flat" && llmStructure && llmEnabled && units.length > 1) {
|
|
3603
|
+
try {
|
|
3604
|
+
const found = await llmHeadings(budget, units, {
|
|
3605
|
+
windowChars: options.structureWindowChars ?? 24e3,
|
|
3606
|
+
maxCalls: Math.max(1, options.maxStructureCalls ?? 12)
|
|
3607
|
+
});
|
|
3608
|
+
if (found.length >= Math.min(minHeadings, 2)) {
|
|
3609
|
+
structure = "llm";
|
|
3610
|
+
headings = found;
|
|
3611
|
+
} else {
|
|
3612
|
+
fallbackReason = "LLM found no usable structure";
|
|
3613
|
+
}
|
|
3614
|
+
} catch (error) {
|
|
3615
|
+
fallbackReason = error instanceof TreeBudgetExceededError ? `LLM structuring stopped: ${error.message}` : `LLM structuring failed: ${error.message}`;
|
|
3616
|
+
}
|
|
3617
|
+
}
|
|
3618
|
+
const draft = structure === "flat" ? units.length === 0 ? [] : groupNodes(units, 0, units.length - 1, groupSize, maxChildren, maxDepth, running) : nestHeadings(headings, units.length);
|
|
3619
|
+
const shaped = shape(draft, units, 1, {
|
|
3620
|
+
maxDepth,
|
|
3621
|
+
maxLeafUnits,
|
|
3622
|
+
groupSize,
|
|
3623
|
+
maxChildren,
|
|
3624
|
+
running
|
|
3625
|
+
});
|
|
3626
|
+
let counter = 0;
|
|
3627
|
+
const finalize = /* @__PURE__ */ __name((nodes2) => nodes2.map((node) => {
|
|
3628
|
+
counter += 1;
|
|
3629
|
+
const id = String(counter).padStart(4, "0");
|
|
3630
|
+
const pageStart = units[node.unitStart]?.page;
|
|
3631
|
+
const pageEnd = units[node.unitEnd]?.page;
|
|
3632
|
+
return {
|
|
3633
|
+
id,
|
|
3634
|
+
title: node.title,
|
|
3635
|
+
unitStart: node.unitStart,
|
|
3636
|
+
unitEnd: node.unitEnd,
|
|
3637
|
+
...pageStart !== void 0 ? {
|
|
3638
|
+
pageStart
|
|
3639
|
+
} : {},
|
|
3640
|
+
...pageEnd !== void 0 ? {
|
|
3641
|
+
pageEnd
|
|
3642
|
+
} : {},
|
|
3643
|
+
...node.keyItems !== void 0 ? {
|
|
3644
|
+
keyItems: node.keyItems
|
|
3645
|
+
} : {},
|
|
3646
|
+
contentHash: stableHash(`${node.title}\0${unitHashes.slice(node.unitStart, node.unitEnd + 1).join(",")}`),
|
|
3647
|
+
children: finalize(node.children)
|
|
3648
|
+
};
|
|
3649
|
+
}), "finalize");
|
|
3650
|
+
const nodes = finalize(shaped);
|
|
3651
|
+
const previousSummaries = /* @__PURE__ */ new Map();
|
|
3652
|
+
if (options.previous !== void 0) {
|
|
3653
|
+
for (const { node } of walkTree(options.previous.nodes)) {
|
|
3654
|
+
if (node.summary !== void 0 && (node.summarySource === "llm" || node.summarySource === "reused")) {
|
|
3655
|
+
previousSummaries.set(node.contentHash, node.summary);
|
|
3656
|
+
}
|
|
3657
|
+
}
|
|
3658
|
+
}
|
|
3659
|
+
const breadthFirst = [];
|
|
3660
|
+
for (let queue = [
|
|
3661
|
+
...nodes
|
|
3662
|
+
]; queue.length > 0; ) {
|
|
3663
|
+
const node = queue.shift();
|
|
3664
|
+
breadthFirst.push(node);
|
|
3665
|
+
queue = queue.concat(node.children);
|
|
3666
|
+
}
|
|
3667
|
+
let reusedSummaries = 0;
|
|
3668
|
+
const pending = [];
|
|
3669
|
+
for (const node of breadthFirst) {
|
|
3670
|
+
const reused = previousSummaries.get(node.contentHash);
|
|
3671
|
+
const own = ownTexts(units, node).join("\n\n");
|
|
3672
|
+
if (reused !== void 0) {
|
|
3673
|
+
node.summary = reused;
|
|
3674
|
+
node.summarySource = "reused";
|
|
3675
|
+
reusedSummaries += 1;
|
|
3676
|
+
} else if (own.trim().length < summaryMinChars && node.children.length === 0) {
|
|
3677
|
+
const text = collapse(own, summaryMinChars);
|
|
3678
|
+
if (text) {
|
|
3679
|
+
node.summary = text;
|
|
3680
|
+
node.summarySource = "text";
|
|
3681
|
+
}
|
|
3682
|
+
} else if (summaries && llmEnabled) {
|
|
3683
|
+
pending.push(node);
|
|
3684
|
+
} else {
|
|
3685
|
+
const text = collapse(own, 300);
|
|
3686
|
+
if (text) {
|
|
3687
|
+
node.summary = text;
|
|
3688
|
+
node.summarySource = "text";
|
|
3689
|
+
}
|
|
3690
|
+
}
|
|
3691
|
+
}
|
|
3692
|
+
let summariesCut = false;
|
|
3693
|
+
await mapLimit(pending, Math.max(1, options.concurrency ?? 4), async (node) => {
|
|
3694
|
+
const own = ownTexts(units, node);
|
|
3695
|
+
const childTitles = [
|
|
3696
|
+
...node.children.map((child) => child.title),
|
|
3697
|
+
...node.keyItems ?? []
|
|
3698
|
+
].map((title) => `- ${title}`).join("\n");
|
|
3699
|
+
const excerpt = excerptOf(own, summaryInputChars);
|
|
3700
|
+
const prompt = `Summarize this section of ${input.title ? `"${input.title}"` : "a document"} in 2-4 sentences (at most 80 words). Name the specific topics, requirements, clauses, figures and defined terms it covers, so a reader can tell whether it answers their question.
|
|
3701
|
+
|
|
3702
|
+
Section title: ${node.title}
|
|
3703
|
+
${childTitles ? `Subsections:
|
|
3704
|
+
${childTitles}
|
|
3705
|
+
` : ""}
|
|
3706
|
+
Section text:
|
|
3707
|
+
${excerpt}
|
|
3708
|
+
|
|
3709
|
+
Reply with the summary only.`;
|
|
3710
|
+
if (!summariesCut && budget.fits(SUMMARY_SYSTEM + prompt, 300)) {
|
|
3711
|
+
try {
|
|
3712
|
+
const reply = await budget.call({
|
|
3713
|
+
kind: "summarize",
|
|
3714
|
+
title: node.title,
|
|
3715
|
+
text: excerpt
|
|
3716
|
+
}, SUMMARY_SYSTEM, prompt, 300);
|
|
3717
|
+
const json2 = parseJsonReply(reply);
|
|
3718
|
+
const summary = typeof json2?.summary === "string" ? json2.summary : reply;
|
|
3719
|
+
if (summary.trim()) {
|
|
3720
|
+
node.summary = collapse(summary, 800);
|
|
3721
|
+
node.summarySource = "llm";
|
|
3722
|
+
return;
|
|
3723
|
+
}
|
|
3724
|
+
} catch (error) {
|
|
3725
|
+
if (error instanceof TreeBudgetExceededError) {
|
|
3726
|
+
summariesCut = true;
|
|
3727
|
+
}
|
|
3728
|
+
fallbackReason ??= `summary failed: ${error.message}`;
|
|
3729
|
+
}
|
|
3730
|
+
} else {
|
|
3731
|
+
summariesCut = true;
|
|
3732
|
+
}
|
|
3733
|
+
const text = collapse(own.join("\n\n"), 300);
|
|
3734
|
+
if (text) {
|
|
3735
|
+
node.summary = text;
|
|
3736
|
+
node.summarySource = "text";
|
|
3737
|
+
}
|
|
3738
|
+
});
|
|
3739
|
+
if (summariesCut) {
|
|
3740
|
+
fallbackReason ??= "LLM budget exhausted before every node was summarized";
|
|
3741
|
+
}
|
|
3742
|
+
let description;
|
|
3743
|
+
const outlineForDescription = nodes.slice(0, 40).map((node) => ({
|
|
3744
|
+
title: node.title,
|
|
3745
|
+
...node.summary !== void 0 ? {
|
|
3746
|
+
summary: node.summary
|
|
3747
|
+
} : {}
|
|
3748
|
+
}));
|
|
3749
|
+
if (describe && llmEnabled && nodes.length > 0) {
|
|
3750
|
+
const prompt = `Write a one-paragraph description (at most 80 words) of the document ${input.title ? `"${input.title}"` : ""} from its table of contents below, so a reader can tell which questions it can answer.
|
|
3751
|
+
|
|
3752
|
+
${outlineForDescription.map((entry) => `- ${entry.title}${entry.summary ? `: ${entry.summary}` : ""}`).join("\n")}
|
|
3753
|
+
|
|
3754
|
+
Reply with the description only.`;
|
|
3755
|
+
if (budget.fits(SUMMARY_SYSTEM + prompt, 300)) {
|
|
3756
|
+
try {
|
|
3757
|
+
description = collapse(await budget.call({
|
|
3758
|
+
kind: "describe",
|
|
3759
|
+
...input.title !== void 0 ? {
|
|
3760
|
+
title: input.title
|
|
3761
|
+
} : {},
|
|
3762
|
+
outline: outlineForDescription
|
|
3763
|
+
}, SUMMARY_SYSTEM, prompt, 300), 1e3);
|
|
3764
|
+
} catch (error) {
|
|
3765
|
+
fallbackReason ??= `description failed: ${error.message}`;
|
|
3766
|
+
}
|
|
3767
|
+
}
|
|
3768
|
+
}
|
|
3769
|
+
if (description === void 0 && options.previous?.description !== void 0 && reusedSummaries > 0) {
|
|
3770
|
+
description = options.previous.description;
|
|
3771
|
+
}
|
|
3772
|
+
const tree = {
|
|
3773
|
+
documentId: options.documentId,
|
|
3774
|
+
version: 1,
|
|
3775
|
+
...input.title !== void 0 ? {
|
|
3776
|
+
title: input.title
|
|
3777
|
+
} : {},
|
|
3778
|
+
...description ? {
|
|
3779
|
+
description
|
|
3780
|
+
} : {},
|
|
3781
|
+
structure,
|
|
3782
|
+
fingerprint,
|
|
3783
|
+
unitCount: units.length,
|
|
3784
|
+
nodes,
|
|
3785
|
+
stats: {
|
|
3786
|
+
llmCalls: budget.calls,
|
|
3787
|
+
inputTokens: budget.inputTokens,
|
|
3788
|
+
outputTokens: budget.outputTokens,
|
|
3789
|
+
durationMs: Date.now() - started,
|
|
3790
|
+
...fallbackReason !== void 0 ? {
|
|
3791
|
+
fallbackReason
|
|
3792
|
+
} : {},
|
|
3793
|
+
budgetExhausted: budget.exhausted,
|
|
3794
|
+
reusedSummaries
|
|
3795
|
+
},
|
|
3796
|
+
...options.metadata !== void 0 ? {
|
|
3797
|
+
metadata: options.metadata
|
|
3798
|
+
} : {},
|
|
3799
|
+
...options.source !== void 0 ? {
|
|
3800
|
+
source: options.source
|
|
3801
|
+
} : {},
|
|
3802
|
+
builtAt: (/* @__PURE__ */ new Date()).toISOString()
|
|
3803
|
+
};
|
|
3804
|
+
return {
|
|
3805
|
+
tree,
|
|
3806
|
+
units,
|
|
3807
|
+
unchanged: false
|
|
3808
|
+
};
|
|
3809
|
+
}
|
|
3810
|
+
__name(buildDocumentTree, "buildDocumentTree");
|
|
3811
|
+
|
|
3812
|
+
// src/tree/keyword-tree-llm.ts
|
|
3813
|
+
var STOP = new Set("the and for are was were with that this from which what when where who whom how does did has have had not but its into under over than then there their about any all can may must shall should would could will our your they them his her per each other such only also more most".split(" "));
|
|
3814
|
+
function terms(text) {
|
|
3815
|
+
return new Set((text.toLowerCase().match(/[\p{L}\p{N}][\p{L}\p{N}.\-]*[\p{L}\p{N}]|[\p{L}\p{N}]{3,}/gu) ?? []).filter((term) => term.length >= 3 && !STOP.has(term)));
|
|
3816
|
+
}
|
|
3817
|
+
__name(terms, "terms");
|
|
3818
|
+
function overlap(query, text) {
|
|
3819
|
+
let score = 0;
|
|
3820
|
+
for (const term of terms(text)) {
|
|
3821
|
+
if (query.has(term)) {
|
|
3822
|
+
score += 1;
|
|
3823
|
+
}
|
|
3824
|
+
}
|
|
3825
|
+
return score;
|
|
3826
|
+
}
|
|
3827
|
+
__name(overlap, "overlap");
|
|
3828
|
+
function keywordTreeLlm() {
|
|
3829
|
+
return async (request) => {
|
|
3830
|
+
const { task } = request;
|
|
3831
|
+
let text;
|
|
3832
|
+
switch (task.kind) {
|
|
3833
|
+
case "navigate": {
|
|
3834
|
+
const query = terms(task.query);
|
|
3835
|
+
const ranked = task.candidates.map((candidate2, order) => ({
|
|
3836
|
+
candidate: candidate2,
|
|
3837
|
+
order,
|
|
3838
|
+
score: overlap(query, `${candidate2.title} ${candidate2.summary ?? ""}`)
|
|
3839
|
+
})).filter((entry) => entry.score > 0).sort((a, b) => b.score - a.score || (b.candidate.depth ?? 0) - (a.candidate.depth ?? 0) || a.order - b.order).slice(0, task.maxNodes).map((entry) => entry.candidate);
|
|
3840
|
+
text = task.mode === "single-pass" ? JSON.stringify({
|
|
3841
|
+
thinking: "keyword overlap",
|
|
3842
|
+
read: ranked.map((c) => c.id)
|
|
3843
|
+
}) : JSON.stringify({
|
|
3844
|
+
thinking: "keyword overlap",
|
|
3845
|
+
expand: ranked.filter((c) => c.hasChildren).map((c) => c.id),
|
|
3846
|
+
read: ranked.filter((c) => !c.hasChildren).map((c) => c.id)
|
|
3847
|
+
});
|
|
3848
|
+
break;
|
|
3849
|
+
}
|
|
3850
|
+
case "select": {
|
|
3851
|
+
const query = terms(task.query);
|
|
3852
|
+
const ranked = task.documents.map((document) => ({
|
|
3853
|
+
id: document.id,
|
|
3854
|
+
score: overlap(query, `${document.title ?? ""} ${document.description ?? ""}`)
|
|
3855
|
+
})).filter((entry) => entry.score > 0).sort((a, b) => b.score - a.score).slice(0, task.maxDocuments).map((entry) => entry.id);
|
|
3856
|
+
text = JSON.stringify({
|
|
3857
|
+
thinking: "keyword overlap",
|
|
3858
|
+
documents: ranked
|
|
3859
|
+
});
|
|
3860
|
+
break;
|
|
3861
|
+
}
|
|
3862
|
+
case "summarize": {
|
|
3863
|
+
const flat = task.text.replace(/\s+/g, " ").trim();
|
|
3864
|
+
const sentence = /^.*?[.!?](\s|$)/.exec(flat)?.[0] ?? flat;
|
|
3865
|
+
text = `${task.title}: ${sentence.slice(0, 200).trim()}`;
|
|
3866
|
+
break;
|
|
3867
|
+
}
|
|
3868
|
+
case "describe":
|
|
3869
|
+
text = `${task.title ?? "Document"} covering ${task.outline.map((entry) => entry.title).join(", ")}.`;
|
|
3870
|
+
break;
|
|
3871
|
+
case "structure": {
|
|
3872
|
+
const sections = [];
|
|
3873
|
+
for (const unit of task.units) {
|
|
3874
|
+
for (const line of unit.text.split("\n")) {
|
|
3875
|
+
const trimmed = line.trim();
|
|
3876
|
+
const markdown = /^(#{1,6})\s+(.+)$/.exec(trimmed);
|
|
3877
|
+
if (markdown?.[1] !== void 0 && markdown[2] !== void 0) {
|
|
3878
|
+
sections.push({
|
|
3879
|
+
title: markdown[2],
|
|
3880
|
+
level: markdown[1].length,
|
|
3881
|
+
unit: unit.unit
|
|
3882
|
+
});
|
|
3883
|
+
} else if (/^[A-Z][A-Z0-9 ,.'&-]{3,80}$/.test(trimmed) && /[A-Z]{3}/.test(trimmed)) {
|
|
3884
|
+
sections.push({
|
|
3885
|
+
title: trimmed,
|
|
3886
|
+
level: 1,
|
|
3887
|
+
unit: unit.unit
|
|
3888
|
+
});
|
|
3889
|
+
}
|
|
3890
|
+
}
|
|
3891
|
+
}
|
|
3892
|
+
text = JSON.stringify({
|
|
3893
|
+
sections
|
|
3894
|
+
});
|
|
3895
|
+
break;
|
|
3896
|
+
}
|
|
3897
|
+
}
|
|
3898
|
+
return {
|
|
3899
|
+
text,
|
|
3900
|
+
usage: {
|
|
3901
|
+
inputTokens: Math.ceil(request.prompt.length / 4),
|
|
3902
|
+
outputTokens: Math.ceil(text.length / 4)
|
|
3903
|
+
}
|
|
3904
|
+
};
|
|
3905
|
+
};
|
|
3906
|
+
}
|
|
3907
|
+
__name(keywordTreeLlm, "keywordTreeLlm");
|
|
3908
|
+
|
|
3909
|
+
// src/tree/store.ts
|
|
3910
|
+
function headerOf(tree) {
|
|
3911
|
+
return {
|
|
3912
|
+
documentId: tree.documentId,
|
|
3913
|
+
...tree.title !== void 0 ? {
|
|
3914
|
+
title: tree.title
|
|
3915
|
+
} : {},
|
|
3916
|
+
...tree.description !== void 0 ? {
|
|
3917
|
+
description: tree.description
|
|
3918
|
+
} : {},
|
|
3919
|
+
structure: tree.structure,
|
|
3920
|
+
fingerprint: tree.fingerprint,
|
|
3921
|
+
unitCount: tree.unitCount,
|
|
3922
|
+
...tree.metadata !== void 0 ? {
|
|
3923
|
+
metadata: tree.metadata
|
|
3924
|
+
} : {},
|
|
3925
|
+
...tree.source !== void 0 ? {
|
|
3926
|
+
source: tree.source
|
|
3927
|
+
} : {},
|
|
3928
|
+
builtAt: tree.builtAt
|
|
3929
|
+
};
|
|
3930
|
+
}
|
|
3931
|
+
__name(headerOf, "headerOf");
|
|
3932
|
+
function passes(tree, filter) {
|
|
3933
|
+
if (filter === void 0 || Object.keys(filter).length === 0) {
|
|
3934
|
+
return true;
|
|
3935
|
+
}
|
|
3936
|
+
if (filterMatchesNothing(filter)) {
|
|
3937
|
+
return false;
|
|
3938
|
+
}
|
|
3939
|
+
return matchesFilter(tree.metadata, filter);
|
|
3940
|
+
}
|
|
3941
|
+
__name(passes, "passes");
|
|
3942
|
+
var MemoryDocumentTreeStore = class {
|
|
3943
|
+
static {
|
|
3944
|
+
__name(this, "MemoryDocumentTreeStore");
|
|
3945
|
+
}
|
|
3946
|
+
trees = /* @__PURE__ */ new Map();
|
|
3947
|
+
async put(tree, units) {
|
|
3948
|
+
this.trees.set(tree.documentId, {
|
|
3949
|
+
tree: structuredClone(tree),
|
|
3950
|
+
units: units.map((unit) => ({
|
|
3951
|
+
...unit
|
|
3952
|
+
}))
|
|
3953
|
+
});
|
|
3954
|
+
}
|
|
3955
|
+
async get(documentId, filter) {
|
|
3956
|
+
const entry = this.trees.get(documentId);
|
|
3957
|
+
return entry !== void 0 && passes(entry.tree, filter) ? structuredClone(entry.tree) : void 0;
|
|
3958
|
+
}
|
|
3959
|
+
async getMany(documentIds, filter) {
|
|
3960
|
+
const trees = [];
|
|
3961
|
+
for (const id of new Set(documentIds)) {
|
|
3962
|
+
const tree = await this.get(id, filter);
|
|
3963
|
+
if (tree !== void 0) {
|
|
3964
|
+
trees.push(tree);
|
|
3965
|
+
}
|
|
3966
|
+
}
|
|
3967
|
+
return trees;
|
|
3968
|
+
}
|
|
3969
|
+
async list(options = {}) {
|
|
3970
|
+
const ids = options.documentIds !== void 0 ? new Set(options.documentIds) : void 0;
|
|
3971
|
+
const headers = [
|
|
3972
|
+
...this.trees.values()
|
|
3973
|
+
].filter(({ tree }) => (ids === void 0 || ids.has(tree.documentId)) && passes(tree, options.filter)).map(({ tree }) => headerOf(tree)).sort((a, b) => a.documentId.localeCompare(b.documentId));
|
|
3974
|
+
return options.limit !== void 0 ? headers.slice(0, options.limit) : headers;
|
|
3975
|
+
}
|
|
3976
|
+
async readUnits(documentId, start, end) {
|
|
3977
|
+
const units = this.trees.get(documentId)?.units ?? [];
|
|
3978
|
+
return units.slice(Math.max(0, start), Math.max(0, end + 1)).map((unit) => ({
|
|
3979
|
+
...unit
|
|
3980
|
+
}));
|
|
3981
|
+
}
|
|
3982
|
+
async remove(documentId) {
|
|
3983
|
+
this.trees.delete(documentId);
|
|
3984
|
+
}
|
|
3985
|
+
};
|
|
3986
|
+
|
|
3987
|
+
// src/tree/pg-tree-store.ts
|
|
3988
|
+
var IDENTIFIER = /^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?$/;
|
|
3989
|
+
function json(value) {
|
|
3990
|
+
if (value === null || value === void 0) {
|
|
3991
|
+
return void 0;
|
|
3992
|
+
}
|
|
3993
|
+
return typeof value === "string" ? JSON.parse(value) : value;
|
|
3994
|
+
}
|
|
3995
|
+
__name(json, "json");
|
|
3996
|
+
var PgDocumentTreeStore = class {
|
|
3997
|
+
static {
|
|
3998
|
+
__name(this, "PgDocumentTreeStore");
|
|
3999
|
+
}
|
|
4000
|
+
client;
|
|
4001
|
+
table;
|
|
4002
|
+
unitsTable;
|
|
4003
|
+
insertBatchSize;
|
|
4004
|
+
constructor(client, options = {}) {
|
|
4005
|
+
this.client = toPgClient(client);
|
|
4006
|
+
this.table = options.table ?? "agent_rag_trees";
|
|
4007
|
+
this.unitsTable = options.unitsTable ?? `${this.table}_units`;
|
|
4008
|
+
for (const name of [
|
|
4009
|
+
this.table,
|
|
4010
|
+
this.unitsTable
|
|
4011
|
+
]) {
|
|
4012
|
+
if (!IDENTIFIER.test(name)) {
|
|
4013
|
+
throw new Error(`PgDocumentTreeStore: invalid table name ${JSON.stringify(name)}`);
|
|
4014
|
+
}
|
|
4015
|
+
}
|
|
4016
|
+
this.insertBatchSize = Math.max(1, Math.floor(options.insertBatchSize ?? 200));
|
|
4017
|
+
}
|
|
4018
|
+
/** The idempotent DDL {@link PgDocumentTreeStore.ensureSchema} runs, in order — for a migration. */
|
|
4019
|
+
schemaStatements() {
|
|
4020
|
+
const index = this.table.replace(".", "_");
|
|
4021
|
+
return [
|
|
4022
|
+
`CREATE TABLE IF NOT EXISTS ${this.table} (
|
|
4023
|
+
document_id TEXT PRIMARY KEY,
|
|
4024
|
+
title TEXT,
|
|
4025
|
+
description TEXT,
|
|
4026
|
+
structure TEXT NOT NULL,
|
|
4027
|
+
fingerprint TEXT NOT NULL,
|
|
4028
|
+
unit_count INTEGER NOT NULL,
|
|
4029
|
+
source TEXT,
|
|
4030
|
+
metadata JSONB,
|
|
4031
|
+
tree JSONB NOT NULL,
|
|
4032
|
+
built_at TIMESTAMPTZ NOT NULL
|
|
4033
|
+
)`,
|
|
4034
|
+
`CREATE INDEX IF NOT EXISTS ${index}_metadata_idx ON ${this.table} USING gin (metadata)`,
|
|
4035
|
+
`CREATE TABLE IF NOT EXISTS ${this.unitsTable} (
|
|
4036
|
+
document_id TEXT NOT NULL,
|
|
4037
|
+
idx INTEGER NOT NULL,
|
|
4038
|
+
page INTEGER,
|
|
4039
|
+
text TEXT NOT NULL,
|
|
4040
|
+
PRIMARY KEY (document_id, idx)
|
|
4041
|
+
)`
|
|
4042
|
+
];
|
|
4043
|
+
}
|
|
4044
|
+
async ensureSchema() {
|
|
4045
|
+
for (const statement of this.schemaStatements()) {
|
|
4046
|
+
await this.client.query(statement);
|
|
4047
|
+
}
|
|
4048
|
+
}
|
|
4049
|
+
/**
|
|
4050
|
+
* The SQL conditions a metadata `filter` compiles to (ANDed), appending bindings to `params`.
|
|
4051
|
+
* Override to map a filter onto your own columns. Return `['false']` to deny.
|
|
4052
|
+
*/
|
|
4053
|
+
whereConditions(filter, params) {
|
|
4054
|
+
return buildWhere(filter, params);
|
|
4055
|
+
}
|
|
4056
|
+
/** Replace the tree and its units in one transaction (when the client has one). */
|
|
4057
|
+
async put(tree, units) {
|
|
4058
|
+
const clean2 = stripNulBytes(tree);
|
|
4059
|
+
const { metadata, ...rest } = clean2;
|
|
4060
|
+
const write = /* @__PURE__ */ __name(async (client) => {
|
|
4061
|
+
await client.query(`DELETE FROM ${this.unitsTable} WHERE document_id = $1`, [
|
|
4062
|
+
clean2.documentId
|
|
4063
|
+
]);
|
|
4064
|
+
for (let start = 0; start < units.length; start += this.insertBatchSize) {
|
|
4065
|
+
const batch = units.slice(start, start + this.insertBatchSize);
|
|
4066
|
+
const params = [];
|
|
4067
|
+
const rows = batch.map((unit) => {
|
|
4068
|
+
params.push(clean2.documentId, unit.index, unit.page ?? null, stripNulBytes(unit.text));
|
|
4069
|
+
const at = params.length;
|
|
4070
|
+
return `($${at - 3}, $${at - 2}, $${at - 1}, $${at})`;
|
|
4071
|
+
});
|
|
4072
|
+
await client.query(`INSERT INTO ${this.unitsTable} (document_id, idx, page, text) VALUES ${rows.join(", ")}`, params);
|
|
4073
|
+
}
|
|
4074
|
+
await client.query(`INSERT INTO ${this.table}
|
|
4075
|
+
(document_id, title, description, structure, fingerprint, unit_count, source, metadata, tree, built_at)
|
|
4076
|
+
VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9::jsonb, $10)
|
|
4077
|
+
ON CONFLICT (document_id) DO UPDATE SET
|
|
4078
|
+
title = EXCLUDED.title, description = EXCLUDED.description, structure = EXCLUDED.structure,
|
|
4079
|
+
fingerprint = EXCLUDED.fingerprint, unit_count = EXCLUDED.unit_count, source = EXCLUDED.source,
|
|
4080
|
+
metadata = EXCLUDED.metadata, tree = EXCLUDED.tree, built_at = EXCLUDED.built_at`, [
|
|
4081
|
+
clean2.documentId,
|
|
4082
|
+
clean2.title ?? null,
|
|
4083
|
+
clean2.description ?? null,
|
|
4084
|
+
clean2.structure,
|
|
4085
|
+
clean2.fingerprint,
|
|
4086
|
+
clean2.unitCount,
|
|
4087
|
+
clean2.source ?? null,
|
|
4088
|
+
metadata === void 0 ? null : JSON.stringify(metadata),
|
|
4089
|
+
JSON.stringify(rest),
|
|
4090
|
+
clean2.builtAt
|
|
4091
|
+
]);
|
|
4092
|
+
}, "write");
|
|
4093
|
+
const transaction = this.client.transaction?.bind(this.client);
|
|
4094
|
+
await (transaction !== void 0 ? transaction(write) : write(this.client));
|
|
4095
|
+
}
|
|
4096
|
+
async get(documentId, filter) {
|
|
4097
|
+
return (await this.getMany([
|
|
4098
|
+
documentId
|
|
4099
|
+
], filter))[0];
|
|
4100
|
+
}
|
|
4101
|
+
async getMany(documentIds, filter) {
|
|
4102
|
+
if (documentIds.length === 0) {
|
|
4103
|
+
return [];
|
|
4104
|
+
}
|
|
4105
|
+
const params = [
|
|
4106
|
+
[
|
|
4107
|
+
...new Set(documentIds)
|
|
4108
|
+
]
|
|
4109
|
+
];
|
|
4110
|
+
const conditions = [
|
|
4111
|
+
"document_id = ANY($1::text[])",
|
|
4112
|
+
...this.whereConditions(filter, params)
|
|
4113
|
+
];
|
|
4114
|
+
const rows = await this.client.query(`SELECT document_id, metadata, tree FROM ${this.table} WHERE ${conditions.join(" AND ")}`, params);
|
|
4115
|
+
return rows.map((row) => {
|
|
4116
|
+
const tree = json(row.tree);
|
|
4117
|
+
const metadata = json(row.metadata);
|
|
4118
|
+
return {
|
|
4119
|
+
...tree,
|
|
4120
|
+
...metadata !== void 0 ? {
|
|
4121
|
+
metadata
|
|
4122
|
+
} : {}
|
|
4123
|
+
};
|
|
4124
|
+
});
|
|
4125
|
+
}
|
|
4126
|
+
async list(options = {}) {
|
|
4127
|
+
const params = [];
|
|
4128
|
+
const conditions = this.whereConditions(options.filter, params);
|
|
4129
|
+
if (options.documentIds !== void 0) {
|
|
4130
|
+
params.push([
|
|
4131
|
+
...options.documentIds
|
|
4132
|
+
]);
|
|
4133
|
+
conditions.push(`document_id = ANY($${params.length}::text[])`);
|
|
4134
|
+
}
|
|
4135
|
+
let limit = "";
|
|
4136
|
+
if (options.limit !== void 0) {
|
|
4137
|
+
params.push(Math.max(0, Math.floor(options.limit)));
|
|
4138
|
+
limit = `LIMIT $${params.length}`;
|
|
4139
|
+
}
|
|
4140
|
+
const rows = await this.client.query(`SELECT document_id, title, description, structure, fingerprint, unit_count, source, metadata, built_at
|
|
4141
|
+
FROM ${this.table}
|
|
4142
|
+
${conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : ""}
|
|
4143
|
+
ORDER BY document_id ${limit}`, params);
|
|
4144
|
+
return rows.map((row) => {
|
|
4145
|
+
const metadata = json(row.metadata);
|
|
4146
|
+
return headerOf({
|
|
4147
|
+
documentId: row.document_id,
|
|
4148
|
+
version: 1,
|
|
4149
|
+
...row.title !== null ? {
|
|
4150
|
+
title: row.title
|
|
4151
|
+
} : {},
|
|
4152
|
+
...row.description !== null ? {
|
|
4153
|
+
description: row.description
|
|
4154
|
+
} : {},
|
|
4155
|
+
structure: row.structure,
|
|
4156
|
+
fingerprint: row.fingerprint,
|
|
4157
|
+
unitCount: Number(row.unit_count),
|
|
4158
|
+
nodes: [],
|
|
4159
|
+
stats: {
|
|
4160
|
+
llmCalls: 0,
|
|
4161
|
+
inputTokens: 0,
|
|
4162
|
+
outputTokens: 0,
|
|
4163
|
+
durationMs: 0,
|
|
4164
|
+
budgetExhausted: false,
|
|
4165
|
+
reusedSummaries: 0
|
|
4166
|
+
},
|
|
4167
|
+
...metadata !== void 0 ? {
|
|
4168
|
+
metadata
|
|
4169
|
+
} : {},
|
|
4170
|
+
...row.source !== null ? {
|
|
4171
|
+
source: row.source
|
|
4172
|
+
} : {},
|
|
4173
|
+
builtAt: row.built_at instanceof Date ? row.built_at.toISOString() : String(row.built_at)
|
|
4174
|
+
});
|
|
4175
|
+
});
|
|
4176
|
+
}
|
|
4177
|
+
async readUnits(documentId, start, end) {
|
|
4178
|
+
const rows = await this.client.query(`SELECT idx, page, text FROM ${this.unitsTable}
|
|
4179
|
+
WHERE document_id = $1 AND idx BETWEEN $2 AND $3 ORDER BY idx`, [
|
|
4180
|
+
documentId,
|
|
4181
|
+
Math.max(0, start),
|
|
4182
|
+
end
|
|
4183
|
+
]);
|
|
4184
|
+
return rows.map((row) => ({
|
|
4185
|
+
index: Number(row.idx),
|
|
4186
|
+
text: row.text,
|
|
4187
|
+
...row.page !== null ? {
|
|
4188
|
+
page: Number(row.page)
|
|
4189
|
+
} : {}
|
|
4190
|
+
}));
|
|
4191
|
+
}
|
|
4192
|
+
async remove(documentId) {
|
|
4193
|
+
const work = /* @__PURE__ */ __name(async (client) => {
|
|
4194
|
+
await client.query(`DELETE FROM ${this.unitsTable} WHERE document_id = $1`, [
|
|
4195
|
+
documentId
|
|
4196
|
+
]);
|
|
4197
|
+
await client.query(`DELETE FROM ${this.table} WHERE document_id = $1`, [
|
|
4198
|
+
documentId
|
|
4199
|
+
]);
|
|
4200
|
+
}, "work");
|
|
4201
|
+
const transaction = this.client.transaction?.bind(this.client);
|
|
4202
|
+
await (transaction !== void 0 ? transaction(work) : work(this.client));
|
|
4203
|
+
}
|
|
4204
|
+
};
|
|
4205
|
+
|
|
4206
|
+
// src/tree/index-document.ts
|
|
4207
|
+
async function indexDocumentTree(input, options) {
|
|
4208
|
+
const { store, minUnits = 20, minChars = 6e4, ...build } = options;
|
|
4209
|
+
const { units } = toUnits(input);
|
|
4210
|
+
const chars = units.reduce((total, unit) => total + unit.text.length, 0);
|
|
4211
|
+
if (units.length < minUnits && chars < minChars) {
|
|
4212
|
+
await store.remove(options.documentId);
|
|
4213
|
+
return {
|
|
4214
|
+
status: "skipped",
|
|
4215
|
+
reason: `${units.length} units / ${chars} chars is below the threshold (${minUnits} units or ${minChars} chars)`
|
|
4216
|
+
};
|
|
4217
|
+
}
|
|
4218
|
+
const previous = await store.get(options.documentId);
|
|
4219
|
+
const built = await buildDocumentTree(input, {
|
|
4220
|
+
...build,
|
|
4221
|
+
...previous !== void 0 ? {
|
|
4222
|
+
previous
|
|
4223
|
+
} : {}
|
|
4224
|
+
});
|
|
4225
|
+
if (built.unchanged) {
|
|
4226
|
+
const metadataMoved = JSON.stringify(previous?.metadata ?? null) !== JSON.stringify(built.tree.metadata ?? null) || previous?.source !== built.tree.source;
|
|
4227
|
+
if (metadataMoved) {
|
|
4228
|
+
await store.put(built.tree, built.units);
|
|
4229
|
+
}
|
|
4230
|
+
return {
|
|
4231
|
+
status: "unchanged",
|
|
4232
|
+
tree: built.tree
|
|
4233
|
+
};
|
|
4234
|
+
}
|
|
4235
|
+
await store.put(built.tree, built.units);
|
|
4236
|
+
return {
|
|
4237
|
+
status: "built",
|
|
4238
|
+
tree: built.tree
|
|
4239
|
+
};
|
|
4240
|
+
}
|
|
4241
|
+
__name(indexDocumentTree, "indexDocumentTree");
|
|
4242
|
+
|
|
4243
|
+
// src/tree/navigate.ts
|
|
4244
|
+
var NAVIGATE_SYSTEM = "You navigate the table of contents of a document to find where a question is answered. You answer with JSON only.";
|
|
4245
|
+
var DEFAULT_BUDGET2 = {
|
|
4246
|
+
maxCalls: 6,
|
|
4247
|
+
maxInputTokens: 6e4,
|
|
4248
|
+
maxOutputTokens: 4e3,
|
|
4249
|
+
timeoutMs: 6e4
|
|
4250
|
+
};
|
|
4251
|
+
function pages(node) {
|
|
4252
|
+
if (node.pageStart === void 0) {
|
|
4253
|
+
return void 0;
|
|
4254
|
+
}
|
|
4255
|
+
return node.pageEnd === void 0 || node.pageEnd === node.pageStart ? `p. ${node.pageStart}` : `pp. ${node.pageStart}-${node.pageEnd}`;
|
|
4256
|
+
}
|
|
4257
|
+
__name(pages, "pages");
|
|
4258
|
+
function candidate(node, summaryChars, depth) {
|
|
4259
|
+
const shownPages = pages(node);
|
|
4260
|
+
let summary = node.summary;
|
|
4261
|
+
if (summary !== void 0 && summary.length > summaryChars) {
|
|
4262
|
+
summary = `${summary.slice(0, Math.max(0, summaryChars - 1))}\u2026`;
|
|
4263
|
+
}
|
|
4264
|
+
if (node.keyItems !== void 0 && node.keyItems.length > 0 && summaryChars > 0) {
|
|
4265
|
+
summary = `${summary ?? ""}${summary ? " " : ""}Includes: ${node.keyItems.slice(0, 12).join("; ")}`;
|
|
4266
|
+
}
|
|
4267
|
+
return {
|
|
4268
|
+
id: node.id,
|
|
4269
|
+
title: node.title,
|
|
4270
|
+
...summary && summaryChars > 0 ? {
|
|
4271
|
+
summary
|
|
4272
|
+
} : {},
|
|
4273
|
+
...shownPages !== void 0 ? {
|
|
4274
|
+
pages: shownPages
|
|
4275
|
+
} : {},
|
|
4276
|
+
...depth !== void 0 ? {
|
|
4277
|
+
depth
|
|
4278
|
+
} : {
|
|
4279
|
+
hasChildren: node.children.length > 0
|
|
4280
|
+
}
|
|
4281
|
+
};
|
|
4282
|
+
}
|
|
4283
|
+
__name(candidate, "candidate");
|
|
4284
|
+
function render(candidates) {
|
|
4285
|
+
return candidates.map((entry) => {
|
|
4286
|
+
const indent = " ".repeat(entry.depth ?? 0);
|
|
4287
|
+
const meta = [
|
|
4288
|
+
entry.pages,
|
|
4289
|
+
entry.hasChildren ? "has subsections" : void 0
|
|
4290
|
+
].filter(Boolean).join(", ");
|
|
4291
|
+
return `${indent}[${entry.id}] ${entry.title}${meta ? ` (${meta})` : ""}${entry.summary ? ` \u2014 ${entry.summary}` : ""}`;
|
|
4292
|
+
}).join("\n");
|
|
4293
|
+
}
|
|
4294
|
+
__name(render, "render");
|
|
4295
|
+
function fit(nodes, maxChars) {
|
|
4296
|
+
let last;
|
|
4297
|
+
for (const summaryChars of [
|
|
4298
|
+
400,
|
|
4299
|
+
160,
|
|
4300
|
+
0
|
|
4301
|
+
]) {
|
|
4302
|
+
const candidates = nodes.map(({ node, depth }) => candidate(node, summaryChars, depth));
|
|
4303
|
+
const text = render(candidates);
|
|
4304
|
+
last = {
|
|
4305
|
+
candidates,
|
|
4306
|
+
text
|
|
4307
|
+
};
|
|
4308
|
+
if (text.length <= maxChars) {
|
|
4309
|
+
break;
|
|
4310
|
+
}
|
|
4311
|
+
}
|
|
4312
|
+
return last;
|
|
4313
|
+
}
|
|
4314
|
+
__name(fit, "fit");
|
|
4315
|
+
var TreeNavigationRetriever = class {
|
|
4316
|
+
static {
|
|
4317
|
+
__name(this, "TreeNavigationRetriever");
|
|
4318
|
+
}
|
|
4319
|
+
options;
|
|
4320
|
+
store;
|
|
4321
|
+
llm;
|
|
4322
|
+
constructor(options) {
|
|
4323
|
+
this.options = options;
|
|
4324
|
+
this.store = options.store;
|
|
4325
|
+
this.llm = options.llm;
|
|
4326
|
+
}
|
|
4327
|
+
describeRetrieval() {
|
|
4328
|
+
return {
|
|
4329
|
+
retriever: "tree-navigation"
|
|
4330
|
+
};
|
|
4331
|
+
}
|
|
4332
|
+
async retrieve(query, options = {}) {
|
|
4333
|
+
const result = await this.search(query, options.filter !== void 0 ? {
|
|
4334
|
+
filter: options.filter
|
|
4335
|
+
} : {});
|
|
4336
|
+
return result.passages.slice(0, options.topK ?? result.passages.length);
|
|
4337
|
+
}
|
|
4338
|
+
/** {@link retrieve} with the audit trail: which documents, which nodes, why. */
|
|
4339
|
+
async search(query, options = {}) {
|
|
4340
|
+
const maxDocuments = Math.max(1, this.options.maxDocuments ?? 3);
|
|
4341
|
+
const maxCandidates = Math.max(maxDocuments, this.options.maxSelectionCandidates ?? 200);
|
|
4342
|
+
const headers = await this.store.list({
|
|
4343
|
+
...options.filter !== void 0 ? {
|
|
4344
|
+
filter: options.filter
|
|
4345
|
+
} : {},
|
|
4346
|
+
limit: maxCandidates
|
|
4347
|
+
});
|
|
4348
|
+
if (headers.length <= maxDocuments) {
|
|
4349
|
+
return this.navigate(query, headers.map((header) => header.documentId), options);
|
|
4350
|
+
}
|
|
4351
|
+
const budget = new BudgetedLlm(this.llm, {
|
|
4352
|
+
...DEFAULT_BUDGET2,
|
|
4353
|
+
...this.options.budget,
|
|
4354
|
+
...options.budget
|
|
4355
|
+
}, options.signal);
|
|
4356
|
+
const documents = headers.map((header) => ({
|
|
4357
|
+
id: header.documentId,
|
|
4358
|
+
...header.title !== void 0 ? {
|
|
4359
|
+
title: header.title
|
|
4360
|
+
} : {},
|
|
4361
|
+
...header.description !== void 0 ? {
|
|
4362
|
+
description: header.description.slice(0, 600)
|
|
4363
|
+
} : {}
|
|
4364
|
+
}));
|
|
4365
|
+
const prompt = `Select the documents that may contain the answer to the question, most relevant first, at most ${maxDocuments}.
|
|
4366
|
+
|
|
4367
|
+
Question: ${query}
|
|
4368
|
+
|
|
4369
|
+
Documents:
|
|
4370
|
+
${documents.map((document) => `[${document.id}] ${document.title ?? ""}${document.description ? ` \u2014 ${document.description}` : ""}`).join("\n")}
|
|
4371
|
+
|
|
4372
|
+
Reply with JSON only: {"thinking": "<why>", "documents": ["<document id>", ...]} \u2014 an empty list if none is relevant.`;
|
|
4373
|
+
const known = new Set(documents.map((document) => document.id));
|
|
4374
|
+
try {
|
|
4375
|
+
const reply = parseJsonReply(await budget.call({
|
|
4376
|
+
kind: "select",
|
|
4377
|
+
query,
|
|
4378
|
+
documents,
|
|
4379
|
+
maxDocuments
|
|
4380
|
+
}, NAVIGATE_SYSTEM, prompt, 800));
|
|
4381
|
+
const selected = [
|
|
4382
|
+
...new Set(stringList(reply?.documents ?? reply?.answer).filter((id) => known.has(id)))
|
|
4383
|
+
].slice(0, maxDocuments);
|
|
4384
|
+
const result = await this.navigate(query, selected, options);
|
|
4385
|
+
const reasoning = typeof reply?.thinking === "string" ? reply.thinking.slice(0, 2e3) : void 0;
|
|
4386
|
+
return {
|
|
4387
|
+
...result,
|
|
4388
|
+
selection: {
|
|
4389
|
+
candidates: documents.length,
|
|
4390
|
+
selected,
|
|
4391
|
+
...reasoning !== void 0 ? {
|
|
4392
|
+
reasoning
|
|
4393
|
+
} : {}
|
|
4394
|
+
}
|
|
4395
|
+
};
|
|
4396
|
+
} catch (error) {
|
|
4397
|
+
return {
|
|
4398
|
+
passages: [],
|
|
4399
|
+
documents: [],
|
|
4400
|
+
selection: {
|
|
4401
|
+
candidates: documents.length,
|
|
4402
|
+
selected: [],
|
|
4403
|
+
error: error.message
|
|
4404
|
+
}
|
|
4405
|
+
};
|
|
4406
|
+
}
|
|
4407
|
+
}
|
|
4408
|
+
/**
|
|
4409
|
+
* Navigate the given documents (in parallel, up to `concurrency`) and return their passages in
|
|
4410
|
+
* the order given, within `maxContextChars`. Documents without a tree, or whose tree fails
|
|
4411
|
+
* `filter`, yield no passages and `stoppedBy: 'not-found'`.
|
|
4412
|
+
*/
|
|
4413
|
+
async navigate(query, documentIds, options = {}) {
|
|
4414
|
+
const ids = [
|
|
4415
|
+
...new Set(documentIds)
|
|
4416
|
+
];
|
|
4417
|
+
const trees = new Map((await this.store.getMany(ids, options.filter)).map((tree) => [
|
|
4418
|
+
tree.documentId,
|
|
4419
|
+
tree
|
|
4420
|
+
]));
|
|
4421
|
+
const navigations = await mapLimit(ids, Math.max(1, this.options.concurrency ?? 3), async (id) => {
|
|
4422
|
+
const tree = trees.get(id);
|
|
4423
|
+
if (tree === void 0) {
|
|
4424
|
+
return {
|
|
4425
|
+
navigation: {
|
|
4426
|
+
documentId: id,
|
|
4427
|
+
mode: "single-pass",
|
|
4428
|
+
steps: [],
|
|
4429
|
+
nodes: [],
|
|
4430
|
+
llmCalls: 0,
|
|
4431
|
+
inputTokens: 0,
|
|
4432
|
+
outputTokens: 0,
|
|
4433
|
+
stoppedBy: "not-found"
|
|
4434
|
+
},
|
|
4435
|
+
passages: []
|
|
4436
|
+
};
|
|
4437
|
+
}
|
|
4438
|
+
const navigation = await this.walk(query, tree, options);
|
|
4439
|
+
return {
|
|
4440
|
+
navigation,
|
|
4441
|
+
passages: await this.passagesOf(tree, navigation)
|
|
4442
|
+
};
|
|
4443
|
+
});
|
|
4444
|
+
const maxContextChars = this.options.maxContextChars ?? 48e3;
|
|
4445
|
+
let used = 0;
|
|
4446
|
+
const passages = [];
|
|
4447
|
+
for (const entry of navigations) {
|
|
4448
|
+
for (const passage of entry.passages) {
|
|
4449
|
+
if (used >= maxContextChars) {
|
|
4450
|
+
break;
|
|
4451
|
+
}
|
|
4452
|
+
const text = passage.text.slice(0, maxContextChars - used);
|
|
4453
|
+
used += text.length;
|
|
4454
|
+
passages.push({
|
|
4455
|
+
...passage,
|
|
4456
|
+
text
|
|
4457
|
+
});
|
|
4458
|
+
}
|
|
4459
|
+
}
|
|
4460
|
+
return {
|
|
4461
|
+
passages,
|
|
4462
|
+
documents: navigations.map((entry) => entry.navigation)
|
|
4463
|
+
};
|
|
4464
|
+
}
|
|
4465
|
+
async walk(query, tree, options) {
|
|
4466
|
+
const budget = new BudgetedLlm(this.llm, {
|
|
4467
|
+
...DEFAULT_BUDGET2,
|
|
4468
|
+
...this.options.budget,
|
|
4469
|
+
...options.budget
|
|
4470
|
+
}, options.signal);
|
|
4471
|
+
const maxNodes = Math.max(1, this.options.maxNodes ?? 5);
|
|
4472
|
+
const singlePassMaxChars = this.options.singlePassMaxChars ?? 4e4;
|
|
4473
|
+
const index = indexTree(tree);
|
|
4474
|
+
const whole = [
|
|
4475
|
+
...walkTree(tree.nodes)
|
|
4476
|
+
].map(({ node, path }) => ({
|
|
4477
|
+
node,
|
|
4478
|
+
depth: path.length
|
|
4479
|
+
}));
|
|
4480
|
+
const outline = fit(whole, Number.POSITIVE_INFINITY);
|
|
4481
|
+
const title = tree.title ?? tree.source ?? tree.documentId;
|
|
4482
|
+
const steps = [];
|
|
4483
|
+
const chosen = [];
|
|
4484
|
+
let stoppedBy;
|
|
4485
|
+
let error;
|
|
4486
|
+
const mode = outline.text.length <= singlePassMaxChars ? "single-pass" : "beam";
|
|
4487
|
+
const stop = /* @__PURE__ */ __name((caught) => {
|
|
4488
|
+
if (caught instanceof TreeBudgetExceededError) {
|
|
4489
|
+
stoppedBy = caught.limit === "timeout" ? "timeout" : "budget";
|
|
4490
|
+
} else {
|
|
4491
|
+
stoppedBy = "error";
|
|
4492
|
+
}
|
|
4493
|
+
error = caught.message;
|
|
4494
|
+
}, "stop");
|
|
4495
|
+
if (mode === "single-pass") {
|
|
4496
|
+
const prompt = `Find the sections of the document that are likely to contain the answer to the question. Prefer the most specific (deepest) sections; list at most ${maxNodes}, most relevant first.
|
|
4497
|
+
|
|
4498
|
+
Question: ${query}
|
|
4499
|
+
|
|
4500
|
+
Document: ${title}${tree.description ? `
|
|
4501
|
+
About: ${tree.description}` : ""}
|
|
4502
|
+
Table of contents ([id] title (pages) \u2014 summary):
|
|
4503
|
+
${outline.text}
|
|
4504
|
+
|
|
4505
|
+
Reply with JSON only: {"thinking": "<which sections answer and why>", "read": ["<id>", ...]} \u2014 an empty list if none does.`;
|
|
4506
|
+
try {
|
|
4507
|
+
const reply = requireJson(await budget.call({
|
|
4508
|
+
kind: "navigate",
|
|
4509
|
+
query,
|
|
4510
|
+
mode,
|
|
4511
|
+
candidates: outline.candidates,
|
|
4512
|
+
maxNodes
|
|
4513
|
+
}, NAVIGATE_SYSTEM, prompt, 1500));
|
|
4514
|
+
const read = [
|
|
4515
|
+
...new Set(stringList(reply?.read ?? reply?.node_list ?? reply?.nodes))
|
|
4516
|
+
].map((id) => resolveId(index, id)).filter((id) => id !== void 0).slice(0, maxNodes);
|
|
4517
|
+
chosen.push(...read);
|
|
4518
|
+
steps.push({
|
|
4519
|
+
step: 1,
|
|
4520
|
+
shown: outline.candidates.map((entry) => entry.id),
|
|
4521
|
+
expanded: [],
|
|
4522
|
+
read,
|
|
4523
|
+
...reasoningOf(reply)
|
|
4524
|
+
});
|
|
4525
|
+
} catch (caught) {
|
|
4526
|
+
stop(caught);
|
|
4527
|
+
}
|
|
4528
|
+
} else {
|
|
4529
|
+
const beamWidth = Math.max(1, this.options.beamWidth ?? 4);
|
|
4530
|
+
const maxSteps = Math.max(1, this.options.maxSteps ?? 4);
|
|
4531
|
+
let frontier = tree.nodes;
|
|
4532
|
+
let pendingExpand = [];
|
|
4533
|
+
for (let step = 1; frontier.length > 0 && chosen.length < maxNodes; step++) {
|
|
4534
|
+
if (step > maxSteps) {
|
|
4535
|
+
stoppedBy = "maxSteps";
|
|
4536
|
+
break;
|
|
4537
|
+
}
|
|
4538
|
+
const shown = fit(frontier.map((node) => ({
|
|
4539
|
+
node
|
|
4540
|
+
})), singlePassMaxChars);
|
|
4541
|
+
const path = index.get(frontier[0]?.id ?? "")?.path ?? [];
|
|
4542
|
+
const where = path.length > 0 ? `
|
|
4543
|
+
You are inside: ${path.map((node) => node.title).join(" > ")}` : "";
|
|
4544
|
+
const prompt = `Find where the question is answered in the document by walking its table of contents.
|
|
4545
|
+
For each relevant section below, either "expand" it (open its subsections, only for sections that have subsections) or "read" it (its text answers the question). Choose at most ${beamWidth} in total, most relevant first; choose none if nothing here is relevant.
|
|
4546
|
+
|
|
4547
|
+
Question: ${query}
|
|
4548
|
+
|
|
4549
|
+
Document: ${title}${tree.description ? `
|
|
4550
|
+
About: ${tree.description}` : ""}${where}
|
|
4551
|
+
Sections ([id] title (pages) \u2014 summary):
|
|
4552
|
+
${shown.text}
|
|
4553
|
+
|
|
4554
|
+
Reply with JSON only: {"thinking": "<why>", "expand": ["<id>", ...], "read": ["<id>", ...]}`;
|
|
4555
|
+
let reply;
|
|
4556
|
+
try {
|
|
4557
|
+
reply = requireJson(await budget.call({
|
|
4558
|
+
kind: "navigate",
|
|
4559
|
+
query,
|
|
4560
|
+
mode,
|
|
4561
|
+
candidates: shown.candidates,
|
|
4562
|
+
maxNodes: beamWidth
|
|
4563
|
+
}, NAVIGATE_SYSTEM, prompt, 1200));
|
|
4564
|
+
} catch (caught) {
|
|
4565
|
+
stop(caught);
|
|
4566
|
+
break;
|
|
4567
|
+
}
|
|
4568
|
+
const onFrontier = new Set(frontier.map((node) => node.id));
|
|
4569
|
+
const pick = /* @__PURE__ */ __name((value) => [
|
|
4570
|
+
...new Set(stringList(value))
|
|
4571
|
+
].map((id) => resolveId(index, id)).filter((id) => id !== void 0 && onFrontier.has(id)), "pick");
|
|
4572
|
+
let expand = pick(reply?.expand);
|
|
4573
|
+
let read = pick(reply?.read);
|
|
4574
|
+
read = [
|
|
4575
|
+
...read,
|
|
4576
|
+
...expand.filter((id) => (index.get(id)?.node.children.length ?? 0) === 0)
|
|
4577
|
+
];
|
|
4578
|
+
expand = expand.filter((id) => (index.get(id)?.node.children.length ?? 0) > 0);
|
|
4579
|
+
const limited = [
|
|
4580
|
+
...expand,
|
|
4581
|
+
...read
|
|
4582
|
+
].slice(0, beamWidth);
|
|
4583
|
+
expand = expand.filter((id) => limited.includes(id));
|
|
4584
|
+
read = [
|
|
4585
|
+
...new Set(read.filter((id) => limited.includes(id)))
|
|
4586
|
+
];
|
|
4587
|
+
steps.push({
|
|
4588
|
+
step,
|
|
4589
|
+
shown: shown.candidates.map((entry) => entry.id),
|
|
4590
|
+
expanded: expand,
|
|
4591
|
+
read,
|
|
4592
|
+
...reasoningOf(reply)
|
|
4593
|
+
});
|
|
4594
|
+
chosen.push(...read);
|
|
4595
|
+
pendingExpand = expand;
|
|
4596
|
+
frontier = expand.flatMap((id) => index.get(id)?.node.children ?? []);
|
|
4597
|
+
if (frontier.length === 0) {
|
|
4598
|
+
pendingExpand = [];
|
|
4599
|
+
}
|
|
4600
|
+
}
|
|
4601
|
+
if (stoppedBy !== void 0) {
|
|
4602
|
+
chosen.push(...pendingExpand);
|
|
4603
|
+
}
|
|
4604
|
+
}
|
|
4605
|
+
const unique = [
|
|
4606
|
+
...new Set(chosen)
|
|
4607
|
+
];
|
|
4608
|
+
const specific = unique.filter((id) => !unique.some((other) => other !== id && index.get(other)?.path.some((node) => node.id === id)));
|
|
4609
|
+
const nodes = specific.slice(0, maxNodes).flatMap((id) => {
|
|
4610
|
+
const entry = index.get(id);
|
|
4611
|
+
if (entry === void 0) {
|
|
4612
|
+
return [];
|
|
4613
|
+
}
|
|
4614
|
+
const { node, path } = entry;
|
|
4615
|
+
return [
|
|
4616
|
+
{
|
|
4617
|
+
nodeId: node.id,
|
|
4618
|
+
title: node.title,
|
|
4619
|
+
path: [
|
|
4620
|
+
...path,
|
|
4621
|
+
node
|
|
4622
|
+
].map((step) => ({
|
|
4623
|
+
id: step.id,
|
|
4624
|
+
title: step.title
|
|
4625
|
+
})),
|
|
4626
|
+
...node.pageStart !== void 0 ? {
|
|
4627
|
+
pageStart: node.pageStart
|
|
4628
|
+
} : {},
|
|
4629
|
+
...node.pageEnd !== void 0 ? {
|
|
4630
|
+
pageEnd: node.pageEnd
|
|
4631
|
+
} : {},
|
|
4632
|
+
unitStart: node.unitStart,
|
|
4633
|
+
unitEnd: node.unitEnd
|
|
4634
|
+
}
|
|
4635
|
+
];
|
|
4636
|
+
});
|
|
4637
|
+
return {
|
|
4638
|
+
documentId: tree.documentId,
|
|
4639
|
+
...tree.title !== void 0 ? {
|
|
4640
|
+
title: tree.title
|
|
4641
|
+
} : {},
|
|
4642
|
+
mode,
|
|
4643
|
+
steps,
|
|
4644
|
+
nodes,
|
|
4645
|
+
llmCalls: budget.calls,
|
|
4646
|
+
inputTokens: budget.inputTokens,
|
|
4647
|
+
outputTokens: budget.outputTokens,
|
|
4648
|
+
...stoppedBy !== void 0 ? {
|
|
4649
|
+
stoppedBy
|
|
4650
|
+
} : {},
|
|
4651
|
+
...error !== void 0 ? {
|
|
4652
|
+
error
|
|
4653
|
+
} : {}
|
|
4654
|
+
};
|
|
4655
|
+
}
|
|
4656
|
+
async passagesOf(tree, navigation) {
|
|
4657
|
+
const maxPassageChars = this.options.maxPassageChars ?? 12e3;
|
|
4658
|
+
const maxUnits = Math.max(1, this.options.maxUnitsPerNode ?? 25);
|
|
4659
|
+
const reasoning = navigation.steps.map((step) => step.reasoning).filter((text) => text !== void 0).join("\n");
|
|
4660
|
+
const passages = [];
|
|
4661
|
+
for (const [rank, node] of navigation.nodes.entries()) {
|
|
4662
|
+
const units = await this.store.readUnits(tree.documentId, node.unitStart, Math.min(node.unitEnd, node.unitStart + maxUnits - 1));
|
|
4663
|
+
let text = units.map((unit) => unit.text).join("\n\n");
|
|
4664
|
+
if (text.length > maxPassageChars) {
|
|
4665
|
+
text = `${text.slice(0, maxPassageChars - 1)}\u2026`;
|
|
4666
|
+
}
|
|
4667
|
+
passages.push({
|
|
4668
|
+
id: `${tree.documentId}#node:${node.nodeId}`,
|
|
4669
|
+
text,
|
|
4670
|
+
// Rank within the document's choice; not a similarity.
|
|
4671
|
+
score: 1 / (rank + 1),
|
|
4672
|
+
source: tree.source ?? tree.title ?? tree.documentId,
|
|
4673
|
+
metadata: {
|
|
4674
|
+
...tree.metadata,
|
|
4675
|
+
documentId: tree.documentId,
|
|
4676
|
+
nodeId: node.nodeId,
|
|
4677
|
+
title: node.title,
|
|
4678
|
+
...node.pageStart !== void 0 ? {
|
|
4679
|
+
pageStart: node.pageStart
|
|
4680
|
+
} : {},
|
|
4681
|
+
...node.pageEnd !== void 0 ? {
|
|
4682
|
+
pageEnd: node.pageEnd
|
|
4683
|
+
} : {},
|
|
4684
|
+
path: node.path,
|
|
4685
|
+
retrieval: "tree-navigation",
|
|
4686
|
+
...reasoning ? {
|
|
4687
|
+
reasoning: reasoning.slice(0, 2e3)
|
|
4688
|
+
} : {}
|
|
4689
|
+
}
|
|
4690
|
+
});
|
|
4691
|
+
}
|
|
4692
|
+
return passages;
|
|
4693
|
+
}
|
|
4694
|
+
};
|
|
4695
|
+
function resolveId(index, raw) {
|
|
4696
|
+
const id = raw.trim().replace(/^\[|\]$/g, "");
|
|
4697
|
+
if (index.has(id)) {
|
|
4698
|
+
return id;
|
|
4699
|
+
}
|
|
4700
|
+
if (/^\d+$/.test(id)) {
|
|
4701
|
+
const padded = id.padStart(4, "0");
|
|
4702
|
+
if (index.has(padded)) {
|
|
4703
|
+
return padded;
|
|
4704
|
+
}
|
|
4705
|
+
}
|
|
4706
|
+
return void 0;
|
|
4707
|
+
}
|
|
4708
|
+
__name(resolveId, "resolveId");
|
|
4709
|
+
function requireJson(reply) {
|
|
4710
|
+
const json2 = parseJsonReply(reply);
|
|
4711
|
+
if (json2 === void 0) {
|
|
4712
|
+
throw new Error(`navigation reply was not JSON: ${JSON.stringify(reply.slice(0, 120))}`);
|
|
4713
|
+
}
|
|
4714
|
+
return json2;
|
|
4715
|
+
}
|
|
4716
|
+
__name(requireJson, "requireJson");
|
|
4717
|
+
function reasoningOf(reply) {
|
|
4718
|
+
const thinking = reply?.thinking ?? reply?.reasoning;
|
|
4719
|
+
return typeof thinking === "string" && thinking.trim() ? {
|
|
4720
|
+
reasoning: thinking.slice(0, 2e3)
|
|
4721
|
+
} : {};
|
|
4722
|
+
}
|
|
4723
|
+
__name(reasoningOf, "reasoningOf");
|
|
4724
|
+
|
|
4725
|
+
// src/tree/two-stage.ts
|
|
4726
|
+
function defaultDocumentIdOf(passage) {
|
|
4727
|
+
const fromMetadata = passage.metadata?.documentId;
|
|
4728
|
+
return typeof fromMetadata === "string" ? fromMetadata : documentIdOf(passage.id);
|
|
4729
|
+
}
|
|
4730
|
+
__name(defaultDocumentIdOf, "defaultDocumentIdOf");
|
|
4731
|
+
var TwoStageRetriever = class {
|
|
4732
|
+
static {
|
|
4733
|
+
__name(this, "TwoStageRetriever");
|
|
4734
|
+
}
|
|
4735
|
+
firstStage;
|
|
4736
|
+
navigator;
|
|
4737
|
+
options;
|
|
4738
|
+
constructor(firstStage, navigator, options = {}) {
|
|
4739
|
+
this.firstStage = firstStage;
|
|
4740
|
+
this.navigator = navigator;
|
|
4741
|
+
this.options = options;
|
|
4742
|
+
}
|
|
4743
|
+
describeRetrieval() {
|
|
4744
|
+
return {
|
|
4745
|
+
retriever: "two-stage"
|
|
4746
|
+
};
|
|
4747
|
+
}
|
|
4748
|
+
async retrieve(query, options = {}) {
|
|
4749
|
+
return (await this.search(query, options)).passages;
|
|
4750
|
+
}
|
|
4751
|
+
/** {@link retrieve} with the navigations that were run. */
|
|
4752
|
+
async search(query, options = {}) {
|
|
4753
|
+
const topK = options.topK ?? 5;
|
|
4754
|
+
const docOf = this.options.documentIdOf ?? defaultDocumentIdOf;
|
|
4755
|
+
const first = await this.firstStage.retrieve(query, {
|
|
4756
|
+
topK: Math.max(topK, this.options.fetchTopK ?? 20),
|
|
4757
|
+
...options.filter !== void 0 ? {
|
|
4758
|
+
filter: options.filter
|
|
4759
|
+
} : {}
|
|
4760
|
+
});
|
|
4761
|
+
const ranked = [];
|
|
4762
|
+
for (const passage of first) {
|
|
4763
|
+
const id = docOf(passage);
|
|
4764
|
+
if (!ranked.includes(id)) {
|
|
4765
|
+
ranked.push(id);
|
|
4766
|
+
}
|
|
4767
|
+
}
|
|
4768
|
+
const considered = ranked.slice(0, Math.max(1, this.options.considerDocuments ?? 3));
|
|
4769
|
+
const withTrees = new Set((await this.navigator.store.list({
|
|
4770
|
+
documentIds: considered,
|
|
4771
|
+
...options.filter !== void 0 ? {
|
|
4772
|
+
filter: options.filter
|
|
4773
|
+
} : {}
|
|
4774
|
+
})).map((header) => header.documentId));
|
|
4775
|
+
const toNavigate = considered.filter((id) => withTrees.has(id)).slice(0, Math.max(0, this.options.navigateDocuments ?? 2));
|
|
4776
|
+
const navigated = /* @__PURE__ */ new Map();
|
|
4777
|
+
let navigations = [];
|
|
4778
|
+
if (toNavigate.length > 0) {
|
|
4779
|
+
const result = await this.navigator.navigate(query, toNavigate, {
|
|
4780
|
+
...options.filter !== void 0 ? {
|
|
4781
|
+
filter: options.filter
|
|
4782
|
+
} : {},
|
|
4783
|
+
...this.options.budget !== void 0 ? {
|
|
4784
|
+
budget: this.options.budget
|
|
4785
|
+
} : {}
|
|
4786
|
+
});
|
|
4787
|
+
navigations = result.documents;
|
|
4788
|
+
for (const passage of result.passages) {
|
|
4789
|
+
const id = String(passage.metadata?.documentId ?? "");
|
|
4790
|
+
navigated.set(id, [
|
|
4791
|
+
...navigated.get(id) ?? [],
|
|
4792
|
+
passage
|
|
4793
|
+
]);
|
|
4794
|
+
}
|
|
4795
|
+
for (const navigation of navigations) {
|
|
4796
|
+
this.options.onNavigation?.(navigation);
|
|
4797
|
+
}
|
|
4798
|
+
}
|
|
4799
|
+
const merged = [];
|
|
4800
|
+
const emitted = /* @__PURE__ */ new Set();
|
|
4801
|
+
for (const passage of first) {
|
|
4802
|
+
const id = docOf(passage);
|
|
4803
|
+
const replacement = navigated.get(id);
|
|
4804
|
+
if (replacement === void 0 || replacement.length === 0) {
|
|
4805
|
+
merged.push(passage);
|
|
4806
|
+
continue;
|
|
4807
|
+
}
|
|
4808
|
+
if (!emitted.has(id)) {
|
|
4809
|
+
emitted.add(id);
|
|
4810
|
+
merged.push(...replacement);
|
|
4811
|
+
}
|
|
4812
|
+
if (this.options.keepFirstStageChunks) {
|
|
4813
|
+
merged.push(passage);
|
|
4814
|
+
}
|
|
4815
|
+
}
|
|
4816
|
+
return {
|
|
4817
|
+
passages: merged.slice(0, topK),
|
|
4818
|
+
navigations
|
|
4819
|
+
};
|
|
4820
|
+
}
|
|
4821
|
+
};
|
|
4822
|
+
|
|
4823
|
+
// src/tree/navigate-tool.ts
|
|
4824
|
+
function isInput(value) {
|
|
4825
|
+
return typeof value === "object" && value !== null && typeof value.documentId === "string" && typeof value.question === "string";
|
|
4826
|
+
}
|
|
4827
|
+
__name(isInput, "isInput");
|
|
4828
|
+
var inputSchema = {
|
|
4829
|
+
"~standard": {
|
|
4830
|
+
version: 1,
|
|
4831
|
+
vendor: "nestjs-agent-rag",
|
|
4832
|
+
validate(value) {
|
|
4833
|
+
return isInput(value) ? {
|
|
4834
|
+
value: {
|
|
4835
|
+
documentId: value.documentId,
|
|
4836
|
+
question: value.question
|
|
4837
|
+
}
|
|
4838
|
+
} : {
|
|
4839
|
+
issues: [
|
|
4840
|
+
{
|
|
4841
|
+
message: "navigate_document expects { documentId: string, question: string }"
|
|
4842
|
+
}
|
|
4843
|
+
]
|
|
4844
|
+
};
|
|
4845
|
+
}
|
|
4846
|
+
}
|
|
4847
|
+
};
|
|
4848
|
+
function createNavigateDocumentTool(navigator, options = {}) {
|
|
4849
|
+
return {
|
|
4850
|
+
spec: {
|
|
4851
|
+
name: options.name ?? "navigate_document",
|
|
4852
|
+
kind: "read",
|
|
4853
|
+
description: options.description ?? "Read inside one long document (regulation, contract, filing, manual) by navigating its table of contents to the sections that answer a question. Use it after a search returned a long document: pass that document id and the question. Returns the matching sections with titles and page ranges for citation.",
|
|
4854
|
+
inputSchema
|
|
4855
|
+
},
|
|
4856
|
+
handler: {
|
|
4857
|
+
async execute(input, ctx) {
|
|
4858
|
+
if (!isInput(input)) {
|
|
4859
|
+
throw new Error("navigate_document expects { documentId: string, question: string }");
|
|
4860
|
+
}
|
|
4861
|
+
const filter = options.filter?.(ctx);
|
|
4862
|
+
const navigateOptions = {
|
|
4863
|
+
...filter !== void 0 ? {
|
|
4864
|
+
filter
|
|
4865
|
+
} : {},
|
|
4866
|
+
...options.budget !== void 0 ? {
|
|
4867
|
+
budget: options.budget
|
|
4868
|
+
} : {}
|
|
4869
|
+
};
|
|
4870
|
+
let documentId = input.documentId.trim();
|
|
4871
|
+
let result = await navigator.navigate(input.question, [
|
|
4872
|
+
documentId
|
|
4873
|
+
], navigateOptions);
|
|
4874
|
+
let navigation = result.documents[0];
|
|
4875
|
+
if (navigation?.stoppedBy === "not-found") {
|
|
4876
|
+
const bare = documentIdOf(documentId.replace(/#node:\d+$/, ""));
|
|
4877
|
+
if (bare !== documentId) {
|
|
4878
|
+
documentId = bare;
|
|
4879
|
+
result = await navigator.navigate(input.question, [
|
|
4880
|
+
documentId
|
|
4881
|
+
], navigateOptions);
|
|
4882
|
+
navigation = result.documents[0];
|
|
4883
|
+
}
|
|
4884
|
+
}
|
|
4885
|
+
if (navigation === void 0 || navigation.stoppedBy === "not-found") {
|
|
4886
|
+
return {
|
|
4887
|
+
documentId,
|
|
4888
|
+
found: false,
|
|
4889
|
+
passages: [],
|
|
4890
|
+
nodes: [],
|
|
4891
|
+
steps: [],
|
|
4892
|
+
note: "This document has no navigable table of contents (it is short, not indexed, or not accessible). Use the search passages instead."
|
|
4893
|
+
};
|
|
4894
|
+
}
|
|
4895
|
+
return {
|
|
4896
|
+
documentId,
|
|
4897
|
+
found: true,
|
|
4898
|
+
passages: result.passages,
|
|
4899
|
+
nodes: navigation.nodes,
|
|
4900
|
+
steps: navigation.steps,
|
|
4901
|
+
...navigation.stoppedBy !== void 0 ? {
|
|
4902
|
+
stoppedBy: navigation.stoppedBy
|
|
4903
|
+
} : {},
|
|
4904
|
+
...navigation.nodes.length === 0 ? {
|
|
4905
|
+
note: "No section of this document looked relevant to the question."
|
|
4906
|
+
} : {}
|
|
4907
|
+
};
|
|
4908
|
+
}
|
|
4909
|
+
}
|
|
4910
|
+
};
|
|
4911
|
+
}
|
|
4912
|
+
__name(createNavigateDocumentTool, "createNavigateDocumentTool");
|
|
2608
4913
|
export {
|
|
2609
4914
|
DEFAULT_STOP_WORDS,
|
|
2610
4915
|
EmbeddingRetriever,
|
|
@@ -2614,7 +4919,9 @@ export {
|
|
|
2614
4919
|
HybridRetriever,
|
|
2615
4920
|
KeywordRetriever,
|
|
2616
4921
|
LexicalRetriever,
|
|
4922
|
+
MemoryDocumentTreeStore,
|
|
2617
4923
|
MemoryVectorStore,
|
|
4924
|
+
PgDocumentTreeStore,
|
|
2618
4925
|
PgLexicalVectorStore,
|
|
2619
4926
|
PgVectorStore,
|
|
2620
4927
|
RAG_DIAGNOSTIC_LIB,
|
|
@@ -2623,26 +4930,44 @@ export {
|
|
|
2623
4930
|
RedisVectorSchemaMismatchError,
|
|
2624
4931
|
RedisVectorStore,
|
|
2625
4932
|
RerankingRetriever,
|
|
4933
|
+
TreeBudgetExceededError,
|
|
4934
|
+
TreeNavigationRetriever,
|
|
4935
|
+
TwoStageRetriever,
|
|
2626
4936
|
UnsafeRemovalError,
|
|
2627
4937
|
anyTermTsquery,
|
|
2628
4938
|
applyMetadataPatch,
|
|
4939
|
+
buildDocumentTree,
|
|
4940
|
+
cachedTreeLlm,
|
|
2629
4941
|
chunkDocuments,
|
|
2630
4942
|
chunkIndexOf,
|
|
2631
4943
|
chunkText,
|
|
4944
|
+
createNavigateDocumentTool,
|
|
2632
4945
|
createRetrievalTool,
|
|
2633
4946
|
describeRetrieval,
|
|
2634
4947
|
describeSharedSource,
|
|
2635
4948
|
describeSource,
|
|
2636
4949
|
describesRetrieval,
|
|
4950
|
+
detectHeadings,
|
|
2637
4951
|
documentIdOf,
|
|
2638
4952
|
emitRagRetrieval,
|
|
4953
|
+
estimateTokens,
|
|
2639
4954
|
hasSearchSyntax,
|
|
4955
|
+
indexDocumentTree,
|
|
4956
|
+
indexTree,
|
|
2640
4957
|
ingestChunks,
|
|
2641
4958
|
ingestDocuments,
|
|
2642
4959
|
instrumentRetriever,
|
|
4960
|
+
isBatchTooLarge,
|
|
2643
4961
|
isLexicalVectorStore,
|
|
2644
4962
|
keywordTerms,
|
|
4963
|
+
keywordTreeLlm,
|
|
4964
|
+
openAiChatTreeLlm,
|
|
2645
4965
|
openAiEmbeddings,
|
|
2646
|
-
|
|
4966
|
+
parseJsonReply,
|
|
4967
|
+
spanText,
|
|
4968
|
+
stableHash,
|
|
4969
|
+
toPgClient,
|
|
4970
|
+
treeLlmFromModelProvider,
|
|
4971
|
+
walkTree
|
|
2647
4972
|
};
|
|
2648
4973
|
//# sourceMappingURL=index.js.map
|