diff --git a/workflowai/core/domain/run.py b/workflowai/core/domain/run.py index 0b4d522..62f15af 100644 --- a/workflowai/core/domain/run.py +++ b/workflowai/core/domain/run.py @@ -110,7 +110,8 @@ def format_output(self) -> str: # Format the output string output: list[str] = [] # In case of partial validation, it is possible that the output is an empty model - if dumped_output := self.output.model_dump(): + if dumped_output := self.output.model_dump(mode="json"): + # Use model_dump_json which handles datetime serialization correctly output += [ "\nOutput:", "=" * 50, diff --git a/workflowai/core/domain/run_test.py b/workflowai/core/domain/run_test.py index 0a3b413..e85d408 100644 --- a/workflowai/core/domain/run_test.py +++ b/workflowai/core/domain/run_test.py @@ -1,3 +1,4 @@ +from datetime import datetime, timezone from unittest.mock import Mock, patch import pytest @@ -17,6 +18,11 @@ class _TestOutput(BaseModel): message: str +class _TestOutputWithDatetime(BaseModel): + timestamp: datetime + message: str + + @pytest.fixture def mock_agent() -> Mock: mock = Mock(spec=_AgentBase) @@ -169,6 +175,29 @@ def test_format_output_tool_call_requests(self): URL: https://workflowai.hello/_/agents/agent-id/runs/run-id""" ) + def test_format_output_with_datetime(self): + """Test that datetimes in the output model are correctly serialized to ISO strings.""" + test_dt = datetime(2024, 1, 1, 12, 30, 0, tzinfo=timezone.utc) + run = Run[_TestOutputWithDatetime]( + id="run-dt-id", + agent_id="agent-dt-id", + schema_id=2, + output=_TestOutputWithDatetime(timestamp=test_dt, message="datetime test"), + duration_seconds=0.5, + cost_usd=0.0001, + ) + + expected_json_part = '{\n "timestamp": "2024-01-01T12:30:00Z",\n "message": "datetime test"\n}' + expected = f"""\nOutput: +================================================== +{expected_json_part} +================================================== +Cost: $ 0.00010 +Latency: 0.50s +URL: https://workflowai.hello/_/agents/agent-dt-id/runs/run-dt-id""" + + assert run.format_output() == expected + class TestRunURL: # The @patch decorator from unittest.mock temporarily replaces the value of an attribute