From e2210f9d6ae73cc7f558c65a1c1cf5928a7f6cfb Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Sat, 8 Aug 2026 12:47:12 +0000 Subject: [PATCH] fix(eval): persist plan-execute token usage for cost/tokens reporting PlanExecuteRunner already meters real token usage via _TokenMeter and writes it onto the OTel span, but never passed it to persist_trajectory(). StepResult has no token field, so _from_plan_execute() in metrics.py always built OpsMetrics with tokens_in=tokens_out=0, and _estimate_cost short-circuited to None. Every plan-execute evaluation report showed zero tokens and no cost, while every SDK-based runner reported correctly because their Trajectory dataclass carries per-turn token fields. Thread the meter's totals through persist_trajectory() as optional keyword arguments (omitted from the record when not given, so other runners' persisted shape is unchanged), and have metrics_from_trajectory() read them back into _from_plan_execute() for tokens_in/tokens_out and the cost estimate. Fixes #499 Signed-off-by: Amir Fathi --- src/agent/plan_execute/runner.py | 2 ++ src/agent/tests/test_runner.py | 32 +++++++++++++++++++++ src/evaluation/metrics.py | 18 +++++++++--- src/evaluation/tests/test_metrics.py | 23 +++++++++++++++ src/observability/persistence.py | 11 +++++++ src/observability/tests/test_persistence.py | 23 +++++++++++++++ 6 files changed, 105 insertions(+), 4 deletions(-) diff --git a/src/agent/plan_execute/runner.py b/src/agent/plan_execute/runner.py index 445f47b09..f630e9917 100644 --- a/src/agent/plan_execute/runner.py +++ b/src/agent/plan_execute/runner.py @@ -181,5 +181,7 @@ async def run(self, question: str) -> OrchestratorResult: question=question, answer=answer or "", trajectory=trajectory, + tokens_in=self._meter.input_tokens, + tokens_out=self._meter.output_tokens, ) return result diff --git a/src/agent/tests/test_runner.py b/src/agent/tests/test_runner.py index 65090a93f..2fb884c11 100644 --- a/src/agent/tests/test_runner.py +++ b/src/agent/tests/test_runner.py @@ -179,6 +179,38 @@ async def test_orchestrator_accumulates_token_usage_across_llm_calls(): assert runner._meter.output_tokens == 100 +@pytest.mark.anyio +async def test_orchestrator_persists_token_usage_alongside_trajectory( + monkeypatch, tmp_path +): + """The meter's totals must reach the persisted record, not just the span. + + Regression for the plan-execute runner's tracked usage never reaching + persist_trajectory(): metrics built from the persisted file (offline + evaluation) always saw tokens_in=tokens_out=0 for this runner. + """ + from observability import set_run_context + + monkeypatch.setenv("AGENT_TRAJECTORY_DIR", str(tmp_path)) + set_run_context(run_id="run-usage") + + llm = _UsageReportingLLM( + [ + (_TWO_STEP_PLAN, 100, 50), + (_STEP1_ARGS, 20, 5), + (_STEP2_ARGS, 30, 5), + (_FINAL_ANSWER, 200, 40), + ] + ) + runner = PlanExecuteRunner(llm) + with _patch_mcp()[0], _patch_mcp()[1]: + await runner.run("Q") + + record = json.loads((tmp_path / "run-usage.json").read_text()) + assert record["tokens_in"] == runner._meter.input_tokens == 350 + assert record["tokens_out"] == runner._meter.output_tokens == 100 + + @pytest.mark.anyio async def test_orchestrator_no_tool_returns_expected_output(sequential_llm): """A step with tool=none returns expected_output without any MCP or LLM call.""" diff --git a/src/evaluation/metrics.py b/src/evaluation/metrics.py index 9450c3100..c3ef2882e 100644 --- a/src/evaluation/metrics.py +++ b/src/evaluation/metrics.py @@ -31,7 +31,12 @@ def metrics_from_trajectory(record: PersistedTrajectory) -> OpsMetrics: if isinstance(traj, dict) and "turns" in traj: return _from_sdk_trajectory(traj, record.model) if isinstance(traj, list): - return _from_plan_execute(traj, record.model) + return _from_plan_execute( + traj, + record.model, + tokens_in=getattr(record, "tokens_in", None) or 0, + tokens_out=getattr(record, "tokens_out", None) or 0, + ) return OpsMetrics() @@ -128,10 +133,13 @@ def _usage_from_raw_events(events: list[Any]) -> tuple[int, int]: return input_tokens or sdk_input_tokens, output_tokens or sdk_output_tokens -def _from_plan_execute(steps: list[Any], model: str) -> OpsMetrics: +def _from_plan_execute( + steps: list[Any], model: str, tokens_in: int = 0, tokens_out: int = 0 +) -> OpsMetrics: # plan-execute persists ``list[StepResult]``; the dataclass exposes # ``server`` / ``tool`` / ``response`` fields but no per-step token - # counts, so we surface what is available and leave the rest at zero. + # counts, so the run-level totals the runner threaded onto the record + # (``tokens_in``/``tokens_out``) are the only source for these. tool_names = [ s.get("tool") for s in steps @@ -141,7 +149,9 @@ def _from_plan_execute(steps: list[Any], model: str) -> OpsMetrics: turn_count=len(steps), tool_call_count=len(tool_names), unique_tools=sorted(set(tool_names)), - est_cost_usd=_estimate_cost(model, 0, 0), + tokens_in=tokens_in, + tokens_out=tokens_out, + est_cost_usd=_estimate_cost(model, tokens_in, tokens_out), ) diff --git a/src/evaluation/tests/test_metrics.py b/src/evaluation/tests/test_metrics.py index d1bce4838..30b71fd33 100644 --- a/src/evaluation/tests/test_metrics.py +++ b/src/evaluation/tests/test_metrics.py @@ -103,6 +103,29 @@ def test_plan_execute_list_trajectory(self, make_persisted_record): assert m.turn_count == 3 assert m.tool_call_count == 3 assert m.unique_tools == ["assets", "sites"] + # No run-level tokens_in/tokens_out on the record (pre-fix persisted + # files, or a runner that never set them): stay at zero, no cost. + assert m.tokens_in == 0 + assert m.tokens_out == 0 + assert m.est_cost_usd is None + + def test_plan_execute_list_trajectory_reads_run_level_tokens( + self, make_persisted_record + ): + rec = PersistedTrajectory.from_raw( + make_persisted_record( + model="gpt-4o", + trajectory=[ + {"step_number": 1, "task": "t", "server": "iot", "tool": "sites", "response": "ok"}, + ], + tokens_in=1000, + tokens_out=500, + ) + ) + m = metrics_from_trajectory(rec) + assert m.tokens_in == 1000 + assert m.tokens_out == 500 + assert m.est_cost_usd == round((1000 * 2.5 + 500 * 10.0) / 1_000_000, 6) class TestAggregateOps: diff --git a/src/observability/persistence.py b/src/observability/persistence.py index a47825a1d..1e1e010ba 100644 --- a/src/observability/persistence.py +++ b/src/observability/persistence.py @@ -43,6 +43,8 @@ def persist_trajectory( question: str, answer: str, trajectory: Any, + tokens_in: int | None = None, + tokens_out: int | None = None, ) -> Path | None: """Write a per-run evaluation record when ``AGENT_TRAJECTORY_DIR`` is set. @@ -50,6 +52,11 @@ def persist_trajectory( :func:`agent_run_span`, so CLI-level wiring doesn't have to touch the runner's public signature. + ``tokens_in`` / ``tokens_out`` are for runners (like plan-execute) whose + trajectory shape has no per-turn token fields of its own, so the totals + have to be threaded through separately. Omitted when ``None`` so callers + that don't pass them keep the record's existing shape. + Returns the output path, or ``None`` when persistence is disabled. """ dir_env = os.environ.get(_TRAJECTORY_DIR_ENV) @@ -77,6 +84,10 @@ def persist_trajectory( "answer": answer, "trajectory": _serialize_trajectory(trajectory), } + if tokens_in is not None: + record["tokens_in"] = tokens_in + if tokens_out is not None: + record["tokens_out"] = tokens_out try: out_path.write_text(json.dumps(record, indent=2, default=str), encoding="utf-8") except OSError: diff --git a/src/observability/tests/test_persistence.py b/src/observability/tests/test_persistence.py index 555ab465e..5e040ba36 100644 --- a/src/observability/tests/test_persistence.py +++ b/src/observability/tests/test_persistence.py @@ -116,6 +116,29 @@ class _FakeStep: assert record["trajectory"] == [ {"step_number": 1, "task": "do thing", "success": True} ] + assert "tokens_in" not in record + assert "tokens_out" not in record + + +def test_persist_includes_run_level_tokens_when_given(monkeypatch, tmp_path: Path): + """plan-execute's trajectory has no per-turn token fields, so the runner + passes the meter's totals separately; they must land on the record.""" + monkeypatch.setenv("AGENT_TRAJECTORY_DIR", str(tmp_path)) + set_run_context(run_id="r4") + + out = persist_trajectory( + runner_name="plan-execute", + model="watsonx/model", + question="q", + answer="a", + trajectory=[], + tokens_in=42, + tokens_out=17, + ) + + record = json.loads(out.read_text()) + assert record["tokens_in"] == 42 + assert record["tokens_out"] == 17 def test_persist_skips_when_no_run_id(monkeypatch, tmp_path: Path, caplog):