diff --git a/langchain/src/vectorstores/memory.ts b/langchain/src/vectorstores/memory.ts index fc77840a77ab..d6210a95f1cc 100644 --- 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; + 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 index 92ad7839ff48..02b893aa0cfd 100644 --- 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,