Repository: deepset-ai/haystack

Issue:
Add `max_retries` and `timeout` params to all `AzureOpenAI` classes
**Is your feature request related to a problem? Please describe.**

Currently all `OpenAI` related classes (e.g. `OpenAIDocumentEmbedder`, `OpenAIChatGenerator`) can be initialised by setting `max_retries` and `timeout` params.

The corresponding `AzureOpenAI` don't always have the same params.

**Describe the solution you'd like**

It would be nice to have these params in the `AzureOpenAI` classes

**Describe alternatives you've considered**

Subclass `AzureOpenAI` and create custom components.

**Additional context**

cc @anakin87 :)

Discussion:
coveralls: ## Pull Request Test Coverage Report for [Build 9840411542](https://coveralls.io/builds/68512733)


### Details

* **0** of **0**   changed or added relevant lines in **0** files are covered.
* **23** unchanged lines in **1** file lost coverage.
* Overall coverage increased (+**0.002%**) to **90.009%**

---


|  Files with Coverage Reduction | New Missed Lines | % |
| :-----|--------------|--: |
| [components/embedders/azure_document_embedder.py](https://coveralls.io/builds/68512733/source?filename=components%2Fembedders%2Fazure_document_embedder.py#L97) | 23 | 53.62% |
<!-- | **Total:** | **23** |  | -->

|  Totals | [![Coverage Status](https://coveralls.io/builds/68512733/badge)](https://coveralls.io/builds/68512733) |
| :-- | --: |
| Change from base [Build 9839651196](https://coveralls.io/builds/68511553): |  0.002% |
| Covered Lines: | 6775 |
| Relevant Lines: | 7527 |

---
##### 💛  - [Coveralls](https://coveralls.io)

davidsbatista: thanks again 👍🏽 

Patch:
diff --git a/haystack/components/embedders/azure_document_embedder.py b/haystack/components/embedders/azure_document_embedder.py
index 54880878..50720d07 100644
--- a/haystack/components/embedders/azure_document_embedder.py
+++ b/haystack/components/embedders/azure_document_embedder.py
@@ -33,7 +33,7 @@ class AzureOpenAIDocumentEmbedder:
     ```
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913 (too-many-arguments)
         self,
         azure_endpoint: Optional[str] = None,
         api_version: Optional[str] = "2023-05-15",
@@ -48,6 +48,8 @@ class AzureOpenAIDocumentEmbedder:
         progress_bar: bool = True,
         meta_fields_to_embed: Optional[List[str]] = None,
         embedding_separator: str = "\n",
+        timeout: Optional[float] = None,
+        max_retries: Optional[int] = None,
     ):
         """
         Create an AzureOpenAIDocumentEmbedder component.
@@ -84,6 +86,10 @@ class AzureOpenAIDocumentEmbedder:
             List of meta fields that will be embedded along with the Document text.
         :param embedding_separator:
             Separator used to concatenate the meta fields to the Document text.
+        :param timeout: The timeout in seconds to be passed to the underlying `AzureOpenAI` client, if not set it is
+            inferred from the `OPENAI_TIMEOUT` environment variable or set to 30.
+        :param max_retries: Maximum retries to establish a connection with AzureOpenAI if it returns an internal error,
+            if not set it is inferred from the `OPENAI_MAX_RETRIES` environment variable or set to 5.
         """
         # if not provided as a parameter, azure_endpoint is read from the env var AZURE_OPENAI_ENDPOINT
         azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
@@ -106,6 +112,8 @@ class AzureOpenAIDocumentEmbedder:
         self.progress_bar = progress_bar
         self.meta_fields_to_embed = meta_fields_to_embed or []
         self.embedding_separator = embedding_separator
+        self.timeout = timeout or float(os.environ.get("OPENAI_TIMEOUT", 30.0))
+        self.max_retries = max_retries or int(os.environ.get("OPENAI_MAX_RETRIES", 5))
 
         self._client = AzureOpenAI(
             api_version=api_version,
@@ -114,6 +122,8 @@ class AzureOpenAIDocumentEmbedder:
             api_key=api_key.resolve_value() if api_key is not None else None,
             azure_ad_token=azure_ad_token.resolve_value() if azure_ad_token is not None else None,
             organization=organization,
+            timeout=self.timeout,
+            max_retries=self.max_retries,
         )
 
     def _get_telemetry_data(self) -> Dict[str, Any]:
@@ -144,6 +154,8 @@ class AzureOpenAIDocumentEmbedder:
             embedding_separator=self.embedding_separator,
             api_key=self.api_key.to_dict() if self.api_key is not None else None,
             azure_ad_token=self.azure_ad_token.to_dict() if self.azure_ad_token is not None else None,
+            timeout=self.timeout,
+            max_retries=self.max_retries,
         )
 
     @classmethod
diff --git a/haystack/components/generators/chat/azure.py b/haystack/components/generators/chat/azure.py
index 06da6095..9ef4311e 100644
--- a/haystack/components/generators/chat/azure.py
+++ b/haystack/components/generators/chat/azure.py
@@ -81,6 +81,7 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
         organization: Optional[str] = None,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
         timeout: Optional[float] = None,
+        max_retries: Optional[int] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """
@@ -95,6 +96,10 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
         [production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization).
         :param streaming_callback: A callback function that is called when a new token is received from the stream.
             The callback function accepts StreamingChunk as an argument.
+        :param timeout: The timeout in seconds to be passed to the underlying `AzureOpenAI` client, if not set it is
+            inferred from the `OPENAI_TIMEOUT` environment variable or set to 30.
+        :param max_retries: Maximum retries to establish a connection with AzureOpenAI if it returns an internal error,
+            if not set it is inferred from the `OPENAI_MAX_RETRIES` environment variable or set to 5.
         :param generation_kwargs: Other parameters to use for the model. These parameters are all sent directly to
             the OpenAI endpoint. See OpenAI [documentation](https://platform.openai.com/docs/api-reference/chat) for
             more details.
@@ -140,7 +145,8 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
         self.azure_deployment = azure_deployment
         self.organization = organization
         self.model = azure_deployment or "gpt-35-turbo"
-        self.timeout = timeout
+        self.timeout = timeout or float(os.environ.get("OPENAI_TIMEOUT", 30.0))
+        self.max_retries = max_retries or int(os.environ.get("OPENAI_MAX_RETRIES", 5))
 
         self.client = AzureOpenAI(
             api_version=api_version,
@@ -149,6 +155,8 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
             api_key=api_key.resolve_value() if api_key is not None else None,
             azure_ad_token=azure_ad_token.resolve_value() if azure_ad_token is not None else None,
             organization=organization,
+            timeout=self.timeout,
+            max_retries=self.max_retries,
         )
 
     def to_dict(self) -> Dict[str, Any]:
@@ -168,6 +176,7 @@ class AzureOpenAIChatGenerator(OpenAIChatGenerator):
             streaming_callback=callback_name,
             generation_kwargs=self.generation_kwargs,
             timeout=self.timeout,
+            max_retries=self.max_retries,
             api_key=self.api_key.to_dict() if self.api_key is not None else None,
             azure_ad_token=self.azure_ad_token.to_dict() if self.azure_ad_token is not None else None,
         )
diff --git a/releasenotes/notes/max_retries-for-AzureOpenAIChatGenerator-9e49b4c7bec5c72b.yaml b/releasenotes/notes/max_retries-for-AzureOpenAIChatGenerator-9e49b4c7bec5c72b.yaml
new file mode 100644
index 00000000..888a53a8
--- /dev/null
+++ b/releasenotes/notes/max_retries-for-AzureOpenAIChatGenerator-9e49b4c7bec5c72b.yaml
@@ -0,0 +1,4 @@
+---
+enhancements:
+  - |
+    Add `max_retries` and `timeout` parameters to the AzureOpenAIChatGenerator initializations.
diff --git a/releasenotes/notes/retries-and-timeout-for-AzureOpenAIDocumentEmbedder-006fd84204942e43.yaml b/releasenotes/notes/retries-and-timeout-for-AzureOpenAIDocumentEmbedder-006fd84204942e43.yaml
new file mode 100644
index 00000000..a4687dcb
--- /dev/null
+++ b/releasenotes/notes/retries-and-timeout-for-AzureOpenAIDocumentEmbedder-006fd84204942e43.yaml
@@ -0,0 +1,4 @@
+---
+enhancements:
+  - |
+    Add `max_retries`, `timeout` parameters to the `AzureOpenAIDocumentEmbedder` initialization.


Test patch:
diff --git a/test/components/embedders/test_azure_document_embedder.py b/test/components/embedders/test_azure_document_embedder.py
index 05f04bea..f9663e61 100644
--- a/test/components/embedders/test_azure_document_embedder.py
+++ b/test/components/embedders/test_azure_document_embedder.py
@@ -43,6 +43,8 @@ class TestAzureOpenAIDocumentEmbedder:
                 "progress_bar": True,
                 "meta_fields_to_embed": [],
                 "embedding_separator": "\n",
+                "max_retries": 5,
+                "timeout": 30.0,
             },
         }
 
diff --git a/test/components/generators/chat/test_azure.py b/test/components/generators/chat/test_azure.py
index 80ffb956..80a290a6 100644
--- a/test/components/generators/chat/test_azure.py
+++ b/test/components/generators/chat/test_azure.py
@@ -55,7 +55,8 @@ class TestOpenAIChatGenerator:
                 "organization": None,
                 "streaming_callback": None,
                 "generation_kwargs": {},
-                "timeout": None,
+                "timeout": 30.0,
+                "max_retries": 5,
             },
         }
 
@@ -66,6 +67,7 @@ class TestOpenAIChatGenerator:
             azure_ad_token=Secret.from_env_var("ENV_VAR1", strict=False),
             azure_endpoint="some-non-existing-endpoint",
             timeout=2.5,
+            max_retries=10,
             generation_kwargs={"max_tokens": 10, "some_test_param": "test-params"},
         )
         data = component.to_dict()
@@ -80,6 +82,7 @@ class TestOpenAIChatGenerator:
                 "organization": None,
                 "streaming_callback": None,
                 "timeout": 2.5,
+                "max_retries": 10,
                 "generation_kwargs": {"max_tokens": 10, "some_test_param": "test-params"},
             },
         }