models - openai - async client

Author: pgrayy

src/strands/models/openai.py

@@ -61,7 +61,7 @@ def __init__(self, client_args: Optional[dict[str, Any]] = None, **model_config:
         logger.debug("config=<%s> | initializing", self.config)
 
         client_args = client_args or {}
-        self.client = openai.OpenAI(**client_args)
+        self.client = openai.AsyncOpenAI(**client_args)
 
     @override
     def update_config(self, **model_config: Unpack[OpenAIConfig]) -> None:  # type: ignore[override]
@@ -91,14 +91,14 @@ async def stream(self, request: dict[str, Any]) -> AsyncGenerator[dict[str, Any]
         Returns:
             An iterable of response events from the OpenAI model.
         """
-        response = self.client.chat.completions.create(**request)
+        response = await self.client.chat.completions.create(**request)
 
         yield {"chunk_type": "message_start"}
         yield {"chunk_type": "content_start", "data_type": "text"}
 
         tool_calls: dict[int, list[Any]] = {}
 
-        for event in response:
+        async for event in response:
             # Defensive: skip events with empty or missing choices
             if not getattr(event, "choices", None):
                 continue
@@ -133,7 +133,7 @@ async def stream(self, request: dict[str, Any]) -> AsyncGenerator[dict[str, Any]
         yield {"chunk_type": "message_stop", "data": choice.finish_reason}
 
         # Skip remaining events as we don't have use for anything except the final usage payload
-        for event in response:
+        async for event in response:
             _ = event
 
         yield {"chunk_type": "metadata", "data": event.usage}
@@ -151,7 +151,7 @@ async def structured_output(
         Yields:
             Model events with the last being the structured output.
         """
-        response: ParsedChatCompletion = self.client.beta.chat.completions.parse(  # type: ignore
+        response: ParsedChatCompletion = await self.client.beta.chat.completions.parse(  # type: ignore
             model=self.get_config()["model_id"],
             messages=super().format_request(prompt)["messages"],
             response_format=output_model,

tests-integ/test_model_openai.py

@@ -12,7 +12,7 @@
 from strands.models.openai import OpenAIModel
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def model():
     return OpenAIModel(
         model_id="gpt-4o",
@@ -22,7 +22,7 @@ def model():
     )
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def tools():
     @strands.tool
     def tool_time() -> str:
@@ -35,36 +35,65 @@ def tool_weather() -> str:
     return [tool_time, tool_weather]
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def agent(model, tools):
     return Agent(model=model, tools=tools)
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def test_image_path(request):
     return request.config.rootpath / "tests-integ" / "test_image.png"
 
 
-def test_agent(agent):
+@pytest.fixture(scope="module")
+def weather():
+    class Weather(BaseModel):
+        """Extracts the time and weather from the user's message with the exact strings."""
+
+        time: str
+        weather: str
+
+    return Weather(time="12:00", weather="sunny")
+
+
+def test_agent_invoke(agent):
     result = agent("What is the time and weather in New York?")
     text = result.message["content"][0]["text"].lower()
 
     assert all(string in text for string in ["12:00", "sunny"])
 
 
-def test_structured_output(model):
-    class Weather(BaseModel):
-        """Extracts the time and weather from the user's message with the exact strings."""
+@pytest.mark.asyncio
+async def test_agent_invoke_async(agent):
+    result = await agent.invoke_async("What is the time and weather in New York?")
+    text = result.message["content"][0]["text"].lower()
 
-        time: str
-        weather: str
+    assert all(string in text for string in ["12:00", "sunny"])
+
+
+@pytest.mark.asyncio
+async def test_agent_stream_async(agent):
+    stream = agent.stream_async("What is the time and weather in New York?")
+    async for event in stream:
+        _ = event
+
+    result = event["result"]
+    text = result.message["content"][0]["text"].lower()
+
+    assert all(string in text for string in ["12:00", "sunny"])
+
+
+def test_structured_output(agent, weather):
+    tru_weather = agent.structured_output(type(weather), "The time is 12:00 and the weather is sunny")
+    exp_weather = weather
+    assert tru_weather == exp_weather
 
-    agent = Agent(model=model)
 
-    result = agent.structured_output(Weather, "The time is 12:00 and the weather is sunny")
-    assert isinstance(result, Weather)
-    assert result.time == "12:00"
-    assert result.weather == "sunny"
+@pytest.mark.asyncio
+async def test_structured_output_async(agent, weather):
+    tru_weather = await agent.structured_output_async(type(weather), "The time is 12:00 and the weather is sunny")
+    exp_weather = weather
+    assert tru_weather == exp_weather
 
 
 def test_tool_returning_images(model, test_image_path):

tests/strands/models/test_openai.py

@@ -9,7 +9,7 @@
 
 @pytest.fixture
 def openai_client_cls():
-    with unittest.mock.patch.object(strands.models.openai.openai, "OpenAI") as mock_client_cls:
+    with unittest.mock.patch.object(strands.models.openai.openai, "AsyncOpenAI") as mock_client_cls:
         yield mock_client_cls
 
 
@@ -70,7 +70,7 @@ def test_update_config(model, model_id):
 
 
 @pytest.mark.asyncio
-async def test_stream(openai_client, model, alist):
+async def test_stream(openai_client, model, agenerator, alist):
     mock_tool_call_1_part_1 = unittest.mock.Mock(index=0)
     mock_tool_call_2_part_1 = unittest.mock.Mock(index=1)
     mock_delta_1 = unittest.mock.Mock(
@@ -102,8 +102,8 @@ async def test_stream(openai_client, model, alist):
     mock_event_5 = unittest.mock.Mock(choices=[unittest.mock.Mock(finish_reason="tool_calls", delta=mock_delta_5)])
     mock_event_6 = unittest.mock.Mock()
 
-    openai_client.chat.completions.create.return_value = iter(
-        [mock_event_1, mock_event_2, mock_event_3, mock_event_4, mock_event_5, mock_event_6]
+    openai_client.chat.completions.create = unittest.mock.AsyncMock(
+        return_value=agenerator([mock_event_1, mock_event_2, mock_event_3, mock_event_4, mock_event_5, mock_event_6])
     )
 
     request = {"model": "m1", "messages": [{"role": "user", "content": [{"type": "text", "text": "calculate 2+2"}]}]}
@@ -133,7 +133,7 @@ async def test_stream(openai_client, model, alist):
 
 
 @pytest.mark.asyncio
-async def test_stream_empty(openai_client, model, alist):
+async def test_stream_empty(openai_client, model, agenerator, alist):
     mock_delta = unittest.mock.Mock(content=None, tool_calls=None, reasoning_content=None)
     mock_usage = unittest.mock.Mock(prompt_tokens=0, completion_tokens=0, total_tokens=0)
 
@@ -142,7 +142,9 @@ async def test_stream_empty(openai_client, model, alist):
     mock_event_3 = unittest.mock.Mock()
     mock_event_4 = unittest.mock.Mock(usage=mock_usage)
 
-    openai_client.chat.completions.create.return_value = iter([mock_event_1, mock_event_2, mock_event_3, mock_event_4])
+    openai_client.chat.completions.create = unittest.mock.AsyncMock(
+        return_value=agenerator([mock_event_1, mock_event_2, mock_event_3, mock_event_4]),
+    )
 
     request = {"model": "m1", "messages": [{"role": "user", "content": []}]}
     response = model.stream(request)
@@ -161,7 +163,7 @@ async def test_stream_empty(openai_client, model, alist):
 
 
 @pytest.mark.asyncio
-async def test_stream_with_empty_choices(openai_client, model, alist):
+async def test_stream_with_empty_choices(openai_client, model, agenerator, alist):
     mock_delta = unittest.mock.Mock(content="content", tool_calls=None, reasoning_content=None)
     mock_usage = unittest.mock.Mock(prompt_tokens=10, completion_tokens=20, total_tokens=30)
 
@@ -180,8 +182,8 @@ async def test_stream_with_empty_choices(openai_client, model, alist):
     # Final event with usage info
     mock_event_5 = unittest.mock.Mock(usage=mock_usage)
 
-    openai_client.chat.completions.create.return_value = iter(
-        [mock_event_1, mock_event_2, mock_event_3, mock_event_4, mock_event_5]
+    openai_client.chat.completions.create = unittest.mock.AsyncMock(
+        return_value=agenerator([mock_event_1, mock_event_2, mock_event_3, mock_event_4, mock_event_5])
     )
 
     request = {"model": "m1", "messages": [{"role": "user", "content": ["test"]}]}
@@ -212,7 +214,7 @@ async def test_structured_output(openai_client, model, test_output_model_cls, al
     mock_response = unittest.mock.Mock()
     mock_response.choices = [mock_choice]
 
-    openai_client.beta.chat.completions.parse.return_value = mock_response
+    openai_client.beta.chat.completions.parse = unittest.mock.AsyncMock(return_value=mock_response)
 
     stream = model.structured_output(test_output_model_cls, messages)
     events = await alist(stream)