langchain[minor]: Support storing and retrieving Document id field in MemoryVectorStore (#6572)

Co-authored-by: dschenkelman <[email protected]>
Co-authored-by: Brace Sproul <[email protected]>

Author: 3 people

langchain/src/vectorstores/memory.ts

@@ -17,6 +17,7 @@ interface MemoryVector {
   embedding: number[];
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   metadata: Record<string, any>;
+  id?: string;
 }
 
 /**
@@ -186,6 +187,7 @@ export class MemoryVectorStore extends VectorStore {
       content: documents[idx].pageContent,
       embedding,
       metadata: documents[idx].metadata,
+      id: documents[idx].id,
     }));
 
     this.memoryVectors = this.memoryVectors.concat(memoryVectors);
@@ -204,6 +206,7 @@ export class MemoryVectorStore extends VectorStore {
       const doc = new Document({
         metadata: memoryVector.metadata,
         pageContent: memoryVector.content,
+        id: memoryVector.id,
       });
       return filter(doc);
     };
@@ -215,6 +218,7 @@ export class MemoryVectorStore extends VectorStore {
         metadata: vector.metadata,
         content: vector.content,
         embedding: vector.embedding,
+        id: vector.id,
       }))
       .sort((a, b) => (a.similarity > b.similarity ? -1 : 0))
       .slice(0, k);
@@ -240,6 +244,7 @@ export class MemoryVectorStore extends VectorStore {
       new Document({
         metadata: search.metadata,
         pageContent: search.content,
+        id: search.id,
       }),
       search.similarity,
     ]);
@@ -273,6 +278,7 @@ export class MemoryVectorStore extends VectorStore {
         new Document({
           metadata: searches[idx].metadata,
           pageContent: searches[idx].content,
+          id: searches[idx].id,
         })
     );
   }

langchain/src/vectorstores/tests/memory.test.ts

@@ -30,6 +30,44 @@ test("MemoryVectorStore with external ids", async () => {
   ]);
 });
 
+test("MemoryVectorStore stores and retrieves document IDs", async () => {
+  const embeddings = new SyntheticEmbeddings({
+    vectorSize: 1536,
+  });
+  const store = new MemoryVectorStore(embeddings);
+
+  const filterFunc = (doc: DocumentInterface): boolean => {
+    const { metadata } = doc;
+    if (metadata.namespace <= 2) {
+      return true;
+    }
+    return false;
+  };
+
+  const retriever = store.asRetriever({
+    k: 2,
+    filter: filterFunc,
+  });
+
+  expect(retriever).toBeDefined();
+
+  await retriever.addDocuments([
+    { pageContent: "hello", metadata: { namespace: 1 }, id: "1" },
+    { pageContent: "hello", metadata: { namespace: 2 }, id: "2" },
+    { pageContent: "hello", metadata: { namespace: 3 }, id: "3" },
+    { pageContent: "hello", metadata: { namespace: 4 }, id: "4" },
+  ]);
+
+  const results = await retriever.getRelevantDocuments("hello");
+
+  expect(results).toHaveLength(2);
+
+  expect(results).toEqual([
+    new Document({ metadata: { namespace: 1 }, pageContent: "hello", id: "1" }),
+    new Document({ metadata: { namespace: 2 }, pageContent: "hello", id: "2" }),
+  ]);
+});
+
 test("MemoryVectorStore as retriever can filter metadata", async () => {
   const embeddings = new SyntheticEmbeddings({
     vectorSize: 1536,