fix(mongodb): mongodb storage mget must return undefined when key not…

… exist (langchain-ai#6445) Co-authored-by: Roman Tyshyk <[email protected]> Co-authored-by: jacoblee93 <[email protected]>
CarterMorris · Nov 10, 2024 · 01a9472 · 01a9472
1 parent 24b40fb
commit 01a9472
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 17 deletions.
diff --git a/libs/langchain-mongodb/src/storage.ts b/libs/langchain-mongodb/src/storage.ts
@@ -35,7 +35,7 @@ export interface MongoDBStoreInput {
  * const store = new MongoDBStore({
  *   collection,
  * });
- * 
+ *
  * const docs = [
  *   [uuidv4(), "Dogs are tough."],
  *   [uuidv4(), "Cats are tough."],
@@ -97,11 +97,18 @@ export class MongoDBStore extends BaseStore<string, Uint8Array> {
       .toArray();
 
     const encoder = new TextEncoder();
-    return retrievedValues.map((value) => {
-      if (!("value" in value)) {
+    const valueMap = new Map(
+      retrievedValues.map((item) => [item[this.primaryKey], item])
+    );
+
+    return prefixedKeys.map((prefixedKey) => {
+      const value = valueMap.get(prefixedKey);
+
+      if (!value) {
         return undefined;
       }
-      if (value === undefined || value === null) {
+
+      if (!("value" in value)) {
         return undefined;
       } else if (typeof value.value === "object") {
         return encoder.encode(JSON.stringify(value.value));

diff --git a/libs/langchain-mongodb/src/tests/storage.int.test.ts b/libs/langchain-mongodb/src/tests/storage.int.test.ts
@@ -43,18 +43,22 @@ test("MongoDBStore can set and retrieve", async () => {
       encoder.encode(doc[1]),
     ]);
     await store.mset(docsAsKVPairs);
-    const retrievedDocs = (await store.mget(docs.map((doc) => doc[0]))).flatMap(
-      (doc) => {
-        if (doc !== undefined) {
-          const decodedDoc = decoder.decode(doc);
-          const parsedDoc = JSON.parse(decodedDoc);
-          return [parsedDoc];
-        }
-        return [];
-      }
-    );
-
-    expect(retrievedDocs.sort()).toEqual(docs.map((doc) => doc[1]).sort());
+
+    const keysToRetrieve = docs.map((doc) => doc[0]);
+    keysToRetrieve.unshift("nonexistent_key_0");
+    keysToRetrieve.push("nonexistent_key_3");
+
+    const retrievedDocs = await store.mget(keysToRetrieve);
+    expect(retrievedDocs.length).toBe(keysToRetrieve.length);
+    // Check that the first item is undefined (nonexistent_key_0)
+    expect(retrievedDocs[0]).toBeUndefined();
+
+    // Check that the second and third items match the original docs
+    expect(decoder.decode(retrievedDocs[1])).toBe(docs[0][1]);
+    expect(decoder.decode(retrievedDocs[2])).toBe(docs[1][1]);
+
+    // Check that the last item is undefined (nonexistent_key_1)
+    expect(retrievedDocs[retrievedDocs.length - 1]).toBeUndefined();
   } finally {
     const keys = store.yieldKeys();
     const yieldedKeys = [];
@@ -104,7 +108,9 @@ test("MongoDBStore can delete", async () => {
 
     const retrievedDocs = await store.mget(docs.map((doc) => doc[0]));
 
-    expect(retrievedDocs.length).toBe(0);
+    expect(retrievedDocs.length).toBe(2);
+    const everyValueUndefined = retrievedDocs.every((v) => v === undefined);
+    expect(everyValueUndefined).toBe(true);
   } finally {
     const keys = store.yieldKeys();
     const yieldedKeys = [];