fix: Astra - fix types + add py.typed (#2011)

Author: anakin87

.github/workflows/astra.yml

@@ -50,11 +50,9 @@ jobs:
       - name: Install Hatch
         run: pip install --upgrade hatch
 
-    # TODO: Once this integration is properly typed, use hatch run test:types
-    # https://github.com/deepset-ai/haystack-core-integrations/issues/1771
       - name: Lint
         if: matrix.python-version == '3.9' && runner.os == 'Linux'
-        run: hatch run fmt-check && hatch run lint:typing
+        run: hatch run fmt-check && hatch run test:types
 
       - name: Generate docs
         if: matrix.python-version == '3.9' && runner.os == 'Linux'

integrations/astra/pyproject.toml

@@ -64,17 +64,14 @@ unit = 'pytest -m "not integration" {args:tests}'
 integration = 'pytest -m "integration" {args:tests}'
 all = 'pytest {args:tests}'
 cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
-types = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = """mypy -p haystack_integrations.document_stores.astra \
+-p haystack_integrations.components.retrievers.astra {args}"""
 
-# TODO: remove lint environment once this integration is properly typed
-# test environment should be used instead
-# https://github.com/deepset-ai/haystack-core-integrations/issues/1771
-[tool.hatch.envs.lint]
-installer = "uv"
-detached = true
-dependencies = ["pip", "mypy>=1.0.0", "ruff>=0.0.243"]
-[tool.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 [tool.hatch.metadata]
 allow-direct-references = true
@@ -156,16 +153,4 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-markers = ["unit: unit tests", "integration: integration tests"]
-
-[[tool.mypy.overrides]]
-module = [
-  "astra_client.*",
-  "astrapy.*",
-  "pydantic.*",
-  "haystack.*",
-  "haystack_integrations.*",
-  "pytest.*",
-  "openpyxl.*",
-]
-ignore_missing_imports = true
+markers = ["integration: integration tests"]

integrations/astra/src/haystack_integrations/components/retrievers/astra/retriever.py

@@ -58,7 +58,12 @@ def __init__(
             raise Exception(message)
 
     @component.output_types(documents=List[Document])
-    def run(self, query_embedding: List[float], filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
+    def run(
+        self,
+        query_embedding: List[float],
+        filters: Optional[Dict[str, Any]] = None,
+        top_k: Optional[int] = None,
+    ) -> Dict[str, List[Document]]:
         """Retrieve documents from the AstraDocumentStore.
 
         :param query_embedding: floats representing the query embedding

integrations/astra/src/haystack_integrations/components/retrievers/py.typed

@@ -278,7 +278,7 @@ def batch_generator(chunks, batch_size):
 
         return formatted_docs
 
-    def insert(self, documents: List[Dict]):
+    def insert(self, documents: List[Dict]) -> List[str]:
         """
         Insert documents into the Astra index.
 
@@ -290,7 +290,7 @@ def insert(self, documents: List[Dict]):
 
         return inserted_ids
 
-    def update_document(self, document: Dict, id_key: str):
+    def update_document(self, document: Dict, id_key: str) -> bool:
         """
         Update a document in the Astra index.
 

integrations/astra/src/haystack_integrations/document_stores/astra/astra_client.py

@@ -9,7 +9,7 @@
 from haystack.document_stores.types import DuplicatePolicy
 from haystack.utils import Secret, deserialize_secrets_inplace
 
-from .astra_client import AstraClient
+from .astra_client import AstraClient, QueryResponse
 from .errors import AstraDocumentStoreFilterError
 from .filters import _convert_filters
 
@@ -150,7 +150,7 @@ def write_documents(
         self,
         documents: List[Document],
         policy: DuplicatePolicy = DuplicatePolicy.NONE,
-    ):
+    ) -> int:
         """
         Indexes documents for later queries.
 
@@ -176,7 +176,7 @@ def write_documents(
 
         batch_size = MAX_BATCH_SIZE
 
-        def _convert_input_document(document: Union[dict, Document]):
+        def _convert_input_document(document: Union[dict, Document]) -> Dict[str, Any]:
             if isinstance(document, Document):
                 document_dict = document.to_dict(flatten=False)
             elif isinstance(document, dict):
@@ -217,7 +217,7 @@ def _convert_input_document(document: Union[dict, Document]):
         documents_to_write = [_convert_input_document(doc) for doc in documents]
 
         duplicate_documents = []
-        new_documents: List[Document] = []
+        new_documents: List[Dict] = []
         i = 0
         while i < len(documents_to_write):
             doc = documents_to_write[i]
@@ -238,7 +238,7 @@ def _convert_input_document(document: Union[dict, Document]):
         if policy == DuplicatePolicy.SKIP:
             if len(new_documents) > 0:
                 for batch in _batches(new_documents, batch_size):
-                    inserted_ids = self.index.insert(batch)  # type: ignore
+                    inserted_ids = self.index.insert(batch)
                     insertion_counter += len(inserted_ids)
                     logger.info(f"write_documents inserted documents with id {inserted_ids}")
             else:
@@ -247,7 +247,7 @@ def _convert_input_document(document: Union[dict, Document]):
         elif policy == DuplicatePolicy.OVERWRITE:
             if len(new_documents) > 0:
                 for batch in _batches(new_documents, batch_size):
-                    inserted_ids = self.index.insert(batch)  # type: ignore
+                    inserted_ids = self.index.insert(batch)
                     insertion_counter += len(inserted_ids)
                     logger.info(f"write_documents inserted documents with id {inserted_ids}")
             else:
@@ -256,7 +256,7 @@ def _convert_input_document(document: Union[dict, Document]):
             if len(duplicate_documents) > 0:
                 updated_ids = []
                 for duplicate_doc in duplicate_documents:
-                    updated = self.index.update_document(duplicate_doc, "_id")  # type: ignore
+                    updated = self.index.update_document(duplicate_doc, "_id")
                     if updated:
                         updated_ids.append(duplicate_doc["_id"])
                 insertion_counter = insertion_counter + len(updated_ids)
@@ -267,7 +267,7 @@ def _convert_input_document(document: Union[dict, Document]):
         elif policy == DuplicatePolicy.FAIL:
             if len(new_documents) > 0:
                 for batch in _batches(new_documents, batch_size):
-                    inserted_ids = self.index.insert(batch)  # type: ignore
+                    inserted_ids = self.index.insert(batch)
                     insertion_counter = insertion_counter + len(inserted_ids)
                     logger.info(f"write_documents inserted documents with id {inserted_ids}")
             else:
@@ -326,15 +326,18 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc
         return documents
 
     @staticmethod
-    def _get_result_to_documents(results) -> List[Document]:
+    def _get_result_to_documents(results: QueryResponse) -> List[Document]:
         documents = []
         for match in results.matches:
+            metadata = match.metadata
+            blob = metadata.pop("blob", None) if metadata else None
+            meta = metadata.pop("meta", {}) if metadata else {}
             document = Document(
                 content=match.text,
                 id=match.document_id,
                 embedding=match.values,
-                blob=match.metadata.pop("blob", None),
-                meta=match.metadata.pop("meta", None),
+                blob=blob,
+                meta=meta,
                 score=match.score,
             )
             documents.append(document)