Merge branch 'main' into guillaume/new-missing-token-user-message

Author: yannbu

workflowai/core/client/_models.py

@@ -167,6 +167,7 @@ class CreateAgentResponse(BaseModel):
     id: str
     schema_id: int
     uid: int = 0
+    tenant_uid: int = 0
 
 
 class ModelMetadata(BaseModel):

workflowai/core/client/_types.py

@@ -57,6 +57,9 @@ class AgentInterface(_BaseObject, Generic[AgentInputContra, AgentOutput], Protoc
     @property
     def agent_uid(self) -> int: ...
 
+    @property
+    def tenant_uid(self) -> int: ...
+
     async def run(
         self,
         agent_input: AgentInputContra,

workflowai/core/client/agent.py

@@ -270,6 +270,7 @@ async def register(self):
         )
         self.schema_id = res.schema_id
         self.agent_uid = res.uid
+        self.tenant_uid = res.tenant_uid
         return res.schema_id
 
     @classmethod

workflowai/core/client/agent_test.py

@@ -232,6 +232,7 @@ async def test_auto_register(self, httpx_mock: HTTPXMock, agent_no_schema: Agent
                 "id": "123",
                 "schema_id": 2,
                 "uid": 123,
+                "tenant_uid": 1234,
             },
         )
         run_response = fixtures_json("task_run.json")
@@ -243,7 +244,7 @@ async def test_auto_register(self, httpx_mock: HTTPXMock, agent_no_schema: Agent
         out = await agent_no_schema.run(HelloTaskInput(name="Alice"))
         assert out.id == "8f635b73-f403-47ee-bff9-18320616c6cc"
         assert agent_no_schema.agent_uid == 123
-
+        assert agent_no_schema.tenant_uid == 1234
         run_response["id"] = "8f635b73-f403-47ee-bff9-18320616c6cc"
         # Try and run again
         httpx_mock.add_response(

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,

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