@dudousxd/nestjs-agent-rag 0.13.0 → 0.14.0

package/dist/index.js

@@ -77,6 +77,16 @@ function chunkByRecord(text, chunkSize, separator) {
 __name(chunkByRecord, "chunkByRecord");
 
 // src/vector-store.ts
+async function replaceDocumentChunks(store, documentId, records) {
+  const atomic = store;
+  if (typeof atomic.replaceDocument === "function") {
+    await atomic.replaceDocument(documentId, records);
+    return;
+  }
+  await store.remove(documentId);
+  await store.upsert(records);
+}
+__name(replaceDocumentChunks, "replaceDocumentChunks");
 function isLexicalVectorStore(store) {
   return typeof store.searchText === "function";
 }
@@ -137,12 +147,19 @@ function applyMetadataPatch(metadata, patch) {
 __name(applyMetadataPatch, "applyMetadataPatch");
 
 // src/filter.ts
+function isMissingValue(value) {
+  return value === void 0 || Array.isArray(value) && value.includes(void 0);
+}
+__name(isMissingValue, "isMissingValue");
 function matchesFilter(metadata, filter) {
   const entries = Object.entries(filter);
   if (metadata === void 0) {
     return entries.length === 0;
   }
   return entries.every(([key, expected]) => {
+    if (isMissingValue(expected)) {
+      return false;
+    }
     if (Array.isArray(expected)) {
       const actual = metadata[key];
       return Array.isArray(actual) ? actual.some((value) => expected.includes(value)) : expected.includes(actual);
@@ -155,9 +172,18 @@ function filterMatchesNothing(filter) {
   if (filter === void 0) {
     return false;
   }
-  return Object.values(filter).some((value) => Array.isArray(value) && value.length === 0);
+  return Object.values(filter).some((value) => isMissingValue(value) || Array.isArray(value) && value.length === 0);
 }
 __name(filterMatchesNothing, "filterMatchesNothing");
+function undefinedFilterKeys(filter) {
+  return Object.entries(filter).filter(([, value]) => isMissingValue(value)).map(([key]) => key);
+}
+__name(undefinedFilterKeys, "undefinedFilterKeys");
+function undefinedFilterMessage(keys) {
+  const named = keys.map((key) => `"${key}"`).join(", ");
+  return `removeWhere() refuses a filter with no value for ${named}: the scope was probably never set.`;
+}
+__name(undefinedFilterMessage, "undefinedFilterMessage");
 
 // src/memory-vector-store.ts
 var MemoryVectorStore = class {
@@ -294,6 +320,10 @@ var MemoryVectorStore = class {
     if (Object.keys(filter).length === 0) {
       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()).");
     }
+    const missing = undefinedFilterKeys(filter);
+    if (missing.length > 0) {
+      throw new UnsafeRemovalError("undefined-value", undefinedFilterMessage(missing));
+    }
     if (filterMatchesNothing(filter)) {
       return 0;
     }
@@ -563,6 +593,7 @@ var PgVectorStore = class {
   dimensions;
   /** The widths with a partial HNSW index, for a mixed-dimension table. */
   indexedDimensions;
+  vectorType;
   nullableEmbeddings;
   iterativeScan;
   efSearch;
@@ -586,6 +617,10 @@ var PgVectorStore = class {
         ...new Set(dimensions)
       ];
     }
+    this.vectorType = options.vectorType ?? "vector";
+    if (this.vectorType === "halfvec" && this.dimensions === void 0) {
+      throw new Error('PgVectorStore: vectorType "halfvec" needs a single width, not a list');
+    }
     this.nullableEmbeddings = options.nullableEmbeddings ?? false;
     this.iterativeScan = options.iterativeScan;
     this.efSearch = options.efSearch;
@@ -604,7 +639,7 @@ var PgVectorStore = class {
   */
   schemaStatements() {
     const mixed = this.dimensions === void 0;
-    const column = mixed ? "vector" : `vector(${this.dimensions})`;
+    const column = mixed ? "vector" : `${this.vectorType}(${this.dimensions})`;
     const statements = [
       "CREATE EXTENSION IF NOT EXISTS vector",
       `CREATE TABLE IF NOT EXISTS ${this.table} (
@@ -623,7 +658,7 @@ var PgVectorStore = class {
       }
     } else {
       statements.push(`CREATE INDEX IF NOT EXISTS ${this.table}_embedding_idx
-        ON ${this.table} USING hnsw (embedding vector_cosine_ops)`);
+        ON ${this.table} USING hnsw (embedding ${this.vectorType}_cosine_ops)`);
     }
     return statements;
   }
@@ -639,6 +674,23 @@ var PgVectorStore = class {
   * {@link PgVectorStoreOptions.nullableEmbeddings}).
   */
   async upsert(records) {
+    await this.upsertWith(this.client, records);
+  }
+  /**
+  * Swaps a document's chunks in one transaction, so a crash can't lose it and a search sees the old
+  * version or the new one. A client without `transaction` falls back to two statements.
+  */
+  async replaceDocument(documentId, records) {
+    const swap = /* @__PURE__ */ __name(async (client) => {
+      await client.query(`DELETE FROM ${this.table} WHERE ${DOCUMENT_ID_FROM_CHUNK} = $1`, [
+        stripNulBytes(documentId)
+      ]);
+      await this.upsertWith(client, records);
+    }, "swap");
+    const transaction = this.client.transaction?.bind(this.client);
+    await (transaction === void 0 ? swap(this.client) : transaction(swap));
+  }
+  async upsertWith(client, records) {
     const byId = /* @__PURE__ */ new Map();
     for (const record of records) {
       const id = stripNulBytes(record.id);
@@ -660,9 +712,9 @@ var PgVectorStore = class {
         const metadata = record.metadata !== void 0 ? stripNulBytes(record.metadata) : void 0;
         const base = params.length;
         params.push(record.id, text, source ?? null, metadata !== void 0 ? JSON.stringify(metadata) : null, record.embedding.length > 0 ? toVectorLiteral(record.embedding) : null);
-        return `($${base + 1}, $${base + 2}, $${base + 3}, $${base + 4}::jsonb, $${base + 5}::vector)`;
+        return `($${base + 1}, $${base + 2}, $${base + 3}, $${base + 4}::jsonb, $${base + 5}::${this.vectorType})`;
       });
-      await this.client.query(`INSERT INTO ${this.table} (id, text, source, metadata, embedding)
+      await client.query(`INSERT INTO ${this.table} (id, text, source, metadata, embedding)
          VALUES ${rows.join(", ")}
          ON CONFLICT (id) DO UPDATE SET
            text = EXCLUDED.text,
@@ -786,6 +838,10 @@ var PgVectorStore = class {
     if (Object.keys(filter).length === 0) {
       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()).");
     }
+    const missing = undefinedFilterKeys(filter);
+    if (missing.length > 0) {
+      throw new UnsafeRemovalError("undefined-value", undefinedFilterMessage(missing));
+    }
     if (filterMatchesNothing(filter)) {
       return 0;
     }
@@ -817,7 +873,7 @@ var PgVectorStore = class {
       options.topK
     ];
     const conditions = this.whereConditions(options.filter, params);
-    let vectorType = "vector";
+    let vectorType = this.vectorType;
     if (this.dimensions === void 0) {
       const width = embedding.length;
       if (this.indexedDimensions.includes(width)) {
@@ -827,7 +883,7 @@ var PgVectorStore = class {
     } else if (this.nullableEmbeddings) {
       conditions.push("embedding IS NOT NULL");
     }
-    const column = vectorType === "vector" ? "embedding" : `(embedding::${vectorType})`;
+    const column = vectorType === this.vectorType ? "embedding" : `(embedding::${vectorType})`;
     const sql = `SELECT id, text, source, metadata, 1 - (${column} <=> $1::${vectorType}) AS score
        FROM ${this.table}
        ${conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : ""}
@@ -1014,6 +1070,10 @@ function buildWhere(filter, params) {
   const clauses = [];
   const scalar = {};
   for (const [key, value] of Object.entries(stripped)) {
+    if (value === void 0 || Array.isArray(value) && value.includes(void 0)) {
+      clauses.push("false");
+      continue;
+    }
     if (Array.isArray(value)) {
       if (value.length === 0) {
         clauses.push("false");
@@ -1555,6 +1615,10 @@ var RedisVectorStore = class {
     if (keys.length === 0) {
       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()).");
     }
+    const missing = undefinedFilterKeys(filter);
+    if (missing.length > 0) {
+      throw new UnsafeRemovalError("undefined-value", undefinedFilterMessage(missing));
+    }
     const unindexed = keys.filter((key) => !this.filterableFields.includes(key));
     if (unindexed.length > 0) {
       const named = unindexed.map((key) => `"${key}"`).join(", ");
@@ -2105,7 +2169,7 @@ var KeywordRetriever = class {
   }
 };
 function tokenize(text) {
-  return text.toLowerCase().match(/[a-z0-9]+/g) ?? [];
+  return text.normalize("NFKD").replace(/\p{M}/gu, "").toLowerCase().match(/[\p{L}\p{N}_]+/gu) ?? [];
 }
 __name(tokenize, "tokenize");
 
@@ -2426,13 +2490,19 @@ function createRetrievalTool(retriever, options = {}) {
       inputSchema: querySchema
     },
     handler: {
-      async execute(input) {
+      async execute(input, ctx) {
         if (!isQueryInput(input)) {
           throw new Error("retrieval tool expects { query: string }");
         }
-        const passages = await instrumented.retrieve(input.query, topK !== void 0 ? {
-          topK
-        } : {});
+        const filter = options.filter?.(ctx);
+        const passages = await instrumented.retrieve(input.query, {
+          ...topK !== void 0 ? {
+            topK
+          } : {},
+          ...filter !== void 0 ? {
+            filter
+          } : {}
+        });
         return {
           passages
         };
@@ -2472,7 +2542,7 @@ function openAiEmbeddings(options) {
   const maxInputChars = options.maxInputChars ?? 24e3;
   const timeoutMs = options.timeoutMs ?? 12e4;
   const doFetch = options.fetch ?? fetch;
-  const request = /* @__PURE__ */ __name(async (input) => {
+  const request = /* @__PURE__ */ __name(async (input, count) => {
     const response = await doFetch(url, {
       method: "POST",
       headers: {
@@ -2505,19 +2575,20 @@ function openAiEmbeddings(options) {
     const tokens = json2?.usage?.total_tokens ?? json2?.usage?.prompt_tokens;
     if (typeof tokens === "number" && Number.isFinite(tokens)) {
       options.onUsage?.(tokens);
+      count(tokens);
     }
     return items.map((item) => item.embedding);
   }, "request");
-  const requestSplitting = /* @__PURE__ */ __name(async (input) => {
+  const requestSplitting = /* @__PURE__ */ __name(async (input, count) => {
     if (input.length > batchSize) {
       const vectors = [];
       for (let start = 0; start < input.length; start += batchSize) {
-        vectors.push(...await requestSplitting(input.slice(start, start + batchSize)));
+        vectors.push(...await requestSplitting(input.slice(start, start + batchSize), count));
       }
       return vectors;
     }
     try {
-      return await request(input);
+      return await request(input, count);
     } catch (error) {
       if (!isBatchTooLarge(error) || input.length <= 1) {
         throw error;
@@ -2529,20 +2600,37 @@ function openAiEmbeddings(options) {
         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.`);
       }
       return [
-        ...await requestSplitting(input.slice(0, half)),
-        ...await requestSplitting(input.slice(half))
+        ...await requestSplitting(input.slice(0, half), count),
+        ...await requestSplitting(input.slice(half), count)
       ];
     }
   }, "requestSplitting");
+  const embedAll = /* @__PURE__ */ __name(async (texts, count) => {
+    const vectors = [];
+    for (let start = 0; start < texts.length; ) {
+      const input = texts.slice(start, start + batchSize).map((text) => (text.length > maxInputChars ? text.slice(0, maxInputChars) : text) || " ");
+      vectors.push(...await requestSplitting(input, count));
+      start += input.length;
+    }
+    return vectors;
+  }, "embedAll");
   return {
-    async embed(texts) {
-      const vectors = [];
-      for (let start = 0; start < texts.length; ) {
-        const input = texts.slice(start, start + batchSize).map((text) => (text.length > maxInputChars ? text.slice(0, maxInputChars) : text) || " ");
-        vectors.push(...await requestSplitting(input));
-        start += input.length;
-      }
-      return vectors;
+    embed: /* @__PURE__ */ __name((texts) => embedAll(texts, () => {
+    }), "embed"),
+    async embedWithUsage(texts) {
+      let inputTokens;
+      const vectors = await embedAll(texts, (tokens) => {
+        inputTokens = (inputTokens ?? 0) + tokens;
+      });
+      return {
+        vectors,
+        ...inputTokens !== void 0 ? {
+          usage: {
+            inputTokens,
+            modelId: options.model
+          }
+        } : {}
+      };
     }
   };
 }
@@ -2647,6 +2735,206 @@ function itemsOf(body) {
 }
 __name(itemsOf, "itemsOf");
 
+// src/embedding-usage.ts
+async function embedCountingUsage(provider, texts) {
+  const reporting = provider;
+  return typeof reporting.embedWithUsage === "function" ? reporting.embedWithUsage(texts) : {
+    vectors: await provider.embed(texts)
+  };
+}
+__name(embedCountingUsage, "embedCountingUsage");
+
+// src/rag-index.ts
+var RAG_ENTITY_KEY = "ragEntity";
+var RagScopeRequiredError = class extends Error {
+  static {
+    __name(this, "RagScopeRequiredError");
+  }
+  model;
+  key;
+  constructor(model, key) {
+    super(`${model} needs scope.${key}: its chunks are isolated by it`), this.model = model, this.key = key;
+    this.name = "RagScopeRequiredError";
+  }
+};
+var ANY_OF = /* @__PURE__ */ Symbol.for("@dudousxd/nestjs-agent-rag/anyOf");
+function anyOf(values) {
+  return {
+    [ANY_OF]: values
+  };
+}
+__name(anyOf, "anyOf");
+function isScopeValue(value) {
+  return typeof value === "number" ? Number.isFinite(value) : typeof value === "string" && value !== "";
+}
+__name(isScopeValue, "isScopeValue");
+function isAnyOf(value) {
+  return typeof value === "object" && value !== null && ANY_OF in value;
+}
+__name(isAnyOf, "isAnyOf");
+var RagIndex = class {
+  static {
+    __name(this, "RagIndex");
+  }
+  definition;
+  deps;
+  embedder;
+  constructor(definition, deps) {
+    this.definition = definition;
+    this.deps = deps;
+    const { onUsage } = deps;
+    this.embedder = onUsage === void 0 ? deps.embedder : {
+      embed: /* @__PURE__ */ __name(async (texts) => {
+        const { vectors, usage } = await embedCountingUsage(deps.embedder, texts);
+        if (usage !== void 0) {
+          await onUsage(usage);
+        }
+        return vectors;
+      }, "embed")
+    };
+  }
+  /** The id this row's chunks are stored under: `<name>:<id>`. */
+  documentId(id) {
+    if (!isScopeValue(id)) {
+      throw new Error(`${this.definition.name}: a row needs an id to be indexed or removed`);
+    }
+    return `${this.definition.name}:${String(id)}`;
+  }
+  /** A `Retriever` locked to this scope, for `HybridRetriever`, `createRetrievalTool` and so on. */
+  scoped(scope) {
+    return new FilteredRetriever(this.baseRetriever(), this.scopeFilter(scope));
+  }
+  async search(query, options) {
+    const { scope, filter, topK } = options;
+    return this.scoped(scope).retrieve(query, {
+      ...topK !== void 0 ? {
+        topK
+      } : {},
+      ...filter !== void 0 ? {
+        filter
+      } : {}
+    });
+  }
+  /**
+  * Replaces the row's chunks and returns how many were written. Embeddings come first: if the
+  * provider fails, the previous chunks are still there. The swap is one transaction on Postgres; on a
+  * store without `replaceDocument` a crash mid-swap loses the document until the next `index`.
+  */
+  async index(row) {
+    const { definition, deps } = this;
+    const documentId = this.documentId(definition.id(row));
+    const scope = this.scopeFilter(Object.fromEntries(Object.entries(definition.scope).map(([key, read]) => [
+      key,
+      read(row)
+    ])));
+    const text = definition.text(row)?.trim();
+    if (!text) {
+      await deps.store.remove(documentId);
+      return 0;
+    }
+    const pieces = chunkText(text, definition.chunk).slice(0, definition.maxChunks);
+    const embeddings = await this.embedder.embed(pieces);
+    const source = definition.source?.(row) ?? void 0;
+    const metadata = {
+      ...definition.metadata?.(row),
+      ...scope
+    };
+    await replaceDocumentChunks(deps.store, documentId, pieces.map((piece, index) => ({
+      id: `${documentId}#${index}`,
+      text: piece,
+      embedding: embeddings[index] ?? [],
+      ...source !== void 0 ? {
+        source
+      } : {},
+      metadata
+    })));
+    return pieces.length;
+  }
+  /** Takes the id itself once the row is gone. */
+  remove(rowOrId) {
+    const id = typeof rowOrId === "string" || typeof rowOrId === "number" ? rowOrId : this.definition.id(rowOrId);
+    return this.deps.store.remove(this.documentId(id));
+  }
+  /**
+  * Diffs the index against the ids that exist, and removes what no hook could: `ON DELETE CASCADE`,
+  * raw SQL, a crash after a commit. `liveIds` must be complete for the part being compared, because
+  * anything left out is removed. Use `dryRun` to look first.
+  */
+  async reconcile(liveIds, options = {}) {
+    const filter = {
+      [RAG_ENTITY_KEY]: this.definition.name
+    };
+    for (const [key, value] of Object.entries(options.scope ?? {})) {
+      if (!isScopeValue(value)) {
+        throw new RagScopeRequiredError(this.definition.name, key);
+      }
+      filter[key] = String(value);
+    }
+    const live = /* @__PURE__ */ new Set();
+    for await (const id of liveIds) {
+      live.add(this.documentId(id));
+    }
+    const indexed = await this.deps.store.listDocumentIds(filter);
+    const orphans = indexed.filter((id) => !live.has(id));
+    if (options.dryRun !== true) {
+      await this.deps.store.removeMany(orphans);
+    }
+    const known = new Set(indexed);
+    return {
+      orphans,
+      missing: [
+        ...live
+      ].filter((id) => !known.has(id))
+    };
+  }
+  /** A `search_*` tool whose scope comes from the context. Pass it to `provideAgentTool`. */
+  tool(options) {
+    return createRetrievalTool(this.baseRetriever(), {
+      name: options.name,
+      ...options.description !== void 0 ? {
+        description: options.description
+      } : {},
+      ...options.topK !== void 0 ? {
+        topK: options.topK
+      } : {},
+      filter: /* @__PURE__ */ __name((ctx) => this.scopeFilter(options.scope(ctx)), "filter")
+    });
+  }
+  baseRetriever() {
+    const { minScore } = this.definition;
+    return this.deps.retriever ?? new EmbeddingRetriever(this.embedder, this.deps.store, minScore !== void 0 ? {
+      minScore
+    } : {});
+  }
+  /** The filter a scope compiles to. Throws on anything but a usable value, before any paid call. */
+  scopeFilter(scope) {
+    const given = scope ?? {};
+    const filter = {
+      [RAG_ENTITY_KEY]: this.definition.name
+    };
+    for (const key of Object.keys(this.definition.scope)) {
+      const value = given[key];
+      if (isAnyOf(value) && value[ANY_OF].every(isScopeValue)) {
+        filter[key] = value[ANY_OF].map(String);
+      } else if (isScopeValue(value)) {
+        filter[key] = String(value);
+      } else {
+        throw new RagScopeRequiredError(this.definition.name, key);
+      }
+    }
+    return filter;
+  }
+};
+function ragFilterableFields(...scopeKeys) {
+  return [
+    .../* @__PURE__ */ new Set([
+      RAG_ENTITY_KEY,
+      ...scopeKeys
+    ])
+  ];
+}
+__name(ragFilterableFields, "ragFilterableFields");
+
 // src/tree/types.ts
 function* walkTree(nodes, path = []) {
   for (const node of nodes) {
@@ -4925,8 +5213,11 @@ export {
   PgLexicalVectorStore,
   PgVectorStore,
   RAG_DIAGNOSTIC_LIB,
+  RAG_ENTITY_KEY,
   RAG_RETRIEVAL_CHANNEL,
   RAG_RETRIEVAL_EVENT,
+  RagIndex,
+  RagScopeRequiredError,
   RedisVectorSchemaMismatchError,
   RedisVectorStore,
   RerankingRetriever,
@@ -4934,6 +5225,7 @@ export {
   TreeNavigationRetriever,
   TwoStageRetriever,
   UnsafeRemovalError,
+  anyOf,
   anyTermTsquery,
   applyMetadataPatch,
   buildDocumentTree,
@@ -4949,6 +5241,7 @@ export {
   describesRetrieval,
   detectHeadings,
   documentIdOf,
+  embedCountingUsage,
   emitRagRetrieval,
   estimateTokens,
   hasSearchSyntax,
@@ -4964,6 +5257,8 @@ export {
   openAiChatTreeLlm,
   openAiEmbeddings,
   parseJsonReply,
+  ragFilterableFields,
+  replaceDocumentChunks,
   spanText,
   stableHash,
   toPgClient,