langchain[patch]: Issue #2756 Add Qdrant custom payload on documents …

…to query them by filter (#3431) * Add optional custom payload param and 1 test * Add Custom Payload in Documents * Resolve comments in PR * Add document changes to langchain-core * Test because all yarn test cases are passing * Update tests and fix yarn lint * Remove payload from Document and add object[] * Remove object[] and replace with objects +comments * Update add document types --------- Co-authored-by: dom_ <[email protected]> Co-authored-by: jacoblee93 <[email protected]>
langchain-ai · Dec 13, 2023 · f7391c8 · f7391c8 · vercel · Dec 13, 2023
1 parent baa7e2e
commit f7391c8
Show file tree

Hide file tree

Showing 2 changed files with 208 additions and 6 deletions.
diff --git a/libs/langchain-community/src/vectorstores/qdrant.ts b/libs/langchain-community/src/vectorstores/qdrant.ts
@@ -1,7 +1,6 @@
 import { QdrantClient } from "@qdrant/js-client-rest";
 import type { Schemas as QdrantSchemas } from "@qdrant/js-client-rest";
 import { v4 as uuid } from "uuid";
-
 import { Embeddings } from "@langchain/core/embeddings";
 import { VectorStore } from "@langchain/core/vectorstores";
 import { Document } from "@langchain/core/documents";
@@ -19,8 +18,15 @@ export interface QdrantLibArgs {
   apiKey?: string;
   collectionName?: string;
   collectionConfig?: QdrantSchemas["CreateCollection"];
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  customPayload?: Record<string, any>[];
 }
 
+export type QdrantAddDocumentOptions = {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  customPayload: Record<string, any>[];
+};
+
 /**
  * Type for the response returned by a search operation in the Qdrant
  * database. It includes the score and payload (metadata and content) for
@@ -84,13 +90,18 @@ export class QdrantVectorStore extends VectorStore {
    * from the documents using the `Embeddings` instance and then adds the
    * vectors to the database.
    * @param documents Array of `Document` instances to be added to the Qdrant database.
+   * @param documentOptions Optional `QdrantAddDocumentOptions` which has a list of JSON objects for extra querying
    * @returns Promise that resolves when the documents have been added to the database.
    */
-  async addDocuments(documents: Document[]): Promise<void> {
+  async addDocuments(
+    documents: Document[],
+    documentOptions?: QdrantAddDocumentOptions
+  ): Promise<void> {
     const texts = documents.map(({ pageContent }) => pageContent);
     await this.addVectors(
       await this.embeddings.embedDocuments(texts),
-      documents
+      documents,
+      documentOptions
     );
   }
 
@@ -100,9 +111,14 @@ export class QdrantVectorStore extends VectorStore {
    * database.
    * @param vectors Array of vectors to be added to the Qdrant database.
    * @param documents Array of `Document` instances associated with the vectors.
+   * @param documentOptions Optional `QdrantAddDocumentOptions` which has a list of JSON objects for extra querying
    * @returns Promise that resolves when the vectors have been added to the database.
    */
-  async addVectors(vectors: number[][], documents: Document[]): Promise<void> {
+  async addVectors(
+    vectors: number[][],
+    documents: Document[],
+    documentOptions?: QdrantAddDocumentOptions
+  ): Promise<void> {
     if (vectors.length === 0) {
       return;
     }
@@ -115,6 +131,7 @@ export class QdrantVectorStore extends VectorStore {
       payload: {
         content: documents[idx].pageContent,
         metadata: documents[idx].metadata,
+        customPayload: documentOptions?.customPayload[idx],
       },
     }));
 
@@ -238,7 +255,14 @@ export class QdrantVectorStore extends VectorStore {
     dbConfig: QdrantLibArgs
   ): Promise<QdrantVectorStore> {
     const instance = new this(embeddings, dbConfig);
-    await instance.addDocuments(docs);
+    if (dbConfig.customPayload) {
+      const documentOptions = {
+        customPayload: dbConfig?.customPayload,
+      };
+      await instance.addDocuments(docs, documentOptions);
+    } else {
+      await instance.addDocuments(docs);
+    }
     return instance;
   }
 

diff --git a/libs/langchain-community/src/vectorstores/tests/qdrant.test.ts b/libs/langchain-community/src/vectorstores/tests/qdrant.test.ts
@@ -14,7 +14,10 @@ test("QdrantVectorStore works", async () => {
 
   const embeddings = new FakeEmbeddings();
 
-  const store = new QdrantVectorStore(embeddings, { client: client as any });
+  const store = new QdrantVectorStore(embeddings, {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    client: client as any,
+  });
 
   expect(store).toBeDefined();
 
@@ -31,3 +34,178 @@ test("QdrantVectorStore works", async () => {
 
   expect(results).toHaveLength(0);
 });
+
+test("QdrantVectorStore adds vectors with custom payload", async () => {
+  // Mock Qdrant client
+  const client = {
+    upsert: jest.fn(),
+    search: jest.fn<any>().mockResolvedValue([]),
+    getCollections: jest.fn<any>().mockResolvedValue({ collections: [] }),
+    createCollection: jest.fn(),
+  };
+
+  // Mock embeddings
+  const embeddings = new FakeEmbeddings();
+
+  // Create QdrantVectorStore instance with the mock client
+  const qdrantVectorStore = new QdrantVectorStore(embeddings, {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    client: client as any,
+  });
+
+  // Define a custom payload
+  const customPayload = {
+    customPayload: [
+      {
+        customField1: "value1",
+        customField2: "value2",
+      },
+    ],
+  };
+
+  // Add documents with custom payload
+  await qdrantVectorStore.addDocuments(
+    [
+      {
+        pageContent: "hello",
+        metadata: {},
+      },
+    ],
+    customPayload
+  );
+
+  // Verify that the Qdrant client's upsert method was called with the correct arguments
+  expect(client.upsert).toHaveBeenCalledTimes(1);
+  expect(client.upsert).toHaveBeenCalledWith("documents", {
+    wait: true,
+    points: [
+      expect.objectContaining({
+        payload: expect.objectContaining({
+          content: "hello",
+          metadata: {},
+          customPayload: customPayload.customPayload[0],
+        }),
+      }),
+    ],
+  });
+});
+
+test("QdrantVectorStore adds vectors with multiple custom payload", async () => {
+  // Mock Qdrant client
+  const client = {
+    upsert: jest.fn(),
+    search: jest.fn<any>().mockResolvedValue([]),
+    getCollections: jest.fn<any>().mockResolvedValue({ collections: [] }),
+    createCollection: jest.fn(),
+  };
+
+  // Mock embeddings
+  const embeddings = new FakeEmbeddings();
+
+  // Create QdrantVectorStore instance with the mock client
+  const qdrantVectorStore = new QdrantVectorStore(embeddings, {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    client: client as any,
+  });
+
+  // Define a custom payload
+  const customPayload = {
+    customPayload: [
+      {
+        customField1: "value1",
+        customField2: "value2",
+      },
+      {
+        customField3: "value3",
+      },
+    ],
+  };
+
+  // Add documents with custom payload
+  await qdrantVectorStore.addDocuments(
+    [
+      {
+        pageContent: "hello",
+        metadata: {},
+      },
+      {
+        pageContent: "Goodbye",
+        metadata: {},
+      },
+      {
+        pageContent: "D01",
+        metadata: {},
+      },
+    ],
+    customPayload
+  );
+
+  // Verify that the Qdrant client's upsert method was called with the correct arguments
+  expect(client.upsert).toHaveBeenCalledTimes(1);
+  expect(client.upsert).toHaveBeenCalledWith("documents", {
+    wait: true,
+    points: [
+      expect.objectContaining({
+        payload: expect.objectContaining({
+          content: "hello",
+          metadata: {},
+          customPayload: customPayload.customPayload[0],
+        }),
+      }),
+      expect.objectContaining({
+        payload: expect.objectContaining({
+          content: "Goodbye",
+          metadata: {},
+          customPayload: customPayload.customPayload[1],
+        }),
+      }),
+      expect.objectContaining({
+        payload: expect.objectContaining({
+          content: "D01",
+          metadata: {},
+        }),
+      }),
+    ],
+  });
+});
+
+test("QdrantVectorStore adds vectors with no custom payload", async () => {
+  // Mock Qdrant client
+  const client = {
+    upsert: jest.fn(),
+    search: jest.fn<any>().mockResolvedValue([]),
+    getCollections: jest.fn<any>().mockResolvedValue({ collections: [] }),
+    createCollection: jest.fn(),
+  };
+
+  // Mock embeddings
+  const embeddings = new FakeEmbeddings();
+
+  // Create QdrantVectorStore instance with the mock client
+  const qdrantVectorStore = new QdrantVectorStore(embeddings, {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    client: client as any,
+  });
+
+  // Add documents with custom payload
+  await qdrantVectorStore.addDocuments([
+    {
+      pageContent: "hello",
+      metadata: {},
+    },
+  ]);
+
+  // Verify that the Qdrant client's upsert method was called with the correct arguments
+  expect(client.upsert).toHaveBeenCalledTimes(1);
+  expect(client.upsert).toHaveBeenCalledWith("documents", {
+    wait: true,
+    points: [
+      expect.objectContaining({
+        payload: expect.objectContaining({
+          content: "hello",
+          metadata: {},
+        }),
+      }),
+    ],
+  });
+});