From cd8a5b98fee37fefe712207fadb98a5c03fbb9da Mon Sep 17 00:00:00 2001
From: Nitanshu Vashistha <nitanshu.vzard@gmail.com>
Date: Tue, 9 Jul 2024 13:26:46 +0530
Subject: [PATCH] feat: Configure max_retries & timeout for
 AzureOpenAITextEmbedder (#7993)

max_retries: if not set is read from the OPENAI_MAX_RETRIES
env variable or set to 5.

timeout: if not set is read from the OPENAI_TIMEOUT
env variable or set to 30.

Signed-off-by: Nitanshu Vashistha <nitanshu.vzard@gmail.com>
---
 haystack/components/embedders/azure_text_embedder.py | 12 ++++++++++++
 ...for-AzureOpenAITextEmbedder-b02b760e9d6f28aa.yaml |  4 ++++
 .../components/embedders/test_azure_text_embedder.py |  2 ++
 3 files changed, 18 insertions(+)
 create mode 100644 releasenotes/notes/timeout-and-retries-for-AzureOpenAITextEmbedder-b02b760e9d6f28aa.yaml

diff --git a/haystack/components/embedders/azure_text_embedder.py b/haystack/components/embedders/azure_text_embedder.py
index b1ac40fec..fd0cecf74 100644
--- a/haystack/components/embedders/azure_text_embedder.py
+++ b/haystack/components/embedders/azure_text_embedder.py
@@ -41,6 +41,8 @@ class AzureOpenAITextEmbedder:
         api_key: Optional[Secret] = Secret.from_env_var("AZURE_OPENAI_API_KEY", strict=False),
         azure_ad_token: Optional[Secret] = Secret.from_env_var("AZURE_OPENAI_AD_TOKEN", strict=False),
         organization: Optional[str] = None,
+        timeout: Optional[float] = None,
+        max_retries: Optional[int] = None,
         prefix: str = "",
         suffix: str = "",
     ):
@@ -67,6 +69,10 @@ class AzureOpenAITextEmbedder:
             The Organization ID. See OpenAI's
             [production best practices](https://platform.openai.com/docs/guides/production-best-practices/setting-up-your-organization)
             for more information.
+        :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 prefix:
             A string to add at the beginning of each text.
         :param suffix:
@@ -90,6 +96,8 @@ class AzureOpenAITextEmbedder:
         self.azure_deployment = azure_deployment
         self.dimensions = dimensions
         self.organization = organization
+        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.prefix = prefix
         self.suffix = suffix
 
@@ -100,6 +108,8 @@ class AzureOpenAITextEmbedder:
             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]:
@@ -126,6 +136,8 @@ class AzureOpenAITextEmbedder:
             suffix=self.suffix,
             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/releasenotes/notes/timeout-and-retries-for-AzureOpenAITextEmbedder-b02b760e9d6f28aa.yaml b/releasenotes/notes/timeout-and-retries-for-AzureOpenAITextEmbedder-b02b760e9d6f28aa.yaml
new file mode 100644
index 000000000..0f30a4738
--- /dev/null
+++ b/releasenotes/notes/timeout-and-retries-for-AzureOpenAITextEmbedder-b02b760e9d6f28aa.yaml
@@ -0,0 +1,4 @@
+---
+enhancements:
+  - |
+    Add `max_retries` and `timeout` parameters to the AzureOpenAITextEmbedder initializations.
diff --git a/test/components/embedders/test_azure_text_embedder.py b/test/components/embedders/test_azure_text_embedder.py
index 77170779b..6a567b8fd 100644
--- a/test/components/embedders/test_azure_text_embedder.py
+++ b/test/components/embedders/test_azure_text_embedder.py
@@ -34,6 +34,8 @@ class TestAzureOpenAITextEmbedder:
                 "organization": None,
                 "azure_endpoint": "https://example-resource.azure.openai.com/",
                 "api_version": "2023-05-15",
+                "max_retries": 5,
+                "timeout": 30.0,
                 "prefix": "",
                 "suffix": "",
             },