Add more type hints

Author: sjrl

integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py

@@ -1,6 +1,6 @@
 import inspect
 from itertools import islice
-from typing import Any, AsyncGenerator, ClassVar, Dict, Generator, List, Optional, Set, Union
+from typing import Any, AsyncGenerator, ClassVar, Dict, Generator, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import qdrant_client
@@ -21,6 +21,7 @@
     convert_haystack_documents_to_qdrant_points,
     convert_id,
     convert_qdrant_point_to_haystack_document,
+    QdrantPoint
 )
 from .filters import convert_filters_to_qdrant
 
@@ -34,7 +35,7 @@ class QdrantStoreError(DocumentStoreError):
 FilterType = Dict[str, Union[Dict[str, Any], List[Any], str, int, float, bool]]
 
 
-def get_batches_from_generator(iterable, n):
+def get_batches_from_generator(iterable: List, n: int) -> Generator:
     """
     Batch elements of an iterable into fixed-length chunks or blocks.
     """
@@ -127,7 +128,7 @@ def __init__(
         write_batch_size: int = 100,
         scroll_size: int = 10_000,
         payload_fields_to_index: Optional[List[dict]] = None,
-    ):
+    ) -> None:
         """
         :param location:
             If `memory` - use in-memory Qdrant instance.
@@ -164,7 +165,7 @@ def __init__(
             Dimension of the embeddings.
         :param on_disk:
             Whether to store the collection on disk.
-        :param use_sparse_embedding:
+        :param use_sparse_embeddings:
             If set to `True`, enables support for sparse embeddings.
         :param sparse_idf:
             If set to `True`, computes the Inverse Document Frequency (IDF) when using sparse embeddings.
@@ -257,7 +258,7 @@ def __init__(
         self.write_batch_size = write_batch_size
         self.scroll_size = scroll_size
 
-    def _initialize_client(self):
+    def _initialize_client(self) -> None:
         if self._client is None:
             client_params = self._prepare_client_params()
             self._client = qdrant_client.QdrantClient(**client_params)
@@ -273,7 +274,7 @@ def _initialize_client(self):
                 self.payload_fields_to_index,
             )
 
-    async def _initialize_async_client(self):
+    async def _initialize_async_client(self) -> None:
         """
         Returns the asynchronous Qdrant client, initializing it if necessary.
         """
@@ -627,8 +628,6 @@ def get_documents_by_id(
 
         :param ids:
             A list of document IDs to retrieve.
-        :param index:
-            The name of the index to retrieve documents from.
         :returns:
             A list of documents.
         """
@@ -660,8 +659,6 @@ async def get_documents_by_id_async(
 
         :param ids:
             A list of document IDs to retrieve.
-        :param index:
-            The name of the index to retrieve documents from.
         :returns:
             A list of documents.
         """
@@ -1209,7 +1206,7 @@ def get_distance(self, similarity: str) -> rest.Distance:
             )
             raise QdrantStoreError(msg) from ke
 
-    def _create_payload_index(self, collection_name: str, payload_fields_to_index: Optional[List[dict]] = None):
+    def _create_payload_index(self, collection_name: str, payload_fields_to_index: Optional[List[dict]] = None) -> None:
         """
         Create payload index for the collection if payload_fields_to_index is provided
         See: https://qdrant.tech/documentation/concepts/indexing/#payload-index
@@ -1228,7 +1225,7 @@ def _create_payload_index(self, collection_name: str, payload_fields_to_index: O
 
     async def _create_payload_index_async(
         self, collection_name: str, payload_fields_to_index: Optional[List[dict]] = None
-    ):
+    ) -> None:
         """
         Asynchronously create payload index for the collection if payload_fields_to_index is provided
         See: https://qdrant.tech/documentation/concepts/indexing/#payload-index
@@ -1256,7 +1253,7 @@ def _set_up_collection(
         sparse_idf: bool,
         on_disk: bool = False,
         payload_fields_to_index: Optional[List[dict]] = None,
-    ):
+    ) -> None:
         """
         Sets up the Qdrant collection with the specified parameters.
         :param collection_name:
@@ -1312,7 +1309,7 @@ async def _set_up_collection_async(
         sparse_idf: bool,
         on_disk: bool = False,
         payload_fields_to_index: Optional[List[dict]] = None,
-    ):
+    ) -> None:
         """
         Asynchronously sets up the Qdrant collection with the specified parameters.
         :param collection_name:
@@ -1366,7 +1363,7 @@ def recreate_collection(
         on_disk: Optional[bool] = None,
         use_sparse_embeddings: Optional[bool] = None,
         sparse_idf: bool = False,
-    ):
+    ) -> None:
         """
         Recreates the Qdrant collection with the specified parameters.
 
@@ -1409,7 +1406,7 @@ async def recreate_collection_async(
         on_disk: Optional[bool] = None,
         use_sparse_embeddings: Optional[bool] = None,
         sparse_idf: bool = False,
-    ):
+    ) -> None:
         """
         Asynchronously recreates the Qdrant collection with the specified parameters.
 
@@ -1448,7 +1445,7 @@ def _handle_duplicate_documents(
         self,
         documents: List[Document],
         policy: DuplicatePolicy = None,
-    ):
+    ) -> List[Document]:
         """
         Checks whether any of the passed documents is already existing in the chosen index and returns a list of
         documents that are not in the index yet.
@@ -1475,7 +1472,7 @@ async def _handle_duplicate_documents_async(
         self,
         documents: List[Document],
         policy: DuplicatePolicy = None,
-    ):
+    ) -> List[Document]:
         """
         Asynchronously checks whether any of the passed documents is already existing
         in the chosen index and returns a list of
@@ -1520,7 +1517,7 @@ def _drop_duplicate_documents(self, documents: List[Document]) -> List[Document]
 
         return _documents
 
-    def _prepare_collection_params(self):
+    def _prepare_collection_params(self) -> Dict[str, Any]:
         """
         Prepares the common parameters for collection creation.
         """
@@ -1536,7 +1533,7 @@ def _prepare_collection_params(self):
             "init_from": self.init_from,
         }
 
-    def _prepare_client_params(self):
+    def _prepare_client_params(self) -> Dict[str, Any]:
         """
         Prepares the common parameters for client initialization.
 
@@ -1564,7 +1561,7 @@ def _prepare_collection_config(
         on_disk: Optional[bool] = None,
         use_sparse_embeddings: Optional[bool] = None,
         sparse_idf: bool = False,
-    ):
+    ) -> Tuple[Dict[str, rest.VectorParams], Optional[Dict[str, rest.SparseVectorParams]]]:
         """
         Prepares the configuration for creating or recreating a Qdrant collection.
 
@@ -1594,9 +1591,12 @@ def _prepare_collection_config(
 
         return vectors_config, sparse_vectors_config
 
-    def _validate_filters(self, filters: Optional[Union[Dict[str, Any], rest.Filter]] = None):
+    def _validate_filters(self, filters: Optional[Union[Dict[str, Any], rest.Filter]] = None) -> None:
         """
         Validates the filters provided for querying.
+
+        :param filters: Filters to validate. Can be a dictionary or an instance of `qdrant_client.http.models.Filter`.
+        :raises ValueError: If the filters are not in the correct format or syntax.
         """
         if filters and not isinstance(filters, dict) and not isinstance(filters, rest.Filter):
             msg = "Filter must be a dictionary or an instance of `qdrant_client.http.models.Filter`"
@@ -1606,7 +1606,7 @@ def _validate_filters(self, filters: Optional[Union[Dict[str, Any], rest.Filter]
             msg = "Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
             raise ValueError(msg)
 
-    def _process_query_point_results(self, results, scale_score: bool = False):
+    def _process_query_point_results(self, results: List[QdrantPoint], scale_score: bool = False) -> List[Document]:
         """
         Processes query results from Qdrant.
         """
@@ -1626,7 +1626,7 @@ def _process_query_point_results(self, results, scale_score: bool = False):
 
         return documents
 
-    def _process_group_results(self, groups):
+    def _process_group_results(self, groups: List[rest.PointGroup]) -> List[Document]:
         """
         Processes grouped query results from Qdrant.
 
@@ -1646,7 +1646,7 @@ def _validate_collection_compatibility(
         collection_info,
         distance,
         embedding_dim: int,
-    ):
+    ) -> None:
         """
         Validates that an existing collection is compatible with the current configuration.
         """