langchain[patch]: Adds metadata to vectorstore retriever memory (#5093)

* adds metadata to vectorstore retriever memory

* extends vectorstore retriever memory docs

* fixes example filter type

* Fix lint

---------

Co-authored-by: jacoblee93 <jacoblee93@gmail.com>

examples/src/memory/vector_store.ts

@@ -67,3 +67,57 @@ console.log({ res3 });
 /*
 { res3: { text: ' Your name is Perry.' } }
 */
+
+// Sometimes we might want to save metadata along with the conversation snippets
+const memoryWithMetadata = new VectorStoreRetrieverMemory({
+  vectorStoreRetriever: vectorStore.asRetriever(
+    1,
+    (doc) => doc.metadata?.userId === "1"
+  ),
+  memoryKey: "history",
+  metadata: { userId: "1", groupId: "42" },
+});
+
+await memoryWithMetadata.saveContext(
+  { input: "Community is my favorite TV Show" },
+  { output: "6 seasons and a movie!" }
+);
+
+console.log(
+  await memoryWithMetadata.loadMemoryVariables({
+    prompt: "what show should i watch? ",
+  })
+);
+/*
+{ history: 'input: Community is my favorite TV Show\noutput: 6 seasons and a movie!' }
+*/
+
+// If we have a retriever whose filter does not match our metadata, our previous messages won't appear
+const memoryWithoutMatchingMetadata = new VectorStoreRetrieverMemory({
+  vectorStoreRetriever: vectorStore.asRetriever(
+    1,
+    (doc) => doc.metadata?.userId === "2"
+  ),
+  memoryKey: "history",
+});
+
+// There are no messages saved for userId 2
+console.log(
+  await memoryWithoutMatchingMetadata.loadMemoryVariables({
+    prompt: "what show should i watch? ",
+  })
+);
+/*
+{ history: '' }
+*/
+
+// If we need the metadata to be dynamic, we can pass a function instead
+const memoryWithMetadataFunction = new VectorStoreRetrieverMemory({
+  vectorStoreRetriever: vectorStore.asRetriever(1),
+  memoryKey: "history",
+  metadata: (inputValues, _outputValues) => ({
+    firstWord: inputValues?.input.split(" ")[0], // First word of the input
+    createdAt: new Date().toLocaleDateString(), // Date when the message was saved
+    userId: "1", // Hardcoded userId
+  }),
+});

langchain/src/memory/tests/vector_store_memory.int.test.ts

@@ -32,3 +32,33 @@ test("Test vector store memory return docs", async () => {
   const result2 = await memory.loadMemoryVariables({ input: "foo" });
   expect(result2).toStrictEqual({ memory: expectedResult });
 });
+
+test("Test vector store memory metadata object", async () => {
+  const vectorStore = new MemoryVectorStore(new OpenAIEmbeddings());
+  const memory = new VectorStoreRetrieverMemory({
+    vectorStoreRetriever: vectorStore.asRetriever(),
+    metadata: { foo: "bar" },
+  });
+
+  await memory.saveContext({ foo: "bar" }, { bar: "foo" });
+  vectorStore.memoryVectors.forEach((vector) => {
+    expect(vector.metadata).toStrictEqual({ foo: "bar" });
+  });
+});
+
+test("Test vector store memory metadata function", async () => {
+  const vectorStore = new MemoryVectorStore(new OpenAIEmbeddings());
+
+  const memory = new VectorStoreRetrieverMemory({
+    vectorStoreRetriever: vectorStore.asRetriever(),
+    metadata: (inputValues, outputValues) => ({
+      foo: `${inputValues?.foo} ${inputValues?.foo}`, // "bar bar"
+      bar: `${outputValues?.bar} ${outputValues?.bar}`, // "foo foo"
+    }),
+  });
+
+  await memory.saveContext({ foo: "bar" }, { bar: "foo" });
+  vectorStore.memoryVectors.forEach((vector) => {
+    expect(vector.metadata).toStrictEqual({ foo: "bar bar", bar: "foo foo" });
+  });
+});

langchain/src/memory/vector_store.ts

@@ -9,6 +9,12 @@ import {
 } from "@langchain/core/memory";
 import { formatDocumentsAsString } from "../util/document.js";
 
+type Metadata = Record<string, unknown>;
+type MetadataFunction = (
+  inputValues?: InputValues,
+  outputValues?: OutputValues
+) => Metadata;
+
 /**
  * Interface for the parameters required to initialize a
  * VectorStoreRetrieverMemory instance.
@@ -19,6 +25,10 @@ export interface VectorStoreRetrieverMemoryParams {
   outputKey?: string;
   memoryKey?: string;
   returnDocs?: boolean;
+  /**
+   * Metadata to be added to the document when saving context.
+   */
+  metadata?: Metadata | MetadataFunction;
 }
 
 /**
@@ -64,12 +74,15 @@ export class VectorStoreRetrieverMemory
 
   returnDocs: boolean;
 
+  metadata?: Metadata | MetadataFunction;
+
   constructor(fields: VectorStoreRetrieverMemoryParams) {
     super();
     this.vectorStoreRetriever = fields.vectorStoreRetriever;
     this.inputKey = fields.inputKey;
     this.memoryKey = fields.memoryKey ?? "memory";
     this.returnDocs = fields.returnDocs ?? false;
+    this.metadata = fields.metadata;
   }
 
   get memoryKeys(): string[] {
@@ -105,13 +118,17 @@ export class VectorStoreRetrieverMemory
     inputValues: InputValues,
     outputValues: OutputValues
   ): Promise<void> {
+    const metadata =
+      typeof this.metadata === "function"
+        ? this.metadata(inputValues, outputValues)
+        : this.metadata;
     const text = Object.entries(inputValues)
       .filter(([k]) => k !== this.memoryKey)
       .concat(Object.entries(outputValues))
       .map(([k, v]) => `${k}: ${v}`)
       .join("\n");
     await this.vectorStoreRetriever.addDocuments([
-      new Document({ pageContent: text }),
+      new Document({ pageContent: text, metadata }),
     ]);
   }
 }