feat: qdrant search params (#1911)

.changeset/silver-taxis-call.md

@@ -0,0 +1,8 @@
+---
+"@llamaindex/doc": patch
+"@llamaindex/examples": patch
+"@llamaindex/core": patch
+"@llamaindex/qdrant": patch
+---
+
+Add functionality for search params when querying Qdrant vector store.

apps/next/src/content/docs/llamaindex/modules/data/stores/vector_stores/qdrant.mdx

@@ -88,7 +88,7 @@ async function main() {
 
   const response = await queryEngine.query({
     query: "What did the author do in college?",
-  });
+  }); // Additional filters and params can be passed as options
 
   // Output response
   console.log(response.toString());

examples/storage/qdrantdb/README.md

@@ -5,7 +5,11 @@ How to run `examples/qdrantdb/preFilters.ts`:
 Add your OpenAI API Key into a file called `.env` in the parent folder of this directory. It should look like this:
 
 ```
-OPEN_API_KEY=sk-you-key
+OPENAI_API_KEY=sk-you-key
 ```
 
 Now, open a new terminal window and inside `examples`, run `npx tsx qdrantdb/preFilters.ts`.
+
+## Notes
+
+- You should have a Qdrant instance running locally.

examples/storage/qdrantdb/searchParams.ts

@@ -0,0 +1,60 @@
+import { QdrantVectorStore } from "@llamaindex/qdrant";
+import * as dotenv from "dotenv";
+import {
+  Document,
+  MetadataMode,
+  NodeWithScore,
+  Settings,
+  VectorStoreIndex,
+  storageContextFromDefaults,
+} from "llamaindex";
+
+// Update callback manager
+Settings.callbackManager.on("retrieve-end", (event) => {
+  const { nodes } = event.detail;
+  console.log(
+    "The retrieved nodes are:",
+    nodes.map((node: NodeWithScore) => node.node.getContent(MetadataMode.NONE)),
+  );
+});
+
+dotenv.config();
+
+const collectionName = "dog_colors";
+const qdrantUrl = "http://127.0.0.1:6333";
+
+async function main() {
+  try {
+    const vectorStore = new QdrantVectorStore({
+      url: qdrantUrl,
+      collectionName,
+    });
+    const ctx = await storageContextFromDefaults({ vectorStore });
+
+    const docs = [
+      new Document({
+        text: "The dog is brown",
+      }),
+    ];
+
+    const index = await VectorStoreIndex.fromDocuments(docs, {
+      storageContext: ctx,
+    });
+
+    const queryEngine = index.asQueryEngine({
+      customParams: {
+        hnsw_ef: 10,
+        exact: true,
+        indexed_only: true,
+      },
+    });
+    const response = await queryEngine.query({
+      query: "What is the color of the dog?",
+    });
+    console.log(response.toString());
+  } catch (error) {
+    console.error(error);
+  }
+}
+
+void main();

packages/core/src/retriever/index.ts

@@ -8,6 +8,7 @@ import { BaseNode, IndexNode, type NodeWithScore, ObjectType } from "../schema";
 export type RetrieveParams = {
   query: MessageContent;
   preFilters?: unknown;
+  customParams?: unknown;
 };
 
 export type RetrieveStartEvent = {

packages/core/src/vector-store/index.ts

@@ -83,7 +83,7 @@ export interface VectorStoreInfo {
   contentInfo: string;
 }
 
-export interface VectorStoreQuery {
+export interface VectorStoreQuery<T = unknown> {
   queryEmbedding?: number[];
   similarityTopK: number;
   docIds?: string[];
@@ -92,6 +92,7 @@ export interface VectorStoreQuery {
   alpha?: number;
   filters?: MetadataFilters | undefined;
   mmrThreshold?: number;
+  customParams?: T | undefined;
 }
 
 // Supported types of vector stores (for each modality)
@@ -103,7 +104,7 @@ export type VectorStoreBaseParams = {
   embeddingModel?: BaseEmbedding | undefined;
 };
 
-export abstract class BaseVectorStore<Client = unknown> {
+export abstract class BaseVectorStore<Client = unknown, T = unknown> {
   embedModel: BaseEmbedding;
   abstract storesText: boolean;
   isEmbeddingQuery?: boolean;
@@ -111,7 +112,7 @@ export abstract class BaseVectorStore<Client = unknown> {
   abstract add(embeddingResults: BaseNode[]): Promise<string[]>;
   abstract delete(refDocId: string, deleteOptions?: object): Promise<void>;
   abstract query(
-    query: VectorStoreQuery,
+    query: VectorStoreQuery<T>,
     options?: object,
   ): Promise<VectorStoreQueryResult>;
 

packages/llamaindex/src/indices/vectorStore/index.ts

@@ -65,6 +65,7 @@ export type VectorIndexChatEngineOptions = {
   retriever?: BaseRetriever;
   similarityTopK?: number;
   preFilters?: MetadataFilters;
+  customParams?: unknown;
 } & Omit<ContextChatEngineOptions, "retriever">;
 
 /**
@@ -285,6 +286,7 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
     retriever?: BaseRetriever;
     responseSynthesizer?: BaseSynthesizer;
     preFilters?: MetadataFilters;
+    customParams?: unknown;
     nodePostprocessors?: BaseNodePostprocessor[];
     similarityTopK?: number;
   }): RetrieverQueryEngine {
@@ -292,11 +294,17 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
       retriever,
       responseSynthesizer,
       preFilters,
+      customParams,
       nodePostprocessors,
       similarityTopK,
     } = options ?? {};
     return new RetrieverQueryEngine(
-      retriever ?? this.asRetriever({ similarityTopK, filters: preFilters }),
+      retriever ??
+        this.asRetriever({
+          similarityTopK,
+          filters: preFilters,
+          customParams: customParams ?? undefined,
+        }),
       responseSynthesizer,
       nodePostprocessors,
     );
@@ -312,11 +320,17 @@ export class VectorStoreIndex extends BaseIndex<IndexDict> {
       retriever,
       similarityTopK,
       preFilters,
+      customParams,
       ...contextChatEngineOptions
     } = options;
     return new ContextChatEngine({
       retriever:
-        retriever ?? this.asRetriever({ similarityTopK, filters: preFilters }),
+        retriever ??
+        this.asRetriever({
+          similarityTopK,
+          filters: preFilters,
+          customParams: customParams ?? undefined,
+        }),
       ...contextChatEngineOptions,
     });
   }
@@ -407,6 +421,7 @@ export type VectorIndexRetrieverOptions = {
   index: VectorStoreIndex;
   filters?: MetadataFilters | undefined;
   mode?: VectorStoreQueryMode;
+  customParams?: unknown | undefined;
 } & (
   | {
       topK?: TopKMap | undefined;
@@ -422,7 +437,7 @@ export class VectorIndexRetriever extends BaseRetriever {
 
   filters?: MetadataFilters | undefined;
   queryMode?: VectorStoreQueryMode | undefined;
-
+  customParams?: unknown | undefined;
   constructor(options: VectorIndexRetrieverOptions) {
     super();
     this.index = options.index;
@@ -440,6 +455,7 @@ export class VectorIndexRetriever extends BaseRetriever {
       };
     }
     this.filters = options.filters;
+    this.customParams = options.customParams;
   }
 
   /**
@@ -468,6 +484,7 @@ export class VectorIndexRetriever extends BaseRetriever {
     type: ModalityType,
     vectorStore: BaseVectorStore,
     filters?: MetadataFilters,
+    customParams?: unknown,
   ): Promise<NodeWithScore[]> {
     // convert string message to multi-modal format
 
@@ -491,6 +508,7 @@ export class VectorIndexRetriever extends BaseRetriever {
           mode: this.queryMode ?? VectorStoreQueryMode.DEFAULT,
           similarityTopK: this.topK[type]!,
           filters: this.filters ?? filters ?? undefined,
+          customParams: this.customParams ?? customParams ?? undefined,
         });
         nodes = nodes.concat(this.buildNodeListFromQueryResult(result));
       }

packages/providers/storage/qdrant/src/QdrantVectorStore.ts

@@ -18,6 +18,7 @@ import { QdrantClient } from "@qdrant/js-client-rest";
 
 type QdrantFilter = Schemas["Filter"];
 type QdrantMustConditions = QdrantFilter["must"];
+type QdrantSearchParams = Schemas["SearchParams"];
 
 type PointStruct = {
   id: string;
@@ -268,19 +269,24 @@ export class QdrantVectorStore extends BaseVectorStore {
   /**
    * Queries the vector store for the closest matching data to the query embeddings.
    * @param query The VectorStoreQuery to be used
-   * @param options Required by VectorStore interface.  Currently ignored.
+   * @param options Required by VectorStore interface.
    * @returns Zero or more Document instances with data from the vector store.
    */
   async query(
-    query: VectorStoreQuery,
+    query: VectorStoreQuery<QdrantSearchParams | undefined>,
     options?: object,
   ): Promise<VectorStoreQueryResult> {
     const qdrantFilters =
       options && "qdrant_filters" in options
         ? options.qdrant_filters
         : undefined;
+    const qdrantSearchParams =
+      options && "qdrant_search_params" in options
+        ? options.qdrant_search_params
+        : undefined;
 
     let queryFilters: QdrantFilter | undefined;
+    let searchParams: QdrantSearchParams | undefined;
 
     if (!query.queryEmbedding) {
       throw new Error("No query embedding provided");
@@ -292,10 +298,17 @@ export class QdrantVectorStore extends BaseVectorStore {
       queryFilters = buildQueryFilter(query);
     }
 
+    if (qdrantSearchParams) {
+      searchParams = qdrantSearchParams;
+    } else {
+      searchParams = buildSearchParams(query);
+    }
+
     const result = (await this.db.search(this.collectionName, {
       vector: query.queryEmbedding,
       limit: query.similarityTopK,
       ...(queryFilters && { filter: queryFilters }),
+      ...(searchParams && { params: searchParams }),
     })) as Array<QuerySearchResult>;
 
     return this.parseToQueryResult(result);
@@ -325,6 +338,18 @@ function buildQueryFilter(query: VectorStoreQuery): QdrantFilter | undefined {
   return { must: mustConditions };
 }
 
+function buildSearchParams(
+  query: VectorStoreQuery<QdrantSearchParams | undefined>,
+): QdrantSearchParams | undefined {
+  if (!query.docIds && !query.queryStr && !query.customParams) return undefined;
+
+  if (query.customParams) {
+    return query.customParams;
+  }
+
+  return undefined;
+}
+
 /**
  * Converts metadata filters to Qdrant-compatible filters
  * @param subFilters The metadata filters to be converted

packages/providers/storage/qdrant/tests/QdrantVectorStore.test.ts

@@ -138,5 +138,35 @@ describe("QdrantVectorStore", () => {
         expect(searchResult.similarities).toEqual([0.1]);
       });
     });
+
+    describe("[QdrantVectorStore] search with params", () => {
+      it("should search in the vector store", async () => {
+        mockQdrantClient.search.mockResolvedValue([
+          {
+            id: "1",
+            score: 0.1,
+            version: 1,
+            payload: { _node_content: JSON.stringify({ text: "hello world" }) },
+          },
+        ]);
+
+        const searchResult = await store.query(
+          {
+            queryEmbedding: [0.1, 0.2],
+            similarityTopK: 1,
+            mode: VectorStoreQueryMode.DEFAULT,
+          },
+          {
+            customParams: {
+              hnsw_ef: 10,
+            },
+          },
+        );
+
+        expect(mockQdrantClient.search).toHaveBeenCalled();
+        expect(searchResult.ids).toEqual(["1"]);
+        expect(searchResult.similarities).toEqual([0.1]);
+      });
+    });
   });
 });