@dudousxd/nestjs-agent-rag 0.14.1 → 0.14.3

package/dist/index.js

@@ -191,6 +191,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
@@ -206,7 +207,30 @@ var MemoryVectorStore = class {
       this.records.set(record.id, record);
     }
   }
+  /** A replace and a removal of one document take turns, in the order they were called. */
+  async replaceDocument(documentId, records) {
+    await this.inTurn(documentId, async () => {
+      this.deleteChunks(documentId);
+      await this.upsert(records);
+    });
+  }
   async remove(documentId) {
+    await this.inTurn(documentId, async () => this.deleteChunks(documentId));
+  }
+  async inTurn(documentId, work) {
+    const previous = this.replacing.get(documentId) ?? Promise.resolve();
+    const run = previous.catch(() => {
+    }).then(work);
+    this.replacing.set(documentId, run);
+    try {
+      await run;
+    } finally {
+      if (this.replacing.get(documentId) === run) {
+        this.replacing.delete(documentId);
+      }
+    }
+  }
+  deleteChunks(documentId) {
     for (const id of this.records.keys()) {
       if (documentIdOf(id) === documentId) {
         this.records.delete(id);
@@ -303,12 +327,9 @@ var MemoryVectorStore = class {
     if (documentIds.length === 0) {
       return;
     }
-    const wanted = new Set(documentIds);
-    for (const id of this.records.keys()) {
-      if (wanted.has(documentIdOf(id))) {
-        this.records.delete(id);
-      }
-    }
+    await Promise.all([
+      ...new Set(documentIds)
+    ].map((documentId) => this.remove(documentId)));
   }
   /**
   * See {@link VectorStore.removeWhere}. The empty-array short-circuit below is redundant
@@ -327,6 +348,7 @@ var MemoryVectorStore = class {
     if (filterMatchesNothing(filter)) {
       return 0;
     }
+    await Promise.allSettled(this.replacing.values());
     let removed = 0;
     for (const [id, record] of this.records) {
       if (matchesFilter(record.metadata, filter)) {
@@ -1121,6 +1143,19 @@ 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 SWAP_CHUNKS_SCRIPT = `
+if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end
+local stale = tonumber(ARGV[2])
+for i = 2, stale + 1 do redis.call('DEL', KEYS[i]) end
+local arg = 3
+for i = stale + 2, #KEYS do
+  local fields = tonumber(ARGV[arg])
+  redis.call('DEL', KEYS[i])
+  redis.call('HSET', KEYS[i], unpack(ARGV, arg + 1, arg + fields))
+  arg = arg + fields + 1
+end
+return 1`;
+var SWAP_MAX_CHUNKS = 256;
 var RedisVectorSchemaMismatchError = class extends Error {
   static {
     __name(this, "RedisVectorSchemaMismatchError");
@@ -1258,42 +1293,56 @@ var RedisVectorStore = class {
   }
   async upsert(records) {
     for (const record of records) {
-      const args = [
+      await this.client.sendCommand([
         "HSET",
         `${this.prefix}${record.id}`,
-        "text",
-        record.text,
-        "embedding",
-        encodeVector(record.embedding)
-      ];
-      if (record.source !== void 0) {
-        args.push("source", record.source);
-      }
-      if (record.metadata !== void 0) {
-        args.push("metadata_json", JSON.stringify(record.metadata));
-        for (const field of this.filterableFields) {
-          const value = record.metadata[field];
-          if (value !== void 0) {
-            args.push(`meta_${field}`, encodeTag(value));
-          }
+        ...this.chunkFields(record)
+      ]);
+    }
+  }
+  chunkFields(record) {
+    const fields = [
+      "text",
+      record.text,
+      "embedding",
+      encodeVector(record.embedding)
+    ];
+    if (record.source !== void 0) {
+      fields.push("source", record.source);
+    }
+    if (record.metadata !== void 0) {
+      fields.push("metadata_json", JSON.stringify(record.metadata));
+      for (const field of this.filterableFields) {
+        const value = record.metadata[field];
+        if (value !== void 0) {
+          fields.push(`meta_${field}`, encodeTag(value));
         }
       }
-      await this.client.sendCommand(args);
     }
+    return fields;
   }
   /**
-  * 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.
+  * Replaces a document's chunks, one writer at a time per document. The swap is one script, so a search
+  * sees the old version or the new one, and a crash leaves the old one. The keys to drop are found
+  * first, under the lock, because a script can't scan the keyspace.
+  *
+  * A document past {@link SWAP_MAX_CHUNKS} would block Redis for too long that way. Its new chunks go in
+  * first and the old ones come out after, so a search may briefly see both, and a crash leaves the
+  * document searchable for the next call to repair.
   */
   async replaceDocument(documentId, records) {
-    await this.withDocumentLock(documentId, async (assertHeld) => {
-      await this.upsert(records);
-      assertHeld();
+    await this.withDocumentLocks([
+      documentId
+    ], async (assertHeld, token) => {
       const current = new Set(records.map((record) => `${this.prefix}${record.id}`));
       const stale = (await this.chunkKeys(documentId)).filter((key) => !current.has(key));
       assertHeld();
+      if (records.length + stale.length <= SWAP_MAX_CHUNKS) {
+        await this.swapChunks(documentId, token, stale, records);
+        return;
+      }
+      await this.upsert(records);
+      assertHeld();
       if (stale.length > 0) {
         await this.client.sendCommand([
           "DEL",
@@ -1302,69 +1351,110 @@ var RedisVectorStore = class {
       }
     });
   }
+  async swapChunks(documentId, token, stale, records) {
+    const keys = [
+      this.lockKey(documentId),
+      ...stale,
+      ...records.map((record) => `${this.prefix}${record.id}`)
+    ];
+    const args = [
+      token,
+      String(stale.length)
+    ];
+    for (const record of records) {
+      const fields = this.chunkFields(record);
+      args.push(String(fields.length), ...fields);
+    }
+    const swapped = await this.client.sendCommand([
+      "EVAL",
+      SWAP_CHUNKS_SCRIPT,
+      String(keys.length),
+      ...keys,
+      ...args
+    ]);
+    if (Number(swapped) === 0) {
+      throw lostLockError(documentId);
+    }
+  }
+  lockKey(documentId) {
+    return `${this.index}:replace-lock:${documentId}`;
+  }
   /**
-  * 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.
+  * A lease on each document, so a crashed writer can't block one forever. They are renewed while the
+  * work runs, which keeps a large document from outliving them.
+  *
+  * Several documents are locked in sorted order, without duplicates. Every other caller either holds
+  * one lock or takes them in that same order, so no two calls can wait on each other.
   */
-  async withDocumentLock(documentId, work) {
-    const key = `${this.index}:replace-lock:${documentId}`;
+  async withDocumentLocks(documentIds, work) {
     const token = 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 replace document "${documentId}"`);
-      }
-      await new Promise((resolve) => setTimeout(resolve, delay));
-      delay = Math.min(delay * 2, 250);
-    }
-    let lost = false;
+    const held = [];
+    let lostDocument;
     const assertHeld = /* @__PURE__ */ __name(() => {
-      if (lost) {
-        throw new Error(`RedisVectorStore: lost the lock while replacing document "${documentId}"; call again to repair it`);
+      if (lostDocument !== void 0) {
+        throw lostLockError(lostDocument);
       }
     }, "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;
+      for (const key of held) {
+        void Promise.resolve(this.client.sendCommand([
+          "EVAL",
+          RENEW_LOCK_SCRIPT,
+          "1",
+          key,
+          token,
+          String(REPLACE_LOCK_TTL_MS)
+        ])).then(
+          (renewed) => {
+            if (Number(renewed) === 0) {
+              lostDocument = key.slice(this.lockKey("").length);
+            }
+          },
+          // A failed renewal isn't a lost lock yet: the lease may still have time left.
+          () => {
           }
-        },
-        // 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);
+      for (const documentId of [
+        ...new Set(documentIds)
+      ].sort()) {
+        await this.acquireLock(documentId, token);
+        held.push(this.lockKey(documentId));
+      }
+      const result = await work(assertHeld, token);
       assertHeld();
       return result;
     } finally {
       clearInterval(renewal);
-      await Promise.resolve(this.client.sendCommand([
+      await Promise.all(held.map((key) => Promise.resolve(this.client.sendCommand([
         "EVAL",
         RELEASE_LOCK_SCRIPT,
         "1",
         key,
         token
       ])).catch(() => {
-      });
+      })));
+    }
+  }
+  async acquireLock(documentId, token) {
+    const deadline = Date.now() + REPLACE_LOCK_WAIT_MS;
+    let delay = 10;
+    while (await this.client.sendCommand([
+      "SET",
+      this.lockKey(documentId),
+      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);
     }
   }
   /**
@@ -1388,13 +1478,17 @@ 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.withDocumentLocks([
+      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`
@@ -1651,44 +1745,51 @@ var RedisVectorStore = class {
   *
   * It deletes by key, exactly as `remove` does, so it is correct on chunks written by any version of
   * this library — there is no field it needs the corpus to already carry.
+  *
+  * It holds every document's lock for the whole scan, so it waits for a `replaceDocument` in progress
+  * and one can't start partway through. That adds two commands per document to the one scan.
   */
   async removeMany(documentIds) {
     if (documentIds.length === 0) {
       return;
     }
-    const wanted = new Set(documentIds);
-    let pending = documentIds.map((id) => `${this.prefix}${id}`);
-    let cursor = "0";
-    do {
-      const reply = await this.client.sendCommand([
-        "SCAN",
-        cursor,
-        "MATCH",
-        `${escapeGlob(this.prefix)}*`,
-        "COUNT",
-        String(SCAN_PAGE)
-      ]);
-      const page = parseScanReply(reply);
-      cursor = page.cursor;
-      for (const key of page.keys) {
-        if (wanted.has(documentIdOf(stripPrefix(key, this.prefix)))) {
-          pending.push(key);
+    await this.withDocumentLocks(documentIds, async (assertHeld) => {
+      const wanted = new Set(documentIds);
+      let pending = documentIds.map((id) => `${this.prefix}${id}`);
+      let cursor = "0";
+      do {
+        const reply = await this.client.sendCommand([
+          "SCAN",
+          cursor,
+          "MATCH",
+          `${escapeGlob(this.prefix)}*`,
+          "COUNT",
+          String(SCAN_PAGE)
+        ]);
+        const page = parseScanReply(reply);
+        cursor = page.cursor;
+        for (const key of page.keys) {
+          if (wanted.has(documentIdOf(stripPrefix(key, this.prefix)))) {
+            pending.push(key);
+          }
         }
-      }
-      if (pending.length >= DELETE_BATCH) {
+        if (pending.length >= DELETE_BATCH) {
+          assertHeld();
+          await this.client.sendCommand([
+            "DEL",
+            ...pending
+          ]);
+          pending = [];
+        }
+      } while (cursor !== "0");
+      if (pending.length > 0) {
+        assertHeld();
         await this.client.sendCommand([
           "DEL",
           ...pending
         ]);
-        pending = [];
       }
-    } while (cursor !== "0");
-    if (pending.length > 0) {
-      await this.client.sendCommand([
-        "DEL",
-        ...pending
-      ]);
-    }
+    });
   }
   /**
   * Delete every chunk matching `filter`, via `FT.SEARCH … NOCONTENT` + `DEL` — no keyspace scan at
@@ -1713,6 +1814,10 @@ var RedisVectorStore = class {
   * from the index, so paging forward would skip their replacements. If a page comes back non-empty
   * but `DEL` reports nothing removed (keys expired between query and delete, or the index is stale),
   * the loop stops rather than spinning on results it can never consume.
+  *
+  * The documents aren't known until a page comes back, so each page locks the documents it touches, one
+  * page at a time, and asks again with `INKEYS` once it holds them. A `replaceDocument` that finished in
+  * between may have dropped those chunks, or rewritten them to no longer match.
   */
   async removeWhere(filter) {
     const keys = Object.keys(filter);
@@ -1734,6 +1839,7 @@ var RedisVectorStore = class {
     }
     const query = buildFilter(filter);
     let removed = 0;
+    let emptyPageBefore = false;
     for (; ; ) {
       const reply = await this.client.sendCommand([
         "FT.SEARCH",
@@ -1750,14 +1856,46 @@ var RedisVectorStore = class {
       if (page.length === 0) {
         break;
       }
-      const deleted = Number(await this.client.sendCommand([
-        "DEL",
-        ...page
-      ]));
-      if (!Number.isFinite(deleted) || deleted === 0) {
+      const documents = page.map((key) => documentIdOf(stripPrefix(key, this.prefix)));
+      const outcome = await this.withDocumentLocks(documents, async (assertHeld) => {
+        const confirmed = parseSearchKeys(await this.client.sendCommand([
+          "FT.SEARCH",
+          this.index,
+          query,
+          "NOCONTENT",
+          "INKEYS",
+          String(page.length),
+          ...page,
+          "LIMIT",
+          "0",
+          String(DELETE_BATCH),
+          "DIALECT",
+          "2"
+        ]));
+        assertHeld();
+        return confirmed.length === 0 ? {
+          confirmed: 0,
+          deleted: 0
+        } : {
+          confirmed: confirmed.length,
+          deleted: Number(await this.client.sendCommand([
+            "DEL",
+            ...confirmed
+          ]))
+        };
+      });
+      if (outcome.confirmed === 0) {
+        if (emptyPageBefore) {
+          break;
+        }
+        emptyPageBefore = true;
+        continue;
+      }
+      emptyPageBefore = false;
+      if (!Number.isFinite(outcome.deleted) || outcome.deleted === 0) {
         break;
       }
-      removed += deleted;
+      removed += outcome.deleted;
     }
     return removed;
   }
@@ -1855,6 +1993,10 @@ var RedisVectorStore = class {
     return parseSearchReply(reply, this.prefix, "lexical");
   }
 };
+function lostLockError(documentId) {
+  return new Error(`RedisVectorStore: lost the lock while changing document "${documentId}"; call again to repair it`);
+}
+__name(lostLockError, "lostLockError");
 var SCAN_PAGE = 1e3;
 var DELETE_BATCH = 512;
 function stripPrefix(key, prefix) {
@@ -2921,9 +3063,10 @@ 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. 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`.
+  * provider fails, the previous chunks are still there. The swap is one transaction on Postgres and one
+  * script on Redis, where a document too large for that 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;