fix(db): record trajectory execution time

tests/db/test_trajectory.py

@@ -1,11 +1,21 @@
-# from utu.agents.common import TaskRecorder
-from utu.agents import SimpleAgent
-from utu.db import DBService, TrajectoryModel
+from types import SimpleNamespace
 
+from utu.db import TrajectoryModel
 
-async def test_traj_model():
-    """Test TrajectoryModel. The recorded trajectory should be saved to db and can be visualized."""
-    agent = SimpleAgent(config="simple/base")
-    task_recorder = await agent.run("hello")
-    trajectory = TrajectoryModel.from_task_recorder(task_recorder)
-    DBService.add(trajectory)
+
+def test_from_task_recorder_records_elapsed_time(monkeypatch):
+    monkeypatch.setattr("utu.db.trajectory_model.time.time", lambda: 15.0)
+    recorder = SimpleNamespace(
+        trace_id="trace-1",
+        task="hello",
+        input="",
+        final_output="done",
+        trajectories=[],
+        started_at=10.0,
+    )
+
+    trajectory = TrajectoryModel.from_task_recorder(recorder)
+
+    assert trajectory.time_cost == 5.0
+    assert trajectory.d_input == "hello"
+    assert trajectory.d_output == "done"

utu/agents/common.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 import traceback
 from collections.abc import AsyncIterator
 from dataclasses import asdict, dataclass, field
@@ -103,6 +104,7 @@ class TaskRecorder(DataClassWithStreamEvents):
 
     # additional infos
     additional_infos: dict = field(default_factory=dict)
+    started_at: float = field(default_factory=time.time)
 
     def to_input_list(self) -> list[TResponseInputItem]:
         return self.get_run_result().to_input_list()

utu/db/trajectory_model.py

@@ -1,4 +1,5 @@
 import json
+import time
 from typing import TYPE_CHECKING
 
 from sqlmodel import Field, SQLModel
@@ -23,11 +24,16 @@ class TrajectoryModel(SQLModel, table=True):
     def from_task_recorder(cls, task_recorder: "TaskRecorder") -> "TrajectoryModel":
         # if isinstance(task_recorder, TaskRecorder):
         d_input = getattr(task_recorder, "task", "") or getattr(task_recorder, "input", "")
+        started_at = getattr(task_recorder, "started_at", None)
+        time_cost = None
+        if isinstance(started_at, (int, float)):
+            time_cost = max(0.0, time.time() - started_at)
+
         return cls(
             trace_id=task_recorder.trace_id,
             trace_url="",
             d_input=d_input,
             d_output=task_recorder.final_output,
             trajectories=json.dumps(task_recorder.trajectories, ensure_ascii=False),
-            time_cost=-1,
+            time_cost=time_cost,
         )