@dudousxd/nestjs-agent-rag 0.13.0 → 0.14.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/README.md CHANGED
@@ -197,6 +197,9 @@ await store.ensureSchema(); // CREATE EXTENSION vector + table + HNSW cosine ind
197
197
  `store.schemaStatements()` returns the same DDL as strings, for a migration instead of a boot-time
198
198
  side effect. Other options (all off by default, so existing tables behave exactly as before):
199
199
 
200
+ - **`vectorType`** (default `vector`) — `halfvec` stores 2 bytes per dimension and its HNSW index goes
201
+ up to 4000 dimensions, where `vector` stops at 2000. Needs a single width. Converting an existing
202
+ table is a migration of yours.
200
203
  - **`upsertBatchSize`** (default 100) — `upsert` writes multi-row `INSERT … ON CONFLICT` statements;
201
204
  duplicate ids in one call collapse to the last occurrence. NUL bytes are stripped from every
202
205
  text/jsonb binding (Postgres rejects them).
@@ -439,6 +442,69 @@ Had `PgVectorStore` simply grown `searchText`, every `isLexicalVectorStore(store
439
442
  started using it against an unindexed table on upgrade. Constructing `PgLexicalVectorStore` is the
440
443
  opt-in; the index is a migration you run and watch (`schemaStatements()` prints it).
441
444
 
445
+ ## RAG for a model: `RagRepository`
446
+
447
+ One class per model, as `nestjs-filter` does with filters. The entity or Drizzle table doesn't change.
448
+
449
+ ```ts
450
+ import { RagModule, RagRepository } from '@dudousxd/nestjs-agent-rag/nestjs';
451
+
452
+ // Once in the app
453
+ RagModule.forRoot({ store, embedder, onUsage: (usage) => meter.add(usage) });
454
+
455
+ @Injectable()
456
+ export class DocumentsRag extends RagRepository<ProjectDocument> {
457
+ protected readonly name = 'project_documents';
458
+ readonly scope = { projectId: (d: ProjectDocument) => d.project.id };
459
+
460
+ id(d: ProjectDocument) { return d.id; }
461
+ protected text(d: ProjectDocument) { return d.textContent; }
462
+ }
463
+
464
+ // In the feature module
465
+ RagModule.forFeature([DocumentsRag]);
466
+
467
+ // Injected by type
468
+ await documentsRag.index(doc);
469
+ await documentsRag.search('question', { scope: { projectId } });
470
+ await documentsRag.remove(doc);
471
+ ```
472
+
473
+ With Drizzle, only the row type changes: `type Row = typeof documents.$inferSelect`, then
474
+ `RagRepository<Row>` and `scope = { projectId: (d: Row) => d.projectId }`.
475
+
476
+ **The scope can't be skipped.** Leaving out `scope`, returning a nullable column, or calling `search`
477
+ without it doesn't compile. At runtime a missing, empty or malformed value throws
478
+ `RagScopeRequiredError` before any embedding is paid for, and `filter` can't override it. Values are
479
+ compared as strings. To search several, use `anyOf(ids)`, where an empty list matches nothing; a plain
480
+ array is refused so a request can't widen a search. A corpus shared by everyone says `scope = {}`.
481
+
482
+ **Saving doesn't index.** Embedding is slow and paid for, so call `index` from a job or wherever the
483
+ text changed. It embeds first, so a provider failure keeps the old chunks. On `PgVectorStore` the swap is
484
+ one transaction; on stores without `replaceDocument`, it is two statements.
485
+
486
+ **Deleting.** With Drizzle, use `.returning()` and call `remove`. With MikroORM, add
487
+ `provideRagDeleteSync(ProjectDocument, DocumentsRag)` from `@dudousxd/nestjs-agent-store-mikro-orm` to
488
+ your providers; it removes after the commit and sees only `em.remove` (on MikroORM 6, not inside a
489
+ nested transaction). For everything no hook sees
490
+ (`ON DELETE CASCADE`, `nativeDelete`, raw SQL), run `reconcile`:
491
+
492
+ ```ts
493
+ await documentsRag.reconcile(allLiveIds, { dryRun: true }); // { orphans, missing }
494
+ await documentsRag.reconcile(liveIdsOfProject, { scope: { projectId } });
495
+ ```
496
+
497
+ `liveIds` must be complete for the part compared: whatever it leaves out is removed.
498
+
499
+ **Agents.** `documentsRag.tool({ name, scope: (ctx) => ({ projectId: ctx.actor.tenantRef as string }) })`
500
+ gives a tool whose scope comes from the context, for `provideAgentTool`. `scoped(scope)` gives a
501
+ `Retriever` for composing. `retrieval: { mode: 'inject' }` passes no actor, so it can't scope per
502
+ request: use the tool for multi-tenant agents.
503
+
504
+ On Redis, declare the keys up front: `filterableFields: ragFilterableFields('projectId')`.
505
+
506
+ Without Nest, `RagIndex` from the main entry does the same with a plain definition.
507
+
442
508
  ## Tree navigation retrieval (PageIndex-style)
443
509
 
444
510
  Chunk search ranks passages; on a **long structured document** (a regulation, an RFP, a contract, a
@@ -592,6 +658,7 @@ await store.removeMany(['doc-a', 'doc-b']); // N documents, one pas
592
658
  |---|---|
593
659
  | `{ collection: 'handbook' }` | deletes that scope, returns the chunk count |
594
660
  | `{ audience: [] }` (or any empty array) | **deletes nothing**, returns `0` — the empty array is this package's deny primitive, and a deny must never be read as "no filter" |
661
+ | `{ tenantId: undefined }` | throws `UnsafeRemovalError` with reason `undefined-value` — a missing value is a bug, and on every read method it matches nothing |
595
662
  | `{}` | throws `UnsafeRemovalError` — an empty filter is far more often a filter built wrong than a deliberate request to wipe the store |
596
663
  | `{ notFilterable: 'x' }` (Redis) | throws `UnsafeRemovalError` — the key has no `meta_*` TAG to match on |
597
664
 
package/dist/index.cjs CHANGED
@@ -19,8 +19,8 @@ var __copyProps = (to, from, except, desc) => {
19
19
  var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
20
20
 
21
21
  // src/index.ts
22
- var index_exports = {};
23
- __export(index_exports, {
22
+ var src_exports = {};
23
+ __export(src_exports, {
24
24
  DEFAULT_STOP_WORDS: () => DEFAULT_STOP_WORDS,
25
25
  EmbeddingRetriever: () => EmbeddingRetriever,
26
26
  FilteredRetriever: () => FilteredRetriever,
@@ -35,8 +35,11 @@ __export(index_exports, {
35
35
  PgLexicalVectorStore: () => PgLexicalVectorStore,
36
36
  PgVectorStore: () => PgVectorStore,
37
37
  RAG_DIAGNOSTIC_LIB: () => RAG_DIAGNOSTIC_LIB,
38
+ RAG_ENTITY_KEY: () => RAG_ENTITY_KEY,
38
39
  RAG_RETRIEVAL_CHANNEL: () => RAG_RETRIEVAL_CHANNEL,
39
40
  RAG_RETRIEVAL_EVENT: () => RAG_RETRIEVAL_EVENT,
41
+ RagIndex: () => RagIndex,
42
+ RagScopeRequiredError: () => RagScopeRequiredError,
40
43
  RedisVectorSchemaMismatchError: () => RedisVectorSchemaMismatchError,
41
44
  RedisVectorStore: () => RedisVectorStore,
42
45
  RerankingRetriever: () => RerankingRetriever,
@@ -44,6 +47,7 @@ __export(index_exports, {
44
47
  TreeNavigationRetriever: () => TreeNavigationRetriever,
45
48
  TwoStageRetriever: () => TwoStageRetriever,
46
49
  UnsafeRemovalError: () => UnsafeRemovalError,
50
+ anyOf: () => anyOf,
47
51
  anyTermTsquery: () => anyTermTsquery,
48
52
  applyMetadataPatch: () => applyMetadataPatch,
49
53
  buildDocumentTree: () => buildDocumentTree,
@@ -59,6 +63,7 @@ __export(index_exports, {
59
63
  describesRetrieval: () => describesRetrieval,
60
64
  detectHeadings: () => detectHeadings,
61
65
  documentIdOf: () => documentIdOf,
66
+ embedCountingUsage: () => embedCountingUsage,
62
67
  emitRagRetrieval: () => emitRagRetrieval,
63
68
  estimateTokens: () => estimateTokens,
64
69
  hasSearchSyntax: () => hasSearchSyntax,
@@ -74,13 +79,15 @@ __export(index_exports, {
74
79
  openAiChatTreeLlm: () => openAiChatTreeLlm,
75
80
  openAiEmbeddings: () => openAiEmbeddings,
76
81
  parseJsonReply: () => parseJsonReply,
82
+ ragFilterableFields: () => ragFilterableFields,
83
+ replaceDocumentChunks: () => replaceDocumentChunks,
77
84
  spanText: () => spanText,
78
85
  stableHash: () => stableHash,
79
86
  toPgClient: () => toPgClient,
80
87
  treeLlmFromModelProvider: () => treeLlmFromModelProvider,
81
88
  walkTree: () => walkTree
82
89
  });
83
- module.exports = __toCommonJS(index_exports);
90
+ module.exports = __toCommonJS(src_exports);
84
91
 
85
92
  // src/chunk.ts
86
93
  function chunkText(text, options = {}) {
@@ -158,6 +165,16 @@ function chunkByRecord(text, chunkSize, separator) {
158
165
  __name(chunkByRecord, "chunkByRecord");
159
166
 
160
167
  // src/vector-store.ts
168
+ async function replaceDocumentChunks(store, documentId, records) {
169
+ const atomic = store;
170
+ if (typeof atomic.replaceDocument === "function") {
171
+ await atomic.replaceDocument(documentId, records);
172
+ return;
173
+ }
174
+ await store.remove(documentId);
175
+ await store.upsert(records);
176
+ }
177
+ __name(replaceDocumentChunks, "replaceDocumentChunks");
161
178
  function isLexicalVectorStore(store) {
162
179
  return typeof store.searchText === "function";
163
180
  }
@@ -218,12 +235,19 @@ function applyMetadataPatch(metadata, patch) {
218
235
  __name(applyMetadataPatch, "applyMetadataPatch");
219
236
 
220
237
  // src/filter.ts
238
+ function isMissingValue(value) {
239
+ return value === void 0 || Array.isArray(value) && value.includes(void 0);
240
+ }
241
+ __name(isMissingValue, "isMissingValue");
221
242
  function matchesFilter(metadata, filter) {
222
243
  const entries = Object.entries(filter);
223
244
  if (metadata === void 0) {
224
245
  return entries.length === 0;
225
246
  }
226
247
  return entries.every(([key, expected]) => {
248
+ if (isMissingValue(expected)) {
249
+ return false;
250
+ }
227
251
  if (Array.isArray(expected)) {
228
252
  const actual = metadata[key];
229
253
  return Array.isArray(actual) ? actual.some((value) => expected.includes(value)) : expected.includes(actual);
@@ -236,9 +260,18 @@ function filterMatchesNothing(filter) {
236
260
  if (filter === void 0) {
237
261
  return false;
238
262
  }
239
- return Object.values(filter).some((value) => Array.isArray(value) && value.length === 0);
263
+ return Object.values(filter).some((value) => isMissingValue(value) || Array.isArray(value) && value.length === 0);
240
264
  }
241
265
  __name(filterMatchesNothing, "filterMatchesNothing");
266
+ function undefinedFilterKeys(filter) {
267
+ return Object.entries(filter).filter(([, value]) => isMissingValue(value)).map(([key]) => key);
268
+ }
269
+ __name(undefinedFilterKeys, "undefinedFilterKeys");
270
+ function undefinedFilterMessage(keys) {
271
+ const named = keys.map((key) => `"${key}"`).join(", ");
272
+ return `removeWhere() refuses a filter with no value for ${named}: the scope was probably never set.`;
273
+ }
274
+ __name(undefinedFilterMessage, "undefinedFilterMessage");
242
275
 
243
276
  // src/memory-vector-store.ts
244
277
  var MemoryVectorStore = class {
@@ -375,6 +408,10 @@ var MemoryVectorStore = class {
375
408
  if (Object.keys(filter).length === 0) {
376
409
  throw new UnsafeRemovalError("empty-filter", "removeWhere() refuses an empty filter: it would delete every chunk in the store. Pass a filter that scopes the removal, or delete deliberately with removeMany(await store.listDocumentIds()).");
377
410
  }
411
+ const missing = undefinedFilterKeys(filter);
412
+ if (missing.length > 0) {
413
+ throw new UnsafeRemovalError("undefined-value", undefinedFilterMessage(missing));
414
+ }
378
415
  if (filterMatchesNothing(filter)) {
379
416
  return 0;
380
417
  }
@@ -644,6 +681,7 @@ var PgVectorStore = class {
644
681
  dimensions;
645
682
  /** The widths with a partial HNSW index, for a mixed-dimension table. */
646
683
  indexedDimensions;
684
+ vectorType;
647
685
  nullableEmbeddings;
648
686
  iterativeScan;
649
687
  efSearch;
@@ -667,6 +705,10 @@ var PgVectorStore = class {
667
705
  ...new Set(dimensions)
668
706
  ];
669
707
  }
708
+ this.vectorType = options.vectorType ?? "vector";
709
+ if (this.vectorType === "halfvec" && this.dimensions === void 0) {
710
+ throw new Error('PgVectorStore: vectorType "halfvec" needs a single width, not a list');
711
+ }
670
712
  this.nullableEmbeddings = options.nullableEmbeddings ?? false;
671
713
  this.iterativeScan = options.iterativeScan;
672
714
  this.efSearch = options.efSearch;
@@ -685,7 +727,7 @@ var PgVectorStore = class {
685
727
  */
686
728
  schemaStatements() {
687
729
  const mixed = this.dimensions === void 0;
688
- const column = mixed ? "vector" : `vector(${this.dimensions})`;
730
+ const column = mixed ? "vector" : `${this.vectorType}(${this.dimensions})`;
689
731
  const statements = [
690
732
  "CREATE EXTENSION IF NOT EXISTS vector",
691
733
  `CREATE TABLE IF NOT EXISTS ${this.table} (
@@ -704,7 +746,7 @@ var PgVectorStore = class {
704
746
  }
705
747
  } else {
706
748
  statements.push(`CREATE INDEX IF NOT EXISTS ${this.table}_embedding_idx
707
- ON ${this.table} USING hnsw (embedding vector_cosine_ops)`);
749
+ ON ${this.table} USING hnsw (embedding ${this.vectorType}_cosine_ops)`);
708
750
  }
709
751
  return statements;
710
752
  }
@@ -720,6 +762,23 @@ var PgVectorStore = class {
720
762
  * {@link PgVectorStoreOptions.nullableEmbeddings}).
721
763
  */
722
764
  async upsert(records) {
765
+ await this.upsertWith(this.client, records);
766
+ }
767
+ /**
768
+ * Swaps a document's chunks in one transaction, so a crash can't lose it and a search sees the old
769
+ * version or the new one. A client without `transaction` falls back to two statements.
770
+ */
771
+ async replaceDocument(documentId, records) {
772
+ const swap = /* @__PURE__ */ __name(async (client) => {
773
+ await client.query(`DELETE FROM ${this.table} WHERE ${DOCUMENT_ID_FROM_CHUNK} = $1`, [
774
+ stripNulBytes(documentId)
775
+ ]);
776
+ await this.upsertWith(client, records);
777
+ }, "swap");
778
+ const transaction = this.client.transaction?.bind(this.client);
779
+ await (transaction === void 0 ? swap(this.client) : transaction(swap));
780
+ }
781
+ async upsertWith(client, records) {
723
782
  const byId = /* @__PURE__ */ new Map();
724
783
  for (const record of records) {
725
784
  const id = stripNulBytes(record.id);
@@ -741,9 +800,9 @@ var PgVectorStore = class {
741
800
  const metadata = record.metadata !== void 0 ? stripNulBytes(record.metadata) : void 0;
742
801
  const base = params.length;
743
802
  params.push(record.id, text, source ?? null, metadata !== void 0 ? JSON.stringify(metadata) : null, record.embedding.length > 0 ? toVectorLiteral(record.embedding) : null);
744
- return `($${base + 1}, $${base + 2}, $${base + 3}, $${base + 4}::jsonb, $${base + 5}::vector)`;
803
+ return `($${base + 1}, $${base + 2}, $${base + 3}, $${base + 4}::jsonb, $${base + 5}::${this.vectorType})`;
745
804
  });
746
- await this.client.query(`INSERT INTO ${this.table} (id, text, source, metadata, embedding)
805
+ await client.query(`INSERT INTO ${this.table} (id, text, source, metadata, embedding)
747
806
  VALUES ${rows.join(", ")}
748
807
  ON CONFLICT (id) DO UPDATE SET
749
808
  text = EXCLUDED.text,
@@ -867,6 +926,10 @@ var PgVectorStore = class {
867
926
  if (Object.keys(filter).length === 0) {
868
927
  throw new UnsafeRemovalError("empty-filter", "removeWhere() refuses an empty filter: it would delete every chunk in the table. Pass a filter that scopes the removal, or delete deliberately with removeMany(await store.listDocumentIds()).");
869
928
  }
929
+ const missing = undefinedFilterKeys(filter);
930
+ if (missing.length > 0) {
931
+ throw new UnsafeRemovalError("undefined-value", undefinedFilterMessage(missing));
932
+ }
870
933
  if (filterMatchesNothing(filter)) {
871
934
  return 0;
872
935
  }
@@ -898,7 +961,7 @@ var PgVectorStore = class {
898
961
  options.topK
899
962
  ];
900
963
  const conditions = this.whereConditions(options.filter, params);
901
- let vectorType = "vector";
964
+ let vectorType = this.vectorType;
902
965
  if (this.dimensions === void 0) {
903
966
  const width = embedding.length;
904
967
  if (this.indexedDimensions.includes(width)) {
@@ -908,7 +971,7 @@ var PgVectorStore = class {
908
971
  } else if (this.nullableEmbeddings) {
909
972
  conditions.push("embedding IS NOT NULL");
910
973
  }
911
- const column = vectorType === "vector" ? "embedding" : `(embedding::${vectorType})`;
974
+ const column = vectorType === this.vectorType ? "embedding" : `(embedding::${vectorType})`;
912
975
  const sql = `SELECT id, text, source, metadata, 1 - (${column} <=> $1::${vectorType}) AS score
913
976
  FROM ${this.table}
914
977
  ${conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : ""}
@@ -1095,6 +1158,10 @@ function buildWhere(filter, params) {
1095
1158
  const clauses = [];
1096
1159
  const scalar = {};
1097
1160
  for (const [key, value] of Object.entries(stripped)) {
1161
+ if (value === void 0 || Array.isArray(value) && value.includes(void 0)) {
1162
+ clauses.push("false");
1163
+ continue;
1164
+ }
1098
1165
  if (Array.isArray(value)) {
1099
1166
  if (value.length === 0) {
1100
1167
  clauses.push("false");
@@ -1636,6 +1703,10 @@ var RedisVectorStore = class {
1636
1703
  if (keys.length === 0) {
1637
1704
  throw new UnsafeRemovalError("empty-filter", "removeWhere() refuses an empty filter: it would delete every chunk in the index. Pass a filter that scopes the removal, or delete deliberately with removeMany(await store.listDocumentIds()).");
1638
1705
  }
1706
+ const missing = undefinedFilterKeys(filter);
1707
+ if (missing.length > 0) {
1708
+ throw new UnsafeRemovalError("undefined-value", undefinedFilterMessage(missing));
1709
+ }
1639
1710
  const unindexed = keys.filter((key) => !this.filterableFields.includes(key));
1640
1711
  if (unindexed.length > 0) {
1641
1712
  const named = unindexed.map((key) => `"${key}"`).join(", ");
@@ -2186,7 +2257,7 @@ var KeywordRetriever = class {
2186
2257
  }
2187
2258
  };
2188
2259
  function tokenize(text) {
2189
- return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
2260
+ return text.normalize("NFKD").replace(/\p{M}/gu, "").toLowerCase().match(/[\p{L}\p{N}_]+/gu) ?? [];
2190
2261
  }
2191
2262
  __name(tokenize, "tokenize");
2192
2263
 
@@ -2507,13 +2578,19 @@ function createRetrievalTool(retriever, options = {}) {
2507
2578
  inputSchema: querySchema
2508
2579
  },
2509
2580
  handler: {
2510
- async execute(input) {
2581
+ async execute(input, ctx) {
2511
2582
  if (!isQueryInput(input)) {
2512
2583
  throw new Error("retrieval tool expects { query: string }");
2513
2584
  }
2514
- const passages = await instrumented.retrieve(input.query, topK !== void 0 ? {
2515
- topK
2516
- } : {});
2585
+ const filter = options.filter?.(ctx);
2586
+ const passages = await instrumented.retrieve(input.query, {
2587
+ ...topK !== void 0 ? {
2588
+ topK
2589
+ } : {},
2590
+ ...filter !== void 0 ? {
2591
+ filter
2592
+ } : {}
2593
+ });
2517
2594
  return {
2518
2595
  passages
2519
2596
  };
@@ -2553,7 +2630,7 @@ function openAiEmbeddings(options) {
2553
2630
  const maxInputChars = options.maxInputChars ?? 24e3;
2554
2631
  const timeoutMs = options.timeoutMs ?? 12e4;
2555
2632
  const doFetch = options.fetch ?? fetch;
2556
- const request = /* @__PURE__ */ __name(async (input) => {
2633
+ const request = /* @__PURE__ */ __name(async (input, count) => {
2557
2634
  const response = await doFetch(url, {
2558
2635
  method: "POST",
2559
2636
  headers: {
@@ -2586,19 +2663,20 @@ function openAiEmbeddings(options) {
2586
2663
  const tokens = json2?.usage?.total_tokens ?? json2?.usage?.prompt_tokens;
2587
2664
  if (typeof tokens === "number" && Number.isFinite(tokens)) {
2588
2665
  options.onUsage?.(tokens);
2666
+ count(tokens);
2589
2667
  }
2590
2668
  return items.map((item) => item.embedding);
2591
2669
  }, "request");
2592
- const requestSplitting = /* @__PURE__ */ __name(async (input) => {
2670
+ const requestSplitting = /* @__PURE__ */ __name(async (input, count) => {
2593
2671
  if (input.length > batchSize) {
2594
2672
  const vectors = [];
2595
2673
  for (let start = 0; start < input.length; start += batchSize) {
2596
- vectors.push(...await requestSplitting(input.slice(start, start + batchSize)));
2674
+ vectors.push(...await requestSplitting(input.slice(start, start + batchSize), count));
2597
2675
  }
2598
2676
  return vectors;
2599
2677
  }
2600
2678
  try {
2601
- return await request(input);
2679
+ return await request(input, count);
2602
2680
  } catch (error) {
2603
2681
  if (!isBatchTooLarge(error) || input.length <= 1) {
2604
2682
  throw error;
@@ -2610,20 +2688,37 @@ function openAiEmbeddings(options) {
2610
2688
  options.onWarn?.(`embeddings server ${url} refused a batch of ${input.length} inputs (${error.message}); retrying in batches of ${half}. Raise the server's batch limit (TEI: --max-client-batch-size) or lower batchSize.`);
2611
2689
  }
2612
2690
  return [
2613
- ...await requestSplitting(input.slice(0, half)),
2614
- ...await requestSplitting(input.slice(half))
2691
+ ...await requestSplitting(input.slice(0, half), count),
2692
+ ...await requestSplitting(input.slice(half), count)
2615
2693
  ];
2616
2694
  }
2617
2695
  }, "requestSplitting");
2696
+ const embedAll = /* @__PURE__ */ __name(async (texts, count) => {
2697
+ const vectors = [];
2698
+ for (let start = 0; start < texts.length; ) {
2699
+ const input = texts.slice(start, start + batchSize).map((text) => (text.length > maxInputChars ? text.slice(0, maxInputChars) : text) || " ");
2700
+ vectors.push(...await requestSplitting(input, count));
2701
+ start += input.length;
2702
+ }
2703
+ return vectors;
2704
+ }, "embedAll");
2618
2705
  return {
2619
- async embed(texts) {
2620
- const vectors = [];
2621
- for (let start = 0; start < texts.length; ) {
2622
- const input = texts.slice(start, start + batchSize).map((text) => (text.length > maxInputChars ? text.slice(0, maxInputChars) : text) || " ");
2623
- vectors.push(...await requestSplitting(input));
2624
- start += input.length;
2625
- }
2626
- return vectors;
2706
+ embed: /* @__PURE__ */ __name((texts) => embedAll(texts, () => {
2707
+ }), "embed"),
2708
+ async embedWithUsage(texts) {
2709
+ let inputTokens;
2710
+ const vectors = await embedAll(texts, (tokens) => {
2711
+ inputTokens = (inputTokens ?? 0) + tokens;
2712
+ });
2713
+ return {
2714
+ vectors,
2715
+ ...inputTokens !== void 0 ? {
2716
+ usage: {
2717
+ inputTokens,
2718
+ modelId: options.model
2719
+ }
2720
+ } : {}
2721
+ };
2627
2722
  }
2628
2723
  };
2629
2724
  }
@@ -2728,6 +2823,206 @@ function itemsOf(body) {
2728
2823
  }
2729
2824
  __name(itemsOf, "itemsOf");
2730
2825
 
2826
+ // src/embedding-usage.ts
2827
+ async function embedCountingUsage(provider, texts) {
2828
+ const reporting = provider;
2829
+ return typeof reporting.embedWithUsage === "function" ? reporting.embedWithUsage(texts) : {
2830
+ vectors: await provider.embed(texts)
2831
+ };
2832
+ }
2833
+ __name(embedCountingUsage, "embedCountingUsage");
2834
+
2835
+ // src/rag-index.ts
2836
+ var RAG_ENTITY_KEY = "ragEntity";
2837
+ var RagScopeRequiredError = class extends Error {
2838
+ static {
2839
+ __name(this, "RagScopeRequiredError");
2840
+ }
2841
+ model;
2842
+ key;
2843
+ constructor(model, key) {
2844
+ super(`${model} needs scope.${key}: its chunks are isolated by it`), this.model = model, this.key = key;
2845
+ this.name = "RagScopeRequiredError";
2846
+ }
2847
+ };
2848
+ var ANY_OF = /* @__PURE__ */ Symbol.for("@dudousxd/nestjs-agent-rag/anyOf");
2849
+ function anyOf(values) {
2850
+ return {
2851
+ [ANY_OF]: values
2852
+ };
2853
+ }
2854
+ __name(anyOf, "anyOf");
2855
+ function isScopeValue(value) {
2856
+ return typeof value === "number" ? Number.isFinite(value) : typeof value === "string" && value !== "";
2857
+ }
2858
+ __name(isScopeValue, "isScopeValue");
2859
+ function isAnyOf(value) {
2860
+ return typeof value === "object" && value !== null && ANY_OF in value;
2861
+ }
2862
+ __name(isAnyOf, "isAnyOf");
2863
+ var RagIndex = class {
2864
+ static {
2865
+ __name(this, "RagIndex");
2866
+ }
2867
+ definition;
2868
+ deps;
2869
+ embedder;
2870
+ constructor(definition, deps) {
2871
+ this.definition = definition;
2872
+ this.deps = deps;
2873
+ const { onUsage } = deps;
2874
+ this.embedder = onUsage === void 0 ? deps.embedder : {
2875
+ embed: /* @__PURE__ */ __name(async (texts) => {
2876
+ const { vectors, usage } = await embedCountingUsage(deps.embedder, texts);
2877
+ if (usage !== void 0) {
2878
+ await onUsage(usage);
2879
+ }
2880
+ return vectors;
2881
+ }, "embed")
2882
+ };
2883
+ }
2884
+ /** The id this row's chunks are stored under: `<name>:<id>`. */
2885
+ documentId(id) {
2886
+ if (!isScopeValue(id)) {
2887
+ throw new Error(`${this.definition.name}: a row needs an id to be indexed or removed`);
2888
+ }
2889
+ return `${this.definition.name}:${String(id)}`;
2890
+ }
2891
+ /** A `Retriever` locked to this scope, for `HybridRetriever`, `createRetrievalTool` and so on. */
2892
+ scoped(scope) {
2893
+ return new FilteredRetriever(this.baseRetriever(), this.scopeFilter(scope));
2894
+ }
2895
+ async search(query, options) {
2896
+ const { scope, filter, topK } = options;
2897
+ return this.scoped(scope).retrieve(query, {
2898
+ ...topK !== void 0 ? {
2899
+ topK
2900
+ } : {},
2901
+ ...filter !== void 0 ? {
2902
+ filter
2903
+ } : {}
2904
+ });
2905
+ }
2906
+ /**
2907
+ * Replaces the row's chunks and returns how many were written. Embeddings come first: if the
2908
+ * provider fails, the previous chunks are still there. The swap is one transaction on Postgres; on a
2909
+ * store without `replaceDocument` a crash mid-swap loses the document until the next `index`.
2910
+ */
2911
+ async index(row) {
2912
+ const { definition, deps } = this;
2913
+ const documentId = this.documentId(definition.id(row));
2914
+ const scope = this.scopeFilter(Object.fromEntries(Object.entries(definition.scope).map(([key, read]) => [
2915
+ key,
2916
+ read(row)
2917
+ ])));
2918
+ const text = definition.text(row)?.trim();
2919
+ if (!text) {
2920
+ await deps.store.remove(documentId);
2921
+ return 0;
2922
+ }
2923
+ const pieces = chunkText(text, definition.chunk).slice(0, definition.maxChunks);
2924
+ const embeddings = await this.embedder.embed(pieces);
2925
+ const source = definition.source?.(row) ?? void 0;
2926
+ const metadata = {
2927
+ ...definition.metadata?.(row),
2928
+ ...scope
2929
+ };
2930
+ await replaceDocumentChunks(deps.store, documentId, pieces.map((piece, index) => ({
2931
+ id: `${documentId}#${index}`,
2932
+ text: piece,
2933
+ embedding: embeddings[index] ?? [],
2934
+ ...source !== void 0 ? {
2935
+ source
2936
+ } : {},
2937
+ metadata
2938
+ })));
2939
+ return pieces.length;
2940
+ }
2941
+ /** Takes the id itself once the row is gone. */
2942
+ remove(rowOrId) {
2943
+ const id = typeof rowOrId === "string" || typeof rowOrId === "number" ? rowOrId : this.definition.id(rowOrId);
2944
+ return this.deps.store.remove(this.documentId(id));
2945
+ }
2946
+ /**
2947
+ * Diffs the index against the ids that exist, and removes what no hook could: `ON DELETE CASCADE`,
2948
+ * raw SQL, a crash after a commit. `liveIds` must be complete for the part being compared, because
2949
+ * anything left out is removed. Use `dryRun` to look first.
2950
+ */
2951
+ async reconcile(liveIds, options = {}) {
2952
+ const filter = {
2953
+ [RAG_ENTITY_KEY]: this.definition.name
2954
+ };
2955
+ for (const [key, value] of Object.entries(options.scope ?? {})) {
2956
+ if (!isScopeValue(value)) {
2957
+ throw new RagScopeRequiredError(this.definition.name, key);
2958
+ }
2959
+ filter[key] = String(value);
2960
+ }
2961
+ const live = /* @__PURE__ */ new Set();
2962
+ for await (const id of liveIds) {
2963
+ live.add(this.documentId(id));
2964
+ }
2965
+ const indexed = await this.deps.store.listDocumentIds(filter);
2966
+ const orphans = indexed.filter((id) => !live.has(id));
2967
+ if (options.dryRun !== true) {
2968
+ await this.deps.store.removeMany(orphans);
2969
+ }
2970
+ const known = new Set(indexed);
2971
+ return {
2972
+ orphans,
2973
+ missing: [
2974
+ ...live
2975
+ ].filter((id) => !known.has(id))
2976
+ };
2977
+ }
2978
+ /** A `search_*` tool whose scope comes from the context. Pass it to `provideAgentTool`. */
2979
+ tool(options) {
2980
+ return createRetrievalTool(this.baseRetriever(), {
2981
+ name: options.name,
2982
+ ...options.description !== void 0 ? {
2983
+ description: options.description
2984
+ } : {},
2985
+ ...options.topK !== void 0 ? {
2986
+ topK: options.topK
2987
+ } : {},
2988
+ filter: /* @__PURE__ */ __name((ctx) => this.scopeFilter(options.scope(ctx)), "filter")
2989
+ });
2990
+ }
2991
+ baseRetriever() {
2992
+ const { minScore } = this.definition;
2993
+ return this.deps.retriever ?? new EmbeddingRetriever(this.embedder, this.deps.store, minScore !== void 0 ? {
2994
+ minScore
2995
+ } : {});
2996
+ }
2997
+ /** The filter a scope compiles to. Throws on anything but a usable value, before any paid call. */
2998
+ scopeFilter(scope) {
2999
+ const given = scope ?? {};
3000
+ const filter = {
3001
+ [RAG_ENTITY_KEY]: this.definition.name
3002
+ };
3003
+ for (const key of Object.keys(this.definition.scope)) {
3004
+ const value = given[key];
3005
+ if (isAnyOf(value) && value[ANY_OF].every(isScopeValue)) {
3006
+ filter[key] = value[ANY_OF].map(String);
3007
+ } else if (isScopeValue(value)) {
3008
+ filter[key] = String(value);
3009
+ } else {
3010
+ throw new RagScopeRequiredError(this.definition.name, key);
3011
+ }
3012
+ }
3013
+ return filter;
3014
+ }
3015
+ };
3016
+ function ragFilterableFields(...scopeKeys) {
3017
+ return [
3018
+ .../* @__PURE__ */ new Set([
3019
+ RAG_ENTITY_KEY,
3020
+ ...scopeKeys
3021
+ ])
3022
+ ];
3023
+ }
3024
+ __name(ragFilterableFields, "ragFilterableFields");
3025
+
2731
3026
  // src/tree/types.ts
2732
3027
  function* walkTree(nodes, path = []) {
2733
3028
  for (const node of nodes) {
@@ -5007,8 +5302,11 @@ __name(createNavigateDocumentTool, "createNavigateDocumentTool");
5007
5302
  PgLexicalVectorStore,
5008
5303
  PgVectorStore,
5009
5304
  RAG_DIAGNOSTIC_LIB,
5305
+ RAG_ENTITY_KEY,
5010
5306
  RAG_RETRIEVAL_CHANNEL,
5011
5307
  RAG_RETRIEVAL_EVENT,
5308
+ RagIndex,
5309
+ RagScopeRequiredError,
5012
5310
  RedisVectorSchemaMismatchError,
5013
5311
  RedisVectorStore,
5014
5312
  RerankingRetriever,
@@ -5016,6 +5314,7 @@ __name(createNavigateDocumentTool, "createNavigateDocumentTool");
5016
5314
  TreeNavigationRetriever,
5017
5315
  TwoStageRetriever,
5018
5316
  UnsafeRemovalError,
5317
+ anyOf,
5019
5318
  anyTermTsquery,
5020
5319
  applyMetadataPatch,
5021
5320
  buildDocumentTree,
@@ -5031,6 +5330,7 @@ __name(createNavigateDocumentTool, "createNavigateDocumentTool");
5031
5330
  describesRetrieval,
5032
5331
  detectHeadings,
5033
5332
  documentIdOf,
5333
+ embedCountingUsage,
5034
5334
  emitRagRetrieval,
5035
5335
  estimateTokens,
5036
5336
  hasSearchSyntax,
@@ -5046,6 +5346,8 @@ __name(createNavigateDocumentTool, "createNavigateDocumentTool");
5046
5346
  openAiChatTreeLlm,
5047
5347
  openAiEmbeddings,
5048
5348
  parseJsonReply,
5349
+ ragFilterableFields,
5350
+ replaceDocumentChunks,
5049
5351
  spanText,
5050
5352
  stableHash,
5051
5353
  toPgClient,