@dudousxd/nestjs-agent-rag 0.11.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 +109 -4
- package/dist/index.cjs +2539 -68
- package/dist/index.cjs.map +1 -1
- package/dist/index.d.cts +862 -11
- package/dist/index.d.ts +862 -11
- package/dist/index.js +2514 -67
- package/dist/index.js.map +1 -1
- package/package.json +4 -3
package/dist/index.cjs
CHANGED
|
@@ -21,6 +21,7 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
|
|
|
21
21
|
// src/index.ts
|
|
22
22
|
var index_exports = {};
|
|
23
23
|
__export(index_exports, {
|
|
24
|
+
DEFAULT_STOP_WORDS: () => DEFAULT_STOP_WORDS,
|
|
24
25
|
EmbeddingRetriever: () => EmbeddingRetriever,
|
|
25
26
|
FilteredRetriever: () => FilteredRetriever,
|
|
26
27
|
HttpModelError: () => HttpModelError,
|
|
@@ -28,7 +29,9 @@ __export(index_exports, {
|
|
|
28
29
|
HybridRetriever: () => HybridRetriever,
|
|
29
30
|
KeywordRetriever: () => KeywordRetriever,
|
|
30
31
|
LexicalRetriever: () => LexicalRetriever,
|
|
32
|
+
MemoryDocumentTreeStore: () => MemoryDocumentTreeStore,
|
|
31
33
|
MemoryVectorStore: () => MemoryVectorStore,
|
|
34
|
+
PgDocumentTreeStore: () => PgDocumentTreeStore,
|
|
32
35
|
PgLexicalVectorStore: () => PgLexicalVectorStore,
|
|
33
36
|
PgVectorStore: () => PgVectorStore,
|
|
34
37
|
RAG_DIAGNOSTIC_LIB: () => RAG_DIAGNOSTIC_LIB,
|
|
@@ -37,31 +40,52 @@ __export(index_exports, {
|
|
|
37
40
|
RedisVectorSchemaMismatchError: () => RedisVectorSchemaMismatchError,
|
|
38
41
|
RedisVectorStore: () => RedisVectorStore,
|
|
39
42
|
RerankingRetriever: () => RerankingRetriever,
|
|
43
|
+
TreeBudgetExceededError: () => TreeBudgetExceededError,
|
|
44
|
+
TreeNavigationRetriever: () => TreeNavigationRetriever,
|
|
45
|
+
TwoStageRetriever: () => TwoStageRetriever,
|
|
40
46
|
UnsafeRemovalError: () => UnsafeRemovalError,
|
|
47
|
+
anyTermTsquery: () => anyTermTsquery,
|
|
41
48
|
applyMetadataPatch: () => applyMetadataPatch,
|
|
49
|
+
buildDocumentTree: () => buildDocumentTree,
|
|
50
|
+
cachedTreeLlm: () => cachedTreeLlm,
|
|
42
51
|
chunkDocuments: () => chunkDocuments,
|
|
43
52
|
chunkIndexOf: () => chunkIndexOf,
|
|
44
53
|
chunkText: () => chunkText,
|
|
54
|
+
createNavigateDocumentTool: () => createNavigateDocumentTool,
|
|
45
55
|
createRetrievalTool: () => createRetrievalTool,
|
|
46
56
|
describeRetrieval: () => describeRetrieval,
|
|
47
57
|
describeSharedSource: () => describeSharedSource,
|
|
48
58
|
describeSource: () => describeSource,
|
|
49
59
|
describesRetrieval: () => describesRetrieval,
|
|
60
|
+
detectHeadings: () => detectHeadings,
|
|
50
61
|
documentIdOf: () => documentIdOf,
|
|
51
62
|
emitRagRetrieval: () => emitRagRetrieval,
|
|
63
|
+
estimateTokens: () => estimateTokens,
|
|
64
|
+
hasSearchSyntax: () => hasSearchSyntax,
|
|
65
|
+
indexDocumentTree: () => indexDocumentTree,
|
|
66
|
+
indexTree: () => indexTree,
|
|
52
67
|
ingestChunks: () => ingestChunks,
|
|
53
68
|
ingestDocuments: () => ingestDocuments,
|
|
54
69
|
instrumentRetriever: () => instrumentRetriever,
|
|
70
|
+
isBatchTooLarge: () => isBatchTooLarge,
|
|
55
71
|
isLexicalVectorStore: () => isLexicalVectorStore,
|
|
72
|
+
keywordTerms: () => keywordTerms,
|
|
73
|
+
keywordTreeLlm: () => keywordTreeLlm,
|
|
74
|
+
openAiChatTreeLlm: () => openAiChatTreeLlm,
|
|
56
75
|
openAiEmbeddings: () => openAiEmbeddings,
|
|
57
|
-
|
|
76
|
+
parseJsonReply: () => parseJsonReply,
|
|
77
|
+
spanText: () => spanText,
|
|
78
|
+
stableHash: () => stableHash,
|
|
79
|
+
toPgClient: () => toPgClient,
|
|
80
|
+
treeLlmFromModelProvider: () => treeLlmFromModelProvider,
|
|
81
|
+
walkTree: () => walkTree
|
|
58
82
|
});
|
|
59
83
|
module.exports = __toCommonJS(index_exports);
|
|
60
84
|
|
|
61
85
|
// src/chunk.ts
|
|
62
86
|
function chunkText(text, options = {}) {
|
|
63
87
|
const chunkSize = options.chunkSize ?? 800;
|
|
64
|
-
const
|
|
88
|
+
const overlap2 = options.overlap ?? 100;
|
|
65
89
|
const normalized = text.trim();
|
|
66
90
|
if (normalized.length === 0) {
|
|
67
91
|
return [];
|
|
@@ -93,7 +117,7 @@ function chunkText(text, options = {}) {
|
|
|
93
117
|
if (breakAt >= normalized.length) {
|
|
94
118
|
break;
|
|
95
119
|
}
|
|
96
|
-
start = Math.max(breakAt -
|
|
120
|
+
start = Math.max(breakAt - overlap2, start + 1);
|
|
97
121
|
}
|
|
98
122
|
return chunks;
|
|
99
123
|
}
|
|
@@ -120,12 +144,12 @@ function chunkByRecord(text, chunkSize, separator) {
|
|
|
120
144
|
}));
|
|
121
145
|
continue;
|
|
122
146
|
}
|
|
123
|
-
const
|
|
124
|
-
if (
|
|
147
|
+
const candidate2 = current.length > 0 ? `${current}${separator}${trimmed}` : trimmed;
|
|
148
|
+
if (candidate2.length > chunkSize) {
|
|
125
149
|
flush();
|
|
126
150
|
current = trimmed;
|
|
127
151
|
} else {
|
|
128
|
-
current =
|
|
152
|
+
current = candidate2;
|
|
129
153
|
}
|
|
130
154
|
}
|
|
131
155
|
flush();
|
|
@@ -413,6 +437,81 @@ function cosineSimilarity(a, b) {
|
|
|
413
437
|
}
|
|
414
438
|
__name(cosineSimilarity, "cosineSimilarity");
|
|
415
439
|
|
|
440
|
+
// src/lexical-query.ts
|
|
441
|
+
var words = /* @__PURE__ */ __name((s) => new Set(s.split(/\s+/).filter(Boolean)), "words");
|
|
442
|
+
var DEFAULT_STOP_WORDS = {
|
|
443
|
+
english: words(`
|
|
444
|
+
me my myself we our ours ourselves you your yours yourself yourselves he him his himself she
|
|
445
|
+
her hers herself it its itself they them their theirs themselves what which who whom this that
|
|
446
|
+
these those am is are was were be been being have has had having do does did doing an the and
|
|
447
|
+
but if or because as until while of at by for with about against between into through during
|
|
448
|
+
before after above below to from up down in out on off over under again further then once here
|
|
449
|
+
there when where why how all any both each few more most other some such no nor not only own
|
|
450
|
+
same so than too very can will just don should now`),
|
|
451
|
+
portuguese: words(`
|
|
452
|
+
de que do da em um para com n\xE3o uma os no se na por mais as dos como mas ao ele das seu sua ou
|
|
453
|
+
quando muito nos j\xE1 eu tamb\xE9m s\xF3 pelo pela at\xE9 isso ela entre depois sem mesmo aos seus quem
|
|
454
|
+
nas me esse eles voc\xEA essa num nem suas meu \xE0s minha numa pelos elas qual n\xF3s lhe deles essas
|
|
455
|
+
esses pelas este dele tu te voc\xEAs vos lhes meus minhas teu tua teus tuas nosso nossa nossos
|
|
456
|
+
nossas dela delas esta estes estas aquele aquela aqueles aquelas isto aquilo estou est\xE1 estamos
|
|
457
|
+
est\xE3o estive esteve estivemos estiveram estava est\xE1vamos estavam estivera estiv\xE9ramos esteja
|
|
458
|
+
estejamos estejam estivesse estiv\xE9ssemos estivessem estiver estivermos estiverem hei h\xE1 havemos
|
|
459
|
+
h\xE3o houve houvemos houveram houvera houv\xE9ramos haja hajamos hajam houvesse houv\xE9ssemos
|
|
460
|
+
houvessem houver houvermos houverem houverei houver\xE1 houveremos houver\xE3o houveria houver\xEDamos
|
|
461
|
+
houveriam sou somos s\xE3o era \xE9ramos eram fui foi fomos foram fora f\xF4ramos seja sejamos sejam
|
|
462
|
+
fosse f\xF4ssemos fossem for formos forem serei ser\xE1 seremos ser\xE3o seria ser\xEDamos seriam tenho tem
|
|
463
|
+
temos t\xE9m tinha t\xEDnhamos tinham tive teve tivemos tiveram tivera tiv\xE9ramos tenha tenhamos
|
|
464
|
+
tenham tivesse tiv\xE9ssemos tivessem tiver tivermos tiverem terei ter\xE1 teremos ter\xE3o teria
|
|
465
|
+
ter\xEDamos teriam`),
|
|
466
|
+
spanish: words(`
|
|
467
|
+
de la que el en los del se las por un para con no una su al lo como m\xE1s pero sus le ya este s\xED
|
|
468
|
+
porque esta entre cuando muy sin sobre tambi\xE9n me hasta hay donde quien desde todo nos durante
|
|
469
|
+
todos uno les ni contra otros ese eso ante ellos esto m\xED antes algunos qu\xE9 unos yo otro otras
|
|
470
|
+
otra \xE9l tanto esa estos mucho quienes nada muchos cual poco ella estar estas algunas algo
|
|
471
|
+
nosotros mi mis t\xFA te ti tu tus ellas nosotras vosotros vosotras os m\xEDo m\xEDa m\xEDos m\xEDas tuyo
|
|
472
|
+
tuya tuyos tuyas suyo suya suyos suyas nuestro nuestra nuestros nuestras vuestro vuestra
|
|
473
|
+
vuestros vuestras esos esas estoy est\xE1s est\xE1 estamos est\xE1is est\xE1n est\xE9 est\xE9s estemos est\xE9is
|
|
474
|
+
est\xE9n estar\xE9 estar\xE1s estar\xE1 estaremos estar\xE9is estar\xE1n estar\xEDa estar\xEDas estar\xEDamos estar\xEDais
|
|
475
|
+
estar\xEDan estaba estabas est\xE1bamos estabais estaban estuve estuviste estuvo estuvimos
|
|
476
|
+
estuvisteis estuvieron estuviera estuvieras estuvi\xE9ramos estuvierais estuvieran estuviese
|
|
477
|
+
estuvieses estuvi\xE9semos estuvieseis estuviesen estando estad he has ha hemos hab\xE9is han haya
|
|
478
|
+
hayas hayamos hay\xE1is hayan habr\xE9 habr\xE1s habr\xE1 habremos habr\xE9is habr\xE1n habr\xEDa habr\xEDas habr\xEDamos
|
|
479
|
+
habr\xEDais habr\xEDan hab\xEDa hab\xEDas hab\xEDamos hab\xEDais hab\xEDan hube hubiste hubo hubimos hubisteis
|
|
480
|
+
hubieron hubiera hubieras hubi\xE9ramos hubierais hubieran hubiese hubieses hubi\xE9semos hubieseis
|
|
481
|
+
hubiesen habiendo habido habida habidos habidas soy eres es somos sois son sea seas seamos
|
|
482
|
+
se\xE1is sean ser\xE9 ser\xE1s ser\xE1 seremos ser\xE9is ser\xE1n ser\xEDa ser\xEDas ser\xEDamos ser\xEDais ser\xEDan era eras
|
|
483
|
+
\xE9ramos erais eran fui fuiste fue fuimos fuisteis fueron fuera fueras fu\xE9ramos fuerais fueran
|
|
484
|
+
fuese fueses fu\xE9semos fueseis fuesen tengo tienes tiene tenemos ten\xE9is tienen tenga tengas
|
|
485
|
+
tengamos teng\xE1is tengan tendr\xE9 tendr\xE1s tendr\xE1 tendremos tendr\xE9is tendr\xE1n tendr\xEDa tendr\xEDas
|
|
486
|
+
tendr\xEDamos tendr\xEDais tendr\xEDan ten\xEDa ten\xEDas ten\xEDamos ten\xEDais ten\xEDan tuve tuviste tuvo tuvimos
|
|
487
|
+
tuvisteis tuvieron tuviera tuvieras tuvi\xE9ramos tuvierais tuvieran tuviese tuvieses tuvi\xE9semos
|
|
488
|
+
tuvieseis tuviesen teniendo tenido tenida tenidos tenidas tened`)
|
|
489
|
+
};
|
|
490
|
+
var TOKEN = /[\p{L}\p{N}]{2,}/gu;
|
|
491
|
+
function keywordTerms(query, max = 24, stopWords = DEFAULT_STOP_WORDS) {
|
|
492
|
+
const tokens = [
|
|
493
|
+
...new Set(query.toLowerCase().normalize("NFC").match(TOKEN) ?? [])
|
|
494
|
+
];
|
|
495
|
+
const counts = Object.entries(stopWords).map(([, stop2]) => [
|
|
496
|
+
stop2,
|
|
497
|
+
tokens.filter((t) => stop2.has(t)).length
|
|
498
|
+
]);
|
|
499
|
+
const best = Math.max(0, ...counts.map(([, n]) => n));
|
|
500
|
+
if (best === 0) return tokens.slice(0, max);
|
|
501
|
+
const stop = counts.filter(([, n]) => n === best).map(([s]) => s);
|
|
502
|
+
const terms2 = tokens.filter((t) => !stop.some((s) => s.has(t)));
|
|
503
|
+
return (terms2.length > 0 ? terms2 : tokens).slice(0, max);
|
|
504
|
+
}
|
|
505
|
+
__name(keywordTerms, "keywordTerms");
|
|
506
|
+
function hasSearchSyntax(query) {
|
|
507
|
+
return /"[^"]+"/.test(query) || /(^|\s)-[\p{L}\p{N}]/u.test(query);
|
|
508
|
+
}
|
|
509
|
+
__name(hasSearchSyntax, "hasSearchSyntax");
|
|
510
|
+
function anyTermTsquery(terms2) {
|
|
511
|
+
return terms2.map((t) => `'${t.replace(/'/g, "")}'`).join(" | ");
|
|
512
|
+
}
|
|
513
|
+
__name(anyTermTsquery, "anyTermTsquery");
|
|
514
|
+
|
|
416
515
|
// src/pg-client.ts
|
|
417
516
|
var ADAPTED = /* @__PURE__ */ Symbol.for("@dudousxd/nestjs-agent-rag/pg-client");
|
|
418
517
|
function toPgClient(source) {
|
|
@@ -871,6 +970,7 @@ var PgLexicalVectorStore = class extends PgVectorStore {
|
|
|
871
970
|
tsConfig;
|
|
872
971
|
tsColumn;
|
|
873
972
|
anyTermFallback;
|
|
973
|
+
stopWords;
|
|
874
974
|
constructor(client, options = {}) {
|
|
875
975
|
super(client, options);
|
|
876
976
|
const config = options.fullText?.config ?? "simple";
|
|
@@ -884,6 +984,8 @@ var PgLexicalVectorStore = class extends PgVectorStore {
|
|
|
884
984
|
this.tsConfig = config;
|
|
885
985
|
this.tsColumn = column;
|
|
886
986
|
this.anyTermFallback = options.fullText?.anyTermFallback ?? true;
|
|
987
|
+
const stopWords = options.fullText?.stopWords;
|
|
988
|
+
this.stopWords = stopWords === false ? {} : stopWords ?? DEFAULT_STOP_WORDS;
|
|
887
989
|
}
|
|
888
990
|
schemaStatements() {
|
|
889
991
|
const statements = super.schemaStatements();
|
|
@@ -897,19 +999,18 @@ var PgLexicalVectorStore = class extends PgVectorStore {
|
|
|
897
999
|
if (filterMatchesNothing(options.filter) || query.trim() === "") {
|
|
898
1000
|
return [];
|
|
899
1001
|
}
|
|
900
|
-
|
|
901
|
-
|
|
902
|
-
|
|
903
|
-
|
|
904
|
-
|
|
905
|
-
...new Set(query.toLowerCase().match(/[\p{L}\p{N}]{2,}/gu) ?? [])
|
|
906
|
-
].slice(0, 24);
|
|
907
|
-
if (terms.length < 2) {
|
|
908
|
-
return all;
|
|
1002
|
+
if (!this.anyTermFallback || hasSearchSyntax(query)) {
|
|
1003
|
+
const strict = await this.runTextSearch(query, options);
|
|
1004
|
+
if (strict.length > 0 || !this.anyTermFallback) {
|
|
1005
|
+
return strict;
|
|
1006
|
+
}
|
|
909
1007
|
}
|
|
910
|
-
|
|
1008
|
+
const terms2 = keywordTerms(query, 24, this.stopWords);
|
|
1009
|
+
return terms2.length === 0 ? [] : this.runRankedTermSearch(terms2, options);
|
|
911
1010
|
}
|
|
912
|
-
|
|
1011
|
+
/** `websearch_to_tsquery` on the query as written: every word, quoted phrases, `-exclusions`. */
|
|
1012
|
+
async runTextSearch(query, options) {
|
|
1013
|
+
const tsquery = `websearch_to_tsquery('${this.tsConfig}', $1)`;
|
|
913
1014
|
const params = [
|
|
914
1015
|
query,
|
|
915
1016
|
options.topK
|
|
@@ -925,6 +1026,47 @@ var PgLexicalVectorStore = class extends PgVectorStore {
|
|
|
925
1026
|
LIMIT $2`, params);
|
|
926
1027
|
return rows.map(toPassage);
|
|
927
1028
|
}
|
|
1029
|
+
/**
|
|
1030
|
+
* Rows holding any of `terms`, ranked by the sum of the matched terms' IDF among those rows,
|
|
1031
|
+
* `ln(1 + (N - df + 0.5) / (df + 0.5))`, then `ts_rank_cd`. Each (row, term) match is evaluated
|
|
1032
|
+
* once, and `ts_rank_cd` only for the rows ranked within `topK` (ties included): a question full
|
|
1033
|
+
* of common words can match most of the table.
|
|
1034
|
+
*/
|
|
1035
|
+
async runRankedTermSearch(terms2, options) {
|
|
1036
|
+
const any = `to_tsquery('${this.tsConfig}', $1)`;
|
|
1037
|
+
const params = [
|
|
1038
|
+
anyTermTsquery(terms2),
|
|
1039
|
+
options.topK,
|
|
1040
|
+
terms2
|
|
1041
|
+
];
|
|
1042
|
+
const conditions = [
|
|
1043
|
+
`${this.tsvector()} @@ ${any}`,
|
|
1044
|
+
...this.whereConditions(options.filter, params)
|
|
1045
|
+
];
|
|
1046
|
+
const rows = await this.client.query(`WITH __terms AS (
|
|
1047
|
+
SELECT DISTINCT to_tsquery('${this.tsConfig}', quote_literal(t)) AS q
|
|
1048
|
+
FROM unnest($3::text[]) AS t
|
|
1049
|
+
), __cand AS MATERIALIZED (
|
|
1050
|
+
SELECT id, text, source, metadata, ${this.tsvector()} AS __tsv
|
|
1051
|
+
FROM ${this.table}
|
|
1052
|
+
WHERE ${conditions.join(" AND ")}
|
|
1053
|
+
), __hits AS MATERIALIZED (
|
|
1054
|
+
SELECT __cand.id AS __id, __terms.q FROM __cand JOIN __terms ON __cand.__tsv @@ __terms.q
|
|
1055
|
+
), __idf AS (
|
|
1056
|
+
SELECT q, ln(1 + ((SELECT count(*) FROM __cand) - count(*) + 0.5) / (count(*) + 0.5)) AS w
|
|
1057
|
+
FROM __hits GROUP BY q
|
|
1058
|
+
), __scored AS (
|
|
1059
|
+
SELECT __hits.__id, sum(__idf.w) AS s, rank() OVER (ORDER BY sum(__idf.w) DESC) AS r
|
|
1060
|
+
FROM __hits JOIN __idf USING (q) GROUP BY __hits.__id
|
|
1061
|
+
)
|
|
1062
|
+
SELECT __cand.id, __cand.text, __cand.source, __cand.metadata,
|
|
1063
|
+
__scored.s + ts_rank_cd(__cand.__tsv, ${any}, 32) AS score
|
|
1064
|
+
FROM __scored JOIN __cand ON __cand.id = __scored.__id
|
|
1065
|
+
WHERE __scored.r <= $2
|
|
1066
|
+
ORDER BY __scored.s DESC, ts_rank_cd(__cand.__tsv, ${any}, 32) DESC, __cand.id
|
|
1067
|
+
LIMIT $2`, params);
|
|
1068
|
+
return rows.map(toPassage);
|
|
1069
|
+
}
|
|
928
1070
|
/** Must be character-for-character the index expression, or the planner won't use the index. */
|
|
929
1071
|
tsvector() {
|
|
930
1072
|
return this.tsColumn ?? `to_tsvector('${this.tsConfig}'::regconfig, text)`;
|
|
@@ -1599,12 +1741,12 @@ var RedisVectorStore = class {
|
|
|
1599
1741
|
if (filterMatchesNothing(options.filter)) {
|
|
1600
1742
|
return [];
|
|
1601
1743
|
}
|
|
1602
|
-
const
|
|
1603
|
-
if (
|
|
1744
|
+
const terms2 = tokenizeQuery(query);
|
|
1745
|
+
if (terms2.length === 0) {
|
|
1604
1746
|
return [];
|
|
1605
1747
|
}
|
|
1606
1748
|
const filterQuery = buildFilter(options.filter);
|
|
1607
|
-
const textQuery = `@text:(${
|
|
1749
|
+
const textQuery = `@text:(${terms2.join("|")})`;
|
|
1608
1750
|
const reply = await this.client.sendCommand([
|
|
1609
1751
|
"FT.SEARCH",
|
|
1610
1752
|
this.index,
|
|
@@ -1858,8 +2000,8 @@ function describesRetrieval(value) {
|
|
|
1858
2000
|
return typeof value === "object" && value !== null && typeof Reflect.get(value, "describeRetrieval") === "function";
|
|
1859
2001
|
}
|
|
1860
2002
|
__name(describesRetrieval, "describesRetrieval");
|
|
1861
|
-
function describeRetrieval(
|
|
1862
|
-
return describesRetrieval(
|
|
2003
|
+
function describeRetrieval(candidate2) {
|
|
2004
|
+
return describesRetrieval(candidate2) ? candidate2.describeRetrieval() : {};
|
|
1863
2005
|
}
|
|
1864
2006
|
__name(describeRetrieval, "describeRetrieval");
|
|
1865
2007
|
function describeSource(base) {
|
|
@@ -2392,53 +2534,94 @@ var HttpModelError = class extends Error {
|
|
|
2392
2534
|
this.name = "HttpModelError";
|
|
2393
2535
|
}
|
|
2394
2536
|
};
|
|
2537
|
+
function isBatchTooLarge(error) {
|
|
2538
|
+
if (!(error instanceof HttpModelError)) {
|
|
2539
|
+
return false;
|
|
2540
|
+
}
|
|
2541
|
+
if (error.status === 413) {
|
|
2542
|
+
return true;
|
|
2543
|
+
}
|
|
2544
|
+
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);
|
|
2545
|
+
}
|
|
2546
|
+
__name(isBatchTooLarge, "isBatchTooLarge");
|
|
2547
|
+
var learnedBatchLimits = /* @__PURE__ */ new Map();
|
|
2395
2548
|
function openAiEmbeddings(options) {
|
|
2396
2549
|
const url = `${(options.baseUrl ?? "https://api.openai.com/v1").replace(/\/+$/, "")}/embeddings`;
|
|
2397
|
-
const
|
|
2550
|
+
const limitKey = `${url} ${options.model}`;
|
|
2551
|
+
const configured = Math.max(1, Math.floor(options.batchSize ?? 64));
|
|
2552
|
+
let batchSize = Math.min(configured, learnedBatchLimits.get(limitKey) ?? configured);
|
|
2398
2553
|
const maxInputChars = options.maxInputChars ?? 24e3;
|
|
2399
2554
|
const timeoutMs = options.timeoutMs ?? 12e4;
|
|
2400
2555
|
const doFetch = options.fetch ?? fetch;
|
|
2556
|
+
const request = /* @__PURE__ */ __name(async (input) => {
|
|
2557
|
+
const response = await doFetch(url, {
|
|
2558
|
+
method: "POST",
|
|
2559
|
+
headers: {
|
|
2560
|
+
"content-type": "application/json",
|
|
2561
|
+
...options.apiKey ? {
|
|
2562
|
+
authorization: `Bearer ${options.apiKey}`
|
|
2563
|
+
} : {},
|
|
2564
|
+
...options.headers
|
|
2565
|
+
},
|
|
2566
|
+
body: JSON.stringify({
|
|
2567
|
+
model: options.model,
|
|
2568
|
+
input,
|
|
2569
|
+
...options.dimensions !== void 0 ? {
|
|
2570
|
+
dimensions: options.dimensions
|
|
2571
|
+
} : {}
|
|
2572
|
+
}),
|
|
2573
|
+
signal: AbortSignal.timeout(timeoutMs)
|
|
2574
|
+
});
|
|
2575
|
+
const body = await response.text();
|
|
2576
|
+
if (!response.ok) {
|
|
2577
|
+
throw new HttpModelError(response.status, `Embeddings request failed (${response.status}): ${errorMessageOf(body)}`);
|
|
2578
|
+
}
|
|
2579
|
+
const json2 = parseJson(body);
|
|
2580
|
+
const items = [
|
|
2581
|
+
...json2?.data ?? []
|
|
2582
|
+
].sort((a, b) => (a.index ?? 0) - (b.index ?? 0));
|
|
2583
|
+
if (items.length !== input.length || !items.every((item) => Array.isArray(item.embedding))) {
|
|
2584
|
+
throw new HttpModelError(502, `Embeddings returned ${items.length} vectors for ${input.length} inputs`);
|
|
2585
|
+
}
|
|
2586
|
+
const tokens = json2?.usage?.total_tokens ?? json2?.usage?.prompt_tokens;
|
|
2587
|
+
if (typeof tokens === "number" && Number.isFinite(tokens)) {
|
|
2588
|
+
options.onUsage?.(tokens);
|
|
2589
|
+
}
|
|
2590
|
+
return items.map((item) => item.embedding);
|
|
2591
|
+
}, "request");
|
|
2592
|
+
const requestSplitting = /* @__PURE__ */ __name(async (input) => {
|
|
2593
|
+
if (input.length > batchSize) {
|
|
2594
|
+
const vectors = [];
|
|
2595
|
+
for (let start = 0; start < input.length; start += batchSize) {
|
|
2596
|
+
vectors.push(...await requestSplitting(input.slice(start, start + batchSize)));
|
|
2597
|
+
}
|
|
2598
|
+
return vectors;
|
|
2599
|
+
}
|
|
2600
|
+
try {
|
|
2601
|
+
return await request(input);
|
|
2602
|
+
} catch (error) {
|
|
2603
|
+
if (!isBatchTooLarge(error) || input.length <= 1) {
|
|
2604
|
+
throw error;
|
|
2605
|
+
}
|
|
2606
|
+
const half = Math.ceil(input.length / 2);
|
|
2607
|
+
if (half < batchSize) {
|
|
2608
|
+
batchSize = half;
|
|
2609
|
+
learnedBatchLimits.set(limitKey, half);
|
|
2610
|
+
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.`);
|
|
2611
|
+
}
|
|
2612
|
+
return [
|
|
2613
|
+
...await requestSplitting(input.slice(0, half)),
|
|
2614
|
+
...await requestSplitting(input.slice(half))
|
|
2615
|
+
];
|
|
2616
|
+
}
|
|
2617
|
+
}, "requestSplitting");
|
|
2401
2618
|
return {
|
|
2402
2619
|
async embed(texts) {
|
|
2403
2620
|
const vectors = [];
|
|
2404
|
-
for (let start = 0; start < texts.length;
|
|
2621
|
+
for (let start = 0; start < texts.length; ) {
|
|
2405
2622
|
const input = texts.slice(start, start + batchSize).map((text) => (text.length > maxInputChars ? text.slice(0, maxInputChars) : text) || " ");
|
|
2406
|
-
|
|
2407
|
-
|
|
2408
|
-
headers: {
|
|
2409
|
-
"content-type": "application/json",
|
|
2410
|
-
...options.apiKey ? {
|
|
2411
|
-
authorization: `Bearer ${options.apiKey}`
|
|
2412
|
-
} : {},
|
|
2413
|
-
...options.headers
|
|
2414
|
-
},
|
|
2415
|
-
body: JSON.stringify({
|
|
2416
|
-
model: options.model,
|
|
2417
|
-
input,
|
|
2418
|
-
...options.dimensions !== void 0 ? {
|
|
2419
|
-
dimensions: options.dimensions
|
|
2420
|
-
} : {}
|
|
2421
|
-
}),
|
|
2422
|
-
signal: AbortSignal.timeout(timeoutMs)
|
|
2423
|
-
});
|
|
2424
|
-
const body = await response.text();
|
|
2425
|
-
if (!response.ok) {
|
|
2426
|
-
throw new HttpModelError(response.status, `Embeddings request failed (${response.status}): ${errorMessageOf(body)}`);
|
|
2427
|
-
}
|
|
2428
|
-
const json = parseJson(body);
|
|
2429
|
-
const items = [
|
|
2430
|
-
...json?.data ?? []
|
|
2431
|
-
].sort((a, b) => (a.index ?? 0) - (b.index ?? 0));
|
|
2432
|
-
if (items.length !== input.length || !items.every((item) => Array.isArray(item.embedding))) {
|
|
2433
|
-
throw new HttpModelError(502, `Embeddings returned ${items.length} vectors for ${input.length} inputs`);
|
|
2434
|
-
}
|
|
2435
|
-
for (const item of items) {
|
|
2436
|
-
vectors.push(item.embedding);
|
|
2437
|
-
}
|
|
2438
|
-
const tokens = json?.usage?.total_tokens ?? json?.usage?.prompt_tokens;
|
|
2439
|
-
if (typeof tokens === "number" && Number.isFinite(tokens)) {
|
|
2440
|
-
options.onUsage?.(tokens);
|
|
2441
|
-
}
|
|
2623
|
+
vectors.push(...await requestSplitting(input));
|
|
2624
|
+
start += input.length;
|
|
2442
2625
|
}
|
|
2443
2626
|
return vectors;
|
|
2444
2627
|
}
|
|
@@ -2454,8 +2637,8 @@ function parseJson(body) {
|
|
|
2454
2637
|
}
|
|
2455
2638
|
__name(parseJson, "parseJson");
|
|
2456
2639
|
function errorMessageOf(body) {
|
|
2457
|
-
const
|
|
2458
|
-
const message = typeof
|
|
2640
|
+
const json2 = parseJson(body);
|
|
2641
|
+
const message = typeof json2?.error === "string" ? json2.error : json2?.error?.message ?? json2?.message ?? void 0;
|
|
2459
2642
|
return typeof message === "string" ? message : body.slice(0, 300);
|
|
2460
2643
|
}
|
|
2461
2644
|
__name(errorMessageOf, "errorMessageOf");
|
|
@@ -2530,22 +2713,2287 @@ var HttpReranker = class {
|
|
|
2530
2713
|
}
|
|
2531
2714
|
};
|
|
2532
2715
|
function itemsOf(body) {
|
|
2533
|
-
let
|
|
2716
|
+
let json2;
|
|
2534
2717
|
try {
|
|
2535
|
-
|
|
2718
|
+
json2 = JSON.parse(body);
|
|
2536
2719
|
} catch {
|
|
2537
2720
|
return void 0;
|
|
2538
2721
|
}
|
|
2539
|
-
if (Array.isArray(
|
|
2540
|
-
return
|
|
2722
|
+
if (Array.isArray(json2)) {
|
|
2723
|
+
return json2;
|
|
2541
2724
|
}
|
|
2542
|
-
const wrapped =
|
|
2725
|
+
const wrapped = json2;
|
|
2543
2726
|
const items = wrapped?.results ?? wrapped?.data;
|
|
2544
2727
|
return Array.isArray(items) ? items : void 0;
|
|
2545
2728
|
}
|
|
2546
2729
|
__name(itemsOf, "itemsOf");
|
|
2730
|
+
|
|
2731
|
+
// src/tree/types.ts
|
|
2732
|
+
function* walkTree(nodes, path = []) {
|
|
2733
|
+
for (const node of nodes) {
|
|
2734
|
+
yield {
|
|
2735
|
+
node,
|
|
2736
|
+
path
|
|
2737
|
+
};
|
|
2738
|
+
yield* walkTree(node.children, [
|
|
2739
|
+
...path,
|
|
2740
|
+
node
|
|
2741
|
+
]);
|
|
2742
|
+
}
|
|
2743
|
+
}
|
|
2744
|
+
__name(walkTree, "walkTree");
|
|
2745
|
+
function indexTree(tree) {
|
|
2746
|
+
const index = /* @__PURE__ */ new Map();
|
|
2747
|
+
for (const entry of walkTree(tree.nodes)) {
|
|
2748
|
+
index.set(entry.node.id, entry);
|
|
2749
|
+
}
|
|
2750
|
+
return index;
|
|
2751
|
+
}
|
|
2752
|
+
__name(indexTree, "indexTree");
|
|
2753
|
+
|
|
2754
|
+
// src/tree/headings.ts
|
|
2755
|
+
var MARKDOWN = /^(#{1,6})\s+(.+?)\s*#*\s*$/;
|
|
2756
|
+
var KEYWORD_LEVEL = {
|
|
2757
|
+
part: 1,
|
|
2758
|
+
chapter: 1,
|
|
2759
|
+
title: 1,
|
|
2760
|
+
appendix: 1,
|
|
2761
|
+
annex: 1,
|
|
2762
|
+
exhibit: 1,
|
|
2763
|
+
schedule: 1,
|
|
2764
|
+
subpart: 2,
|
|
2765
|
+
article: 2,
|
|
2766
|
+
item: 2,
|
|
2767
|
+
section: 3
|
|
2768
|
+
};
|
|
2769
|
+
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})$/;
|
|
2770
|
+
var NUMBERED = /^(\d{1,3}(?:\.\d{1,4}){1,4})([-–]\d{1,4})?\.?\s+([A-Z[][^\n]{1,100})$/;
|
|
2771
|
+
function clean(title) {
|
|
2772
|
+
return title.replace(/\s+/g, " ").replace(/[\s.:—–-]+$/, "").trim();
|
|
2773
|
+
}
|
|
2774
|
+
__name(clean, "clean");
|
|
2775
|
+
function looksLikeBody(rest) {
|
|
2776
|
+
return /(\d\)?|[%$])\s*$/.test(rest) || rest.split(" ").length > 16;
|
|
2777
|
+
}
|
|
2778
|
+
__name(looksLikeBody, "looksLikeBody");
|
|
2779
|
+
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(" "));
|
|
2780
|
+
function dangling(title) {
|
|
2781
|
+
if (/[,;(]$/.test(title)) {
|
|
2782
|
+
return true;
|
|
2783
|
+
}
|
|
2784
|
+
const last = title.split(/\s+/).pop()?.toLowerCase() ?? "";
|
|
2785
|
+
return DANGLING.has(last);
|
|
2786
|
+
}
|
|
2787
|
+
__name(dangling, "dangling");
|
|
2788
|
+
function detectHeadings(texts, markdownOnly = false) {
|
|
2789
|
+
const headings = [];
|
|
2790
|
+
let inFence = false;
|
|
2791
|
+
let run = [];
|
|
2792
|
+
const dropped = /* @__PURE__ */ new Set();
|
|
2793
|
+
const closeRun = /* @__PURE__ */ __name(() => {
|
|
2794
|
+
if (run.length >= 5) {
|
|
2795
|
+
run.forEach((entry, position) => {
|
|
2796
|
+
const heading = headings[entry.index];
|
|
2797
|
+
const keep = position === run.length - 1 || position === 0 && entry.keyword && heading !== void 0 && heading.level === 1;
|
|
2798
|
+
if (!keep) {
|
|
2799
|
+
dropped.add(entry.index);
|
|
2800
|
+
}
|
|
2801
|
+
});
|
|
2802
|
+
}
|
|
2803
|
+
run = [];
|
|
2804
|
+
}, "closeRun");
|
|
2805
|
+
texts.forEach((text, unit) => {
|
|
2806
|
+
let seenText = false;
|
|
2807
|
+
const lines = text.split("\n").map((raw) => raw.trim());
|
|
2808
|
+
for (let at = 0; at < lines.length; at++) {
|
|
2809
|
+
let line = lines[at] ?? "";
|
|
2810
|
+
if (line.startsWith("```")) {
|
|
2811
|
+
inFence = !inFence;
|
|
2812
|
+
seenText = true;
|
|
2813
|
+
closeRun();
|
|
2814
|
+
continue;
|
|
2815
|
+
}
|
|
2816
|
+
if (line === "") {
|
|
2817
|
+
continue;
|
|
2818
|
+
}
|
|
2819
|
+
if (inFence || line.length > 140) {
|
|
2820
|
+
seenText = true;
|
|
2821
|
+
closeRun();
|
|
2822
|
+
continue;
|
|
2823
|
+
}
|
|
2824
|
+
const atUnitStart = !seenText;
|
|
2825
|
+
seenText = true;
|
|
2826
|
+
const markdown = MARKDOWN.exec(line);
|
|
2827
|
+
if (markdown?.[1] !== void 0 && markdown[2] !== void 0) {
|
|
2828
|
+
closeRun();
|
|
2829
|
+
const title = clean(markdown[2].replace(/\*\*/g, ""));
|
|
2830
|
+
if (title !== "") {
|
|
2831
|
+
headings.push({
|
|
2832
|
+
title,
|
|
2833
|
+
level: markdown[1].length,
|
|
2834
|
+
unit,
|
|
2835
|
+
atUnitStart
|
|
2836
|
+
});
|
|
2837
|
+
}
|
|
2838
|
+
continue;
|
|
2839
|
+
}
|
|
2840
|
+
if (markdownOnly) {
|
|
2841
|
+
continue;
|
|
2842
|
+
}
|
|
2843
|
+
const candidate2 = KEYWORD.test(line) || NUMBERED.test(line);
|
|
2844
|
+
if (candidate2) {
|
|
2845
|
+
for (let joined = 0; joined < 2; joined++) {
|
|
2846
|
+
const next = lines[at + 1];
|
|
2847
|
+
if (!/[a-z]-$/.test(line) || next === void 0 || next === "" || next.length > 80) {
|
|
2848
|
+
break;
|
|
2849
|
+
}
|
|
2850
|
+
line = `${line.slice(0, -1)}${next}`;
|
|
2851
|
+
at += 1;
|
|
2852
|
+
}
|
|
2853
|
+
}
|
|
2854
|
+
const keyword = KEYWORD.exec(line);
|
|
2855
|
+
if (keyword?.[1] !== void 0 && keyword[2] !== void 0) {
|
|
2856
|
+
const rest = keyword[3] ?? "";
|
|
2857
|
+
const allCaps = keyword[1] === keyword[1].toUpperCase();
|
|
2858
|
+
if (!looksLikeBody(rest) && !/^[a-z,;)]/.test(rest) && (rest !== "" || allCaps) && !dangling(clean(line))) {
|
|
2859
|
+
const level = KEYWORD_LEVEL[keyword[1].toLowerCase()] ?? 2;
|
|
2860
|
+
run.push({
|
|
2861
|
+
index: headings.length,
|
|
2862
|
+
keyword: true
|
|
2863
|
+
});
|
|
2864
|
+
headings.push({
|
|
2865
|
+
title: clean(line),
|
|
2866
|
+
level,
|
|
2867
|
+
unit,
|
|
2868
|
+
atUnitStart
|
|
2869
|
+
});
|
|
2870
|
+
continue;
|
|
2871
|
+
}
|
|
2872
|
+
}
|
|
2873
|
+
const numbered = NUMBERED.exec(line);
|
|
2874
|
+
if (numbered?.[1] !== void 0 && numbered[3] !== void 0 && !looksLikeBody(numbered[3]) && !dangling(clean(numbered[3]))) {
|
|
2875
|
+
const level = Math.min(6, numbered[1].split(".").length + 1 + (numbered[2] ? 1 : 0));
|
|
2876
|
+
run.push({
|
|
2877
|
+
index: headings.length,
|
|
2878
|
+
keyword: false
|
|
2879
|
+
});
|
|
2880
|
+
headings.push({
|
|
2881
|
+
title: clean(line),
|
|
2882
|
+
level,
|
|
2883
|
+
unit,
|
|
2884
|
+
atUnitStart
|
|
2885
|
+
});
|
|
2886
|
+
continue;
|
|
2887
|
+
}
|
|
2888
|
+
if (!candidate2 && /[A-Za-z]{3,}/.test(line) && !isPageFurniture(line) && !isTitleCase(line)) {
|
|
2889
|
+
closeRun();
|
|
2890
|
+
}
|
|
2891
|
+
}
|
|
2892
|
+
});
|
|
2893
|
+
closeRun();
|
|
2894
|
+
return dropContentsDuplicates(dropRunningHeaders(headings.filter((_, index) => !dropped.has(index)), texts.length));
|
|
2895
|
+
}
|
|
2896
|
+
__name(detectHeadings, "detectHeadings");
|
|
2897
|
+
function dropContentsDuplicates(headings) {
|
|
2898
|
+
const last = /* @__PURE__ */ new Map();
|
|
2899
|
+
headings.forEach((heading, index) => {
|
|
2900
|
+
if (/\d/.test(heading.title)) {
|
|
2901
|
+
last.set(heading.title.toLowerCase(), index);
|
|
2902
|
+
}
|
|
2903
|
+
});
|
|
2904
|
+
return headings.filter((heading, index) => {
|
|
2905
|
+
const key = heading.title.toLowerCase();
|
|
2906
|
+
return !last.has(key) || last.get(key) === index;
|
|
2907
|
+
});
|
|
2908
|
+
}
|
|
2909
|
+
__name(dropContentsDuplicates, "dropContentsDuplicates");
|
|
2910
|
+
function isTitleCase(line) {
|
|
2911
|
+
if (line.length > 70 || /\.\s*$/.test(line)) {
|
|
2912
|
+
return false;
|
|
2913
|
+
}
|
|
2914
|
+
const words2 = line.split(/\s+/).filter((word) => /^[A-Za-z]{4,}/.test(word));
|
|
2915
|
+
return words2.length > 0 && words2.filter((word) => /^[A-Z]/.test(word)).length / words2.length >= 0.6;
|
|
2916
|
+
}
|
|
2917
|
+
__name(isTitleCase, "isTitleCase");
|
|
2918
|
+
function isPageFurniture(line) {
|
|
2919
|
+
return line.length < 50 && /^[\d\s\-–—.()]*(\d+\s*CFR|Federal Acquisition Regulation|Page \d+)/i.test(line);
|
|
2920
|
+
}
|
|
2921
|
+
__name(isPageFurniture, "isPageFurniture");
|
|
2922
|
+
function runningLines(texts) {
|
|
2923
|
+
const counts = /* @__PURE__ */ new Map();
|
|
2924
|
+
for (const text of texts) {
|
|
2925
|
+
const seen = /* @__PURE__ */ new Set();
|
|
2926
|
+
for (const raw of text.split("\n")) {
|
|
2927
|
+
const key = raw.trim().replace(/\d+/g, "#").toLowerCase();
|
|
2928
|
+
if (key && !seen.has(key)) {
|
|
2929
|
+
seen.add(key);
|
|
2930
|
+
counts.set(key, (counts.get(key) ?? 0) + 1);
|
|
2931
|
+
}
|
|
2932
|
+
}
|
|
2933
|
+
}
|
|
2934
|
+
const threshold = Math.max(3, Math.ceil(texts.length * 0.2));
|
|
2935
|
+
return new Set([
|
|
2936
|
+
...counts
|
|
2937
|
+
].filter(([, count]) => count >= threshold).map(([key]) => key));
|
|
2938
|
+
}
|
|
2939
|
+
__name(runningLines, "runningLines");
|
|
2940
|
+
function dropRunningHeaders(headings, unitCount) {
|
|
2941
|
+
const counts = /* @__PURE__ */ new Map();
|
|
2942
|
+
for (const heading of headings) {
|
|
2943
|
+
const key = heading.title.toLowerCase();
|
|
2944
|
+
counts.set(key, (counts.get(key) ?? 0) + 1);
|
|
2945
|
+
}
|
|
2946
|
+
const seen = /* @__PURE__ */ new Set();
|
|
2947
|
+
const threshold = Math.max(3, Math.ceil(unitCount * 0.2));
|
|
2948
|
+
return headings.filter((heading) => {
|
|
2949
|
+
const key = heading.title.toLowerCase();
|
|
2950
|
+
if ((counts.get(key) ?? 0) >= threshold) {
|
|
2951
|
+
if (seen.has(key)) {
|
|
2952
|
+
return false;
|
|
2953
|
+
}
|
|
2954
|
+
}
|
|
2955
|
+
seen.add(key);
|
|
2956
|
+
return true;
|
|
2957
|
+
});
|
|
2958
|
+
}
|
|
2959
|
+
__name(dropRunningHeaders, "dropRunningHeaders");
|
|
2960
|
+
|
|
2961
|
+
// src/tree/llm.ts
|
|
2962
|
+
var TreeBudgetExceededError = class extends Error {
|
|
2963
|
+
static {
|
|
2964
|
+
__name(this, "TreeBudgetExceededError");
|
|
2965
|
+
}
|
|
2966
|
+
limit;
|
|
2967
|
+
constructor(limit) {
|
|
2968
|
+
super(`tree LLM budget exhausted (${limit})`), this.limit = limit;
|
|
2969
|
+
this.name = "TreeBudgetExceededError";
|
|
2970
|
+
}
|
|
2971
|
+
};
|
|
2972
|
+
function estimateTokens(text) {
|
|
2973
|
+
return Math.ceil(text.length / 4);
|
|
2974
|
+
}
|
|
2975
|
+
__name(estimateTokens, "estimateTokens");
|
|
2976
|
+
var BudgetedLlm = class {
|
|
2977
|
+
static {
|
|
2978
|
+
__name(this, "BudgetedLlm");
|
|
2979
|
+
}
|
|
2980
|
+
llm;
|
|
2981
|
+
budget;
|
|
2982
|
+
outerSignal;
|
|
2983
|
+
calls = 0;
|
|
2984
|
+
inputTokens = 0;
|
|
2985
|
+
outputTokens = 0;
|
|
2986
|
+
exhausted = false;
|
|
2987
|
+
deadline;
|
|
2988
|
+
constructor(llm, budget, outerSignal) {
|
|
2989
|
+
this.llm = llm;
|
|
2990
|
+
this.budget = budget;
|
|
2991
|
+
this.outerSignal = outerSignal;
|
|
2992
|
+
this.deadline = budget.timeoutMs !== void 0 ? Date.now() + budget.timeoutMs : Number.POSITIVE_INFINITY;
|
|
2993
|
+
}
|
|
2994
|
+
/** Time left before the deadline, ms (`Infinity` without one). */
|
|
2995
|
+
remainingMs() {
|
|
2996
|
+
return this.deadline - Date.now();
|
|
2997
|
+
}
|
|
2998
|
+
/** Would a call with this prompt fit? Does not reserve anything; a `false` marks the budget exhausted. */
|
|
2999
|
+
fits(prompt, maxOutputTokens) {
|
|
3000
|
+
const fits = this.limitHit(prompt, maxOutputTokens) === void 0;
|
|
3001
|
+
if (!fits) {
|
|
3002
|
+
this.exhausted = true;
|
|
3003
|
+
}
|
|
3004
|
+
return fits;
|
|
3005
|
+
}
|
|
3006
|
+
limitHit(prompt, maxOutputTokens) {
|
|
3007
|
+
if (this.calls >= this.budget.maxCalls) {
|
|
3008
|
+
return "calls";
|
|
3009
|
+
}
|
|
3010
|
+
if (this.inputTokens + estimateTokens(prompt) > this.budget.maxInputTokens) {
|
|
3011
|
+
return "inputTokens";
|
|
3012
|
+
}
|
|
3013
|
+
if (this.outputTokens + Math.min(maxOutputTokens, 16) > this.budget.maxOutputTokens) {
|
|
3014
|
+
return "outputTokens";
|
|
3015
|
+
}
|
|
3016
|
+
if (this.remainingMs() <= 0 || this.outerSignal?.aborted) {
|
|
3017
|
+
return "timeout";
|
|
3018
|
+
}
|
|
3019
|
+
return void 0;
|
|
3020
|
+
}
|
|
3021
|
+
/**
|
|
3022
|
+
* Make one call, or throw {@link TreeBudgetExceededError} without calling. Provider errors
|
|
3023
|
+
* propagate (after being counted). Never retries.
|
|
3024
|
+
*/
|
|
3025
|
+
async call(task, system, prompt, maxOutputTokens) {
|
|
3026
|
+
const limit = this.limitHit(system + prompt, maxOutputTokens);
|
|
3027
|
+
if (limit !== void 0) {
|
|
3028
|
+
this.exhausted = true;
|
|
3029
|
+
throw new TreeBudgetExceededError(limit);
|
|
3030
|
+
}
|
|
3031
|
+
const outputCap = Math.max(16, Math.min(maxOutputTokens, this.budget.maxOutputTokens - this.outputTokens));
|
|
3032
|
+
const timeoutMs = Math.min(this.budget.callTimeoutMs ?? Number.POSITIVE_INFINITY, this.remainingMs());
|
|
3033
|
+
const reservedInput = estimateTokens(system + prompt);
|
|
3034
|
+
this.inputTokens += reservedInput;
|
|
3035
|
+
this.outputTokens += outputCap;
|
|
3036
|
+
const controller = new AbortController();
|
|
3037
|
+
const onOuterAbort = /* @__PURE__ */ __name(() => controller.abort(this.outerSignal?.reason), "onOuterAbort");
|
|
3038
|
+
this.outerSignal?.addEventListener("abort", onOuterAbort, {
|
|
3039
|
+
once: true
|
|
3040
|
+
});
|
|
3041
|
+
let timer;
|
|
3042
|
+
this.calls += 1;
|
|
3043
|
+
try {
|
|
3044
|
+
const response = await Promise.race([
|
|
3045
|
+
this.llm({
|
|
3046
|
+
task,
|
|
3047
|
+
system,
|
|
3048
|
+
prompt,
|
|
3049
|
+
maxOutputTokens: outputCap,
|
|
3050
|
+
signal: controller.signal
|
|
3051
|
+
}),
|
|
3052
|
+
new Promise((_, reject) => {
|
|
3053
|
+
if (Number.isFinite(timeoutMs)) {
|
|
3054
|
+
timer = setTimeout(() => {
|
|
3055
|
+
controller.abort(new TreeBudgetExceededError("timeout"));
|
|
3056
|
+
reject(new TreeBudgetExceededError("timeout"));
|
|
3057
|
+
}, Math.max(0, timeoutMs));
|
|
3058
|
+
}
|
|
3059
|
+
})
|
|
3060
|
+
]);
|
|
3061
|
+
this.inputTokens += (response.usage?.inputTokens ?? reservedInput) - reservedInput;
|
|
3062
|
+
this.outputTokens += (response.usage?.outputTokens ?? estimateTokens(response.text)) - outputCap;
|
|
3063
|
+
return response.text;
|
|
3064
|
+
} catch (error) {
|
|
3065
|
+
this.outputTokens -= outputCap;
|
|
3066
|
+
if (error instanceof TreeBudgetExceededError) {
|
|
3067
|
+
this.exhausted = true;
|
|
3068
|
+
}
|
|
3069
|
+
throw error;
|
|
3070
|
+
} finally {
|
|
3071
|
+
if (timer !== void 0) {
|
|
3072
|
+
clearTimeout(timer);
|
|
3073
|
+
}
|
|
3074
|
+
this.outerSignal?.removeEventListener("abort", onOuterAbort);
|
|
3075
|
+
}
|
|
3076
|
+
}
|
|
3077
|
+
};
|
|
3078
|
+
function parseJsonReply(text) {
|
|
3079
|
+
const tryParse = /* @__PURE__ */ __name((candidate2) => {
|
|
3080
|
+
try {
|
|
3081
|
+
const value = JSON.parse(candidate2);
|
|
3082
|
+
return value !== null && typeof value === "object" && !Array.isArray(value) ? value : void 0;
|
|
3083
|
+
} catch {
|
|
3084
|
+
return void 0;
|
|
3085
|
+
}
|
|
3086
|
+
}, "tryParse");
|
|
3087
|
+
const trimmed = text.trim();
|
|
3088
|
+
const direct = tryParse(trimmed);
|
|
3089
|
+
if (direct !== void 0) {
|
|
3090
|
+
return direct;
|
|
3091
|
+
}
|
|
3092
|
+
const fenced = /```(?:json)?\s*([\s\S]*?)```/i.exec(trimmed);
|
|
3093
|
+
if (fenced?.[1] !== void 0) {
|
|
3094
|
+
const inner = tryParse(fenced[1].trim());
|
|
3095
|
+
if (inner !== void 0) {
|
|
3096
|
+
return inner;
|
|
3097
|
+
}
|
|
3098
|
+
}
|
|
3099
|
+
const start = trimmed.indexOf("{");
|
|
3100
|
+
const end = trimmed.lastIndexOf("}");
|
|
3101
|
+
if (start >= 0 && end > start) {
|
|
3102
|
+
return tryParse(trimmed.slice(start, end + 1));
|
|
3103
|
+
}
|
|
3104
|
+
return void 0;
|
|
3105
|
+
}
|
|
3106
|
+
__name(parseJsonReply, "parseJsonReply");
|
|
3107
|
+
function stringList(value) {
|
|
3108
|
+
if (!Array.isArray(value)) {
|
|
3109
|
+
return [];
|
|
3110
|
+
}
|
|
3111
|
+
return value.map((item) => typeof item === "number" ? String(item) : item).filter((item) => typeof item === "string");
|
|
3112
|
+
}
|
|
3113
|
+
__name(stringList, "stringList");
|
|
3114
|
+
var NOOP_SINK = {
|
|
3115
|
+
write() {
|
|
3116
|
+
},
|
|
3117
|
+
end() {
|
|
3118
|
+
},
|
|
3119
|
+
fail() {
|
|
3120
|
+
}
|
|
3121
|
+
};
|
|
3122
|
+
function treeLlmFromModelProvider(provider) {
|
|
3123
|
+
return async (request) => {
|
|
3124
|
+
const result = await provider.runTurn({
|
|
3125
|
+
system: request.system,
|
|
3126
|
+
messages: [
|
|
3127
|
+
{
|
|
3128
|
+
role: "user",
|
|
3129
|
+
content: request.prompt
|
|
3130
|
+
}
|
|
3131
|
+
],
|
|
3132
|
+
tools: [],
|
|
3133
|
+
sink: NOOP_SINK,
|
|
3134
|
+
abortSignal: request.signal
|
|
3135
|
+
});
|
|
3136
|
+
return {
|
|
3137
|
+
text: result.text,
|
|
3138
|
+
usage: {
|
|
3139
|
+
inputTokens: result.usage.inputTokens,
|
|
3140
|
+
outputTokens: result.usage.outputTokens
|
|
3141
|
+
}
|
|
3142
|
+
};
|
|
3143
|
+
};
|
|
3144
|
+
}
|
|
3145
|
+
__name(treeLlmFromModelProvider, "treeLlmFromModelProvider");
|
|
3146
|
+
function openAiChatTreeLlm(options) {
|
|
3147
|
+
const url = `${(options.baseUrl ?? "https://api.openai.com/v1").replace(/\/+$/, "")}/chat/completions`;
|
|
3148
|
+
const doFetch = options.fetch ?? fetch;
|
|
3149
|
+
return async (request) => {
|
|
3150
|
+
const json2 = request.task.kind !== "summarize";
|
|
3151
|
+
const response = await doFetch(url, {
|
|
3152
|
+
method: "POST",
|
|
3153
|
+
headers: {
|
|
3154
|
+
"content-type": "application/json",
|
|
3155
|
+
...options.apiKey ? {
|
|
3156
|
+
authorization: `Bearer ${options.apiKey}`
|
|
3157
|
+
} : {},
|
|
3158
|
+
...options.headers
|
|
3159
|
+
},
|
|
3160
|
+
body: JSON.stringify({
|
|
3161
|
+
model: options.model,
|
|
3162
|
+
messages: [
|
|
3163
|
+
{
|
|
3164
|
+
role: "system",
|
|
3165
|
+
content: request.system
|
|
3166
|
+
},
|
|
3167
|
+
{
|
|
3168
|
+
role: "user",
|
|
3169
|
+
content: request.prompt
|
|
3170
|
+
}
|
|
3171
|
+
],
|
|
3172
|
+
temperature: options.temperature ?? 0,
|
|
3173
|
+
max_tokens: request.maxOutputTokens,
|
|
3174
|
+
...options.jsonMode && json2 ? {
|
|
3175
|
+
response_format: {
|
|
3176
|
+
type: "json_object"
|
|
3177
|
+
}
|
|
3178
|
+
} : {},
|
|
3179
|
+
...options.body
|
|
3180
|
+
}),
|
|
3181
|
+
signal: request.signal
|
|
3182
|
+
});
|
|
3183
|
+
const body = await response.text();
|
|
3184
|
+
if (!response.ok) {
|
|
3185
|
+
throw new HttpModelError(response.status, `Chat completion failed (${response.status}): ${errorMessageOf(body)}`);
|
|
3186
|
+
}
|
|
3187
|
+
const parsed = parseJsonReply(body);
|
|
3188
|
+
const content = parsed?.choices?.[0]?.message?.content;
|
|
3189
|
+
if (typeof content !== "string") {
|
|
3190
|
+
throw new HttpModelError(502, "Chat completion returned no message content");
|
|
3191
|
+
}
|
|
3192
|
+
const usage = {};
|
|
3193
|
+
if (typeof parsed?.usage?.prompt_tokens === "number") {
|
|
3194
|
+
usage.inputTokens = parsed.usage.prompt_tokens;
|
|
3195
|
+
}
|
|
3196
|
+
if (typeof parsed?.usage?.completion_tokens === "number") {
|
|
3197
|
+
usage.outputTokens = parsed.usage.completion_tokens;
|
|
3198
|
+
}
|
|
3199
|
+
return {
|
|
3200
|
+
text: content,
|
|
3201
|
+
usage
|
|
3202
|
+
};
|
|
3203
|
+
};
|
|
3204
|
+
}
|
|
3205
|
+
__name(openAiChatTreeLlm, "openAiChatTreeLlm");
|
|
3206
|
+
function cachedTreeLlm(llm, cache = /* @__PURE__ */ new Map()) {
|
|
3207
|
+
return async (request) => {
|
|
3208
|
+
const key = stableHash(`${request.task.kind}\0${request.maxOutputTokens}\0${request.system}\0${request.prompt}`);
|
|
3209
|
+
const hit = await cache.get(key);
|
|
3210
|
+
if (hit !== void 0) {
|
|
3211
|
+
return {
|
|
3212
|
+
text: hit.text,
|
|
3213
|
+
usage: {
|
|
3214
|
+
inputTokens: 0,
|
|
3215
|
+
outputTokens: 0
|
|
3216
|
+
}
|
|
3217
|
+
};
|
|
3218
|
+
}
|
|
3219
|
+
const response = await llm(request);
|
|
3220
|
+
await cache.set(key, response);
|
|
3221
|
+
return response;
|
|
3222
|
+
};
|
|
3223
|
+
}
|
|
3224
|
+
__name(cachedTreeLlm, "cachedTreeLlm");
|
|
3225
|
+
function stableHash(text) {
|
|
3226
|
+
let h1 = 2166136261;
|
|
3227
|
+
let h2 = 16777619 ^ text.length;
|
|
3228
|
+
for (let i = 0; i < text.length; i++) {
|
|
3229
|
+
const c = text.charCodeAt(i);
|
|
3230
|
+
h1 = Math.imul(h1 ^ c, 16777619);
|
|
3231
|
+
h2 = Math.imul(h2 ^ c, 1540483477);
|
|
3232
|
+
h2 ^= h2 >>> 13;
|
|
3233
|
+
}
|
|
3234
|
+
h1 = Math.imul(h1 ^ h1 >>> 16, 2246822507);
|
|
3235
|
+
h2 = Math.imul(h2 ^ h2 >>> 15, 3266489909);
|
|
3236
|
+
return (h1 >>> 0).toString(16).padStart(8, "0") + (h2 >>> 0).toString(16).padStart(8, "0");
|
|
3237
|
+
}
|
|
3238
|
+
__name(stableHash, "stableHash");
|
|
3239
|
+
|
|
3240
|
+
// src/tree/build.ts
|
|
3241
|
+
var DEFAULT_BUDGET = {
|
|
3242
|
+
maxCalls: 60,
|
|
3243
|
+
maxInputTokens: 4e5,
|
|
3244
|
+
maxOutputTokens: 4e4,
|
|
3245
|
+
timeoutMs: 3e5,
|
|
3246
|
+
callTimeoutMs: 6e4
|
|
3247
|
+
};
|
|
3248
|
+
var STRUCTURE_SYSTEM = "You extract the table of contents of a document. You answer with JSON only, no prose.";
|
|
3249
|
+
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.";
|
|
3250
|
+
function toUnits(input) {
|
|
3251
|
+
if (input.pages !== void 0 && input.sections !== void 0) {
|
|
3252
|
+
throw new Error("buildDocumentTree: give `pages` or `sections`, not both");
|
|
3253
|
+
}
|
|
3254
|
+
if (input.sections !== void 0) {
|
|
3255
|
+
const units = input.sections.map((section, index) => ({
|
|
3256
|
+
index,
|
|
3257
|
+
text: section.text,
|
|
3258
|
+
...section.page !== void 0 ? {
|
|
3259
|
+
page: section.page
|
|
3260
|
+
} : {}
|
|
3261
|
+
}));
|
|
3262
|
+
const titled = input.sections.some((section) => section.title?.trim());
|
|
3263
|
+
const declared = titled ? input.sections.flatMap((section, unit) => section.title?.trim() ? [
|
|
3264
|
+
{
|
|
3265
|
+
title: section.title.trim(),
|
|
3266
|
+
level: Math.max(1, Math.min(6, section.level ?? 1)),
|
|
3267
|
+
unit,
|
|
3268
|
+
atUnitStart: true
|
|
3269
|
+
}
|
|
3270
|
+
] : []) : void 0;
|
|
3271
|
+
return {
|
|
3272
|
+
units,
|
|
3273
|
+
declared,
|
|
3274
|
+
kind: "sections"
|
|
3275
|
+
};
|
|
3276
|
+
}
|
|
3277
|
+
const pages2 = input.pages ?? [];
|
|
3278
|
+
return {
|
|
3279
|
+
units: pages2.map((text, index) => ({
|
|
3280
|
+
index,
|
|
3281
|
+
text,
|
|
3282
|
+
page: index + 1
|
|
3283
|
+
})),
|
|
3284
|
+
declared: void 0,
|
|
3285
|
+
kind: "pages"
|
|
3286
|
+
};
|
|
3287
|
+
}
|
|
3288
|
+
__name(toUnits, "toUnits");
|
|
3289
|
+
function outlineHeadings(input, units) {
|
|
3290
|
+
const headings = [];
|
|
3291
|
+
for (const entry of input.outline ?? []) {
|
|
3292
|
+
let unit = entry.unit;
|
|
3293
|
+
if (unit === void 0 && entry.page !== void 0) {
|
|
3294
|
+
unit = units.findIndex((candidate2) => candidate2.page === entry.page);
|
|
3295
|
+
if (unit < 0) {
|
|
3296
|
+
unit = entry.page - 1;
|
|
3297
|
+
}
|
|
3298
|
+
}
|
|
3299
|
+
if (unit === void 0 || unit < 0 || unit >= units.length || !entry.title.trim()) {
|
|
3300
|
+
continue;
|
|
3301
|
+
}
|
|
3302
|
+
headings.push({
|
|
3303
|
+
title: entry.title.trim(),
|
|
3304
|
+
level: Math.max(1, Math.min(6, Math.floor(entry.level))),
|
|
3305
|
+
unit,
|
|
3306
|
+
atUnitStart: false
|
|
3307
|
+
});
|
|
3308
|
+
}
|
|
3309
|
+
return headings.sort((a, b) => a.unit - b.unit);
|
|
3310
|
+
}
|
|
3311
|
+
__name(outlineHeadings, "outlineHeadings");
|
|
3312
|
+
function nestHeadings(headings, unitCount) {
|
|
3313
|
+
const roots = [];
|
|
3314
|
+
const stack = [];
|
|
3315
|
+
headings.forEach((heading, index) => {
|
|
3316
|
+
let end = unitCount - 1;
|
|
3317
|
+
for (let next = index + 1; next < headings.length; next++) {
|
|
3318
|
+
const candidate2 = headings[next];
|
|
3319
|
+
if (candidate2 !== void 0 && candidate2.level <= heading.level) {
|
|
3320
|
+
end = candidate2.atUnitStart ? candidate2.unit - 1 : candidate2.unit;
|
|
3321
|
+
break;
|
|
3322
|
+
}
|
|
3323
|
+
}
|
|
3324
|
+
const node = {
|
|
3325
|
+
title: heading.title,
|
|
3326
|
+
unitStart: heading.unit,
|
|
3327
|
+
unitEnd: Math.max(heading.unit, end),
|
|
3328
|
+
children: []
|
|
3329
|
+
};
|
|
3330
|
+
while (stack.length > 0 && (stack[stack.length - 1]?.level ?? 0) >= heading.level) {
|
|
3331
|
+
stack.pop();
|
|
3332
|
+
}
|
|
3333
|
+
const parent = stack[stack.length - 1];
|
|
3334
|
+
if (parent === void 0) {
|
|
3335
|
+
roots.push(node);
|
|
3336
|
+
} else {
|
|
3337
|
+
parent.node.children.push(node);
|
|
3338
|
+
parent.node.unitEnd = Math.max(parent.node.unitEnd, node.unitEnd);
|
|
3339
|
+
}
|
|
3340
|
+
stack.push({
|
|
3341
|
+
node,
|
|
3342
|
+
level: heading.level
|
|
3343
|
+
});
|
|
3344
|
+
});
|
|
3345
|
+
const first = roots[0];
|
|
3346
|
+
if (first !== void 0 && first.unitStart > 0) {
|
|
3347
|
+
roots.unshift({
|
|
3348
|
+
title: "Front matter",
|
|
3349
|
+
unitStart: 0,
|
|
3350
|
+
unitEnd: first.unitStart - 1,
|
|
3351
|
+
children: []
|
|
3352
|
+
});
|
|
3353
|
+
}
|
|
3354
|
+
return roots;
|
|
3355
|
+
}
|
|
3356
|
+
__name(nestHeadings, "nestHeadings");
|
|
3357
|
+
function firstLine(text, running) {
|
|
3358
|
+
const line = text.split("\n").map((candidate2) => candidate2.trim()).find((candidate2) => /[A-Za-z]{3,}/.test(candidate2) && !running.has(candidate2.replace(/\d+/g, "#").toLowerCase()));
|
|
3359
|
+
return line === void 0 ? "" : line.length > 80 ? `${line.slice(0, 77)}\u2026` : line;
|
|
3360
|
+
}
|
|
3361
|
+
__name(firstLine, "firstLine");
|
|
3362
|
+
function rangeLabel(units, start, end) {
|
|
3363
|
+
const first = units[start];
|
|
3364
|
+
const last = units[end];
|
|
3365
|
+
if (first?.page !== void 0 && last?.page !== void 0) {
|
|
3366
|
+
return first.page === last.page ? `Page ${first.page}` : `Pages ${first.page}\u2013${last.page}`;
|
|
3367
|
+
}
|
|
3368
|
+
return start === end ? `Section ${start + 1}` : `Sections ${start + 1}\u2013${end + 1}`;
|
|
3369
|
+
}
|
|
3370
|
+
__name(rangeLabel, "rangeLabel");
|
|
3371
|
+
function groupNodes(units, start, end, groupSize, maxChildren, levelsLeft, running) {
|
|
3372
|
+
let level = [];
|
|
3373
|
+
for (let from = start; from <= end; from += groupSize) {
|
|
3374
|
+
const to = Math.min(end, from + groupSize - 1);
|
|
3375
|
+
const lead = firstLine(units[from]?.text ?? "", running);
|
|
3376
|
+
level.push({
|
|
3377
|
+
title: `${rangeLabel(units, from, to)}${lead ? `: ${lead}` : ""}`,
|
|
3378
|
+
unitStart: from,
|
|
3379
|
+
unitEnd: to,
|
|
3380
|
+
children: []
|
|
3381
|
+
});
|
|
3382
|
+
}
|
|
3383
|
+
let depth = 1;
|
|
3384
|
+
while (level.length > maxChildren && depth < levelsLeft) {
|
|
3385
|
+
const parents = [];
|
|
3386
|
+
const per = Math.max(maxChildren, Math.ceil(level.length / maxChildren));
|
|
3387
|
+
for (let index = 0; index < level.length; index += per) {
|
|
3388
|
+
const children = level.slice(index, index + per);
|
|
3389
|
+
const from = children[0]?.unitStart ?? start;
|
|
3390
|
+
const to = children[children.length - 1]?.unitEnd ?? end;
|
|
3391
|
+
parents.push({
|
|
3392
|
+
title: rangeLabel(units, from, to),
|
|
3393
|
+
unitStart: from,
|
|
3394
|
+
unitEnd: to,
|
|
3395
|
+
children
|
|
3396
|
+
});
|
|
3397
|
+
}
|
|
3398
|
+
level = parents;
|
|
3399
|
+
depth += 1;
|
|
3400
|
+
}
|
|
3401
|
+
return level;
|
|
3402
|
+
}
|
|
3403
|
+
__name(groupNodes, "groupNodes");
|
|
3404
|
+
function descendantTitles(nodes, into = []) {
|
|
3405
|
+
for (const { node } of walkDraft(nodes)) {
|
|
3406
|
+
if (into.length < 30) {
|
|
3407
|
+
into.push(node.title);
|
|
3408
|
+
}
|
|
3409
|
+
}
|
|
3410
|
+
return into;
|
|
3411
|
+
}
|
|
3412
|
+
__name(descendantTitles, "descendantTitles");
|
|
3413
|
+
function shape(nodes, units, depth, options) {
|
|
3414
|
+
const shaped = nodes.map((node) => {
|
|
3415
|
+
if (depth >= options.maxDepth || node.unitStart === node.unitEnd && node.children.length > 0) {
|
|
3416
|
+
const keyItems = descendantTitles(node.children, [
|
|
3417
|
+
...node.keyItems ?? []
|
|
3418
|
+
]);
|
|
3419
|
+
return {
|
|
3420
|
+
...node,
|
|
3421
|
+
...keyItems.length > 0 ? {
|
|
3422
|
+
keyItems
|
|
3423
|
+
} : {},
|
|
3424
|
+
children: []
|
|
3425
|
+
};
|
|
3426
|
+
}
|
|
3427
|
+
let children = node.children;
|
|
3428
|
+
const only = children[0];
|
|
3429
|
+
if (children.length === 1 && only !== void 0 && only.unitStart === node.unitStart && only.unitEnd === node.unitEnd) {
|
|
3430
|
+
children = only.children;
|
|
3431
|
+
}
|
|
3432
|
+
if (children.length === 0 && node.unitEnd - node.unitStart + 1 > options.maxLeafUnits) {
|
|
3433
|
+
children = groupNodes(units, node.unitStart, node.unitEnd, options.groupSize, options.maxChildren, options.maxDepth - depth, options.running);
|
|
3434
|
+
}
|
|
3435
|
+
return {
|
|
3436
|
+
...node,
|
|
3437
|
+
children: shape(children, units, depth + 1, options)
|
|
3438
|
+
};
|
|
3439
|
+
});
|
|
3440
|
+
const merged = [];
|
|
3441
|
+
for (const node of shaped) {
|
|
3442
|
+
const previous = merged[merged.length - 1];
|
|
3443
|
+
if (previous !== void 0 && previous.children.length === 0 && node.children.length === 0 && previous.unitStart === node.unitStart && previous.unitEnd === node.unitEnd) {
|
|
3444
|
+
previous.keyItems = [
|
|
3445
|
+
...previous.keyItems ?? [],
|
|
3446
|
+
node.title,
|
|
3447
|
+
...node.keyItems ?? []
|
|
3448
|
+
].slice(0, 30);
|
|
3449
|
+
continue;
|
|
3450
|
+
}
|
|
3451
|
+
merged.push(node);
|
|
3452
|
+
}
|
|
3453
|
+
return merged;
|
|
3454
|
+
}
|
|
3455
|
+
__name(shape, "shape");
|
|
3456
|
+
function* walkDraft(nodes) {
|
|
3457
|
+
for (const node of nodes) {
|
|
3458
|
+
yield {
|
|
3459
|
+
node
|
|
3460
|
+
};
|
|
3461
|
+
yield* walkDraft(node.children);
|
|
3462
|
+
}
|
|
3463
|
+
}
|
|
3464
|
+
__name(walkDraft, "walkDraft");
|
|
3465
|
+
function spanText(units, start, end) {
|
|
3466
|
+
const parts = [];
|
|
3467
|
+
for (let index = Math.max(0, start); index <= end && index < units.length; index++) {
|
|
3468
|
+
parts.push(units[index]?.text ?? "");
|
|
3469
|
+
}
|
|
3470
|
+
return parts.join("\n\n");
|
|
3471
|
+
}
|
|
3472
|
+
__name(spanText, "spanText");
|
|
3473
|
+
function ownTexts(units, node) {
|
|
3474
|
+
const first = node.children[0];
|
|
3475
|
+
const end = first === void 0 ? node.unitEnd : Math.max(node.unitStart, first.unitStart - 1);
|
|
3476
|
+
const texts = [];
|
|
3477
|
+
for (let index = node.unitStart; index <= end && index < units.length; index++) {
|
|
3478
|
+
texts.push(units[index]?.text ?? "");
|
|
3479
|
+
}
|
|
3480
|
+
const head = texts[0];
|
|
3481
|
+
if (head !== void 0) {
|
|
3482
|
+
const probe = node.title.slice(0, 30);
|
|
3483
|
+
const at = probe.length >= 4 ? head.indexOf(probe) : -1;
|
|
3484
|
+
if (at > 0) {
|
|
3485
|
+
texts[0] = head.slice(at);
|
|
3486
|
+
}
|
|
3487
|
+
}
|
|
3488
|
+
return texts;
|
|
3489
|
+
}
|
|
3490
|
+
__name(ownTexts, "ownTexts");
|
|
3491
|
+
function excerptOf(texts, max) {
|
|
3492
|
+
const total = texts.reduce((sum, text) => sum + text.length, 0);
|
|
3493
|
+
if (total <= max) {
|
|
3494
|
+
return texts.join("\n\n");
|
|
3495
|
+
}
|
|
3496
|
+
const per = Math.max(400, Math.floor(max / texts.length));
|
|
3497
|
+
const parts = [];
|
|
3498
|
+
let used = 0;
|
|
3499
|
+
for (const text of texts) {
|
|
3500
|
+
if (used >= max) {
|
|
3501
|
+
break;
|
|
3502
|
+
}
|
|
3503
|
+
const part = text.length > per ? `${text.slice(0, per)}\u2026` : text;
|
|
3504
|
+
parts.push(part);
|
|
3505
|
+
used += part.length;
|
|
3506
|
+
}
|
|
3507
|
+
return parts.join("\n\n");
|
|
3508
|
+
}
|
|
3509
|
+
__name(excerptOf, "excerptOf");
|
|
3510
|
+
function collapse(text, max) {
|
|
3511
|
+
const flat = text.replace(/\s+/g, " ").trim();
|
|
3512
|
+
return flat.length > max ? `${flat.slice(0, max - 1)}\u2026` : flat;
|
|
3513
|
+
}
|
|
3514
|
+
__name(collapse, "collapse");
|
|
3515
|
+
async function llmHeadings(llm, units, options) {
|
|
3516
|
+
const perUnit = Math.max(500, Math.floor(options.windowChars / 4));
|
|
3517
|
+
const windows = [];
|
|
3518
|
+
let current = [];
|
|
3519
|
+
let size = 0;
|
|
3520
|
+
for (const unit of units) {
|
|
3521
|
+
const length = Math.min(unit.text.length, perUnit);
|
|
3522
|
+
if (current.length > 0 && size + length > options.windowChars) {
|
|
3523
|
+
windows.push(current);
|
|
3524
|
+
current = [];
|
|
3525
|
+
size = 0;
|
|
3526
|
+
}
|
|
3527
|
+
current.push(unit);
|
|
3528
|
+
size += length;
|
|
3529
|
+
}
|
|
3530
|
+
if (current.length > 0) {
|
|
3531
|
+
windows.push(current);
|
|
3532
|
+
}
|
|
3533
|
+
if (windows.length > options.maxCalls) {
|
|
3534
|
+
throw new Error(`document needs ${windows.length} structuring calls, more than maxStructureCalls (${options.maxCalls})`);
|
|
3535
|
+
}
|
|
3536
|
+
const headings = [];
|
|
3537
|
+
for (const window of windows) {
|
|
3538
|
+
const shown = window.map((unit) => ({
|
|
3539
|
+
unit: unit.index,
|
|
3540
|
+
...unit.page !== void 0 ? {
|
|
3541
|
+
page: unit.page
|
|
3542
|
+
} : {},
|
|
3543
|
+
text: unit.text.length > perUnit ? `${unit.text.slice(0, perUnit)}\u2026` : unit.text
|
|
3544
|
+
}));
|
|
3545
|
+
const body = shown.map((unit) => `<unit ${unit.unit}>
|
|
3546
|
+
${unit.text}
|
|
3547
|
+
</unit ${unit.unit}>`).join("\n");
|
|
3548
|
+
const prompt = `Below is part of a document, split into numbered units (pages or sections).
|
|
3549
|
+
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.
|
|
3550
|
+
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.
|
|
3551
|
+
|
|
3552
|
+
${body}
|
|
3553
|
+
|
|
3554
|
+
Reply with JSON only: {"sections": [{"title": "...", "level": 1, "unit": 0}]}`;
|
|
3555
|
+
const reply = parseJsonReply(await llm.call({
|
|
3556
|
+
kind: "structure",
|
|
3557
|
+
units: shown
|
|
3558
|
+
}, STRUCTURE_SYSTEM, prompt, 2e3));
|
|
3559
|
+
const sections = reply?.sections;
|
|
3560
|
+
if (!Array.isArray(sections)) {
|
|
3561
|
+
throw new Error("structuring reply had no `sections` list");
|
|
3562
|
+
}
|
|
3563
|
+
const first = window[0]?.index ?? 0;
|
|
3564
|
+
const last = window[window.length - 1]?.index ?? first;
|
|
3565
|
+
for (const entry of sections) {
|
|
3566
|
+
if (entry === null || typeof entry !== "object") {
|
|
3567
|
+
continue;
|
|
3568
|
+
}
|
|
3569
|
+
const { title, level, unit } = entry;
|
|
3570
|
+
const at = typeof unit === "number" ? unit : Number(unit);
|
|
3571
|
+
if (typeof title !== "string" || !title.trim() || !Number.isInteger(at) || at < first || at > last) {
|
|
3572
|
+
continue;
|
|
3573
|
+
}
|
|
3574
|
+
const previous = headings[headings.length - 1];
|
|
3575
|
+
if (previous !== void 0 && at < previous.unit) {
|
|
3576
|
+
continue;
|
|
3577
|
+
}
|
|
3578
|
+
headings.push({
|
|
3579
|
+
title: collapse(title, 160),
|
|
3580
|
+
level: Math.max(1, Math.min(6, Math.floor(typeof level === "number" ? level : 1))),
|
|
3581
|
+
unit: at,
|
|
3582
|
+
atUnitStart: false
|
|
3583
|
+
});
|
|
3584
|
+
}
|
|
3585
|
+
}
|
|
3586
|
+
return headings;
|
|
3587
|
+
}
|
|
3588
|
+
__name(llmHeadings, "llmHeadings");
|
|
3589
|
+
async function mapLimit(items, limit, work) {
|
|
3590
|
+
const results = new Array(items.length);
|
|
3591
|
+
let next = 0;
|
|
3592
|
+
const workers = Array.from({
|
|
3593
|
+
length: Math.max(1, Math.min(limit, items.length))
|
|
3594
|
+
}, async () => {
|
|
3595
|
+
while (next < items.length) {
|
|
3596
|
+
const index = next++;
|
|
3597
|
+
results[index] = await work(items[index], index);
|
|
3598
|
+
}
|
|
3599
|
+
});
|
|
3600
|
+
await Promise.all(workers);
|
|
3601
|
+
return results;
|
|
3602
|
+
}
|
|
3603
|
+
__name(mapLimit, "mapLimit");
|
|
3604
|
+
async function buildDocumentTree(input, options) {
|
|
3605
|
+
const started = Date.now();
|
|
3606
|
+
const { units, declared, kind } = toUnits(input);
|
|
3607
|
+
const maxDepth = Math.max(1, options.maxDepth ?? 4);
|
|
3608
|
+
const groupSize = Math.max(1, options.groupSize ?? 5);
|
|
3609
|
+
const maxLeafUnits = Math.max(groupSize, options.maxLeafUnits ?? 10);
|
|
3610
|
+
const maxChildren = Math.max(2, options.maxChildren ?? 12);
|
|
3611
|
+
const minHeadings = Math.max(1, options.minHeadings ?? 3);
|
|
3612
|
+
const llmEnabled = options.llm !== void 0;
|
|
3613
|
+
const summaries = options.summaries ?? llmEnabled;
|
|
3614
|
+
const describe = options.describe ?? llmEnabled;
|
|
3615
|
+
const llmStructure = options.llmStructure ?? llmEnabled;
|
|
3616
|
+
const detect = options.detectHeadings ?? true;
|
|
3617
|
+
const summaryMinChars = options.summaryMinChars ?? 600;
|
|
3618
|
+
const summaryInputChars = options.summaryInputChars ?? 8e3;
|
|
3619
|
+
const unitHashes = units.map((unit) => stableHash(unit.text));
|
|
3620
|
+
const fingerprint = stableHash(JSON.stringify({
|
|
3621
|
+
v: 1,
|
|
3622
|
+
kind,
|
|
3623
|
+
title: input.title ?? null,
|
|
3624
|
+
outline: input.outline ?? null,
|
|
3625
|
+
sections: input.sections?.map((section) => [
|
|
3626
|
+
section.title ?? null,
|
|
3627
|
+
section.level ?? null,
|
|
3628
|
+
section.page ?? null
|
|
3629
|
+
]) ?? null,
|
|
3630
|
+
units: unitHashes,
|
|
3631
|
+
maxDepth,
|
|
3632
|
+
groupSize,
|
|
3633
|
+
maxLeafUnits,
|
|
3634
|
+
maxChildren,
|
|
3635
|
+
minHeadings,
|
|
3636
|
+
detect,
|
|
3637
|
+
summaries,
|
|
3638
|
+
describe,
|
|
3639
|
+
llmStructure: llmStructure && llmEnabled,
|
|
3640
|
+
summaryMinChars,
|
|
3641
|
+
summaryInputChars
|
|
3642
|
+
}));
|
|
3643
|
+
if (options.previous?.fingerprint === fingerprint) {
|
|
3644
|
+
const { metadata: _metadata, source: _source, ...previous } = options.previous;
|
|
3645
|
+
const metadata = options.metadata ?? options.previous.metadata;
|
|
3646
|
+
const source = options.source ?? options.previous.source;
|
|
3647
|
+
return {
|
|
3648
|
+
tree: {
|
|
3649
|
+
...previous,
|
|
3650
|
+
...metadata !== void 0 ? {
|
|
3651
|
+
metadata
|
|
3652
|
+
} : {},
|
|
3653
|
+
...source !== void 0 ? {
|
|
3654
|
+
source
|
|
3655
|
+
} : {}
|
|
3656
|
+
},
|
|
3657
|
+
units,
|
|
3658
|
+
unchanged: true
|
|
3659
|
+
};
|
|
3660
|
+
}
|
|
3661
|
+
const budget = new BudgetedLlm(options.llm ?? (() => Promise.reject(new Error("no llm"))), {
|
|
3662
|
+
...DEFAULT_BUDGET,
|
|
3663
|
+
...options.budget
|
|
3664
|
+
}, options.signal);
|
|
3665
|
+
let fallbackReason;
|
|
3666
|
+
let structure = "flat";
|
|
3667
|
+
let headings = [];
|
|
3668
|
+
const running = runningLines(units.map((unit) => unit.text));
|
|
3669
|
+
const outline = outlineHeadings(input, units);
|
|
3670
|
+
if (outline.length >= Math.min(minHeadings, 2)) {
|
|
3671
|
+
structure = "outline";
|
|
3672
|
+
headings = outline;
|
|
3673
|
+
} else if (declared !== void 0 && declared.length >= 1) {
|
|
3674
|
+
structure = "headings";
|
|
3675
|
+
headings = declared;
|
|
3676
|
+
} else if (detect !== false && units.length > 0) {
|
|
3677
|
+
const found = detectHeadings(units.map((unit) => unit.text), detect === "markdown");
|
|
3678
|
+
if (found.length >= minHeadings) {
|
|
3679
|
+
structure = "headings";
|
|
3680
|
+
headings = found;
|
|
3681
|
+
}
|
|
3682
|
+
}
|
|
3683
|
+
if (structure === "flat" && llmStructure && llmEnabled && units.length > 1) {
|
|
3684
|
+
try {
|
|
3685
|
+
const found = await llmHeadings(budget, units, {
|
|
3686
|
+
windowChars: options.structureWindowChars ?? 24e3,
|
|
3687
|
+
maxCalls: Math.max(1, options.maxStructureCalls ?? 12)
|
|
3688
|
+
});
|
|
3689
|
+
if (found.length >= Math.min(minHeadings, 2)) {
|
|
3690
|
+
structure = "llm";
|
|
3691
|
+
headings = found;
|
|
3692
|
+
} else {
|
|
3693
|
+
fallbackReason = "LLM found no usable structure";
|
|
3694
|
+
}
|
|
3695
|
+
} catch (error) {
|
|
3696
|
+
fallbackReason = error instanceof TreeBudgetExceededError ? `LLM structuring stopped: ${error.message}` : `LLM structuring failed: ${error.message}`;
|
|
3697
|
+
}
|
|
3698
|
+
}
|
|
3699
|
+
const draft = structure === "flat" ? units.length === 0 ? [] : groupNodes(units, 0, units.length - 1, groupSize, maxChildren, maxDepth, running) : nestHeadings(headings, units.length);
|
|
3700
|
+
const shaped = shape(draft, units, 1, {
|
|
3701
|
+
maxDepth,
|
|
3702
|
+
maxLeafUnits,
|
|
3703
|
+
groupSize,
|
|
3704
|
+
maxChildren,
|
|
3705
|
+
running
|
|
3706
|
+
});
|
|
3707
|
+
let counter = 0;
|
|
3708
|
+
const finalize = /* @__PURE__ */ __name((nodes2) => nodes2.map((node) => {
|
|
3709
|
+
counter += 1;
|
|
3710
|
+
const id = String(counter).padStart(4, "0");
|
|
3711
|
+
const pageStart = units[node.unitStart]?.page;
|
|
3712
|
+
const pageEnd = units[node.unitEnd]?.page;
|
|
3713
|
+
return {
|
|
3714
|
+
id,
|
|
3715
|
+
title: node.title,
|
|
3716
|
+
unitStart: node.unitStart,
|
|
3717
|
+
unitEnd: node.unitEnd,
|
|
3718
|
+
...pageStart !== void 0 ? {
|
|
3719
|
+
pageStart
|
|
3720
|
+
} : {},
|
|
3721
|
+
...pageEnd !== void 0 ? {
|
|
3722
|
+
pageEnd
|
|
3723
|
+
} : {},
|
|
3724
|
+
...node.keyItems !== void 0 ? {
|
|
3725
|
+
keyItems: node.keyItems
|
|
3726
|
+
} : {},
|
|
3727
|
+
contentHash: stableHash(`${node.title}\0${unitHashes.slice(node.unitStart, node.unitEnd + 1).join(",")}`),
|
|
3728
|
+
children: finalize(node.children)
|
|
3729
|
+
};
|
|
3730
|
+
}), "finalize");
|
|
3731
|
+
const nodes = finalize(shaped);
|
|
3732
|
+
const previousSummaries = /* @__PURE__ */ new Map();
|
|
3733
|
+
if (options.previous !== void 0) {
|
|
3734
|
+
for (const { node } of walkTree(options.previous.nodes)) {
|
|
3735
|
+
if (node.summary !== void 0 && (node.summarySource === "llm" || node.summarySource === "reused")) {
|
|
3736
|
+
previousSummaries.set(node.contentHash, node.summary);
|
|
3737
|
+
}
|
|
3738
|
+
}
|
|
3739
|
+
}
|
|
3740
|
+
const breadthFirst = [];
|
|
3741
|
+
for (let queue = [
|
|
3742
|
+
...nodes
|
|
3743
|
+
]; queue.length > 0; ) {
|
|
3744
|
+
const node = queue.shift();
|
|
3745
|
+
breadthFirst.push(node);
|
|
3746
|
+
queue = queue.concat(node.children);
|
|
3747
|
+
}
|
|
3748
|
+
let reusedSummaries = 0;
|
|
3749
|
+
const pending = [];
|
|
3750
|
+
for (const node of breadthFirst) {
|
|
3751
|
+
const reused = previousSummaries.get(node.contentHash);
|
|
3752
|
+
const own = ownTexts(units, node).join("\n\n");
|
|
3753
|
+
if (reused !== void 0) {
|
|
3754
|
+
node.summary = reused;
|
|
3755
|
+
node.summarySource = "reused";
|
|
3756
|
+
reusedSummaries += 1;
|
|
3757
|
+
} else if (own.trim().length < summaryMinChars && node.children.length === 0) {
|
|
3758
|
+
const text = collapse(own, summaryMinChars);
|
|
3759
|
+
if (text) {
|
|
3760
|
+
node.summary = text;
|
|
3761
|
+
node.summarySource = "text";
|
|
3762
|
+
}
|
|
3763
|
+
} else if (summaries && llmEnabled) {
|
|
3764
|
+
pending.push(node);
|
|
3765
|
+
} else {
|
|
3766
|
+
const text = collapse(own, 300);
|
|
3767
|
+
if (text) {
|
|
3768
|
+
node.summary = text;
|
|
3769
|
+
node.summarySource = "text";
|
|
3770
|
+
}
|
|
3771
|
+
}
|
|
3772
|
+
}
|
|
3773
|
+
let summariesCut = false;
|
|
3774
|
+
await mapLimit(pending, Math.max(1, options.concurrency ?? 4), async (node) => {
|
|
3775
|
+
const own = ownTexts(units, node);
|
|
3776
|
+
const childTitles = [
|
|
3777
|
+
...node.children.map((child) => child.title),
|
|
3778
|
+
...node.keyItems ?? []
|
|
3779
|
+
].map((title) => `- ${title}`).join("\n");
|
|
3780
|
+
const excerpt = excerptOf(own, summaryInputChars);
|
|
3781
|
+
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.
|
|
3782
|
+
|
|
3783
|
+
Section title: ${node.title}
|
|
3784
|
+
${childTitles ? `Subsections:
|
|
3785
|
+
${childTitles}
|
|
3786
|
+
` : ""}
|
|
3787
|
+
Section text:
|
|
3788
|
+
${excerpt}
|
|
3789
|
+
|
|
3790
|
+
Reply with the summary only.`;
|
|
3791
|
+
if (!summariesCut && budget.fits(SUMMARY_SYSTEM + prompt, 300)) {
|
|
3792
|
+
try {
|
|
3793
|
+
const reply = await budget.call({
|
|
3794
|
+
kind: "summarize",
|
|
3795
|
+
title: node.title,
|
|
3796
|
+
text: excerpt
|
|
3797
|
+
}, SUMMARY_SYSTEM, prompt, 300);
|
|
3798
|
+
const json2 = parseJsonReply(reply);
|
|
3799
|
+
const summary = typeof json2?.summary === "string" ? json2.summary : reply;
|
|
3800
|
+
if (summary.trim()) {
|
|
3801
|
+
node.summary = collapse(summary, 800);
|
|
3802
|
+
node.summarySource = "llm";
|
|
3803
|
+
return;
|
|
3804
|
+
}
|
|
3805
|
+
} catch (error) {
|
|
3806
|
+
if (error instanceof TreeBudgetExceededError) {
|
|
3807
|
+
summariesCut = true;
|
|
3808
|
+
}
|
|
3809
|
+
fallbackReason ??= `summary failed: ${error.message}`;
|
|
3810
|
+
}
|
|
3811
|
+
} else {
|
|
3812
|
+
summariesCut = true;
|
|
3813
|
+
}
|
|
3814
|
+
const text = collapse(own.join("\n\n"), 300);
|
|
3815
|
+
if (text) {
|
|
3816
|
+
node.summary = text;
|
|
3817
|
+
node.summarySource = "text";
|
|
3818
|
+
}
|
|
3819
|
+
});
|
|
3820
|
+
if (summariesCut) {
|
|
3821
|
+
fallbackReason ??= "LLM budget exhausted before every node was summarized";
|
|
3822
|
+
}
|
|
3823
|
+
let description;
|
|
3824
|
+
const outlineForDescription = nodes.slice(0, 40).map((node) => ({
|
|
3825
|
+
title: node.title,
|
|
3826
|
+
...node.summary !== void 0 ? {
|
|
3827
|
+
summary: node.summary
|
|
3828
|
+
} : {}
|
|
3829
|
+
}));
|
|
3830
|
+
if (describe && llmEnabled && nodes.length > 0) {
|
|
3831
|
+
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.
|
|
3832
|
+
|
|
3833
|
+
${outlineForDescription.map((entry) => `- ${entry.title}${entry.summary ? `: ${entry.summary}` : ""}`).join("\n")}
|
|
3834
|
+
|
|
3835
|
+
Reply with the description only.`;
|
|
3836
|
+
if (budget.fits(SUMMARY_SYSTEM + prompt, 300)) {
|
|
3837
|
+
try {
|
|
3838
|
+
description = collapse(await budget.call({
|
|
3839
|
+
kind: "describe",
|
|
3840
|
+
...input.title !== void 0 ? {
|
|
3841
|
+
title: input.title
|
|
3842
|
+
} : {},
|
|
3843
|
+
outline: outlineForDescription
|
|
3844
|
+
}, SUMMARY_SYSTEM, prompt, 300), 1e3);
|
|
3845
|
+
} catch (error) {
|
|
3846
|
+
fallbackReason ??= `description failed: ${error.message}`;
|
|
3847
|
+
}
|
|
3848
|
+
}
|
|
3849
|
+
}
|
|
3850
|
+
if (description === void 0 && options.previous?.description !== void 0 && reusedSummaries > 0) {
|
|
3851
|
+
description = options.previous.description;
|
|
3852
|
+
}
|
|
3853
|
+
const tree = {
|
|
3854
|
+
documentId: options.documentId,
|
|
3855
|
+
version: 1,
|
|
3856
|
+
...input.title !== void 0 ? {
|
|
3857
|
+
title: input.title
|
|
3858
|
+
} : {},
|
|
3859
|
+
...description ? {
|
|
3860
|
+
description
|
|
3861
|
+
} : {},
|
|
3862
|
+
structure,
|
|
3863
|
+
fingerprint,
|
|
3864
|
+
unitCount: units.length,
|
|
3865
|
+
nodes,
|
|
3866
|
+
stats: {
|
|
3867
|
+
llmCalls: budget.calls,
|
|
3868
|
+
inputTokens: budget.inputTokens,
|
|
3869
|
+
outputTokens: budget.outputTokens,
|
|
3870
|
+
durationMs: Date.now() - started,
|
|
3871
|
+
...fallbackReason !== void 0 ? {
|
|
3872
|
+
fallbackReason
|
|
3873
|
+
} : {},
|
|
3874
|
+
budgetExhausted: budget.exhausted,
|
|
3875
|
+
reusedSummaries
|
|
3876
|
+
},
|
|
3877
|
+
...options.metadata !== void 0 ? {
|
|
3878
|
+
metadata: options.metadata
|
|
3879
|
+
} : {},
|
|
3880
|
+
...options.source !== void 0 ? {
|
|
3881
|
+
source: options.source
|
|
3882
|
+
} : {},
|
|
3883
|
+
builtAt: (/* @__PURE__ */ new Date()).toISOString()
|
|
3884
|
+
};
|
|
3885
|
+
return {
|
|
3886
|
+
tree,
|
|
3887
|
+
units,
|
|
3888
|
+
unchanged: false
|
|
3889
|
+
};
|
|
3890
|
+
}
|
|
3891
|
+
__name(buildDocumentTree, "buildDocumentTree");
|
|
3892
|
+
|
|
3893
|
+
// src/tree/keyword-tree-llm.ts
|
|
3894
|
+
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(" "));
|
|
3895
|
+
function terms(text) {
|
|
3896
|
+
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)));
|
|
3897
|
+
}
|
|
3898
|
+
__name(terms, "terms");
|
|
3899
|
+
function overlap(query, text) {
|
|
3900
|
+
let score = 0;
|
|
3901
|
+
for (const term of terms(text)) {
|
|
3902
|
+
if (query.has(term)) {
|
|
3903
|
+
score += 1;
|
|
3904
|
+
}
|
|
3905
|
+
}
|
|
3906
|
+
return score;
|
|
3907
|
+
}
|
|
3908
|
+
__name(overlap, "overlap");
|
|
3909
|
+
function keywordTreeLlm() {
|
|
3910
|
+
return async (request) => {
|
|
3911
|
+
const { task } = request;
|
|
3912
|
+
let text;
|
|
3913
|
+
switch (task.kind) {
|
|
3914
|
+
case "navigate": {
|
|
3915
|
+
const query = terms(task.query);
|
|
3916
|
+
const ranked = task.candidates.map((candidate2, order) => ({
|
|
3917
|
+
candidate: candidate2,
|
|
3918
|
+
order,
|
|
3919
|
+
score: overlap(query, `${candidate2.title} ${candidate2.summary ?? ""}`)
|
|
3920
|
+
})).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);
|
|
3921
|
+
text = task.mode === "single-pass" ? JSON.stringify({
|
|
3922
|
+
thinking: "keyword overlap",
|
|
3923
|
+
read: ranked.map((c) => c.id)
|
|
3924
|
+
}) : JSON.stringify({
|
|
3925
|
+
thinking: "keyword overlap",
|
|
3926
|
+
expand: ranked.filter((c) => c.hasChildren).map((c) => c.id),
|
|
3927
|
+
read: ranked.filter((c) => !c.hasChildren).map((c) => c.id)
|
|
3928
|
+
});
|
|
3929
|
+
break;
|
|
3930
|
+
}
|
|
3931
|
+
case "select": {
|
|
3932
|
+
const query = terms(task.query);
|
|
3933
|
+
const ranked = task.documents.map((document) => ({
|
|
3934
|
+
id: document.id,
|
|
3935
|
+
score: overlap(query, `${document.title ?? ""} ${document.description ?? ""}`)
|
|
3936
|
+
})).filter((entry) => entry.score > 0).sort((a, b) => b.score - a.score).slice(0, task.maxDocuments).map((entry) => entry.id);
|
|
3937
|
+
text = JSON.stringify({
|
|
3938
|
+
thinking: "keyword overlap",
|
|
3939
|
+
documents: ranked
|
|
3940
|
+
});
|
|
3941
|
+
break;
|
|
3942
|
+
}
|
|
3943
|
+
case "summarize": {
|
|
3944
|
+
const flat = task.text.replace(/\s+/g, " ").trim();
|
|
3945
|
+
const sentence = /^.*?[.!?](\s|$)/.exec(flat)?.[0] ?? flat;
|
|
3946
|
+
text = `${task.title}: ${sentence.slice(0, 200).trim()}`;
|
|
3947
|
+
break;
|
|
3948
|
+
}
|
|
3949
|
+
case "describe":
|
|
3950
|
+
text = `${task.title ?? "Document"} covering ${task.outline.map((entry) => entry.title).join(", ")}.`;
|
|
3951
|
+
break;
|
|
3952
|
+
case "structure": {
|
|
3953
|
+
const sections = [];
|
|
3954
|
+
for (const unit of task.units) {
|
|
3955
|
+
for (const line of unit.text.split("\n")) {
|
|
3956
|
+
const trimmed = line.trim();
|
|
3957
|
+
const markdown = /^(#{1,6})\s+(.+)$/.exec(trimmed);
|
|
3958
|
+
if (markdown?.[1] !== void 0 && markdown[2] !== void 0) {
|
|
3959
|
+
sections.push({
|
|
3960
|
+
title: markdown[2],
|
|
3961
|
+
level: markdown[1].length,
|
|
3962
|
+
unit: unit.unit
|
|
3963
|
+
});
|
|
3964
|
+
} else if (/^[A-Z][A-Z0-9 ,.'&-]{3,80}$/.test(trimmed) && /[A-Z]{3}/.test(trimmed)) {
|
|
3965
|
+
sections.push({
|
|
3966
|
+
title: trimmed,
|
|
3967
|
+
level: 1,
|
|
3968
|
+
unit: unit.unit
|
|
3969
|
+
});
|
|
3970
|
+
}
|
|
3971
|
+
}
|
|
3972
|
+
}
|
|
3973
|
+
text = JSON.stringify({
|
|
3974
|
+
sections
|
|
3975
|
+
});
|
|
3976
|
+
break;
|
|
3977
|
+
}
|
|
3978
|
+
}
|
|
3979
|
+
return {
|
|
3980
|
+
text,
|
|
3981
|
+
usage: {
|
|
3982
|
+
inputTokens: Math.ceil(request.prompt.length / 4),
|
|
3983
|
+
outputTokens: Math.ceil(text.length / 4)
|
|
3984
|
+
}
|
|
3985
|
+
};
|
|
3986
|
+
};
|
|
3987
|
+
}
|
|
3988
|
+
__name(keywordTreeLlm, "keywordTreeLlm");
|
|
3989
|
+
|
|
3990
|
+
// src/tree/store.ts
|
|
3991
|
+
function headerOf(tree) {
|
|
3992
|
+
return {
|
|
3993
|
+
documentId: tree.documentId,
|
|
3994
|
+
...tree.title !== void 0 ? {
|
|
3995
|
+
title: tree.title
|
|
3996
|
+
} : {},
|
|
3997
|
+
...tree.description !== void 0 ? {
|
|
3998
|
+
description: tree.description
|
|
3999
|
+
} : {},
|
|
4000
|
+
structure: tree.structure,
|
|
4001
|
+
fingerprint: tree.fingerprint,
|
|
4002
|
+
unitCount: tree.unitCount,
|
|
4003
|
+
...tree.metadata !== void 0 ? {
|
|
4004
|
+
metadata: tree.metadata
|
|
4005
|
+
} : {},
|
|
4006
|
+
...tree.source !== void 0 ? {
|
|
4007
|
+
source: tree.source
|
|
4008
|
+
} : {},
|
|
4009
|
+
builtAt: tree.builtAt
|
|
4010
|
+
};
|
|
4011
|
+
}
|
|
4012
|
+
__name(headerOf, "headerOf");
|
|
4013
|
+
function passes(tree, filter) {
|
|
4014
|
+
if (filter === void 0 || Object.keys(filter).length === 0) {
|
|
4015
|
+
return true;
|
|
4016
|
+
}
|
|
4017
|
+
if (filterMatchesNothing(filter)) {
|
|
4018
|
+
return false;
|
|
4019
|
+
}
|
|
4020
|
+
return matchesFilter(tree.metadata, filter);
|
|
4021
|
+
}
|
|
4022
|
+
__name(passes, "passes");
|
|
4023
|
+
var MemoryDocumentTreeStore = class {
|
|
4024
|
+
static {
|
|
4025
|
+
__name(this, "MemoryDocumentTreeStore");
|
|
4026
|
+
}
|
|
4027
|
+
trees = /* @__PURE__ */ new Map();
|
|
4028
|
+
async put(tree, units) {
|
|
4029
|
+
this.trees.set(tree.documentId, {
|
|
4030
|
+
tree: structuredClone(tree),
|
|
4031
|
+
units: units.map((unit) => ({
|
|
4032
|
+
...unit
|
|
4033
|
+
}))
|
|
4034
|
+
});
|
|
4035
|
+
}
|
|
4036
|
+
async get(documentId, filter) {
|
|
4037
|
+
const entry = this.trees.get(documentId);
|
|
4038
|
+
return entry !== void 0 && passes(entry.tree, filter) ? structuredClone(entry.tree) : void 0;
|
|
4039
|
+
}
|
|
4040
|
+
async getMany(documentIds, filter) {
|
|
4041
|
+
const trees = [];
|
|
4042
|
+
for (const id of new Set(documentIds)) {
|
|
4043
|
+
const tree = await this.get(id, filter);
|
|
4044
|
+
if (tree !== void 0) {
|
|
4045
|
+
trees.push(tree);
|
|
4046
|
+
}
|
|
4047
|
+
}
|
|
4048
|
+
return trees;
|
|
4049
|
+
}
|
|
4050
|
+
async list(options = {}) {
|
|
4051
|
+
const ids = options.documentIds !== void 0 ? new Set(options.documentIds) : void 0;
|
|
4052
|
+
const headers = [
|
|
4053
|
+
...this.trees.values()
|
|
4054
|
+
].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));
|
|
4055
|
+
return options.limit !== void 0 ? headers.slice(0, options.limit) : headers;
|
|
4056
|
+
}
|
|
4057
|
+
async readUnits(documentId, start, end) {
|
|
4058
|
+
const units = this.trees.get(documentId)?.units ?? [];
|
|
4059
|
+
return units.slice(Math.max(0, start), Math.max(0, end + 1)).map((unit) => ({
|
|
4060
|
+
...unit
|
|
4061
|
+
}));
|
|
4062
|
+
}
|
|
4063
|
+
async remove(documentId) {
|
|
4064
|
+
this.trees.delete(documentId);
|
|
4065
|
+
}
|
|
4066
|
+
};
|
|
4067
|
+
|
|
4068
|
+
// src/tree/pg-tree-store.ts
|
|
4069
|
+
var IDENTIFIER = /^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?$/;
|
|
4070
|
+
function json(value) {
|
|
4071
|
+
if (value === null || value === void 0) {
|
|
4072
|
+
return void 0;
|
|
4073
|
+
}
|
|
4074
|
+
return typeof value === "string" ? JSON.parse(value) : value;
|
|
4075
|
+
}
|
|
4076
|
+
__name(json, "json");
|
|
4077
|
+
var PgDocumentTreeStore = class {
|
|
4078
|
+
static {
|
|
4079
|
+
__name(this, "PgDocumentTreeStore");
|
|
4080
|
+
}
|
|
4081
|
+
client;
|
|
4082
|
+
table;
|
|
4083
|
+
unitsTable;
|
|
4084
|
+
insertBatchSize;
|
|
4085
|
+
constructor(client, options = {}) {
|
|
4086
|
+
this.client = toPgClient(client);
|
|
4087
|
+
this.table = options.table ?? "agent_rag_trees";
|
|
4088
|
+
this.unitsTable = options.unitsTable ?? `${this.table}_units`;
|
|
4089
|
+
for (const name of [
|
|
4090
|
+
this.table,
|
|
4091
|
+
this.unitsTable
|
|
4092
|
+
]) {
|
|
4093
|
+
if (!IDENTIFIER.test(name)) {
|
|
4094
|
+
throw new Error(`PgDocumentTreeStore: invalid table name ${JSON.stringify(name)}`);
|
|
4095
|
+
}
|
|
4096
|
+
}
|
|
4097
|
+
this.insertBatchSize = Math.max(1, Math.floor(options.insertBatchSize ?? 200));
|
|
4098
|
+
}
|
|
4099
|
+
/** The idempotent DDL {@link PgDocumentTreeStore.ensureSchema} runs, in order — for a migration. */
|
|
4100
|
+
schemaStatements() {
|
|
4101
|
+
const index = this.table.replace(".", "_");
|
|
4102
|
+
return [
|
|
4103
|
+
`CREATE TABLE IF NOT EXISTS ${this.table} (
|
|
4104
|
+
document_id TEXT PRIMARY KEY,
|
|
4105
|
+
title TEXT,
|
|
4106
|
+
description TEXT,
|
|
4107
|
+
structure TEXT NOT NULL,
|
|
4108
|
+
fingerprint TEXT NOT NULL,
|
|
4109
|
+
unit_count INTEGER NOT NULL,
|
|
4110
|
+
source TEXT,
|
|
4111
|
+
metadata JSONB,
|
|
4112
|
+
tree JSONB NOT NULL,
|
|
4113
|
+
built_at TIMESTAMPTZ NOT NULL
|
|
4114
|
+
)`,
|
|
4115
|
+
`CREATE INDEX IF NOT EXISTS ${index}_metadata_idx ON ${this.table} USING gin (metadata)`,
|
|
4116
|
+
`CREATE TABLE IF NOT EXISTS ${this.unitsTable} (
|
|
4117
|
+
document_id TEXT NOT NULL,
|
|
4118
|
+
idx INTEGER NOT NULL,
|
|
4119
|
+
page INTEGER,
|
|
4120
|
+
text TEXT NOT NULL,
|
|
4121
|
+
PRIMARY KEY (document_id, idx)
|
|
4122
|
+
)`
|
|
4123
|
+
];
|
|
4124
|
+
}
|
|
4125
|
+
async ensureSchema() {
|
|
4126
|
+
for (const statement of this.schemaStatements()) {
|
|
4127
|
+
await this.client.query(statement);
|
|
4128
|
+
}
|
|
4129
|
+
}
|
|
4130
|
+
/**
|
|
4131
|
+
* The SQL conditions a metadata `filter` compiles to (ANDed), appending bindings to `params`.
|
|
4132
|
+
* Override to map a filter onto your own columns. Return `['false']` to deny.
|
|
4133
|
+
*/
|
|
4134
|
+
whereConditions(filter, params) {
|
|
4135
|
+
return buildWhere(filter, params);
|
|
4136
|
+
}
|
|
4137
|
+
/** Replace the tree and its units in one transaction (when the client has one). */
|
|
4138
|
+
async put(tree, units) {
|
|
4139
|
+
const clean2 = stripNulBytes(tree);
|
|
4140
|
+
const { metadata, ...rest } = clean2;
|
|
4141
|
+
const write = /* @__PURE__ */ __name(async (client) => {
|
|
4142
|
+
await client.query(`DELETE FROM ${this.unitsTable} WHERE document_id = $1`, [
|
|
4143
|
+
clean2.documentId
|
|
4144
|
+
]);
|
|
4145
|
+
for (let start = 0; start < units.length; start += this.insertBatchSize) {
|
|
4146
|
+
const batch = units.slice(start, start + this.insertBatchSize);
|
|
4147
|
+
const params = [];
|
|
4148
|
+
const rows = batch.map((unit) => {
|
|
4149
|
+
params.push(clean2.documentId, unit.index, unit.page ?? null, stripNulBytes(unit.text));
|
|
4150
|
+
const at = params.length;
|
|
4151
|
+
return `($${at - 3}, $${at - 2}, $${at - 1}, $${at})`;
|
|
4152
|
+
});
|
|
4153
|
+
await client.query(`INSERT INTO ${this.unitsTable} (document_id, idx, page, text) VALUES ${rows.join(", ")}`, params);
|
|
4154
|
+
}
|
|
4155
|
+
await client.query(`INSERT INTO ${this.table}
|
|
4156
|
+
(document_id, title, description, structure, fingerprint, unit_count, source, metadata, tree, built_at)
|
|
4157
|
+
VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9::jsonb, $10)
|
|
4158
|
+
ON CONFLICT (document_id) DO UPDATE SET
|
|
4159
|
+
title = EXCLUDED.title, description = EXCLUDED.description, structure = EXCLUDED.structure,
|
|
4160
|
+
fingerprint = EXCLUDED.fingerprint, unit_count = EXCLUDED.unit_count, source = EXCLUDED.source,
|
|
4161
|
+
metadata = EXCLUDED.metadata, tree = EXCLUDED.tree, built_at = EXCLUDED.built_at`, [
|
|
4162
|
+
clean2.documentId,
|
|
4163
|
+
clean2.title ?? null,
|
|
4164
|
+
clean2.description ?? null,
|
|
4165
|
+
clean2.structure,
|
|
4166
|
+
clean2.fingerprint,
|
|
4167
|
+
clean2.unitCount,
|
|
4168
|
+
clean2.source ?? null,
|
|
4169
|
+
metadata === void 0 ? null : JSON.stringify(metadata),
|
|
4170
|
+
JSON.stringify(rest),
|
|
4171
|
+
clean2.builtAt
|
|
4172
|
+
]);
|
|
4173
|
+
}, "write");
|
|
4174
|
+
const transaction = this.client.transaction?.bind(this.client);
|
|
4175
|
+
await (transaction !== void 0 ? transaction(write) : write(this.client));
|
|
4176
|
+
}
|
|
4177
|
+
async get(documentId, filter) {
|
|
4178
|
+
return (await this.getMany([
|
|
4179
|
+
documentId
|
|
4180
|
+
], filter))[0];
|
|
4181
|
+
}
|
|
4182
|
+
async getMany(documentIds, filter) {
|
|
4183
|
+
if (documentIds.length === 0) {
|
|
4184
|
+
return [];
|
|
4185
|
+
}
|
|
4186
|
+
const params = [
|
|
4187
|
+
[
|
|
4188
|
+
...new Set(documentIds)
|
|
4189
|
+
]
|
|
4190
|
+
];
|
|
4191
|
+
const conditions = [
|
|
4192
|
+
"document_id = ANY($1::text[])",
|
|
4193
|
+
...this.whereConditions(filter, params)
|
|
4194
|
+
];
|
|
4195
|
+
const rows = await this.client.query(`SELECT document_id, metadata, tree FROM ${this.table} WHERE ${conditions.join(" AND ")}`, params);
|
|
4196
|
+
return rows.map((row) => {
|
|
4197
|
+
const tree = json(row.tree);
|
|
4198
|
+
const metadata = json(row.metadata);
|
|
4199
|
+
return {
|
|
4200
|
+
...tree,
|
|
4201
|
+
...metadata !== void 0 ? {
|
|
4202
|
+
metadata
|
|
4203
|
+
} : {}
|
|
4204
|
+
};
|
|
4205
|
+
});
|
|
4206
|
+
}
|
|
4207
|
+
async list(options = {}) {
|
|
4208
|
+
const params = [];
|
|
4209
|
+
const conditions = this.whereConditions(options.filter, params);
|
|
4210
|
+
if (options.documentIds !== void 0) {
|
|
4211
|
+
params.push([
|
|
4212
|
+
...options.documentIds
|
|
4213
|
+
]);
|
|
4214
|
+
conditions.push(`document_id = ANY($${params.length}::text[])`);
|
|
4215
|
+
}
|
|
4216
|
+
let limit = "";
|
|
4217
|
+
if (options.limit !== void 0) {
|
|
4218
|
+
params.push(Math.max(0, Math.floor(options.limit)));
|
|
4219
|
+
limit = `LIMIT $${params.length}`;
|
|
4220
|
+
}
|
|
4221
|
+
const rows = await this.client.query(`SELECT document_id, title, description, structure, fingerprint, unit_count, source, metadata, built_at
|
|
4222
|
+
FROM ${this.table}
|
|
4223
|
+
${conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : ""}
|
|
4224
|
+
ORDER BY document_id ${limit}`, params);
|
|
4225
|
+
return rows.map((row) => {
|
|
4226
|
+
const metadata = json(row.metadata);
|
|
4227
|
+
return headerOf({
|
|
4228
|
+
documentId: row.document_id,
|
|
4229
|
+
version: 1,
|
|
4230
|
+
...row.title !== null ? {
|
|
4231
|
+
title: row.title
|
|
4232
|
+
} : {},
|
|
4233
|
+
...row.description !== null ? {
|
|
4234
|
+
description: row.description
|
|
4235
|
+
} : {},
|
|
4236
|
+
structure: row.structure,
|
|
4237
|
+
fingerprint: row.fingerprint,
|
|
4238
|
+
unitCount: Number(row.unit_count),
|
|
4239
|
+
nodes: [],
|
|
4240
|
+
stats: {
|
|
4241
|
+
llmCalls: 0,
|
|
4242
|
+
inputTokens: 0,
|
|
4243
|
+
outputTokens: 0,
|
|
4244
|
+
durationMs: 0,
|
|
4245
|
+
budgetExhausted: false,
|
|
4246
|
+
reusedSummaries: 0
|
|
4247
|
+
},
|
|
4248
|
+
...metadata !== void 0 ? {
|
|
4249
|
+
metadata
|
|
4250
|
+
} : {},
|
|
4251
|
+
...row.source !== null ? {
|
|
4252
|
+
source: row.source
|
|
4253
|
+
} : {},
|
|
4254
|
+
builtAt: row.built_at instanceof Date ? row.built_at.toISOString() : String(row.built_at)
|
|
4255
|
+
});
|
|
4256
|
+
});
|
|
4257
|
+
}
|
|
4258
|
+
async readUnits(documentId, start, end) {
|
|
4259
|
+
const rows = await this.client.query(`SELECT idx, page, text FROM ${this.unitsTable}
|
|
4260
|
+
WHERE document_id = $1 AND idx BETWEEN $2 AND $3 ORDER BY idx`, [
|
|
4261
|
+
documentId,
|
|
4262
|
+
Math.max(0, start),
|
|
4263
|
+
end
|
|
4264
|
+
]);
|
|
4265
|
+
return rows.map((row) => ({
|
|
4266
|
+
index: Number(row.idx),
|
|
4267
|
+
text: row.text,
|
|
4268
|
+
...row.page !== null ? {
|
|
4269
|
+
page: Number(row.page)
|
|
4270
|
+
} : {}
|
|
4271
|
+
}));
|
|
4272
|
+
}
|
|
4273
|
+
async remove(documentId) {
|
|
4274
|
+
const work = /* @__PURE__ */ __name(async (client) => {
|
|
4275
|
+
await client.query(`DELETE FROM ${this.unitsTable} WHERE document_id = $1`, [
|
|
4276
|
+
documentId
|
|
4277
|
+
]);
|
|
4278
|
+
await client.query(`DELETE FROM ${this.table} WHERE document_id = $1`, [
|
|
4279
|
+
documentId
|
|
4280
|
+
]);
|
|
4281
|
+
}, "work");
|
|
4282
|
+
const transaction = this.client.transaction?.bind(this.client);
|
|
4283
|
+
await (transaction !== void 0 ? transaction(work) : work(this.client));
|
|
4284
|
+
}
|
|
4285
|
+
};
|
|
4286
|
+
|
|
4287
|
+
// src/tree/index-document.ts
|
|
4288
|
+
async function indexDocumentTree(input, options) {
|
|
4289
|
+
const { store, minUnits = 20, minChars = 6e4, ...build } = options;
|
|
4290
|
+
const { units } = toUnits(input);
|
|
4291
|
+
const chars = units.reduce((total, unit) => total + unit.text.length, 0);
|
|
4292
|
+
if (units.length < minUnits && chars < minChars) {
|
|
4293
|
+
await store.remove(options.documentId);
|
|
4294
|
+
return {
|
|
4295
|
+
status: "skipped",
|
|
4296
|
+
reason: `${units.length} units / ${chars} chars is below the threshold (${minUnits} units or ${minChars} chars)`
|
|
4297
|
+
};
|
|
4298
|
+
}
|
|
4299
|
+
const previous = await store.get(options.documentId);
|
|
4300
|
+
const built = await buildDocumentTree(input, {
|
|
4301
|
+
...build,
|
|
4302
|
+
...previous !== void 0 ? {
|
|
4303
|
+
previous
|
|
4304
|
+
} : {}
|
|
4305
|
+
});
|
|
4306
|
+
if (built.unchanged) {
|
|
4307
|
+
const metadataMoved = JSON.stringify(previous?.metadata ?? null) !== JSON.stringify(built.tree.metadata ?? null) || previous?.source !== built.tree.source;
|
|
4308
|
+
if (metadataMoved) {
|
|
4309
|
+
await store.put(built.tree, built.units);
|
|
4310
|
+
}
|
|
4311
|
+
return {
|
|
4312
|
+
status: "unchanged",
|
|
4313
|
+
tree: built.tree
|
|
4314
|
+
};
|
|
4315
|
+
}
|
|
4316
|
+
await store.put(built.tree, built.units);
|
|
4317
|
+
return {
|
|
4318
|
+
status: "built",
|
|
4319
|
+
tree: built.tree
|
|
4320
|
+
};
|
|
4321
|
+
}
|
|
4322
|
+
__name(indexDocumentTree, "indexDocumentTree");
|
|
4323
|
+
|
|
4324
|
+
// src/tree/navigate.ts
|
|
4325
|
+
var NAVIGATE_SYSTEM = "You navigate the table of contents of a document to find where a question is answered. You answer with JSON only.";
|
|
4326
|
+
var DEFAULT_BUDGET2 = {
|
|
4327
|
+
maxCalls: 6,
|
|
4328
|
+
maxInputTokens: 6e4,
|
|
4329
|
+
maxOutputTokens: 4e3,
|
|
4330
|
+
timeoutMs: 6e4
|
|
4331
|
+
};
|
|
4332
|
+
function pages(node) {
|
|
4333
|
+
if (node.pageStart === void 0) {
|
|
4334
|
+
return void 0;
|
|
4335
|
+
}
|
|
4336
|
+
return node.pageEnd === void 0 || node.pageEnd === node.pageStart ? `p. ${node.pageStart}` : `pp. ${node.pageStart}-${node.pageEnd}`;
|
|
4337
|
+
}
|
|
4338
|
+
__name(pages, "pages");
|
|
4339
|
+
function candidate(node, summaryChars, depth) {
|
|
4340
|
+
const shownPages = pages(node);
|
|
4341
|
+
let summary = node.summary;
|
|
4342
|
+
if (summary !== void 0 && summary.length > summaryChars) {
|
|
4343
|
+
summary = `${summary.slice(0, Math.max(0, summaryChars - 1))}\u2026`;
|
|
4344
|
+
}
|
|
4345
|
+
if (node.keyItems !== void 0 && node.keyItems.length > 0 && summaryChars > 0) {
|
|
4346
|
+
summary = `${summary ?? ""}${summary ? " " : ""}Includes: ${node.keyItems.slice(0, 12).join("; ")}`;
|
|
4347
|
+
}
|
|
4348
|
+
return {
|
|
4349
|
+
id: node.id,
|
|
4350
|
+
title: node.title,
|
|
4351
|
+
...summary && summaryChars > 0 ? {
|
|
4352
|
+
summary
|
|
4353
|
+
} : {},
|
|
4354
|
+
...shownPages !== void 0 ? {
|
|
4355
|
+
pages: shownPages
|
|
4356
|
+
} : {},
|
|
4357
|
+
...depth !== void 0 ? {
|
|
4358
|
+
depth
|
|
4359
|
+
} : {
|
|
4360
|
+
hasChildren: node.children.length > 0
|
|
4361
|
+
}
|
|
4362
|
+
};
|
|
4363
|
+
}
|
|
4364
|
+
__name(candidate, "candidate");
|
|
4365
|
+
function render(candidates) {
|
|
4366
|
+
return candidates.map((entry) => {
|
|
4367
|
+
const indent = " ".repeat(entry.depth ?? 0);
|
|
4368
|
+
const meta = [
|
|
4369
|
+
entry.pages,
|
|
4370
|
+
entry.hasChildren ? "has subsections" : void 0
|
|
4371
|
+
].filter(Boolean).join(", ");
|
|
4372
|
+
return `${indent}[${entry.id}] ${entry.title}${meta ? ` (${meta})` : ""}${entry.summary ? ` \u2014 ${entry.summary}` : ""}`;
|
|
4373
|
+
}).join("\n");
|
|
4374
|
+
}
|
|
4375
|
+
__name(render, "render");
|
|
4376
|
+
function fit(nodes, maxChars) {
|
|
4377
|
+
let last;
|
|
4378
|
+
for (const summaryChars of [
|
|
4379
|
+
400,
|
|
4380
|
+
160,
|
|
4381
|
+
0
|
|
4382
|
+
]) {
|
|
4383
|
+
const candidates = nodes.map(({ node, depth }) => candidate(node, summaryChars, depth));
|
|
4384
|
+
const text = render(candidates);
|
|
4385
|
+
last = {
|
|
4386
|
+
candidates,
|
|
4387
|
+
text
|
|
4388
|
+
};
|
|
4389
|
+
if (text.length <= maxChars) {
|
|
4390
|
+
break;
|
|
4391
|
+
}
|
|
4392
|
+
}
|
|
4393
|
+
return last;
|
|
4394
|
+
}
|
|
4395
|
+
__name(fit, "fit");
|
|
4396
|
+
var TreeNavigationRetriever = class {
|
|
4397
|
+
static {
|
|
4398
|
+
__name(this, "TreeNavigationRetriever");
|
|
4399
|
+
}
|
|
4400
|
+
options;
|
|
4401
|
+
store;
|
|
4402
|
+
llm;
|
|
4403
|
+
constructor(options) {
|
|
4404
|
+
this.options = options;
|
|
4405
|
+
this.store = options.store;
|
|
4406
|
+
this.llm = options.llm;
|
|
4407
|
+
}
|
|
4408
|
+
describeRetrieval() {
|
|
4409
|
+
return {
|
|
4410
|
+
retriever: "tree-navigation"
|
|
4411
|
+
};
|
|
4412
|
+
}
|
|
4413
|
+
async retrieve(query, options = {}) {
|
|
4414
|
+
const result = await this.search(query, options.filter !== void 0 ? {
|
|
4415
|
+
filter: options.filter
|
|
4416
|
+
} : {});
|
|
4417
|
+
return result.passages.slice(0, options.topK ?? result.passages.length);
|
|
4418
|
+
}
|
|
4419
|
+
/** {@link retrieve} with the audit trail: which documents, which nodes, why. */
|
|
4420
|
+
async search(query, options = {}) {
|
|
4421
|
+
const maxDocuments = Math.max(1, this.options.maxDocuments ?? 3);
|
|
4422
|
+
const maxCandidates = Math.max(maxDocuments, this.options.maxSelectionCandidates ?? 200);
|
|
4423
|
+
const headers = await this.store.list({
|
|
4424
|
+
...options.filter !== void 0 ? {
|
|
4425
|
+
filter: options.filter
|
|
4426
|
+
} : {},
|
|
4427
|
+
limit: maxCandidates
|
|
4428
|
+
});
|
|
4429
|
+
if (headers.length <= maxDocuments) {
|
|
4430
|
+
return this.navigate(query, headers.map((header) => header.documentId), options);
|
|
4431
|
+
}
|
|
4432
|
+
const budget = new BudgetedLlm(this.llm, {
|
|
4433
|
+
...DEFAULT_BUDGET2,
|
|
4434
|
+
...this.options.budget,
|
|
4435
|
+
...options.budget
|
|
4436
|
+
}, options.signal);
|
|
4437
|
+
const documents = headers.map((header) => ({
|
|
4438
|
+
id: header.documentId,
|
|
4439
|
+
...header.title !== void 0 ? {
|
|
4440
|
+
title: header.title
|
|
4441
|
+
} : {},
|
|
4442
|
+
...header.description !== void 0 ? {
|
|
4443
|
+
description: header.description.slice(0, 600)
|
|
4444
|
+
} : {}
|
|
4445
|
+
}));
|
|
4446
|
+
const prompt = `Select the documents that may contain the answer to the question, most relevant first, at most ${maxDocuments}.
|
|
4447
|
+
|
|
4448
|
+
Question: ${query}
|
|
4449
|
+
|
|
4450
|
+
Documents:
|
|
4451
|
+
${documents.map((document) => `[${document.id}] ${document.title ?? ""}${document.description ? ` \u2014 ${document.description}` : ""}`).join("\n")}
|
|
4452
|
+
|
|
4453
|
+
Reply with JSON only: {"thinking": "<why>", "documents": ["<document id>", ...]} \u2014 an empty list if none is relevant.`;
|
|
4454
|
+
const known = new Set(documents.map((document) => document.id));
|
|
4455
|
+
try {
|
|
4456
|
+
const reply = parseJsonReply(await budget.call({
|
|
4457
|
+
kind: "select",
|
|
4458
|
+
query,
|
|
4459
|
+
documents,
|
|
4460
|
+
maxDocuments
|
|
4461
|
+
}, NAVIGATE_SYSTEM, prompt, 800));
|
|
4462
|
+
const selected = [
|
|
4463
|
+
...new Set(stringList(reply?.documents ?? reply?.answer).filter((id) => known.has(id)))
|
|
4464
|
+
].slice(0, maxDocuments);
|
|
4465
|
+
const result = await this.navigate(query, selected, options);
|
|
4466
|
+
const reasoning = typeof reply?.thinking === "string" ? reply.thinking.slice(0, 2e3) : void 0;
|
|
4467
|
+
return {
|
|
4468
|
+
...result,
|
|
4469
|
+
selection: {
|
|
4470
|
+
candidates: documents.length,
|
|
4471
|
+
selected,
|
|
4472
|
+
...reasoning !== void 0 ? {
|
|
4473
|
+
reasoning
|
|
4474
|
+
} : {}
|
|
4475
|
+
}
|
|
4476
|
+
};
|
|
4477
|
+
} catch (error) {
|
|
4478
|
+
return {
|
|
4479
|
+
passages: [],
|
|
4480
|
+
documents: [],
|
|
4481
|
+
selection: {
|
|
4482
|
+
candidates: documents.length,
|
|
4483
|
+
selected: [],
|
|
4484
|
+
error: error.message
|
|
4485
|
+
}
|
|
4486
|
+
};
|
|
4487
|
+
}
|
|
4488
|
+
}
|
|
4489
|
+
/**
|
|
4490
|
+
* Navigate the given documents (in parallel, up to `concurrency`) and return their passages in
|
|
4491
|
+
* the order given, within `maxContextChars`. Documents without a tree, or whose tree fails
|
|
4492
|
+
* `filter`, yield no passages and `stoppedBy: 'not-found'`.
|
|
4493
|
+
*/
|
|
4494
|
+
async navigate(query, documentIds, options = {}) {
|
|
4495
|
+
const ids = [
|
|
4496
|
+
...new Set(documentIds)
|
|
4497
|
+
];
|
|
4498
|
+
const trees = new Map((await this.store.getMany(ids, options.filter)).map((tree) => [
|
|
4499
|
+
tree.documentId,
|
|
4500
|
+
tree
|
|
4501
|
+
]));
|
|
4502
|
+
const navigations = await mapLimit(ids, Math.max(1, this.options.concurrency ?? 3), async (id) => {
|
|
4503
|
+
const tree = trees.get(id);
|
|
4504
|
+
if (tree === void 0) {
|
|
4505
|
+
return {
|
|
4506
|
+
navigation: {
|
|
4507
|
+
documentId: id,
|
|
4508
|
+
mode: "single-pass",
|
|
4509
|
+
steps: [],
|
|
4510
|
+
nodes: [],
|
|
4511
|
+
llmCalls: 0,
|
|
4512
|
+
inputTokens: 0,
|
|
4513
|
+
outputTokens: 0,
|
|
4514
|
+
stoppedBy: "not-found"
|
|
4515
|
+
},
|
|
4516
|
+
passages: []
|
|
4517
|
+
};
|
|
4518
|
+
}
|
|
4519
|
+
const navigation = await this.walk(query, tree, options);
|
|
4520
|
+
return {
|
|
4521
|
+
navigation,
|
|
4522
|
+
passages: await this.passagesOf(tree, navigation)
|
|
4523
|
+
};
|
|
4524
|
+
});
|
|
4525
|
+
const maxContextChars = this.options.maxContextChars ?? 48e3;
|
|
4526
|
+
let used = 0;
|
|
4527
|
+
const passages = [];
|
|
4528
|
+
for (const entry of navigations) {
|
|
4529
|
+
for (const passage of entry.passages) {
|
|
4530
|
+
if (used >= maxContextChars) {
|
|
4531
|
+
break;
|
|
4532
|
+
}
|
|
4533
|
+
const text = passage.text.slice(0, maxContextChars - used);
|
|
4534
|
+
used += text.length;
|
|
4535
|
+
passages.push({
|
|
4536
|
+
...passage,
|
|
4537
|
+
text
|
|
4538
|
+
});
|
|
4539
|
+
}
|
|
4540
|
+
}
|
|
4541
|
+
return {
|
|
4542
|
+
passages,
|
|
4543
|
+
documents: navigations.map((entry) => entry.navigation)
|
|
4544
|
+
};
|
|
4545
|
+
}
|
|
4546
|
+
async walk(query, tree, options) {
|
|
4547
|
+
const budget = new BudgetedLlm(this.llm, {
|
|
4548
|
+
...DEFAULT_BUDGET2,
|
|
4549
|
+
...this.options.budget,
|
|
4550
|
+
...options.budget
|
|
4551
|
+
}, options.signal);
|
|
4552
|
+
const maxNodes = Math.max(1, this.options.maxNodes ?? 5);
|
|
4553
|
+
const singlePassMaxChars = this.options.singlePassMaxChars ?? 4e4;
|
|
4554
|
+
const index = indexTree(tree);
|
|
4555
|
+
const whole = [
|
|
4556
|
+
...walkTree(tree.nodes)
|
|
4557
|
+
].map(({ node, path }) => ({
|
|
4558
|
+
node,
|
|
4559
|
+
depth: path.length
|
|
4560
|
+
}));
|
|
4561
|
+
const outline = fit(whole, Number.POSITIVE_INFINITY);
|
|
4562
|
+
const title = tree.title ?? tree.source ?? tree.documentId;
|
|
4563
|
+
const steps = [];
|
|
4564
|
+
const chosen = [];
|
|
4565
|
+
let stoppedBy;
|
|
4566
|
+
let error;
|
|
4567
|
+
const mode = outline.text.length <= singlePassMaxChars ? "single-pass" : "beam";
|
|
4568
|
+
const stop = /* @__PURE__ */ __name((caught) => {
|
|
4569
|
+
if (caught instanceof TreeBudgetExceededError) {
|
|
4570
|
+
stoppedBy = caught.limit === "timeout" ? "timeout" : "budget";
|
|
4571
|
+
} else {
|
|
4572
|
+
stoppedBy = "error";
|
|
4573
|
+
}
|
|
4574
|
+
error = caught.message;
|
|
4575
|
+
}, "stop");
|
|
4576
|
+
if (mode === "single-pass") {
|
|
4577
|
+
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.
|
|
4578
|
+
|
|
4579
|
+
Question: ${query}
|
|
4580
|
+
|
|
4581
|
+
Document: ${title}${tree.description ? `
|
|
4582
|
+
About: ${tree.description}` : ""}
|
|
4583
|
+
Table of contents ([id] title (pages) \u2014 summary):
|
|
4584
|
+
${outline.text}
|
|
4585
|
+
|
|
4586
|
+
Reply with JSON only: {"thinking": "<which sections answer and why>", "read": ["<id>", ...]} \u2014 an empty list if none does.`;
|
|
4587
|
+
try {
|
|
4588
|
+
const reply = requireJson(await budget.call({
|
|
4589
|
+
kind: "navigate",
|
|
4590
|
+
query,
|
|
4591
|
+
mode,
|
|
4592
|
+
candidates: outline.candidates,
|
|
4593
|
+
maxNodes
|
|
4594
|
+
}, NAVIGATE_SYSTEM, prompt, 1500));
|
|
4595
|
+
const read = [
|
|
4596
|
+
...new Set(stringList(reply?.read ?? reply?.node_list ?? reply?.nodes))
|
|
4597
|
+
].map((id) => resolveId(index, id)).filter((id) => id !== void 0).slice(0, maxNodes);
|
|
4598
|
+
chosen.push(...read);
|
|
4599
|
+
steps.push({
|
|
4600
|
+
step: 1,
|
|
4601
|
+
shown: outline.candidates.map((entry) => entry.id),
|
|
4602
|
+
expanded: [],
|
|
4603
|
+
read,
|
|
4604
|
+
...reasoningOf(reply)
|
|
4605
|
+
});
|
|
4606
|
+
} catch (caught) {
|
|
4607
|
+
stop(caught);
|
|
4608
|
+
}
|
|
4609
|
+
} else {
|
|
4610
|
+
const beamWidth = Math.max(1, this.options.beamWidth ?? 4);
|
|
4611
|
+
const maxSteps = Math.max(1, this.options.maxSteps ?? 4);
|
|
4612
|
+
let frontier = tree.nodes;
|
|
4613
|
+
let pendingExpand = [];
|
|
4614
|
+
for (let step = 1; frontier.length > 0 && chosen.length < maxNodes; step++) {
|
|
4615
|
+
if (step > maxSteps) {
|
|
4616
|
+
stoppedBy = "maxSteps";
|
|
4617
|
+
break;
|
|
4618
|
+
}
|
|
4619
|
+
const shown = fit(frontier.map((node) => ({
|
|
4620
|
+
node
|
|
4621
|
+
})), singlePassMaxChars);
|
|
4622
|
+
const path = index.get(frontier[0]?.id ?? "")?.path ?? [];
|
|
4623
|
+
const where = path.length > 0 ? `
|
|
4624
|
+
You are inside: ${path.map((node) => node.title).join(" > ")}` : "";
|
|
4625
|
+
const prompt = `Find where the question is answered in the document by walking its table of contents.
|
|
4626
|
+
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.
|
|
4627
|
+
|
|
4628
|
+
Question: ${query}
|
|
4629
|
+
|
|
4630
|
+
Document: ${title}${tree.description ? `
|
|
4631
|
+
About: ${tree.description}` : ""}${where}
|
|
4632
|
+
Sections ([id] title (pages) \u2014 summary):
|
|
4633
|
+
${shown.text}
|
|
4634
|
+
|
|
4635
|
+
Reply with JSON only: {"thinking": "<why>", "expand": ["<id>", ...], "read": ["<id>", ...]}`;
|
|
4636
|
+
let reply;
|
|
4637
|
+
try {
|
|
4638
|
+
reply = requireJson(await budget.call({
|
|
4639
|
+
kind: "navigate",
|
|
4640
|
+
query,
|
|
4641
|
+
mode,
|
|
4642
|
+
candidates: shown.candidates,
|
|
4643
|
+
maxNodes: beamWidth
|
|
4644
|
+
}, NAVIGATE_SYSTEM, prompt, 1200));
|
|
4645
|
+
} catch (caught) {
|
|
4646
|
+
stop(caught);
|
|
4647
|
+
break;
|
|
4648
|
+
}
|
|
4649
|
+
const onFrontier = new Set(frontier.map((node) => node.id));
|
|
4650
|
+
const pick = /* @__PURE__ */ __name((value) => [
|
|
4651
|
+
...new Set(stringList(value))
|
|
4652
|
+
].map((id) => resolveId(index, id)).filter((id) => id !== void 0 && onFrontier.has(id)), "pick");
|
|
4653
|
+
let expand = pick(reply?.expand);
|
|
4654
|
+
let read = pick(reply?.read);
|
|
4655
|
+
read = [
|
|
4656
|
+
...read,
|
|
4657
|
+
...expand.filter((id) => (index.get(id)?.node.children.length ?? 0) === 0)
|
|
4658
|
+
];
|
|
4659
|
+
expand = expand.filter((id) => (index.get(id)?.node.children.length ?? 0) > 0);
|
|
4660
|
+
const limited = [
|
|
4661
|
+
...expand,
|
|
4662
|
+
...read
|
|
4663
|
+
].slice(0, beamWidth);
|
|
4664
|
+
expand = expand.filter((id) => limited.includes(id));
|
|
4665
|
+
read = [
|
|
4666
|
+
...new Set(read.filter((id) => limited.includes(id)))
|
|
4667
|
+
];
|
|
4668
|
+
steps.push({
|
|
4669
|
+
step,
|
|
4670
|
+
shown: shown.candidates.map((entry) => entry.id),
|
|
4671
|
+
expanded: expand,
|
|
4672
|
+
read,
|
|
4673
|
+
...reasoningOf(reply)
|
|
4674
|
+
});
|
|
4675
|
+
chosen.push(...read);
|
|
4676
|
+
pendingExpand = expand;
|
|
4677
|
+
frontier = expand.flatMap((id) => index.get(id)?.node.children ?? []);
|
|
4678
|
+
if (frontier.length === 0) {
|
|
4679
|
+
pendingExpand = [];
|
|
4680
|
+
}
|
|
4681
|
+
}
|
|
4682
|
+
if (stoppedBy !== void 0) {
|
|
4683
|
+
chosen.push(...pendingExpand);
|
|
4684
|
+
}
|
|
4685
|
+
}
|
|
4686
|
+
const unique = [
|
|
4687
|
+
...new Set(chosen)
|
|
4688
|
+
];
|
|
4689
|
+
const specific = unique.filter((id) => !unique.some((other) => other !== id && index.get(other)?.path.some((node) => node.id === id)));
|
|
4690
|
+
const nodes = specific.slice(0, maxNodes).flatMap((id) => {
|
|
4691
|
+
const entry = index.get(id);
|
|
4692
|
+
if (entry === void 0) {
|
|
4693
|
+
return [];
|
|
4694
|
+
}
|
|
4695
|
+
const { node, path } = entry;
|
|
4696
|
+
return [
|
|
4697
|
+
{
|
|
4698
|
+
nodeId: node.id,
|
|
4699
|
+
title: node.title,
|
|
4700
|
+
path: [
|
|
4701
|
+
...path,
|
|
4702
|
+
node
|
|
4703
|
+
].map((step) => ({
|
|
4704
|
+
id: step.id,
|
|
4705
|
+
title: step.title
|
|
4706
|
+
})),
|
|
4707
|
+
...node.pageStart !== void 0 ? {
|
|
4708
|
+
pageStart: node.pageStart
|
|
4709
|
+
} : {},
|
|
4710
|
+
...node.pageEnd !== void 0 ? {
|
|
4711
|
+
pageEnd: node.pageEnd
|
|
4712
|
+
} : {},
|
|
4713
|
+
unitStart: node.unitStart,
|
|
4714
|
+
unitEnd: node.unitEnd
|
|
4715
|
+
}
|
|
4716
|
+
];
|
|
4717
|
+
});
|
|
4718
|
+
return {
|
|
4719
|
+
documentId: tree.documentId,
|
|
4720
|
+
...tree.title !== void 0 ? {
|
|
4721
|
+
title: tree.title
|
|
4722
|
+
} : {},
|
|
4723
|
+
mode,
|
|
4724
|
+
steps,
|
|
4725
|
+
nodes,
|
|
4726
|
+
llmCalls: budget.calls,
|
|
4727
|
+
inputTokens: budget.inputTokens,
|
|
4728
|
+
outputTokens: budget.outputTokens,
|
|
4729
|
+
...stoppedBy !== void 0 ? {
|
|
4730
|
+
stoppedBy
|
|
4731
|
+
} : {},
|
|
4732
|
+
...error !== void 0 ? {
|
|
4733
|
+
error
|
|
4734
|
+
} : {}
|
|
4735
|
+
};
|
|
4736
|
+
}
|
|
4737
|
+
async passagesOf(tree, navigation) {
|
|
4738
|
+
const maxPassageChars = this.options.maxPassageChars ?? 12e3;
|
|
4739
|
+
const maxUnits = Math.max(1, this.options.maxUnitsPerNode ?? 25);
|
|
4740
|
+
const reasoning = navigation.steps.map((step) => step.reasoning).filter((text) => text !== void 0).join("\n");
|
|
4741
|
+
const passages = [];
|
|
4742
|
+
for (const [rank, node] of navigation.nodes.entries()) {
|
|
4743
|
+
const units = await this.store.readUnits(tree.documentId, node.unitStart, Math.min(node.unitEnd, node.unitStart + maxUnits - 1));
|
|
4744
|
+
let text = units.map((unit) => unit.text).join("\n\n");
|
|
4745
|
+
if (text.length > maxPassageChars) {
|
|
4746
|
+
text = `${text.slice(0, maxPassageChars - 1)}\u2026`;
|
|
4747
|
+
}
|
|
4748
|
+
passages.push({
|
|
4749
|
+
id: `${tree.documentId}#node:${node.nodeId}`,
|
|
4750
|
+
text,
|
|
4751
|
+
// Rank within the document's choice; not a similarity.
|
|
4752
|
+
score: 1 / (rank + 1),
|
|
4753
|
+
source: tree.source ?? tree.title ?? tree.documentId,
|
|
4754
|
+
metadata: {
|
|
4755
|
+
...tree.metadata,
|
|
4756
|
+
documentId: tree.documentId,
|
|
4757
|
+
nodeId: node.nodeId,
|
|
4758
|
+
title: node.title,
|
|
4759
|
+
...node.pageStart !== void 0 ? {
|
|
4760
|
+
pageStart: node.pageStart
|
|
4761
|
+
} : {},
|
|
4762
|
+
...node.pageEnd !== void 0 ? {
|
|
4763
|
+
pageEnd: node.pageEnd
|
|
4764
|
+
} : {},
|
|
4765
|
+
path: node.path,
|
|
4766
|
+
retrieval: "tree-navigation",
|
|
4767
|
+
...reasoning ? {
|
|
4768
|
+
reasoning: reasoning.slice(0, 2e3)
|
|
4769
|
+
} : {}
|
|
4770
|
+
}
|
|
4771
|
+
});
|
|
4772
|
+
}
|
|
4773
|
+
return passages;
|
|
4774
|
+
}
|
|
4775
|
+
};
|
|
4776
|
+
function resolveId(index, raw) {
|
|
4777
|
+
const id = raw.trim().replace(/^\[|\]$/g, "");
|
|
4778
|
+
if (index.has(id)) {
|
|
4779
|
+
return id;
|
|
4780
|
+
}
|
|
4781
|
+
if (/^\d+$/.test(id)) {
|
|
4782
|
+
const padded = id.padStart(4, "0");
|
|
4783
|
+
if (index.has(padded)) {
|
|
4784
|
+
return padded;
|
|
4785
|
+
}
|
|
4786
|
+
}
|
|
4787
|
+
return void 0;
|
|
4788
|
+
}
|
|
4789
|
+
__name(resolveId, "resolveId");
|
|
4790
|
+
function requireJson(reply) {
|
|
4791
|
+
const json2 = parseJsonReply(reply);
|
|
4792
|
+
if (json2 === void 0) {
|
|
4793
|
+
throw new Error(`navigation reply was not JSON: ${JSON.stringify(reply.slice(0, 120))}`);
|
|
4794
|
+
}
|
|
4795
|
+
return json2;
|
|
4796
|
+
}
|
|
4797
|
+
__name(requireJson, "requireJson");
|
|
4798
|
+
function reasoningOf(reply) {
|
|
4799
|
+
const thinking = reply?.thinking ?? reply?.reasoning;
|
|
4800
|
+
return typeof thinking === "string" && thinking.trim() ? {
|
|
4801
|
+
reasoning: thinking.slice(0, 2e3)
|
|
4802
|
+
} : {};
|
|
4803
|
+
}
|
|
4804
|
+
__name(reasoningOf, "reasoningOf");
|
|
4805
|
+
|
|
4806
|
+
// src/tree/two-stage.ts
|
|
4807
|
+
function defaultDocumentIdOf(passage) {
|
|
4808
|
+
const fromMetadata = passage.metadata?.documentId;
|
|
4809
|
+
return typeof fromMetadata === "string" ? fromMetadata : documentIdOf(passage.id);
|
|
4810
|
+
}
|
|
4811
|
+
__name(defaultDocumentIdOf, "defaultDocumentIdOf");
|
|
4812
|
+
var TwoStageRetriever = class {
|
|
4813
|
+
static {
|
|
4814
|
+
__name(this, "TwoStageRetriever");
|
|
4815
|
+
}
|
|
4816
|
+
firstStage;
|
|
4817
|
+
navigator;
|
|
4818
|
+
options;
|
|
4819
|
+
constructor(firstStage, navigator, options = {}) {
|
|
4820
|
+
this.firstStage = firstStage;
|
|
4821
|
+
this.navigator = navigator;
|
|
4822
|
+
this.options = options;
|
|
4823
|
+
}
|
|
4824
|
+
describeRetrieval() {
|
|
4825
|
+
return {
|
|
4826
|
+
retriever: "two-stage"
|
|
4827
|
+
};
|
|
4828
|
+
}
|
|
4829
|
+
async retrieve(query, options = {}) {
|
|
4830
|
+
return (await this.search(query, options)).passages;
|
|
4831
|
+
}
|
|
4832
|
+
/** {@link retrieve} with the navigations that were run. */
|
|
4833
|
+
async search(query, options = {}) {
|
|
4834
|
+
const topK = options.topK ?? 5;
|
|
4835
|
+
const docOf = this.options.documentIdOf ?? defaultDocumentIdOf;
|
|
4836
|
+
const first = await this.firstStage.retrieve(query, {
|
|
4837
|
+
topK: Math.max(topK, this.options.fetchTopK ?? 20),
|
|
4838
|
+
...options.filter !== void 0 ? {
|
|
4839
|
+
filter: options.filter
|
|
4840
|
+
} : {}
|
|
4841
|
+
});
|
|
4842
|
+
const ranked = [];
|
|
4843
|
+
for (const passage of first) {
|
|
4844
|
+
const id = docOf(passage);
|
|
4845
|
+
if (!ranked.includes(id)) {
|
|
4846
|
+
ranked.push(id);
|
|
4847
|
+
}
|
|
4848
|
+
}
|
|
4849
|
+
const considered = ranked.slice(0, Math.max(1, this.options.considerDocuments ?? 3));
|
|
4850
|
+
const withTrees = new Set((await this.navigator.store.list({
|
|
4851
|
+
documentIds: considered,
|
|
4852
|
+
...options.filter !== void 0 ? {
|
|
4853
|
+
filter: options.filter
|
|
4854
|
+
} : {}
|
|
4855
|
+
})).map((header) => header.documentId));
|
|
4856
|
+
const toNavigate = considered.filter((id) => withTrees.has(id)).slice(0, Math.max(0, this.options.navigateDocuments ?? 2));
|
|
4857
|
+
const navigated = /* @__PURE__ */ new Map();
|
|
4858
|
+
let navigations = [];
|
|
4859
|
+
if (toNavigate.length > 0) {
|
|
4860
|
+
const result = await this.navigator.navigate(query, toNavigate, {
|
|
4861
|
+
...options.filter !== void 0 ? {
|
|
4862
|
+
filter: options.filter
|
|
4863
|
+
} : {},
|
|
4864
|
+
...this.options.budget !== void 0 ? {
|
|
4865
|
+
budget: this.options.budget
|
|
4866
|
+
} : {}
|
|
4867
|
+
});
|
|
4868
|
+
navigations = result.documents;
|
|
4869
|
+
for (const passage of result.passages) {
|
|
4870
|
+
const id = String(passage.metadata?.documentId ?? "");
|
|
4871
|
+
navigated.set(id, [
|
|
4872
|
+
...navigated.get(id) ?? [],
|
|
4873
|
+
passage
|
|
4874
|
+
]);
|
|
4875
|
+
}
|
|
4876
|
+
for (const navigation of navigations) {
|
|
4877
|
+
this.options.onNavigation?.(navigation);
|
|
4878
|
+
}
|
|
4879
|
+
}
|
|
4880
|
+
const merged = [];
|
|
4881
|
+
const emitted = /* @__PURE__ */ new Set();
|
|
4882
|
+
for (const passage of first) {
|
|
4883
|
+
const id = docOf(passage);
|
|
4884
|
+
const replacement = navigated.get(id);
|
|
4885
|
+
if (replacement === void 0 || replacement.length === 0) {
|
|
4886
|
+
merged.push(passage);
|
|
4887
|
+
continue;
|
|
4888
|
+
}
|
|
4889
|
+
if (!emitted.has(id)) {
|
|
4890
|
+
emitted.add(id);
|
|
4891
|
+
merged.push(...replacement);
|
|
4892
|
+
}
|
|
4893
|
+
if (this.options.keepFirstStageChunks) {
|
|
4894
|
+
merged.push(passage);
|
|
4895
|
+
}
|
|
4896
|
+
}
|
|
4897
|
+
return {
|
|
4898
|
+
passages: merged.slice(0, topK),
|
|
4899
|
+
navigations
|
|
4900
|
+
};
|
|
4901
|
+
}
|
|
4902
|
+
};
|
|
4903
|
+
|
|
4904
|
+
// src/tree/navigate-tool.ts
|
|
4905
|
+
function isInput(value) {
|
|
4906
|
+
return typeof value === "object" && value !== null && typeof value.documentId === "string" && typeof value.question === "string";
|
|
4907
|
+
}
|
|
4908
|
+
__name(isInput, "isInput");
|
|
4909
|
+
var inputSchema = {
|
|
4910
|
+
"~standard": {
|
|
4911
|
+
version: 1,
|
|
4912
|
+
vendor: "nestjs-agent-rag",
|
|
4913
|
+
validate(value) {
|
|
4914
|
+
return isInput(value) ? {
|
|
4915
|
+
value: {
|
|
4916
|
+
documentId: value.documentId,
|
|
4917
|
+
question: value.question
|
|
4918
|
+
}
|
|
4919
|
+
} : {
|
|
4920
|
+
issues: [
|
|
4921
|
+
{
|
|
4922
|
+
message: "navigate_document expects { documentId: string, question: string }"
|
|
4923
|
+
}
|
|
4924
|
+
]
|
|
4925
|
+
};
|
|
4926
|
+
}
|
|
4927
|
+
}
|
|
4928
|
+
};
|
|
4929
|
+
function createNavigateDocumentTool(navigator, options = {}) {
|
|
4930
|
+
return {
|
|
4931
|
+
spec: {
|
|
4932
|
+
name: options.name ?? "navigate_document",
|
|
4933
|
+
kind: "read",
|
|
4934
|
+
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.",
|
|
4935
|
+
inputSchema
|
|
4936
|
+
},
|
|
4937
|
+
handler: {
|
|
4938
|
+
async execute(input, ctx) {
|
|
4939
|
+
if (!isInput(input)) {
|
|
4940
|
+
throw new Error("navigate_document expects { documentId: string, question: string }");
|
|
4941
|
+
}
|
|
4942
|
+
const filter = options.filter?.(ctx);
|
|
4943
|
+
const navigateOptions = {
|
|
4944
|
+
...filter !== void 0 ? {
|
|
4945
|
+
filter
|
|
4946
|
+
} : {},
|
|
4947
|
+
...options.budget !== void 0 ? {
|
|
4948
|
+
budget: options.budget
|
|
4949
|
+
} : {}
|
|
4950
|
+
};
|
|
4951
|
+
let documentId = input.documentId.trim();
|
|
4952
|
+
let result = await navigator.navigate(input.question, [
|
|
4953
|
+
documentId
|
|
4954
|
+
], navigateOptions);
|
|
4955
|
+
let navigation = result.documents[0];
|
|
4956
|
+
if (navigation?.stoppedBy === "not-found") {
|
|
4957
|
+
const bare = documentIdOf(documentId.replace(/#node:\d+$/, ""));
|
|
4958
|
+
if (bare !== documentId) {
|
|
4959
|
+
documentId = bare;
|
|
4960
|
+
result = await navigator.navigate(input.question, [
|
|
4961
|
+
documentId
|
|
4962
|
+
], navigateOptions);
|
|
4963
|
+
navigation = result.documents[0];
|
|
4964
|
+
}
|
|
4965
|
+
}
|
|
4966
|
+
if (navigation === void 0 || navigation.stoppedBy === "not-found") {
|
|
4967
|
+
return {
|
|
4968
|
+
documentId,
|
|
4969
|
+
found: false,
|
|
4970
|
+
passages: [],
|
|
4971
|
+
nodes: [],
|
|
4972
|
+
steps: [],
|
|
4973
|
+
note: "This document has no navigable table of contents (it is short, not indexed, or not accessible). Use the search passages instead."
|
|
4974
|
+
};
|
|
4975
|
+
}
|
|
4976
|
+
return {
|
|
4977
|
+
documentId,
|
|
4978
|
+
found: true,
|
|
4979
|
+
passages: result.passages,
|
|
4980
|
+
nodes: navigation.nodes,
|
|
4981
|
+
steps: navigation.steps,
|
|
4982
|
+
...navigation.stoppedBy !== void 0 ? {
|
|
4983
|
+
stoppedBy: navigation.stoppedBy
|
|
4984
|
+
} : {},
|
|
4985
|
+
...navigation.nodes.length === 0 ? {
|
|
4986
|
+
note: "No section of this document looked relevant to the question."
|
|
4987
|
+
} : {}
|
|
4988
|
+
};
|
|
4989
|
+
}
|
|
4990
|
+
}
|
|
4991
|
+
};
|
|
4992
|
+
}
|
|
4993
|
+
__name(createNavigateDocumentTool, "createNavigateDocumentTool");
|
|
2547
4994
|
// Annotate the CommonJS export names for ESM import in node:
|
|
2548
4995
|
0 && (module.exports = {
|
|
4996
|
+
DEFAULT_STOP_WORDS,
|
|
2549
4997
|
EmbeddingRetriever,
|
|
2550
4998
|
FilteredRetriever,
|
|
2551
4999
|
HttpModelError,
|
|
@@ -2553,7 +5001,9 @@ __name(itemsOf, "itemsOf");
|
|
|
2553
5001
|
HybridRetriever,
|
|
2554
5002
|
KeywordRetriever,
|
|
2555
5003
|
LexicalRetriever,
|
|
5004
|
+
MemoryDocumentTreeStore,
|
|
2556
5005
|
MemoryVectorStore,
|
|
5006
|
+
PgDocumentTreeStore,
|
|
2557
5007
|
PgLexicalVectorStore,
|
|
2558
5008
|
PgVectorStore,
|
|
2559
5009
|
RAG_DIAGNOSTIC_LIB,
|
|
@@ -2562,23 +5012,44 @@ __name(itemsOf, "itemsOf");
|
|
|
2562
5012
|
RedisVectorSchemaMismatchError,
|
|
2563
5013
|
RedisVectorStore,
|
|
2564
5014
|
RerankingRetriever,
|
|
5015
|
+
TreeBudgetExceededError,
|
|
5016
|
+
TreeNavigationRetriever,
|
|
5017
|
+
TwoStageRetriever,
|
|
2565
5018
|
UnsafeRemovalError,
|
|
5019
|
+
anyTermTsquery,
|
|
2566
5020
|
applyMetadataPatch,
|
|
5021
|
+
buildDocumentTree,
|
|
5022
|
+
cachedTreeLlm,
|
|
2567
5023
|
chunkDocuments,
|
|
2568
5024
|
chunkIndexOf,
|
|
2569
5025
|
chunkText,
|
|
5026
|
+
createNavigateDocumentTool,
|
|
2570
5027
|
createRetrievalTool,
|
|
2571
5028
|
describeRetrieval,
|
|
2572
5029
|
describeSharedSource,
|
|
2573
5030
|
describeSource,
|
|
2574
5031
|
describesRetrieval,
|
|
5032
|
+
detectHeadings,
|
|
2575
5033
|
documentIdOf,
|
|
2576
5034
|
emitRagRetrieval,
|
|
5035
|
+
estimateTokens,
|
|
5036
|
+
hasSearchSyntax,
|
|
5037
|
+
indexDocumentTree,
|
|
5038
|
+
indexTree,
|
|
2577
5039
|
ingestChunks,
|
|
2578
5040
|
ingestDocuments,
|
|
2579
5041
|
instrumentRetriever,
|
|
5042
|
+
isBatchTooLarge,
|
|
2580
5043
|
isLexicalVectorStore,
|
|
5044
|
+
keywordTerms,
|
|
5045
|
+
keywordTreeLlm,
|
|
5046
|
+
openAiChatTreeLlm,
|
|
2581
5047
|
openAiEmbeddings,
|
|
2582
|
-
|
|
5048
|
+
parseJsonReply,
|
|
5049
|
+
spanText,
|
|
5050
|
+
stableHash,
|
|
5051
|
+
toPgClient,
|
|
5052
|
+
treeLlmFromModelProvider,
|
|
5053
|
+
walkTree
|
|
2583
5054
|
});
|
|
2584
5055
|
//# sourceMappingURL=index.cjs.map
|