refactor: Improve type safety in QdrantVectorStore client parameters (#…

…17197)
run-llama · Dec 8, 2024 · 6eb53bc · 6eb53bc
1 parent 3c5ec51
commit 6eb53bc
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/...s/vector_stores/llama-index-vector-stores-qdrant/llama_index/vector_stores/qdrant/base.py b/...s/vector_stores/llama-index-vector-stores-qdrant/llama_index/vector_stores/qdrant/base.py
@@ -9,6 +9,7 @@
 from typing import Any, List, Optional, Tuple, cast
 
 import qdrant_client
+from qdrant_client import QdrantClient, AsyncQdrantClient
 from grpc import RpcError
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.schema import BaseNode, MetadataMode, TextNode
@@ -74,8 +75,8 @@ class QdrantVectorStore(BasePydanticVectorStore):
 
     Args:
         collection_name: (str): name of the Qdrant collection
-        client (Optional[Any]): QdrantClient instance from `qdrant-client` package
-        aclient (Optional[Any]): AsyncQdrantClient instance from `qdrant-client` package
+        client (Optional[QdrantClient]): QdrantClient instance from `qdrant-client` package
+        aclient (Optional[AsyncQdrantClient]): AsyncQdrantClient instance from `qdrant-client` package
         url (Optional[str]): url of the Qdrant instance
         api_key (Optional[str]): API key for authenticating with Qdrant
         batch_size (int): number of points to upload in a single request to Qdrant. Defaults to 64
@@ -120,8 +121,8 @@ class QdrantVectorStore(BasePydanticVectorStore):
     fastembed_sparse_model: Optional[str]
     text_key: Optional[str]
 
-    _client: qdrant_client.QdrantClient = PrivateAttr()
-    _aclient: qdrant_client.AsyncQdrantClient = PrivateAttr()
+    _client: QdrantClient = PrivateAttr()
+    _aclient: AsyncQdrantClient = PrivateAttr()
     _collection_initialized: bool = PrivateAttr()
     _sparse_doc_fn: Optional[SparseEncoderCallable] = PrivateAttr()
     _sparse_query_fn: Optional[SparseEncoderCallable] = PrivateAttr()
@@ -133,8 +134,8 @@ class QdrantVectorStore(BasePydanticVectorStore):
     def __init__(
         self,
         collection_name: str,
-        client: Optional[Any] = None,
-        aclient: Optional[Any] = None,
+        client: Optional[QdrantClient] = None,
+        aclient: Optional[AsyncQdrantClient] = None,
         url: Optional[str] = None,
         api_key: Optional[str] = None,
         batch_size: int = 64,