token usage summary

Author: SentienceDEV

predicate/agents/planner_executor_agent.py

@@ -53,6 +53,97 @@
 from .recovery import RecoveryCheckpoint, RecoveryState
 
 
+# ---------------------------------------------------------------------------
+# Token Usage Tracking
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TokenUsageTotals:
+    """Accumulated token counts for a single role or model."""
+
+    calls: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+    def add(self, resp: LLMResponse) -> None:
+        """Add token counts from an LLM response."""
+        self.calls += 1
+        pt = resp.prompt_tokens if isinstance(resp.prompt_tokens, int) else 0
+        ct = resp.completion_tokens if isinstance(resp.completion_tokens, int) else 0
+        tt = resp.total_tokens if isinstance(resp.total_tokens, int) else (pt + ct)
+        self.prompt_tokens += max(0, int(pt))
+        self.completion_tokens += max(0, int(ct))
+        self.total_tokens += max(0, int(tt))
+
+
+class _TokenUsageCollector:
+    """Collects token usage statistics by role (planner/executor) and model."""
+
+    def __init__(self) -> None:
+        self._by_role: dict[str, TokenUsageTotals] = {}
+        self._by_model: dict[str, TokenUsageTotals] = {}
+
+    def record(self, *, role: str, resp: LLMResponse) -> None:
+        """Record token usage from an LLM response."""
+        self._by_role.setdefault(role, TokenUsageTotals()).add(resp)
+        m = str(resp.model_name or "").strip() or "unknown"
+        self._by_model.setdefault(m, TokenUsageTotals()).add(resp)
+
+    def reset(self) -> None:
+        """Clear all recorded statistics."""
+        self._by_role.clear()
+        self._by_model.clear()
+
+    def summary(self) -> dict[str, Any]:
+        """
+        Get a summary of all token usage.
+
+        Returns:
+            Dictionary with:
+            - total: aggregate counts across all calls
+            - by_role: breakdown by role (planner, executor, replan)
+            - by_model: breakdown by model name
+        """
+        def _sum(items: dict[str, TokenUsageTotals]) -> TokenUsageTotals:
+            out = TokenUsageTotals()
+            for t in items.values():
+                out.calls += t.calls
+                out.prompt_tokens += t.prompt_tokens
+                out.completion_tokens += t.completion_tokens
+                out.total_tokens += t.total_tokens
+            return out
+
+        total = _sum(self._by_role)
+        return {
+            "total": {
+                "calls": total.calls,
+                "prompt_tokens": total.prompt_tokens,
+                "completion_tokens": total.completion_tokens,
+                "total_tokens": total.total_tokens,
+            },
+            "by_role": {
+                k: {
+                    "calls": v.calls,
+                    "prompt_tokens": v.prompt_tokens,
+                    "completion_tokens": v.completion_tokens,
+                    "total_tokens": v.total_tokens,
+                }
+                for k, v in self._by_role.items()
+            },
+            "by_model": {
+                k: {
+                    "calls": v.calls,
+                    "prompt_tokens": v.prompt_tokens,
+                    "completion_tokens": v.completion_tokens,
+                    "total_tokens": v.total_tokens,
+                }
+                for k, v in self._by_model.items()
+            },
+        }
+
+
 # ---------------------------------------------------------------------------
 # IntentHeuristics Protocol
 # ---------------------------------------------------------------------------
@@ -729,6 +820,7 @@ class RunOutcome:
     step_outcomes: list[StepOutcome] = field(default_factory=list)
     total_duration_ms: int = 0
     error: str | None = None
+    token_usage: dict[str, Any] | None = None  # Token usage summary from get_token_stats()
 
 
 # ---------------------------------------------------------------------------
@@ -1323,6 +1415,37 @@ def __init__(
         # Current automation task (for run-level context)
         self._current_task: AutomationTask | None = None
 
+        # Token usage tracking
+        self._token_collector = _TokenUsageCollector()
+
+    def get_token_stats(self) -> dict[str, Any]:
+        """
+        Get token usage statistics for the agent session.
+
+        Returns:
+            Dictionary with:
+            - total: aggregate counts (calls, prompt_tokens, completion_tokens, total_tokens)
+            - by_role: breakdown by role (planner, executor, replan, vision)
+            - by_model: breakdown by model name
+
+        Example:
+            >>> stats = agent.get_token_stats()
+            >>> print(f"Total tokens: {stats['total']['total_tokens']}")
+            >>> print(f"Planner tokens: {stats['by_role'].get('planner', {}).get('total_tokens', 0)}")
+        """
+        return self._token_collector.summary()
+
+    def reset_token_stats(self) -> None:
+        """Reset token usage statistics to zero."""
+        self._token_collector.reset()
+
+    def _record_token_usage(self, role: str, resp: LLMResponse) -> None:
+        """Record token usage from an LLM response."""
+        try:
+            self._token_collector.record(role=role, resp=resp)
+        except Exception:
+            pass  # Don't fail on token tracking errors
+
     def _format_context(self, snap: Snapshot, goal: str) -> str:
         """
         Format snapshot for LLM context.
@@ -2069,6 +2192,7 @@ async def plan(
                 temperature=self.config.planner_temperature,
                 max_new_tokens=max_tokens,
             )
+            self._record_token_usage("planner", resp)
             last_output = resp.content
 
             if self.config.verbose:
@@ -2169,6 +2293,7 @@ async def replan(
                 temperature=self.config.planner_temperature,
                 max_new_tokens=1024,
             )
+            self._record_token_usage("replan", resp)
             last_output = resp.content
 
             try:
@@ -2327,6 +2452,7 @@ async def _scroll_to_find_element(
                     temperature=self.config.executor_temperature,
                     max_new_tokens=self.config.executor_max_tokens,
                 )
+                self._record_token_usage("executor", resp)
                 parsed_action, parsed_args = self._parse_action(resp.content)
 
                 if parsed_action == "CLICK" and parsed_args:
@@ -2398,6 +2524,7 @@ async def _execute_optional_substeps(
                             temperature=self.config.executor_temperature,
                             max_new_tokens=self.config.executor_max_tokens,
                         )
+                        self._record_token_usage("executor", resp)
                         parsed_action, parsed_args = self._parse_action(resp.content)
                         if parsed_action == "CLICK" and parsed_args:
                             element_id = parsed_args[0]
@@ -2917,6 +3044,7 @@ async def _execute_step(
                         temperature=self.config.executor_temperature,
                         max_new_tokens=self.config.executor_max_tokens,
                     )
+                    self._record_token_usage("executor", resp)
                     llm_response = resp.content
 
                     if self.config.verbose:
@@ -3467,6 +3595,7 @@ async def run(
             step_outcomes=step_outcomes,
             total_duration_ms=int((time.time() - start_time) * 1000),
             error=error,
+            token_usage=self.get_token_stats(),
         )
 
         # Emit run end

tests/unit/test_planner_executor_agent.py

@@ -1351,3 +1351,197 @@ def test_modal_config_has_required_fields_for_drawer_dismissal(self) -> None:
         assert hasattr(config, "role_filter")
         assert hasattr(config, "max_attempts")
         assert hasattr(config, "min_new_elements")
+
+
+# ---------------------------------------------------------------------------
+# Test Token Usage Tracking
+# ---------------------------------------------------------------------------
+
+
+class TestTokenUsageTracking:
+    """Tests for token usage tracking in PlannerExecutorAgent."""
+
+    def test_token_usage_totals_add(self) -> None:
+        """TokenUsageTotals should accumulate tokens correctly."""
+        from predicate.agents.planner_executor_agent import TokenUsageTotals
+        from predicate.llm_provider import LLMResponse
+
+        totals = TokenUsageTotals()
+        assert totals.calls == 0
+        assert totals.prompt_tokens == 0
+        assert totals.completion_tokens == 0
+        assert totals.total_tokens == 0
+
+        # Add first response
+        resp1 = LLMResponse(
+            content="test",
+            prompt_tokens=100,
+            completion_tokens=50,
+            total_tokens=150,
+        )
+        totals.add(resp1)
+        assert totals.calls == 1
+        assert totals.prompt_tokens == 100
+        assert totals.completion_tokens == 50
+        assert totals.total_tokens == 150
+
+        # Add second response
+        resp2 = LLMResponse(
+            content="test2",
+            prompt_tokens=200,
+            completion_tokens=75,
+            total_tokens=275,
+        )
+        totals.add(resp2)
+        assert totals.calls == 2
+        assert totals.prompt_tokens == 300
+        assert totals.completion_tokens == 125
+        assert totals.total_tokens == 425
+
+    def test_token_usage_totals_handles_none_values(self) -> None:
+        """TokenUsageTotals should handle None token counts gracefully."""
+        from predicate.agents.planner_executor_agent import TokenUsageTotals
+        from predicate.llm_provider import LLMResponse
+
+        totals = TokenUsageTotals()
+        resp = LLMResponse(
+            content="test",
+            prompt_tokens=None,
+            completion_tokens=None,
+            total_tokens=None,
+        )
+        totals.add(resp)
+        assert totals.calls == 1
+        assert totals.prompt_tokens == 0
+        assert totals.completion_tokens == 0
+        assert totals.total_tokens == 0
+
+    def test_token_usage_collector_records_by_role(self) -> None:
+        """_TokenUsageCollector should track tokens by role."""
+        from predicate.agents.planner_executor_agent import _TokenUsageCollector
+        from predicate.llm_provider import LLMResponse
+
+        collector = _TokenUsageCollector()
+
+        resp_planner = LLMResponse(
+            content="plan",
+            prompt_tokens=500,
+            completion_tokens=200,
+            total_tokens=700,
+            model_name="gpt-4o",
+        )
+        collector.record(role="planner", resp=resp_planner)
+
+        resp_executor = LLMResponse(
+            content="action",
+            prompt_tokens=100,
+            completion_tokens=20,
+            total_tokens=120,
+            model_name="gpt-4o-mini",
+        )
+        collector.record(role="executor", resp=resp_executor)
+
+        summary = collector.summary()
+
+        # Check total
+        assert summary["total"]["calls"] == 2
+        assert summary["total"]["prompt_tokens"] == 600
+        assert summary["total"]["completion_tokens"] == 220
+        assert summary["total"]["total_tokens"] == 820
+
+        # Check by_role
+        assert "planner" in summary["by_role"]
+        assert summary["by_role"]["planner"]["calls"] == 1
+        assert summary["by_role"]["planner"]["total_tokens"] == 700
+
+        assert "executor" in summary["by_role"]
+        assert summary["by_role"]["executor"]["calls"] == 1
+        assert summary["by_role"]["executor"]["total_tokens"] == 120
+
+    def test_token_usage_collector_records_by_model(self) -> None:
+        """_TokenUsageCollector should track tokens by model name."""
+        from predicate.agents.planner_executor_agent import _TokenUsageCollector
+        from predicate.llm_provider import LLMResponse
+
+        collector = _TokenUsageCollector()
+
+        resp1 = LLMResponse(
+            content="test",
+            prompt_tokens=100,
+            completion_tokens=50,
+            total_tokens=150,
+            model_name="gpt-4o",
+        )
+        collector.record(role="planner", resp=resp1)
+
+        resp2 = LLMResponse(
+            content="test",
+            prompt_tokens=50,
+            completion_tokens=25,
+            total_tokens=75,
+            model_name="gpt-4o-mini",
+        )
+        collector.record(role="executor", resp=resp2)
+
+        summary = collector.summary()
+
+        # Check by_model
+        assert "gpt-4o" in summary["by_model"]
+        assert summary["by_model"]["gpt-4o"]["total_tokens"] == 150
+
+        assert "gpt-4o-mini" in summary["by_model"]
+        assert summary["by_model"]["gpt-4o-mini"]["total_tokens"] == 75
+
+    def test_token_usage_collector_reset(self) -> None:
+        """_TokenUsageCollector reset should clear all data."""
+        from predicate.agents.planner_executor_agent import _TokenUsageCollector
+        from predicate.llm_provider import LLMResponse
+
+        collector = _TokenUsageCollector()
+        resp = LLMResponse(
+            content="test",
+            prompt_tokens=100,
+            completion_tokens=50,
+            total_tokens=150,
+        )
+        collector.record(role="planner", resp=resp)
+        assert collector.summary()["total"]["calls"] == 1
+
+        collector.reset()
+        summary = collector.summary()
+        assert summary["total"]["calls"] == 0
+        assert summary["total"]["total_tokens"] == 0
+        assert summary["by_role"] == {}
+        assert summary["by_model"] == {}
+
+    def test_run_outcome_has_token_usage_field(self) -> None:
+        """RunOutcome should have token_usage field."""
+        from predicate.agents.planner_executor_agent import RunOutcome
+
+        outcome = RunOutcome(
+            run_id="test-run",
+            task="test task",
+            success=True,
+            steps_completed=3,
+            steps_total=3,
+            replans_used=0,
+        )
+        # Default should be None
+        assert outcome.token_usage is None
+
+        # Should accept token usage dict
+        outcome_with_tokens = RunOutcome(
+            run_id="test-run",
+            task="test task",
+            success=True,
+            steps_completed=3,
+            steps_total=3,
+            replans_used=0,
+            token_usage={
+                "total": {"calls": 5, "total_tokens": 1000},
+                "by_role": {"planner": {"calls": 2, "total_tokens": 700}},
+                "by_model": {"gpt-4o": {"calls": 2, "total_tokens": 700}},
+            },
+        )
+        assert outcome_with_tokens.token_usage is not None
+        assert outcome_with_tokens.token_usage["total"]["total_tokens"] == 1000