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