fix: fix Azure AI types + add py.typed (#2003)

* fix: fix Azure AI types + add py.typed

* fix

---------

Co-authored-by: David S. Batista <[email protected]>

Author: anakin87

.github/workflows/azure_ai_search.yml

@@ -47,11 +47,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/azure_ai_search/pyproject.toml

@@ -66,17 +66,14 @@ 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.azure_ai_search \
+-p haystack_integrations.components.retrievers.azure_ai_search {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]
-detached = true
-dependencies = ["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
@@ -161,9 +158,5 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-markers = ["unit: unit tests", "integration: integration tests"]
-pythonpath = ["src"]
-
-[[tool.mypy.overrides]]
-module = ["haystack.*", "haystack_integrations.*", "pytest.*", "azure.identity.*", "mypy.*", "azure.core.*", "azure.search.documents.*"]
-ignore_missing_imports = true
+markers = ["integration: integration tests"]
+pythonpath = ["src"]

integrations/azure_ai_search/src/haystack_integrations/components/retrievers/azure_ai_search/bm25_retriever.py

@@ -24,7 +24,7 @@ def __init__(
         filters: Optional[Dict[str, Any]] = None,
         top_k: int = 10,
         filter_policy: Union[str, FilterPolicy] = FilterPolicy.REPLACE,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Create the AzureAISearchBM25Retriever component.
@@ -96,7 +96,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "AzureAISearchBM25Retriever":
         return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
-    def run(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
+    def run(
+        self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None
+    ) -> Dict[str, List[Document]]:
         """Retrieve documents from the AzureAISearchDocumentStore.
 
         :param query: Text of the query.
@@ -111,11 +113,11 @@ def run(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optio
 
         top_k = top_k or self._top_k
         filters = filters or self._filters
-        if filters:
-            applied_filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+
+        normalized_filters = ""
+        applied_filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+        if applied_filters:
             normalized_filters = _normalize_filters(applied_filters)
-        else:
-            normalized_filters = ""
 
         try:
             docs = self._document_store._bm25_retrieval(

integrations/azure_ai_search/src/haystack_integrations/components/retrievers/azure_ai_search/embedding_retriever.py

@@ -24,7 +24,7 @@ def __init__(
         filters: Optional[Dict[str, Any]] = None,
         top_k: int = 10,
         filter_policy: Union[str, FilterPolicy] = FilterPolicy.REPLACE,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Create the AzureAISearchEmbeddingRetriever component.
@@ -93,7 +93,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "AzureAISearchEmbeddingRetriever":
         return default_from_dict(cls, data)
 
     @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 AzureAISearchDocumentStore.
 
         :param query_embedding: A list of floats representing the query embedding.
@@ -107,11 +109,11 @@ def run(self, query_embedding: List[float], filters: Optional[Dict[str, Any]] =
 
         top_k = top_k or self._top_k
         filters = filters or self._filters
-        if filters:
-            applied_filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+
+        normalized_filters = ""
+        applied_filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+        if applied_filters:
             normalized_filters = _normalize_filters(applied_filters)
-        else:
-            normalized_filters = ""
 
         try:
             docs = self._document_store._embedding_retrieval(

integrations/azure_ai_search/src/haystack_integrations/components/retrievers/azure_ai_search/hybrid_retriever.py

@@ -24,7 +24,7 @@ def __init__(
         filters: Optional[Dict[str, Any]] = None,
         top_k: int = 10,
         filter_policy: Union[str, FilterPolicy] = FilterPolicy.REPLACE,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Create the AzureAISearchHybridRetriever component.
@@ -102,7 +102,7 @@ def run(
         query_embedding: List[float],
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
-    ):
+    ) -> Dict[str, List[Document]]:
         """Retrieve documents from the AzureAISearchDocumentStore.
 
         :param query: Text of the query.
@@ -118,11 +118,11 @@ def run(
 
         top_k = top_k or self._top_k
         filters = filters or self._filters
-        if filters:
-            applied_filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+
+        normalized_filters = ""
+        applied_filters = apply_filter_policy(self._filter_policy, self._filters, filters)
+        if applied_filters:
             normalized_filters = _normalize_filters(applied_filters)
-        else:
-            normalized_filters = ""
 
         try:
             docs = self._document_store._hybrid_retrieval(

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

@@ -2,16 +2,16 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 import logging as python_logging
-import os
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 from azure.core.credentials import AzureKeyCredential
 from azure.core.exceptions import ClientAuthenticationError, HttpResponseError, ResourceNotFoundError
 from azure.core.pipeline.policies import UserAgentPolicy
 from azure.identity import DefaultAzureCredential
 from azure.search.documents import SearchClient
 from azure.search.documents.indexes import SearchIndexClient
+from azure.search.documents.indexes._generated._serialization import Model
 from azure.search.documents.indexes.models import (
     CharFilter,
     CorsOptions,
@@ -53,7 +53,7 @@
 }
 
 # Map of expected field names to their corresponding classes
-AZURE_CLASS_MAPPING = {
+AZURE_CLASS_MAPPING: Dict[str, Type[Model]] = {
     "suggesters": SearchSuggester,
     "analyzers": LexicalAnalyzer,
     "tokenizers": LexicalTokenizer,
@@ -94,7 +94,7 @@ def __init__(
         embedding_dimension: int = 768,
         metadata_fields: Optional[Dict[str, Union[SearchField, type]]] = None,
         vector_search_configuration: Optional[VectorSearch] = None,
-        **index_creation_kwargs,
+        **index_creation_kwargs: Any,
     ):
         """
         A document store using [Azure AI Search](https://azure.microsoft.com/products/ai-services/ai-search/)
@@ -133,16 +133,8 @@ def __init__(
 
         For more information on parameters, see the [official Azure AI Search documentation](https://learn.microsoft.com/en-us/azure/search/).
         """
-
-        azure_endpoint = azure_endpoint or os.environ.get("AZURE_AI_SEARCH_ENDPOINT") or None
-        if not azure_endpoint:
-            msg = "Please provide an Azure endpoint or set the environment variable AZURE_AI_SEARCH_ENDPOINT."
-            raise ValueError(msg)
-
-        api_key = api_key or os.environ.get("AZURE_AI_SEARCH_API_KEY") or None
-
-        self._client = None
-        self._index_client = None
+        self._client: Optional[SearchClient] = None
+        self._index_client: Optional[SearchIndexClient] = None
         self._index_fields = []  # type: List[Any]  # stores all fields in the final schema of index
         self._api_key = api_key
         self._azure_endpoint = azure_endpoint
@@ -155,11 +147,8 @@ def __init__(
 
     @property
     def client(self) -> SearchClient:
-        # resolve secrets for authentication
-        resolved_endpoint = (
-            self._azure_endpoint.resolve_value() if isinstance(self._azure_endpoint, Secret) else self._azure_endpoint
-        )
-        resolved_key = self._api_key.resolve_value() if isinstance(self._api_key, Secret) else self._api_key
+        resolved_endpoint = self._azure_endpoint.resolve_value()
+        resolved_key = self._api_key.resolve_value()
 
         credential = AzureKeyCredential(resolved_key) if resolved_key else DefaultAzureCredential()
 
@@ -168,8 +157,9 @@ def client(self) -> SearchClient:
         try:
             if not self._index_client:
                 self._index_client = SearchIndexClient(
-                    resolved_endpoint,
-                    credential,
+                    # resolve_value, with Secret.from_env_var (strict=True), returns a string or raises an error
+                    endpoint=resolved_endpoint,  # type: ignore[arg-type]
+                    credential=credential,
                     user_agent=ua_policy,
                 )
             if not self._index_exists(self._index_name):
@@ -287,7 +277,7 @@ def _deserialize_index_creation_kwargs(cls, data: Dict[str, Any]) -> Any:
         """
         Deserializes the index creation kwargs to the original classes.
         """
-        result = {}
+        result: Dict[str, Union[List[Model], Model]] = {}
         for key, value in data.items():
             if key in AZURE_CLASS_MAPPING:
                 if isinstance(value, list):
@@ -337,7 +327,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "AzureAISearchDocumentStore":
         else:
             data["init_parameters"]["metadata_fields"] = {}
 
-        for key, _value in AZURE_CLASS_MAPPING.items():
+        for key in AZURE_CLASS_MAPPING:
             if key in data["init_parameters"]:
                 param_value = data["init_parameters"].get(key)
                 data["init_parameters"][key] = cls._deserialize_index_creation_kwargs({key: param_value})
@@ -421,7 +411,7 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc
         if filters:
             normalized_filters = _normalize_filters(filters)
             result = self.client.search(filter=normalized_filters)
-            return self._convert_search_result_to_documents(result)
+            return self._convert_search_result_to_documents(list(result))
         else:
             return self.search_documents()
 
@@ -465,7 +455,7 @@ def _index_exists(self, index_name: Optional[str]) -> bool:
             msg = "Index name is required to check if the index exists."
             raise ValueError(msg)
 
-    def _get_raw_documents_by_id(self, document_ids: List[str]):
+    def _get_raw_documents_by_id(self, document_ids: List[str]) -> List[Dict]:
         """
         Retrieves all Azure documents with a matching document_ids from the document store.
 
@@ -499,7 +489,7 @@ def _embedding_retrieval(
         *,
         top_k: int = 10,
         filters: Optional[str] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> List[Document]:
         """
         Retrieves documents that are most similar to the query embedding using a vector similarity metric.
@@ -533,7 +523,7 @@ def _bm25_retrieval(
         query: str,
         top_k: int = 10,
         filters: Optional[str] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> List[Document]:
         """
         Retrieves documents that are most similar to `query`, using the BM25 algorithm.
@@ -566,7 +556,7 @@ def _hybrid_retrieval(
         query_embedding: List[float],
         top_k: int = 10,
         filters: Optional[str] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> List[Document]:
         """
         Retrieves documents similar to query using the vector configuration in the document store and

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

@@ -37,8 +37,6 @@ def document_store(request):
         client.delete_index(index_name)
 
     store = AzureAISearchDocumentStore(
-        api_key=api_key,
-        azure_endpoint=azure_endpoint,
         index_name=index_name,
         create_index=True,
         embedding_dimension=768,