Repository: deepset-ai/haystack

Issue:
add run_async for OpenAIChatGenerator



Discussion:

Patch:
diff --git a/haystack/components/generators/chat/openai.py b/haystack/components/generators/chat/openai.py
index 1b872a0d7b..52f7c9340c 100644
--- a/haystack/components/generators/chat/openai.py
+++ b/haystack/components/generators/chat/openai.py
@@ -5,24 +5,29 @@
 import json
 import os
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
-from openai import OpenAI, Stream
+from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream
 from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
 from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
 
 from haystack import component, default_from_dict, default_to_dict, logging
-from haystack.dataclasses import ChatMessage, StreamingChunk, ToolCall
+from haystack.dataclasses import (
+    AsyncStreamingCallbackT,
+    ChatMessage,
+    StreamingCallbackT,
+    StreamingChunk,
+    SyncStreamingCallbackT,
+    ToolCall,
+    select_streaming_callback,
+)
 from haystack.tools.tool import Tool, _check_duplicate_tool_names, deserialize_tools_inplace
 from haystack.utils import Secret, deserialize_callable, deserialize_secrets_inplace, serialize_callable
 
 logger = logging.getLogger(__name__)
 
 
-StreamingCallbackT = Callable[[StreamingChunk], None]
-
-
 @component
 class OpenAIChatGenerator:
     """
@@ -145,13 +150,16 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         if max_retries is None:
             max_retries = int(os.environ.get("OPENAI_MAX_RETRIES", 5))
 
-        self.client = OpenAI(
-            api_key=api_key.resolve_value(),
-            organization=organization,
-            base_url=api_base_url,
-            timeout=timeout,
-            max_retries=max_retries,
-        )
+        client_args: Dict[str, Any] = {
+            "api_key": api_key.resolve_value(),
+            "organization": organization,
+            "base_url": api_base_url,
+            "timeout": timeout,
+            "max_retries": max_retries,
+        }
+
+        self.client = OpenAI(**client_args)
+        self.async_client = AsyncOpenAI(**client_args)
 
     def _get_telemetry_data(self) -> Dict[str, Any]:
         """
@@ -268,6 +276,82 @@ def run(
 
         return {"replies": completions}
 
+    @component.output_types(replies=List[ChatMessage])
+    async def run_async(
+        self,
+        messages: List[ChatMessage],
+        streaming_callback: Optional[StreamingCallbackT] = None,
+        generation_kwargs: Optional[Dict[str, Any]] = None,
+        *,
+        tools: Optional[List[Tool]] = None,
+        tools_strict: Optional[bool] = None,
+    ):
+        """
+        Asynchronously invokes chat completion based on the provided messages and generation parameters.
+
+        This is the asynchronous version of the `run` method. It has the same parameters and return values
+        but can be used with `await` in async code.
+
+        :param messages:
+            A list of ChatMessage instances representing the input messages.
+        :param streaming_callback:
+            A callback function that is called when a new token is received from the stream.
+            Must be a coroutine.
+        :param generation_kwargs:
+            Additional keyword arguments for text generation. These parameters will
+            override the parameters passed during component initialization.
+            For details on OpenAI API parameters, see [OpenAI documentation](https://platform.openai.com/docs/api-reference/chat/create).
+        :param tools:
+            A list of tools for which the model can prepare calls. If set, it will override the `tools` parameter set
+            during component initialization.
+        :param tools_strict:
+            Whether to enable strict schema adherence for tool calls. If set to `True`, the model will follow exactly
+            the schema provided in the `parameters` field of the tool definition, but this may increase latency.
+            If set, it will override the `tools_strict` parameter set during component initialization.
+
+        :returns:
+            A dictionary with the following key:
+            - `replies`: A list containing the generated responses as ChatMessage instances.
+        """
+        # validate and select the streaming callback
+        streaming_callback = select_streaming_callback(self.streaming_callback, streaming_callback, requires_async=True)  # type: ignore
+
+        if len(messages) == 0:
+            return {"replies": []}
+
+        api_args = self._prepare_api_call(
+            messages=messages,
+            streaming_callback=streaming_callback,
+            generation_kwargs=generation_kwargs,
+            tools=tools,
+            tools_strict=tools_strict,
+        )
+
+        chat_completion: Union[
+            AsyncStream[ChatCompletionChunk], ChatCompletion
+        ] = await self.async_client.chat.completions.create(**api_args)
+
+        is_streaming = isinstance(chat_completion, AsyncStream)
+        assert is_streaming or streaming_callback is None
+
+        if is_streaming:
+            completions = await self._handle_async_stream_response(
+                chat_completion,  # type: ignore
+                streaming_callback,  # type: ignore
+            )
+        else:
+            assert isinstance(chat_completion, ChatCompletion), "Unexpected response type for non-streaming request."
+            completions = [
+                self._convert_chat_completion_to_chat_message(chat_completion, choice)
+                for choice in chat_completion.choices
+            ]
+
+        # before returning, do post-processing of the completions
+        for message in completions:
+            self._check_finish_reason(message.meta)
+
+        return {"replies": completions}
+
     def _prepare_api_call(  # noqa: PLR0913
         self,
         *,
@@ -309,7 +393,7 @@ def _prepare_api_call(  # noqa: PLR0913
             **generation_kwargs,
         }
 
-    def _handle_stream_response(self, chat_completion: Stream, callback: StreamingCallbackT) -> List[ChatMessage]:
+    def _handle_stream_response(self, chat_completion: Stream, callback: SyncStreamingCallbackT) -> List[ChatMessage]:
         chunks: List[StreamingChunk] = []
         chunk = None
 
@@ -322,6 +406,21 @@ def _handle_stream_response(self, chat_completion: Stream, callback: StreamingCa
 
         return [self._convert_streaming_chunks_to_chat_message(chunk, chunks)]
 
+    async def _handle_async_stream_response(
+        self, chat_completion: AsyncStream, callback: AsyncStreamingCallbackT
+    ) -> List[ChatMessage]:
+        chunks: List[StreamingChunk] = []
+        chunk = None
+
+        async for chunk in chat_completion:  # pylint: disable=not-an-iterable
+            assert len(chunk.choices) == 1, "Streaming responses should have only one choice."
+            chunk_delta: StreamingChunk = self._convert_chat_completion_chunk_to_streaming_chunk(chunk)
+            chunks.append(chunk_delta)
+
+            await callback(chunk_delta)
+
+        return [self._convert_streaming_chunks_to_chat_message(chunk, chunks)]
+
     def _check_finish_reason(self, meta: Dict[str, Any]) -> None:
         if meta["finish_reason"] == "length":
             logger.warning(
diff --git a/haystack/dataclasses/__init__.py b/haystack/dataclasses/__init__.py
index 91e8f0408f..17063e12d5 100644
--- a/haystack/dataclasses/__init__.py
+++ b/haystack/dataclasses/__init__.py
@@ -7,7 +7,13 @@
 from haystack.dataclasses.chat_message import ChatMessage, ChatRole, TextContent, ToolCall, ToolCallResult
 from haystack.dataclasses.document import Document
 from haystack.dataclasses.sparse_embedding import SparseEmbedding
-from haystack.dataclasses.streaming_chunk import StreamingChunk
+from haystack.dataclasses.streaming_chunk import (
+    AsyncStreamingCallbackT,
+    StreamingCallbackT,
+    StreamingChunk,
+    SyncStreamingCallbackT,
+    select_streaming_callback,
+)
 
 __all__ = [
     "Document",
@@ -22,4 +28,8 @@
     "TextContent",
     "StreamingChunk",
     "SparseEmbedding",
+    "select_streaming_callback",
+    "StreamingCallbackT",
+    "SyncStreamingCallbackT",
+    "AsyncStreamingCallbackT",
 ]
diff --git a/haystack/dataclasses/streaming_chunk.py b/haystack/dataclasses/streaming_chunk.py
index 455caa09e2..b11ef5f7f2 100644
--- a/haystack/dataclasses/streaming_chunk.py
+++ b/haystack/dataclasses/streaming_chunk.py
@@ -3,7 +3,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from dataclasses import dataclass, field
-from typing import Any, Dict
+from typing import Any, Awaitable, Callable, Dict, Optional, Union
+
+from haystack.utils.asynchronous import is_callable_async_compatible
 
 
 @dataclass
@@ -19,3 +21,41 @@ class StreamingChunk:
 
     content: str
     meta: Dict[str, Any] = field(default_factory=dict, hash=False)
+
+
+SyncStreamingCallbackT = Callable[[StreamingChunk], None]
+AsyncStreamingCallbackT = Callable[[StreamingChunk], Awaitable[None]]
+
+StreamingCallbackT = Union[SyncStreamingCallbackT, AsyncStreamingCallbackT]
+
+
+def select_streaming_callback(
+    init_callback: Optional[StreamingCallbackT], runtime_callback: Optional[StreamingCallbackT], requires_async: bool
+) -> Optional[StreamingCallbackT]:
+    """
+    Picks the correct streaming callback given an optional initial and runtime callback.
+
+    The runtime callback takes precedence over the initial callback.
+
+    :param init_callback:
+        The initial callback.
+    :param runtime_callback:
+        The runtime callback.
+    :param requires_async:
+        Whether the selected callback must be async compatible.
+    :returns:
+        The selected callback.
+    """
+    if init_callback is not None:
+        if requires_async and not is_callable_async_compatible(init_callback):
+            raise ValueError("The init callback must be async compatible.")
+        if not requires_async and is_callable_async_compatible(init_callback):
+            raise ValueError("The init callback cannot be a coroutine.")
+
+    if runtime_callback is not None:
+        if requires_async and not is_callable_async_compatible(runtime_callback):
+            raise ValueError("The runtime callback must be async compatible.")
+        if not requires_async and is_callable_async_compatible(runtime_callback):
+            raise ValueError("The runtime callback cannot be a coroutine.")
+
+    return runtime_callback or init_callback
diff --git a/haystack/utils/asynchronous.py b/haystack/utils/asynchronous.py
new file mode 100644
index 0000000000..e84ecbd039
--- /dev/null
+++ b/haystack/utils/asynchronous.py
@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+from typing import Callable
+
+
+def is_callable_async_compatible(func: Callable) -> bool:
+    """
+    Returns if the given callable is usable inside a component's `run_async` method.
+
+    :param callable:
+        The callable to check.
+    :returns:
+        True if the callable is compatible, False otherwise.
+    """
+    return inspect.iscoroutinefunction(func)
diff --git a/pyproject.toml b/pyproject.toml
index 2c04ceb905..7251e21a81 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -276,6 +276,7 @@ markers = [
   "document_store",
 ]
 log_cli = true
+asyncio_mode = "auto"
 
 [tool.mypy]
 warn_return_any = false
diff --git a/releasenotes/notes/add-run-async-to-openai-chat-generator-3b4724af4fe4707a.yaml b/releasenotes/notes/add-run-async-to-openai-chat-generator-3b4724af4fe4707a.yaml
new file mode 100644
index 0000000000..ad16a844b0
--- /dev/null
+++ b/releasenotes/notes/add-run-async-to-openai-chat-generator-3b4724af4fe4707a.yaml
@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Add `run_async` method to OpenAIChatGenerator. This method internally uses the async version of the OpenAI client
+    to generate chat completions and supports the same parameters as the `run` method. It returns a coroutine
+    that can be awaited.


Test patch:
diff --git a/test/components/generators/chat/test_openai.py b/test/components/generators/chat/test_openai.py
index 9a5b2e8352..561c5bb700 100644
--- a/test/components/generators/chat/test_openai.py
+++ b/test/components/generators/chat/test_openai.py
@@ -9,7 +9,7 @@
 import os
 from datetime import datetime
 
-from openai import OpenAIError
+from openai import AsyncOpenAI, OpenAIError
 from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessage, ChatCompletionMessageToolCall
 from openai.types.chat.chat_completion import Choice
 from openai.types.chat.chat_completion_message_tool_call import Function
@@ -151,6 +151,23 @@ def test_init_with_parameters_and_env_vars(self, monkeypatch):
         assert component.client.timeout == 100.0
         assert component.client.max_retries == 10
 
+    def test_init_should_also_create_async_client_with_same_args(self, monkeypatch):
+        monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
+        component = OpenAIChatGenerator(
+            api_key=Secret.from_token("test-api-key"),
+            api_base_url="test-base-url",
+            organization="test-organization",
+            timeout=30,
+            max_retries=5,
+        )
+
+        assert isinstance(component.async_client, AsyncOpenAI)
+        assert component.async_client.api_key == "test-api-key"
+        assert component.async_client.organization == "test-organization"
+        assert component.async_client.base_url == "test-base-url/"
+        assert component.async_client.timeout == 30
+        assert component.async_client.max_retries == 5
+
     def test_to_dict_default(self, monkeypatch):
         monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
         component = OpenAIChatGenerator()
diff --git a/test/components/generators/chat/test_openai_async.py b/test/components/generators/chat/test_openai_async.py
new file mode 100644
index 0000000000..9d6ae6b67f
--- /dev/null
+++ b/test/components/generators/chat/test_openai_async.py
@@ -0,0 +1,349 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+from unittest.mock import AsyncMock, patch
+
+from openai import AsyncOpenAI, OpenAIError
+import pytest
+from datetime import datetime
+import os
+
+from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, ChatCompletionChunk
+from openai.types.chat.chat_completion import Choice
+from openai.types.chat.chat_completion_message_tool_call import Function
+from openai.types.chat import chat_completion_chunk
+
+from haystack.dataclasses import StreamingChunk
+from haystack.utils.auth import Secret
+from haystack.dataclasses import ChatMessage, ToolCall
+from haystack.tools import Tool
+from haystack.components.generators.chat.openai import OpenAIChatGenerator
+
+
+@pytest.fixture
+def chat_messages():
+    return [
+        ChatMessage.from_system("You are a helpful assistant"),
+        ChatMessage.from_user("What's the capital of France"),
+    ]
+
+
+@pytest.fixture
+def mock_chat_completion_chunk_with_tools(openai_mock_stream_async):
+    """
+    Mock the OpenAI API completion chunk response and reuse it for tests
+    """
+
+    with patch(
+        "openai.resources.chat.completions.AsyncCompletions.create", new_callable=AsyncMock
+    ) as mock_chat_completion_create:
+        completion = ChatCompletionChunk(
+            id="foo",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            choices=[
+                chat_completion_chunk.Choice(
+                    finish_reason="tool_calls",
+                    logprobs=None,
+                    index=0,
+                    delta=chat_completion_chunk.ChoiceDelta(
+                        role="assistant",
+                        tool_calls=[
+                            chat_completion_chunk.ChoiceDeltaToolCall(
+                                index=0,
+                                id="123",
+                                type="function",
+                                function=chat_completion_chunk.ChoiceDeltaToolCallFunction(
+                                    name="weather", arguments='{"city": "Paris"}'
+                                ),
+                            )
+                        ],
+                    ),
+                )
+            ],
+            created=int(datetime.now().timestamp()),
+            usage={"prompt_tokens": 57, "completion_tokens": 40, "total_tokens": 97},
+        )
+        mock_chat_completion_create.return_value = openai_mock_stream_async(completion)
+        yield mock_chat_completion_create
+
+
+@pytest.fixture
+def tools():
+    tool_parameters = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
+    tool = Tool(
+        name="weather",
+        description="useful to determine the weather in a given location",
+        parameters=tool_parameters,
+        function=lambda x: x,
+    )
+    return [tool]
+
+
+class TestOpenAIChatGeneratorAsync:
+    def test_init_should_also_create_async_client_with_same_args(self, monkeypatch):
+        monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
+        component = OpenAIChatGenerator(
+            api_key=Secret.from_token("test-api-key"),
+            api_base_url="test-base-url",
+            organization="test-organization",
+            timeout=30,
+            max_retries=5,
+        )
+
+        assert isinstance(component.async_client, AsyncOpenAI)
+        assert component.async_client.api_key == "test-api-key"
+        assert component.async_client.organization == "test-organization"
+        assert component.async_client.base_url == "test-base-url/"
+        assert component.async_client.timeout == 30
+        assert component.async_client.max_retries == 5
+
+    @pytest.mark.asyncio
+    async def test_run_async(self, chat_messages, openai_mock_async_chat_completion):
+        component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
+        response = await component.run_async(chat_messages)
+
+        # check that the component returns the correct ChatMessage response
+        assert isinstance(response, dict)
+        assert "replies" in response
+        assert isinstance(response["replies"], list)
+        assert len(response["replies"]) == 1
+        assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
+
+    @pytest.mark.asyncio
+    async def test_run_with_params_async(self, chat_messages, openai_mock_async_chat_completion):
+        component = OpenAIChatGenerator(
+            api_key=Secret.from_token("test-api-key"), generation_kwargs={"max_tokens": 10, "temperature": 0.5}
+        )
+        response = await component.run_async(chat_messages)
+
+        # check that the component calls the OpenAI API with the correct parameters
+        _, kwargs = openai_mock_async_chat_completion.call_args
+        assert kwargs["max_tokens"] == 10
+        assert kwargs["temperature"] == 0.5
+
+        # check that the tools are not passed to the OpenAI API (the generator is initialized without tools)
+        assert "tools" not in kwargs
+
+        # check that the component returns the correct response
+        assert isinstance(response, dict)
+        assert "replies" in response
+        assert isinstance(response["replies"], list)
+        assert len(response["replies"]) == 1
+        assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
+
+    @pytest.mark.asyncio
+    async def test_run_with_params_streaming_async(self, chat_messages, openai_mock_async_chat_completion_chunk):
+        streaming_callback_called = False
+
+        async def streaming_callback(chunk: StreamingChunk) -> None:
+            nonlocal streaming_callback_called
+            streaming_callback_called = True
+
+        component = OpenAIChatGenerator(
+            api_key=Secret.from_token("test-api-key"), streaming_callback=streaming_callback
+        )
+        response = await component.run_async(chat_messages)
+
+        # check we called the streaming callback
+        assert streaming_callback_called
+
+        # check that the component still returns the correct response
+        assert isinstance(response, dict)
+        assert "replies" in response
+        assert isinstance(response["replies"], list)
+        assert len(response["replies"]) == 1
+        assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
+        assert "Hello" in response["replies"][0].text  # see openai_mock_chat_completion_chunk
+
+    @pytest.mark.asyncio
+    async def test_run_with_streaming_callback_in_run_method_async(
+        self, chat_messages, openai_mock_async_chat_completion_chunk
+    ):
+        streaming_callback_called = False
+
+        async def streaming_callback(chunk: StreamingChunk) -> None:
+            nonlocal streaming_callback_called
+            streaming_callback_called = True
+
+        component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"))
+        response = await component.run_async(chat_messages, streaming_callback=streaming_callback)
+
+        # check we called the streaming callback
+        assert streaming_callback_called
+
+        # check that the component still returns the correct response
+        assert isinstance(response, dict)
+        assert "replies" in response
+        assert isinstance(response["replies"], list)
+        assert len(response["replies"]) == 1
+        assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
+        assert "Hello" in response["replies"][0].text  # see openai_mock_chat_completion_chunk
+
+    @pytest.mark.asyncio
+    async def test_run_with_tools_async(self, tools):
+        with patch(
+            "openai.resources.chat.completions.AsyncCompletions.create", new_callable=AsyncMock
+        ) as mock_chat_completion_create:
+            completion = ChatCompletion(
+                id="foo",
+                model="gpt-4",
+                object="chat.completion",
+                choices=[
+                    Choice(
+                        finish_reason="tool_calls",
+                        logprobs=None,
+                        index=0,
+                        message=ChatCompletionMessage(
+                            role="assistant",
+                            tool_calls=[
+                                ChatCompletionMessageToolCall(
+                                    id="123",
+                                    type="function",
+                                    function=Function(name="weather", arguments='{"city": "Paris"}'),
+                                )
+                            ],
+                        ),
+                    )
+                ],
+                created=int(datetime.now().timestamp()),
+                usage={"prompt_tokens": 57, "completion_tokens": 40, "total_tokens": 97},
+            )
+
+            mock_chat_completion_create.return_value = completion
+
+            component = OpenAIChatGenerator(api_key=Secret.from_token("test-api-key"), tools=tools, tools_strict=True)
+            response = await component.run_async([ChatMessage.from_user("What's the weather like in Paris?")])
+
+        # ensure that the tools are passed to the OpenAI API
+        assert mock_chat_completion_create.call_args[1]["tools"] == [
+            {"type": "function", "function": {**tools[0].tool_spec, "strict": True}}
+        ]
+
+        assert len(response["replies"]) == 1
+        message = response["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"
+
+    @pytest.mark.asyncio
+    async def test_run_with_tools_streaming_async(self, mock_chat_completion_chunk_with_tools, tools):
+        streaming_callback_called = False
+
+        async def streaming_callback(chunk: StreamingChunk) -> None:
+            nonlocal streaming_callback_called
+            streaming_callback_called = True
+
+        component = OpenAIChatGenerator(
+            api_key=Secret.from_token("test-api-key"), streaming_callback=streaming_callback
+        )
+        chat_messages = [ChatMessage.from_user("What's the weather like in Paris?")]
+        response = await component.run_async(chat_messages, tools=tools)
+
+        # check we called the streaming callback
+        assert streaming_callback_called
+
+        # check that the component still returns the correct response
+        assert isinstance(response, dict)
+        assert "replies" in response
+        assert isinstance(response["replies"], list)
+        assert len(response["replies"]) == 1
+        assert [isinstance(reply, ChatMessage) for reply in response["replies"]]
+
+        message = response["replies"][0]
+
+        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"
+
+    @pytest.mark.skipif(
+        not os.environ.get("OPENAI_API_KEY", None),
+        reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
+    )
+    @pytest.mark.integration
+    @pytest.mark.asyncio
+    async def test_live_run_async(self):
+        chat_messages = [ChatMessage.from_user("What's the capital of France")]
+        component = OpenAIChatGenerator(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.skipif(
+        not os.environ.get("OPENAI_API_KEY", None),
+        reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
+    )
+    @pytest.mark.integration
+    @pytest.mark.asyncio
+    async def test_live_run_wrong_model_async(self, chat_messages):
+        component = OpenAIChatGenerator(model="something-obviously-wrong")
+        with pytest.raises(OpenAIError):
+            await component.run_async(chat_messages)
+
+    @pytest.mark.skipif(
+        not os.environ.get("OPENAI_API_KEY", None),
+        reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
+    )
+    @pytest.mark.integration
+    @pytest.mark.asyncio
+    async def test_live_run_streaming_async(self):
+        counter = 0
+        responses = ""
+
+        async def callback(chunk: StreamingChunk):
+            nonlocal counter
+            nonlocal responses
+            counter += 1
+            responses += chunk.content if chunk.content else ""
+
+        component = OpenAIChatGenerator(streaming_callback=callback)
+        results = await component.run_async([ChatMessage.from_user("What's the capital of France?")])
+
+        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"
+
+        assert counter > 1
+        assert "Paris" in responses
+
+        # check that the completion_start_time is set and valid ISO format
+        assert "completion_start_time" in message.meta
+        assert datetime.fromisoformat(message.meta["completion_start_time"]) < datetime.now()
+
+    @pytest.mark.skipif(
+        not os.environ.get("OPENAI_API_KEY", None),
+        reason="Export an env var called OPENAI_API_KEY containing the OpenAI API key to run this test.",
+    )
+    @pytest.mark.integration
+    @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 = OpenAIChatGenerator(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"
diff --git a/test/components/generators/conftest.py b/test/components/generators/conftest.py
index 4aa931451f..418120ed6e 100644
--- a/test/components/generators/conftest.py
+++ b/test/components/generators/conftest.py
@@ -3,10 +3,10 @@
 # SPDX-License-Identifier: Apache-2.0
 from datetime import datetime
 from typing import Iterator
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, patch, AsyncMock
 
 import pytest
-from openai import Stream
+from openai import AsyncStream, Stream
 from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessage
 from openai.types.chat.chat_completion_chunk import Choice, ChoiceDelta
 from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessage, ChatCompletionMessageToolCall
@@ -41,6 +41,21 @@ def __stream__(self) -> Iterator[ChatCompletionChunk]:
         yield self.mock_chunk
 
 
+class OpenAIAsyncMockStream(AsyncStream[ChatCompletionChunk]):
+    def __init__(self, mock_chunk: ChatCompletionChunk):
+        self.mock_chunk = mock_chunk
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        # Only yield once, then stop iteration
+        if not hasattr(self, "_done"):
+            self._done = True
+            return self.mock_chunk
+        raise StopAsyncIteration
+
+
 @pytest.fixture
 def openai_mock_stream():
     """
@@ -49,6 +64,14 @@ def openai_mock_stream():
     return OpenAIMockStream
 
 
+@pytest.fixture
+def openai_mock_stream_async():
+    """
+    Fixture that returns a function to create AsyncMockStream instances with custom chunks
+    """
+    return OpenAIAsyncMockStream
+
+
 @pytest.fixture
 def openai_mock_chat_completion():
     """
@@ -75,6 +98,34 @@ def openai_mock_chat_completion():
         yield mock_chat_completion_create
 
 
+@pytest.fixture
+async def openai_mock_async_chat_completion():
+    """
+    Mock the OpenAI API completion response and reuse it for async tests
+    """
+    with patch(
+        "openai.resources.chat.completions.AsyncCompletions.create", new_callable=AsyncMock
+    ) as mock_chat_completion_create:
+        completion = ChatCompletion(
+            id="foo",
+            model="gpt-4",
+            object="chat.completion",
+            choices=[
+                {
+                    "finish_reason": "stop",
+                    "logprobs": None,
+                    "index": 0,
+                    "message": {"content": "Hello world!", "role": "assistant"},
+                }
+            ],
+            created=int(datetime.now().timestamp()),
+            usage={"prompt_tokens": 57, "completion_tokens": 40, "total_tokens": 97},
+        )
+
+        mock_chat_completion_create.return_value = completion
+        yield mock_chat_completion_create
+
+
 @pytest.fixture
 def openai_mock_chat_completion_chunk():
     """
@@ -101,3 +152,30 @@ def openai_mock_chat_completion_chunk():
             completion, cast_to=None, response=None, client=None
         )
         yield mock_chat_completion_create
+
+
+@pytest.fixture
+async def openai_mock_async_chat_completion_chunk():
+    """
+    Mock the OpenAI API completion chunk response and reuse it for async tests
+    """
+    with patch(
+        "openai.resources.chat.completions.AsyncCompletions.create", new_callable=AsyncMock
+    ) as mock_chat_completion_create:
+        completion = ChatCompletionChunk(
+            id="foo",
+            model="gpt-4",
+            object="chat.completion.chunk",
+            choices=[
+                chat_completion_chunk.Choice(
+                    finish_reason="stop",
+                    logprobs=None,
+                    index=0,
+                    delta=chat_completion_chunk.ChoiceDelta(content="Hello", role="assistant"),
+                )
+            ],
+            created=int(datetime.now().timestamp()),
+            usage={"prompt_tokens": 57, "completion_tokens": 40, "total_tokens": 97},
+        )
+        mock_chat_completion_create.return_value = OpenAIAsyncMockStream(completion)
+        yield mock_chat_completion_create