feat: Add run async for AzureOpenAIChatGenerator (#8948)

* Add tests for run_async

* Add reno

* Add async client

* Add init test

* Add comment

* Fix test

* Update releasenotes/notes/run-async-azure-54450f0c2495f5c8.yaml

Co-authored-by: Amna Mubashar <[email protected]>

---------

Co-authored-by: Amna Mubashar <[email protected]>

Author: sjrl

haystack/components/generators/chat/azure.py

@@ -6,7 +6,7 @@
 from typing import Any, Callable, Dict, List, Optional
 
 # pylint: disable=import-error
-from openai.lib.azure import AzureOpenAI
+from openai.lib.azure import AsyncAzureOpenAI, AzureOpenAI
 
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.components.generators.chat import OpenAIChatGenerator
@@ -154,17 +154,20 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         self.tools = tools
         self.tools_strict = tools_strict
 
-        self.client = AzureOpenAI(
-            api_version=api_version,
-            azure_endpoint=azure_endpoint,
-            azure_deployment=azure_deployment,
-            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,
-            default_headers=self.default_headers,
-        )
+        client_args: Dict[str, Any] = {
+            "api_version": api_version,
+            "azure_endpoint": azure_endpoint,
+            "azure_deployment": azure_deployment,
+            "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,
+            "default_headers": self.default_headers,
+        }
+
+        self.client = AzureOpenAI(**client_args)
+        self.async_client = AsyncAzureOpenAI(**client_args)
 
     def to_dict(self) -> Dict[str, Any]:
         """

releasenotes/notes/run-async-azure-54450f0c2495f5c8.yaml

@@ -0,0 +1,7 @@
+---
+---
+features:
+  - |
+    Add `run_async` method to `AzureOpenAIChatGenerator`. This method uses `AsyncAzureOpenAI`
+    to generate chat completions and supports the same parameters as the `run` method. It returns a coroutine
+    that can be awaited.

test/components/generators/chat/test_azure.py

@@ -242,3 +242,69 @@ def test_live_run_with_tools(self, tools):
         assert message.meta["finish_reason"] == "tool_calls"
 
     # additional tests intentionally omitted as they are covered by test_openai.py
+
+
+class TestAzureOpenAIChatGeneratorAsync:
+    def test_init_should_also_create_async_client_with_same_args(self, tools):
+        component = AzureOpenAIChatGenerator(
+            api_key=Secret.from_token("test-api-key"),
+            azure_endpoint="some-non-existing-endpoint",
+            streaming_callback=print_streaming_chunk,
+            generation_kwargs={"max_tokens": 10, "some_test_param": "test-params"},
+            tools=tools,
+            tools_strict=True,
+        )
+        assert component.async_client.api_key == "test-api-key"
+        assert component.azure_deployment == "gpt-4o-mini"
+        assert component.streaming_callback is print_streaming_chunk
+        assert component.generation_kwargs == {"max_tokens": 10, "some_test_param": "test-params"}
+        assert component.tools == tools
+        assert component.tools_strict
+
+    @pytest.mark.integration
+    @pytest.mark.skipif(
+        not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
+        reason=(
+            "Please export env variables called AZURE_OPENAI_API_KEY containing "
+            "the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
+            "the Azure OpenAI endpoint URL to run this test."
+        ),
+    )
+    @pytest.mark.asyncio
+    async def test_live_run_async(self):
+        chat_messages = [ChatMessage.from_user("What's the capital of France")]
+        component = AzureOpenAIChatGenerator(generation_kwargs={"n": 1})
+        results = await component.run_async(chat_messages)
+        assert len(results["replies"]) == 1
+        message: ChatMessage = results["replies"][0]
+        assert "Paris" in message.text
+        assert "gpt-4o" in message.meta["model"]
+        assert message.meta["finish_reason"] == "stop"
+
+    @pytest.mark.integration
+    @pytest.mark.skipif(
+        not os.environ.get("AZURE_OPENAI_API_KEY", None) or not os.environ.get("AZURE_OPENAI_ENDPOINT", None),
+        reason=(
+            "Please export env variables called AZURE_OPENAI_API_KEY containing "
+            "the Azure OpenAI key, AZURE_OPENAI_ENDPOINT containing "
+            "the Azure OpenAI endpoint URL to run this test."
+        ),
+    )
+    @pytest.mark.asyncio
+    async def test_live_run_with_tools_async(self, tools):
+        chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
+        component = AzureOpenAIChatGenerator(tools=tools)
+        results = await component.run_async(chat_messages)
+        assert len(results["replies"]) == 1
+        message = results["replies"][0]
+
+        assert not message.texts
+        assert not message.text
+        assert message.tool_calls
+        tool_call = message.tool_call
+        assert isinstance(tool_call, ToolCall)
+        assert tool_call.tool_name == "weather"
+        assert tool_call.arguments == {"city": "Paris"}
+        assert message.meta["finish_reason"] == "tool_calls"
+
+    # additional tests intentionally omitted as they are covered by test_openai.py