fix: keep streaming working for output schemas with field aliases

atomic-agents/atomic_agents/agents/atomic_agent.py

@@ -637,7 +637,9 @@ def run_stream(self, user_input: Optional[InputSchema] = None) -> Generator[Outp
             yield partial_response
 
         if last_response:
-            full_response_content = self.output_schema(**last_response.model_dump())
+            # model_dump() keys the data by field name, so accept field names
+            # for output schemas whose fields have an alias.
+            full_response_content = self.output_schema.model_validate(last_response.model_dump(), by_name=True)
             self.history.add_message(self.assistant_role, full_response_content)
             self._prepare_messages()
             return full_response_content
@@ -707,7 +709,8 @@ async def run_async_stream(self, user_input: Optional[InputSchema] = None) -> As
             yield partial_response
 
         if last_response:
-            full_response_content = self.output_schema(**last_response.model_dump())
+            # See run_stream(): accept field names for aliased output schemas.
+            full_response_content = self.output_schema.model_validate(last_response.model_dump(), by_name=True)
             self.history.add_message(self.assistant_role, full_response_content)
             self._prepare_messages()
 

atomic-agents/tests/agents/test_atomic_agent.py

@@ -512,6 +512,27 @@ def test_run_stream(mock_instructor, mock_history):
     mock_history.add_message.assert_has_calls([call("user", mock_input), call("assistant", mock_output)])
 
 
+class AliasedOutputSchema(BaseIOSchema):
+    """Output schema whose field is set through an alias."""
+
+    chat_message: str = Field(..., alias="chatMessage", description="The reply text")
+
+
+def test_run_stream_with_aliased_output_schema(mock_instructor, mock_history):
+    streamed = AliasedOutputSchema(chatMessage="Test output")
+    mock_instructor.chat.completions.create_partial.return_value = iter([streamed])
+    config = AgentConfig(client=mock_instructor, model="gpt-5-mini", history=mock_history)
+    agent = AtomicAgent[BasicChatInputSchema, AliasedOutputSchema](config)
+
+    stream = agent.run_stream(BasicChatInputSchema(chat_message="Test input"))
+    with pytest.raises(StopIteration) as stop:
+        while True:
+            next(stream)
+
+    assert stop.value.value == streamed
+    mock_history.add_message.assert_called_with("assistant", streamed)
+
+
 @pytest.mark.asyncio
 async def test_run_async(agent_async, mock_history):
     # Create a mock input
@@ -551,6 +572,23 @@ async def test_run_async_stream(agent_async, mock_history):
     mock_history.add_message.assert_any_call("assistant", full_response_content)
 
 
+@pytest.mark.asyncio
+async def test_run_async_stream_with_aliased_output_schema(mock_instructor_async, mock_history):
+    streamed = AliasedOutputSchema(chatMessage="Test output")
+
+    async def mock_create_partial(*args, **kwargs):
+        yield streamed
+
+    mock_instructor_async.chat.completions.create_partial = mock_create_partial
+    config = AgentConfig(client=mock_instructor_async, model="gpt-5-mini", history=mock_history)
+    agent = AtomicAgent[BasicChatInputSchema, AliasedOutputSchema](config)
+
+    responses = [response async for response in agent.run_async_stream(BasicChatInputSchema(chat_message="Test input"))]
+
+    assert responses == [streamed]
+    mock_history.add_message.assert_called_with("assistant", streamed)
+
+
 def test_model_from_chunks_patched():
     class TestPartialModel(PartialBase):
         @classmethod