diff --git a/libs/langchain-qdrant/src/tests/vectorstores.int.test.ts b/libs/langchain-qdrant/src/tests/vectorstores.int.test.ts index 051c012c6dc1..f0a6afca6549 100644 --- a/libs/langchain-qdrant/src/tests/vectorstores.int.test.ts +++ b/libs/langchain-qdrant/src/tests/vectorstores.int.test.ts @@ -25,6 +25,17 @@ describe("QdrantVectorStore testcase", () => { const results = await qdrantVectorStore.similaritySearch(pageContent, 1); expect(results[0]).toEqual(new Document({ metadata: {}, pageContent })); + + expect(qdrantVectorStore.maxMarginalRelevanceSearch).toBeDefined(); + + const mmrResults = await qdrantVectorStore.maxMarginalRelevanceSearch( + pageContent, + { + k: 1, + } + ); + expect(mmrResults.length).toBe(1); + expect(mmrResults[0]).toEqual(new Document({ metadata: {}, pageContent })); }); test("passing client directly with a model that creates embeddings with a different number of dimensions", async () => { diff --git a/libs/langchain-qdrant/src/tests/vectorstores.test.ts b/libs/langchain-qdrant/src/tests/vectorstores.test.ts index 6b738d8bc464..2fb12e767617 100644 --- a/libs/langchain-qdrant/src/tests/vectorstores.test.ts +++ b/libs/langchain-qdrant/src/tests/vectorstores.test.ts @@ -209,3 +209,46 @@ test("QdrantVectorStore adds vectors with no custom payload", async () => { ], }); }); + +test("QdrantVectorStore MMR works", async () => { + const client = { + upsert: jest.fn(), + search: jest.fn().mockResolvedValue([]), + getCollections: jest.fn().mockResolvedValue({ collections: [] }), + createCollection: jest.fn(), + }; + + const embeddings = new FakeEmbeddings(); + + const store = new QdrantVectorStore(embeddings, { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + client: client as any, + }); + + expect(store).toBeDefined(); + + await store.addDocuments([ + { + pageContent: "hello", + metadata: {}, + }, + ]); + + expect(client.upsert).toHaveBeenCalledTimes(1); + + expect(store.maxMarginalRelevanceSearch).toBeDefined(); + + await store.maxMarginalRelevanceSearch("hello", { + k: 10, + fetchK: 7, + }); + + expect(client.search).toHaveBeenCalledTimes(1); + expect(client.search).toHaveBeenCalledWith("documents", { + filter: undefined, + limit: 7, + vector: [0.1, 0.2, 0.3, 0.4], + with_payload: ["metadata", "content"], + with_vector: true, + }); +}); diff --git a/libs/langchain-qdrant/src/vectorstores.ts b/libs/langchain-qdrant/src/vectorstores.ts index e1b3e0919978..80d98e1226b0 100644 --- a/libs/langchain-qdrant/src/vectorstores.ts +++ b/libs/langchain-qdrant/src/vectorstores.ts @@ -2,9 +2,13 @@ import { QdrantClient } from "@qdrant/js-client-rest"; import type { Schemas as QdrantSchemas } from "@qdrant/js-client-rest"; import { v4 as uuid } from "uuid"; import type { EmbeddingsInterface } from "@langchain/core/embeddings"; -import { VectorStore } from "@langchain/core/vectorstores"; +import { + type MaxMarginalRelevanceSearchOptions, + VectorStore, +} from "@langchain/core/vectorstores"; import { Document } from "@langchain/core/documents"; import { getEnvironmentVariable } from "@langchain/core/utils/env"; +import { maximalMarginalRelevance } from "@langchain/core/utils/math"; const CONTENT_KEY = "content"; const METADATA_KEY = "metadata"; @@ -194,6 +198,8 @@ export class QdrantVectorStore extends VectorStore { vector: query, limit: k, filter, + with_payload: [this.metadataPayloadKey, this.contentPayloadKey], + with_vector: false, }); const result: [Document, number][] = ( @@ -210,6 +216,63 @@ export class QdrantVectorStore extends VectorStore { return result; } + /** + * Return documents selected using the maximal marginal relevance. + * Maximal marginal relevance optimizes for similarity to the query AND diversity + * among selected documents. + * + * @param {string} query - Text to look up documents similar to. + * @param {number} options.k - Number of documents to return. + * @param {number} options.fetchK - Number of documents to fetch before passing to the MMR algorithm. Defaults to 20. + * @param {number} options.lambda - Number between 0 and 1 that determines the degree of diversity among the results, + * where 0 corresponds to maximum diversity and 1 to minimum diversity. + * @param {this["FilterType"]} options.filter - Optional filter to apply to the search results. + * + * @returns {Promise} - List of documents selected by maximal marginal relevance. + */ + async maxMarginalRelevanceSearch( + query: string, + options: MaxMarginalRelevanceSearchOptions + ): Promise { + if (!query) { + return []; + } + + const queryEmbedding = await this.embeddings.embedQuery(query); + + await this.ensureCollection(); + + const results = await this.client.search(this.collectionName, { + vector: queryEmbedding, + limit: options?.fetchK ?? 20, + filter: options?.filter, + with_payload: [this.metadataPayloadKey, this.contentPayloadKey], + with_vector: true, + }); + + const embeddingList = results.map((res) => res.vector) as number[][]; + + const mmrIndexes = maximalMarginalRelevance( + queryEmbedding, + embeddingList, + options?.lambda, + options.k + ); + + const topMmrMatches = mmrIndexes.map((idx) => results[idx]); + + const result = (topMmrMatches as QdrantSearchResponse[]).map( + (res) => + new Document({ + // eslint-disable-next-line @typescript-eslint/no-explicit-any + metadata: res.payload[this.metadataPayloadKey] as Record, + pageContent: res.payload[this.contentPayloadKey] as string, + }) + ); + + return result; + } + /** * Method to ensure the existence of a collection in the Qdrant database. * If the collection does not exist, it is created.