Support storing and retrieving Document id field in MemoryVectorStore

langchain-ai · Aug 20, 2024 · 2c74f3f · 2c74f3f
1 parent fc01973
commit 2c74f3f
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/langchain/src/vectorstores/memory.ts b/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,
         })
     );
   }

diff --git a/langchain/src/vectorstores/tests/memory.test.ts b/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,