diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index ebcd289f..0a27b079 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/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() diff --git a/atomic-agents/tests/agents/test_atomic_agent.py b/atomic-agents/tests/agents/test_atomic_agent.py index 3e2782ee..ab598f3a 100644 --- a/atomic-agents/tests/agents/test_atomic_agent.py +++ b/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