diff --git a/libs/langchain-community/src/vectorstores/qdrant.ts b/libs/langchain-community/src/vectorstores/qdrant.ts index 21062e1e86b5..8e72426281dc 100644 --- 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[]; } +export type QdrantAddDocumentOptions = { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + customPayload: Record[]; +}; + /** * 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 { + async addDocuments( + documents: Document[], + documentOptions?: QdrantAddDocumentOptions + ): Promise { 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 { + async addVectors( + vectors: number[][], + documents: Document[], + documentOptions?: QdrantAddDocumentOptions + ): Promise { 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 { 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 index 7e3aaf0eab2b..ee3fa42cf13e 100644 --- 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().mockResolvedValue([]), + getCollections: jest.fn().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().mockResolvedValue([]), + getCollections: jest.fn().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().mockResolvedValue([]), + getCollections: jest.fn().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: {}, + }), + }), + ], + }); +});