Merge branch 'main' into move-github-components

Author: julian-risch

integrations/azure_ai_search/CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## [integrations/azure_ai_search-v2.0.1] - 2025-04-16
+
+### 🚀 Features
+
+- Fix serialization errors in `AzureAIDocumentStore` (#1647)
+
+
 ## [integrations/azure_ai_search-v2.0.0] - 2025-04-03
 
 ### 🐛 Bug Fixes
@@ -21,6 +28,7 @@
 
 ## [integrations/azure_ai_search-v0.1.3] - 2025-02-17
 
+
 ### 🧹 Chores
 
 - Azure AI search - clarify that dataframe is not supported (#1407)

integrations/azure_ai_search/pyproject.toml

@@ -156,6 +156,7 @@ 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.*"]

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

@@ -4,21 +4,30 @@
 import logging as python_logging
 import os
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, ClassVar, Dict, List, Optional
 
 from azure.core.credentials import AzureKeyCredential
 from azure.core.exceptions import ClientAuthenticationError, HttpResponseError, ResourceNotFoundError
 from azure.identity import DefaultAzureCredential
 from azure.search.documents import SearchClient
 from azure.search.documents.indexes import SearchIndexClient
 from azure.search.documents.indexes.models import (
+    CharFilter,
+    CorsOptions,
     HnswAlgorithmConfiguration,
     HnswParameters,
+    LexicalAnalyzer,
+    LexicalTokenizer,
+    ScoringProfile,
     SearchableField,
     SearchField,
     SearchFieldDataType,
     SearchIndex,
+    SearchResourceEncryptionKey,
+    SearchSuggester,
+    SimilarityAlgorithm,
     SimpleField,
+    TokenFilter,
     VectorSearch,
     VectorSearchAlgorithmMetric,
     VectorSearchProfile,
@@ -40,6 +49,19 @@
     datetime: "Edm.DateTimeOffset",
 }
 
+# Map of expected field names to their corresponding classes
+AZURE_CLASS_MAPPING = {
+    "suggesters": SearchSuggester,
+    "analyzers": LexicalAnalyzer,
+    "tokenizers": LexicalTokenizer,
+    "token_filters": TokenFilter,
+    "char_filters": CharFilter,
+    "cors_options": CorsOptions,
+    "similarity_algorithm": SimilarityAlgorithm,
+    "encryption_key": SearchResourceEncryptionKey,
+    "scoring_profiles": ScoringProfile,
+}
+
 DEFAULT_VECTOR_SEARCH = VectorSearch(
     profiles=[
         VectorSearchProfile(name="default-vector-config", algorithm_configuration_name="cosine-algorithm-config")
@@ -60,6 +82,8 @@
 
 
 class AzureAISearchDocumentStore:
+    TYPE_MAP: ClassVar[Dict[str, type]] = {"str": str, "int": int, "float": float, "bool": bool, "datetime": datetime}
+
     def __init__(
         self,
         *,
@@ -138,7 +162,7 @@ def client(self) -> SearchClient:
                     "The index '{idx_name}' does not exist. A new index will be created.",
                     idx_name=self._index_name,
                 )
-                self._create_index(self._index_name)
+                self._create_index()
         except (HttpResponseError, ClientAuthenticationError) as error:
             msg = f"Failed to authenticate with Azure Search: {error}"
             raise AzureAISearchDocumentStoreConfigError(msg) from error
@@ -154,11 +178,9 @@ def client(self) -> SearchClient:
 
         return self._client
 
-    def _create_index(self, index_name: str) -> None:
+    def _create_index(self) -> None:
         """
-        Creates a new search index.
-        :param index_name: Name of the index to create. If None, the index name from the constructor is used.
-        :param kwargs: Optional keyword parameters.
+        Internally creates a new search index.
         """
 
         # default fields to create index based on Haystack Document (id, content, embedding)
@@ -175,19 +197,66 @@ def _create_index(self, index_name: str) -> None:
             ),
         ]
 
-        if not index_name:
-            index_name = self._index_name
         if self._metadata_fields:
             default_fields.extend(self._create_metadata_index_fields(self._metadata_fields))
         index = SearchIndex(
-            name=index_name,
+            name=self._index_name,
             fields=default_fields,
             vector_search=self._vector_search_configuration,
             **self._index_creation_kwargs,
         )
         if self._index_client:
             self._index_client.create_index(index)
 
+    @classmethod
+    def _deserialize_metadata_fields(cls, fields: Optional[Dict[str, str]]) -> Optional[Dict[str, type]]:
+        """Convert string representations back to type objects."""
+        if not fields:
+            return None
+        try:
+            # Use the class-level TYPE_MAP for conversion.
+            ans = {key: cls.TYPE_MAP[value] for key, value in fields.items()}
+            return ans
+        except KeyError as e:
+            msg = f"Unsupported type encountered in metadata_fields: {e}"
+            raise ValueError(msg) from e
+
+    @staticmethod
+    def _serialize_index_creation_kwargs(index_creation_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Serializes the index creation kwargs to a dictionary.
+        This is needed to handle serialization of Azure AI Search classes
+        that are passed in the index creation kwargs.
+        """
+        result = {}
+        for key, value in index_creation_kwargs.items():
+            if hasattr(value, "as_dict"):
+                result[key] = value.as_dict()
+            elif isinstance(value, list) and all(hasattr(item, "as_dict") for item in value):
+                result[key] = [item.as_dict() for item in value]
+            else:
+                result[key] = value
+        return result
+
+    @classmethod
+    def _deserialize_index_creation_kwargs(cls, data: Dict[str, Any]) -> Any:
+        """
+        Deserializes the index creation kwargs to the original classes.
+        """
+        result = {}
+        for key, value in data.items():
+            if key in AZURE_CLASS_MAPPING:
+                if isinstance(value, list):
+                    result[key] = [AZURE_CLASS_MAPPING[key].from_dict(item) for item in value]
+                else:
+                    result[key] = AZURE_CLASS_MAPPING[key].from_dict(value)
+            elif isinstance(value, dict) and hasattr(value, "from_dict"):
+                result[key] = value.from_dict(value)
+            else:
+                result[key] = value
+
+        return result[key]
+
     def to_dict(self) -> Dict[str, Any]:
         # This is not the best solution to serialise this class but is the fastest to implement.
         # Not all kwargs types can be serialised to text so this can fail. We must serialise each
@@ -198,15 +267,21 @@ def to_dict(self) -> Dict[str, Any]:
         :returns:
             Dictionary with serialized data.
         """
+
+        if self._metadata_fields:
+            serialized_metadata = {key: value.__name__ for key, value in self._metadata_fields.items()}
+        else:
+            serialized_metadata = None
+
         return default_to_dict(
             self,
             azure_endpoint=self._azure_endpoint.to_dict() if self._azure_endpoint else None,
             api_key=self._api_key.to_dict() if self._api_key else None,
             index_name=self._index_name,
             embedding_dimension=self._embedding_dimension,
-            metadata_fields=self._metadata_fields,
+            metadata_fields=serialized_metadata,
             vector_search_configuration=self._vector_search_configuration.as_dict(),
-            **self._index_creation_kwargs,
+            **self._serialize_index_creation_kwargs(self._index_creation_kwargs),
         )
 
     @classmethod
@@ -220,6 +295,13 @@ def from_dict(cls, data: Dict[str, Any]) -> "AzureAISearchDocumentStore":
         :returns:
             Deserialized component.
         """
+        if (fields := data["init_parameters"]["metadata_fields"]) is not None:
+            data["init_parameters"]["metadata_fields"] = cls._deserialize_metadata_fields(fields)
+
+        for key, _value in AZURE_CLASS_MAPPING.items():
+            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})
 
         deserialize_secrets_inplace(data["init_parameters"], keys=["api_key", "azure_endpoint"])
         if (vector_search_configuration := data["init_parameters"].get("vector_search_configuration")) is not None:

integrations/azure_ai_search/tests/test_document_store.py

@@ -8,6 +8,7 @@
 from unittest.mock import patch
 
 import pytest
+from azure.search.documents.indexes.models import CustomAnalyzer, SearchResourceEncryptionKey
 from haystack.dataclasses.document import Document
 from haystack.errors import FilterError
 from haystack.testing.document_store import (
@@ -18,10 +19,12 @@
 )
 from haystack.utils.auth import EnvVarSecret, Secret
 
-from haystack_integrations.document_stores.azure_ai_search import DEFAULT_VECTOR_SEARCH, AzureAISearchDocumentStore
+from haystack_integrations.document_stores.azure_ai_search import (
+    DEFAULT_VECTOR_SEARCH,
+    AzureAISearchDocumentStore,
+)
 
 
-@patch("haystack_integrations.document_stores.azure_ai_search.document_store.AzureAISearchDocumentStore")
 def test_to_dict(monkeypatch):
     monkeypatch.setenv("AZURE_AI_SEARCH_API_KEY", "test-api-key")
     monkeypatch.setenv("AZURE_AI_SEARCH_ENDPOINT", "test-endpoint")
@@ -51,7 +54,68 @@ def test_to_dict(monkeypatch):
     }
 
 
-@patch("haystack_integrations.document_stores.azure_ai_search.document_store.AzureAISearchDocumentStore")
+def test_to_dict_with_params(monkeypatch):
+    monkeypatch.setenv("AZURE_AI_SEARCH_API_KEY", "test-api-key")
+    monkeypatch.setenv("AZURE_AI_SEARCH_ENDPOINT", "test-endpoint")
+    encryption_key = SearchResourceEncryptionKey(
+        key_name="my-key",
+        key_version="my-version",
+        vault_uri="my-uri",
+    )
+    analyzer = CustomAnalyzer(
+        name="url-analyze",
+        tokenizer_name="uax_url_email",
+        token_filters=["lowercase"],  # Using token filter name directly as string
+    )
+    document_store = AzureAISearchDocumentStore(
+        index_name="my_index",
+        embedding_dimension=15,
+        metadata_fields={"Title": str, "Pages": int},
+        encryption_key=encryption_key,
+        analyzers=[analyzer],
+    )
+
+    res = document_store.to_dict()
+    assert res == {
+        "type": "haystack_integrations.document_stores.azure_ai_search.document_store.AzureAISearchDocumentStore",
+        "init_parameters": {
+            "azure_endpoint": {"env_vars": ["AZURE_AI_SEARCH_ENDPOINT"], "strict": True, "type": "env_var"},
+            "api_key": {"env_vars": ["AZURE_AI_SEARCH_API_KEY"], "strict": False, "type": "env_var"},
+            "index_name": "my_index",
+            "embedding_dimension": 15,
+            "metadata_fields": {
+                "Title": "str",
+                "Pages": "int",
+            },
+            "encryption_key": {
+                "key_name": "my-key",
+                "key_version": "my-version",
+                "vault_uri": "my-uri",
+            },
+            "analyzers": [
+                {
+                    "name": "url-analyze",
+                    "odata_type": "#Microsoft.Azure.Search.CustomAnalyzer",
+                    "tokenizer_name": "uax_url_email",
+                    "token_filters": ["lowercase"],
+                }
+            ],
+            "vector_search_configuration": {
+                "profiles": [
+                    {"name": "default-vector-config", "algorithm_configuration_name": "cosine-algorithm-config"}
+                ],
+                "algorithms": [
+                    {
+                        "name": "cosine-algorithm-config",
+                        "kind": "hnsw",
+                        "parameters": {"m": 4, "ef_construction": 400, "ef_search": 500, "metric": "cosine"},
+                    }
+                ],
+            },
+        },
+    }
+
+
 def test_from_dict(monkeypatch):
     monkeypatch.setenv("AZURE_AI_SEARCH_API_KEY", "test-api-key")
     monkeypatch.setenv("AZURE_AI_SEARCH_ENDPOINT", "test-endpoint")
@@ -76,6 +140,66 @@ def test_from_dict(monkeypatch):
     assert document_store._vector_search_configuration == DEFAULT_VECTOR_SEARCH
 
 
+def test_from_dict_with_params(monkeypatch):
+    monkeypatch.setenv("AZURE_AI_SEARCH_API_KEY", "test-api-key")
+    monkeypatch.setenv("AZURE_AI_SEARCH_ENDPOINT", "test-endpoint")
+    encryption_key = SearchResourceEncryptionKey(
+        key_name="my-key",
+        key_version="my-version",
+        vault_uri="my-uri",
+    )
+
+    data = {
+        "type": "haystack_integrations.document_stores.azure_ai_search.document_store.AzureAISearchDocumentStore",
+        "init_parameters": {
+            "azure_endpoint": {"env_vars": ["AZURE_AI_SEARCH_ENDPOINT"], "strict": True, "type": "env_var"},
+            "api_key": {"env_vars": ["AZURE_AI_SEARCH_API_KEY"], "strict": False, "type": "env_var"},
+            "index_name": "my_index",
+            "embedding_dimension": 15,
+            "metadata_fields": {
+                "Title": "str",
+                "Pages": "int",
+            },
+            "encryption_key": {
+                "key_name": "my-key",
+                "key_version": "my-version",
+                "vault_uri": "my-uri",
+            },
+            "analyzers": [
+                {
+                    "name": "url-analyze",
+                    "odata_type": "#Microsoft.Azure.Search.CustomAnalyzer",
+                    "tokenizer_name": "uax_url_email",
+                    "token_filters": ["lowercase"],
+                }
+            ],
+            "vector_search_configuration": {
+                "profiles": [
+                    {"name": "default-vector-config", "algorithm_configuration_name": "cosine-algorithm-config"}
+                ],
+                "algorithms": [
+                    {
+                        "name": "cosine-algorithm-config",
+                        "kind": "hnsw",
+                        "parameters": {"m": 4, "ef_construction": 400, "ef_search": 500, "metric": "cosine"},
+                    }
+                ],
+            },
+        },
+    }
+    document_store = AzureAISearchDocumentStore.from_dict(data)
+    assert isinstance(document_store._api_key, EnvVarSecret)
+    assert isinstance(document_store._azure_endpoint, EnvVarSecret)
+    assert document_store._index_name == "my_index"
+    assert document_store._embedding_dimension == 15
+    assert document_store._metadata_fields == {"Title": str, "Pages": int}
+    assert document_store._index_creation_kwargs["encryption_key"] == encryption_key
+    assert document_store._index_creation_kwargs["analyzers"][0].name == "url-analyze"
+    assert document_store._index_creation_kwargs["analyzers"][0].token_filters == ["lowercase"]
+    assert "CustomAnalyzer" in document_store._index_creation_kwargs["analyzers"][0].odata_type
+    assert document_store._vector_search_configuration.as_dict() == DEFAULT_VECTOR_SEARCH.as_dict()
+
+
 @patch("haystack_integrations.document_stores.azure_ai_search.document_store.AzureAISearchDocumentStore")
 def test_init_is_lazy(_mock_azure_search_client):
     AzureAISearchDocumentStore(azure_endpoint=Secret.from_token("test_endpoint"))
@@ -84,7 +208,6 @@ def test_init_is_lazy(_mock_azure_search_client):
 
 @patch("haystack_integrations.document_stores.azure_ai_search.document_store.AzureAISearchDocumentStore")
 def test_init(_mock_azure_search_client):
-
     document_store = AzureAISearchDocumentStore(
         api_key=Secret.from_token("fake-api-key"),
         azure_endpoint=Secret.from_token("fake_endpoint"),