@dudousxd/nestjs-agent-rag 0.14.0 → 0.14.2

package/README.md

@@ -481,7 +481,8 @@ array is refused so a request can't widen a search. A corpus shared by everyone
 
 **Saving doesn't index.** Embedding is slow and paid for, so call `index` from a job or wherever the
 text changed. It embeds first, so a provider failure keeps the old chunks. On `PgVectorStore` the swap is
-one transaction; on stores without `replaceDocument`, it is two statements.
+one transaction. On Redis, a search can briefly see both versions, but concurrent calls wait for each
+other, so they can't leave a mix behind. On stores without `replaceDocument`, it is two statements.
 
 **Deleting.** With Drizzle, use `.returning()` and call `remove`. With MikroORM, add
 `provideRagDeleteSync(ProjectDocument, DocumentsRag)` from `@dudousxd/nestjs-agent-store-mikro-orm` to

package/dist/index.cjs

@@ -279,6 +279,7 @@ var MemoryVectorStore = class {
     __name(this, "MemoryVectorStore");
   }
   records = /* @__PURE__ */ new Map();
+  replacing = /* @__PURE__ */ new Map();
   /**
   * Telemetry self-description. No `collection`: there is no namespace to name — the corpus IS this
   * instance, and two `MemoryVectorStore`s in one process are two corpora with nothing to tell them
@@ -294,6 +295,23 @@ var MemoryVectorStore = class {
       this.records.set(record.id, record);
     }
   }
+  /** Replaces take turns per document. A `remove` during one isn't queued, and may be undone by it. */
+  async replaceDocument(documentId, records) {
+    const previous = this.replacing.get(documentId) ?? Promise.resolve();
+    const run = previous.catch(() => {
+    }).then(async () => {
+      await this.remove(documentId);
+      await this.upsert(records);
+    });
+    this.replacing.set(documentId, run);
+    try {
+      await run;
+    } finally {
+      if (this.replacing.get(documentId) === run) {
+        this.replacing.delete(documentId);
+      }
+    }
+  }
   async remove(documentId) {
     for (const id of this.records.keys()) {
       if (documentIdOf(id) === documentId) {
@@ -767,6 +785,10 @@ var PgVectorStore = class {
   /**
   * 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.
+  *
+  * Concurrent swaps of one document take turns on an advisory lock. Without it the second `DELETE`
+  * can't see the rows the first just committed, so the first's tail survives next to the second's
+  * chunks, and overlapping writes can deadlock.
   */
   async replaceDocument(documentId, records) {
     const swap = /* @__PURE__ */ __name(async (client) => {
@@ -776,7 +798,16 @@ var PgVectorStore = class {
       await this.upsertWith(client, records);
     }, "swap");
     const transaction = this.client.transaction?.bind(this.client);
-    await (transaction === void 0 ? swap(this.client) : transaction(swap));
+    if (transaction === void 0) {
+      await swap(this.client);
+      return;
+    }
+    await transaction(async (client) => {
+      await client.query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))", [
+        `${this.table}:${stripNulBytes(documentId)}`
+      ]);
+      await swap(client);
+    });
   }
   async upsertWith(client, records) {
     const byId = /* @__PURE__ */ new Map();
@@ -1191,6 +1222,11 @@ function toVectorLiteral(embedding) {
 __name(toVectorLiteral, "toVectorLiteral");
 
 // src/redis-vector-store.ts
+var import_node_crypto = require("crypto");
+var REPLACE_LOCK_TTL_MS = 3e4;
+var REPLACE_LOCK_WAIT_MS = 12e4;
+var RENEW_LOCK_SCRIPT = "if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('PEXPIRE', KEYS[1], ARGV[2]) end return 0";
+var RELEASE_LOCK_SCRIPT = "if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('DEL', KEYS[1]) end return 0";
 var RedisVectorSchemaMismatchError = class extends Error {
   static {
     __name(this, "RedisVectorSchemaMismatchError");
@@ -1352,6 +1388,92 @@ var RedisVectorStore = class {
     }
   }
   /**
+  * Replaces a document's chunks, one writer at a time per document. Redis can't swap them in one step
+  * without scanning the keyspace inside a script, so a search may briefly see old and new chunks
+  * together, but concurrent calls can't leave them mixed. The new chunks go in before the old ones
+  * come out, so a crash leaves the document searchable and the next call repairs it.
+  */
+  async replaceDocument(documentId, records) {
+    await this.withDocumentLock(documentId, async (assertHeld) => {
+      await this.upsert(records);
+      assertHeld();
+      const current = new Set(records.map((record) => `${this.prefix}${record.id}`));
+      const stale = (await this.chunkKeys(documentId)).filter((key) => !current.has(key));
+      assertHeld();
+      if (stale.length > 0) {
+        await this.client.sendCommand([
+          "DEL",
+          ...stale
+        ]);
+      }
+    });
+  }
+  /**
+  * A lease, so a crashed writer can't block the document forever. It is renewed while the work runs,
+  * which keeps a large document from outliving it.
+  */
+  async withDocumentLock(documentId, work) {
+    const key = `${this.index}:replace-lock:${documentId}`;
+    const token = (0, import_node_crypto.randomUUID)();
+    const deadline = Date.now() + REPLACE_LOCK_WAIT_MS;
+    let delay = 10;
+    while (await this.client.sendCommand([
+      "SET",
+      key,
+      token,
+      "NX",
+      "PX",
+      String(REPLACE_LOCK_TTL_MS)
+    ]) === null) {
+      if (Date.now() >= deadline) {
+        throw new Error(`RedisVectorStore: timed out after ${REPLACE_LOCK_WAIT_MS}ms waiting to change document "${documentId}"`);
+      }
+      await new Promise((resolve) => setTimeout(resolve, delay));
+      delay = Math.min(delay * 2, 250);
+    }
+    let lost = false;
+    const assertHeld = /* @__PURE__ */ __name(() => {
+      if (lost) {
+        throw new Error(`RedisVectorStore: lost the lock while changing document "${documentId}"; call again to repair it`);
+      }
+    }, "assertHeld");
+    const renewal = setInterval(() => {
+      void Promise.resolve(this.client.sendCommand([
+        "EVAL",
+        RENEW_LOCK_SCRIPT,
+        "1",
+        key,
+        token,
+        String(REPLACE_LOCK_TTL_MS)
+      ])).then(
+        (renewed) => {
+          if (Number(renewed) === 0) {
+            lost = true;
+          }
+        },
+        // A failed renewal isn't a lost lock yet: the lease may still have time left.
+        () => {
+        }
+      );
+    }, REPLACE_LOCK_TTL_MS / 3);
+    renewal.unref?.();
+    try {
+      const result = await work(assertHeld);
+      assertHeld();
+      return result;
+    } finally {
+      clearInterval(renewal);
+      await Promise.resolve(this.client.sendCommand([
+        "EVAL",
+        RELEASE_LOCK_SCRIPT,
+        "1",
+        key,
+        token
+      ])).catch(() => {
+      });
+    }
+  }
+  /**
   * Single-document deletion, by `SCAN … MATCH`.
   *
   * This deliberately stays a keyspace scan even though the rest of this class no longer needs one.
@@ -1372,13 +1494,15 @@ var RedisVectorStore = class {
   * than N scans). Neither depends on a field that older chunks lack, so neither has a migration.
   */
   async remove(documentId) {
-    const keys = await this.chunkKeys(documentId);
-    if (keys.length > 0) {
-      await this.client.sendCommand([
-        "DEL",
-        ...keys
-      ]);
-    }
+    await this.withDocumentLock(documentId, async () => {
+      const keys = await this.chunkKeys(documentId);
+      if (keys.length > 0) {
+        await this.client.sendCommand([
+          "DEL",
+          ...keys
+        ]);
+      }
+    });
   }
   /**
   * Rewrite the metadata of every chunk of `documentId` in place, leaving `text` and `embedding`
@@ -2905,8 +3029,9 @@ var RagIndex = class {
   }
   /**
   * 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`.
+  * provider fails, the previous chunks are still there. The swap is one transaction on Postgres. Redis
+  * adds the new chunks before dropping the old, so a search can briefly see both. On a store without
+  * `replaceDocument`, a crash mid-swap loses the document until the next `index`.
   */
   async index(row) {
     const { definition, deps } = this;