feat: add per-request visualizer token toggle

Co-authored-by: openhands <[email protected]>

Author: enyst

openhands/sdk/conversation/visualizer.py

@@ -1,3 +1,4 @@
+import os
 import re
 from typing import TYPE_CHECKING
 
@@ -16,6 +17,7 @@
 )
 from openhands.sdk.event.base import Event
 from openhands.sdk.event.condenser import Condensation
+from openhands.sdk.llm.utils.metrics import TokenUsage
 
 
 if TYPE_CHECKING:
@@ -47,6 +49,18 @@
     r"\*(.*?)\*": "italic",
 }
 
+
+_PER_REQUEST_TOKENS_ENV = "OPENHANDS_VISUALIZER_PER_REQUEST_TOKENS"
+_TRUTHY_ENV_VALUES = {"1", "true", "yes", "on"}
+
+
+def _env_flag(name: str) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return False
+    return value.strip().lower() in _TRUTHY_ENV_VALUES
+
+
 _PANEL_PADDING = (1, 1)
 
 
@@ -61,6 +75,7 @@ def __init__(
         highlight_regex: dict[str, str] | None = None,
         skip_user_messages: bool = False,
         conversation_stats: "ConversationStats | None" = None,
+        per_request_token_stats: bool | None = None,
     ):
         """Initialize the visualizer.
 
@@ -72,11 +87,20 @@ def __init__(
             skip_user_messages: If True, skip displaying user messages. Useful for
                                 scenarios where user input is not relevant to show.
             conversation_stats: ConversationStats object to display metrics information.
+            per_request_token_stats: Whether to display per-request token usage instead
+                                     of accumulated totals. When None, uses the
+                                     OPENHANDS_VISUALIZER_PER_REQUEST_TOKENS
+                                     environment flag.
         """
         self._console = Console()
         self._skip_user_messages = skip_user_messages
         self._highlight_patterns: dict[str, str] = highlight_regex or {}
         self._conversation_stats = conversation_stats
+        if per_request_token_stats is None:
+            per_request_token_stats = _env_flag(_PER_REQUEST_TOKENS_ENV)
+        self._use_per_request_token_stats = per_request_token_stats
+        self._last_accumulated_usage: TokenUsage | None = None
+        self._last_request_usage: TokenUsage | None = None
 
     def on_event(self, event: Event) -> None:
         """Main event handler that displays events with Rich formatting."""
@@ -227,6 +251,78 @@ def _create_event_panel(self, event: Event) -> Panel | None:
                 expand=True,
             )
 
+    @staticmethod
+    def _clone_usage(usage: TokenUsage) -> TokenUsage:
+        return TokenUsage.model_validate(usage.model_dump())
+
+    def _get_per_request_usage(self, usage: TokenUsage) -> TokenUsage | None:
+        current = self._clone_usage(usage)
+
+        if self._last_accumulated_usage is None:
+            self._last_accumulated_usage = current
+            self._last_request_usage = self._clone_usage(current)
+            return self._last_request_usage
+
+        prev = self._last_accumulated_usage
+        tracked_fields = (
+            "prompt_tokens",
+            "completion_tokens",
+            "cache_read_tokens",
+            "cache_write_tokens",
+            "reasoning_tokens",
+        )
+
+        decreased = any(
+            getattr(current, field) < getattr(prev, field) for field in tracked_fields
+        )
+        if decreased:
+            self._last_accumulated_usage = current
+            self._last_request_usage = self._clone_usage(current)
+            return self._last_request_usage
+
+        changed = any(
+            getattr(current, field) != getattr(prev, field) for field in tracked_fields
+        )
+        if changed:
+            prompt_delta = max(current.prompt_tokens - prev.prompt_tokens, 0)
+            completion_delta = max(
+                current.completion_tokens - prev.completion_tokens, 0
+            )
+            cache_read_delta = max(
+                current.cache_read_tokens - prev.cache_read_tokens, 0
+            )
+            cache_write_delta = max(
+                current.cache_write_tokens - prev.cache_write_tokens, 0
+            )
+            reasoning_delta = max(current.reasoning_tokens - prev.reasoning_tokens, 0)
+            self._last_request_usage = TokenUsage(
+                model=current.model,
+                prompt_tokens=prompt_delta,
+                completion_tokens=completion_delta,
+                cache_read_tokens=cache_read_delta,
+                cache_write_tokens=cache_write_delta,
+                reasoning_tokens=reasoning_delta,
+                context_window=current.context_window,
+                per_turn_token=prompt_delta + completion_delta,
+                response_id=current.response_id or prev.response_id,
+            )
+
+        if self._last_request_usage is None:
+            self._last_request_usage = TokenUsage(
+                model=current.model,
+                prompt_tokens=0,
+                completion_tokens=0,
+                cache_read_tokens=0,
+                cache_write_tokens=0,
+                reasoning_tokens=0,
+                context_window=current.context_window,
+                per_turn_token=0,
+                response_id=current.response_id,
+            )
+
+        self._last_accumulated_usage = current
+        return self._last_request_usage
+
     def _format_metrics_subtitle(self) -> str | None:
         """Format LLM metrics as a visually appealing subtitle string with icons,
         colors, and k/m abbreviations using conversation stats."""
@@ -240,7 +336,12 @@ def _format_metrics_subtitle(self) -> str | None:
         usage = combined_metrics.accumulated_token_usage
         cost = combined_metrics.accumulated_cost or 0.0
 
-        # helper: 1234 -> "1.2K", 1200000 -> "1.2M"
+        display_usage = usage
+        if self._use_per_request_token_stats:
+            per_request_usage = self._get_per_request_usage(usage)
+            if per_request_usage is not None:
+                display_usage = per_request_usage
+
         def abbr(n: int | float) -> str:
             n = int(n or 0)
             if n >= 1_000_000_000:
@@ -253,28 +354,29 @@ def abbr(n: int | float) -> str:
                 return str(n)
             return s.replace(".0", "")
 
-        input_tokens = abbr(usage.prompt_tokens or 0)
-        output_tokens = abbr(usage.completion_tokens or 0)
-
-        # Cache hit rate (prompt + cache)
-        prompt = usage.prompt_tokens or 0
-        cache_read = usage.cache_read_tokens or 0
-        cache_rate = f"{(cache_read / prompt * 100):.2f}%" if prompt > 0 else "N/A"
-        reasoning_tokens = usage.reasoning_tokens or 0
+        prompt_value = display_usage.prompt_tokens or 0
+        output_value = display_usage.completion_tokens or 0
+        cache_prompt = prompt_value
+        cache_read = display_usage.cache_read_tokens or 0
+        cache_rate = (
+            f"{(cache_read / cache_prompt * 100):.2f}%" if cache_prompt > 0 else "N/A"
+        )
+        reasoning_tokens = display_usage.reasoning_tokens or 0
 
-        # Cost
         cost_str = f"{cost:.4f}" if cost > 0 else "$0.00"
 
-        # Build with fixed color scheme
         parts: list[str] = []
-        parts.append(f"[cyan]↑ input {input_tokens}[/cyan]")
+        parts.append(f"[cyan]↑ input {abbr(prompt_value)}[/cyan]")
         parts.append(f"[magenta]cache hit {cache_rate}[/magenta]")
         if reasoning_tokens > 0:
             parts.append(f"[yellow] reasoning {abbr(reasoning_tokens)}[/yellow]")
-        parts.append(f"[blue]↓ output {output_tokens}[/blue]")
+        parts.append(f"[blue]↓ output {abbr(output_value)}[/blue]")
         parts.append(f"[green]$ {cost_str}[/green]")
 
-        return "Tokens: " + " • ".join(parts)
+        label = (
+            "Tokens (last request)" if self._use_per_request_token_stats else "Tokens"
+        )
+        return f"{label}: " + " • ".join(parts)
 
 
 def create_default_visualizer(

tests/sdk/conversation/test_visualizer.py

@@ -340,6 +340,54 @@ def test_metrics_formatting():
     assert "0.0234" in subtitle  # Cost
 
 
+def test_metrics_formatting_per_request_toggle(monkeypatch):
+    """Per-request metrics should be shown when the env flag is enabled."""
+    from openhands.sdk.conversation.conversation_stats import ConversationStats
+    from openhands.sdk.llm.utils.metrics import Metrics
+
+    monkeypatch.setenv("OPENHANDS_VISUALIZER_PER_REQUEST_TOKENS", "true")
+
+    conversation_stats = ConversationStats()
+    metrics = Metrics(model_name="test-model")
+    metrics.add_cost(0.1)
+    metrics.add_token_usage(
+        prompt_tokens=150,
+        completion_tokens=50,
+        cache_read_tokens=30,
+        cache_write_tokens=0,
+        reasoning_tokens=20,
+        context_window=8000,
+        response_id="first",
+    )
+    conversation_stats.service_to_metrics["test_service"] = metrics
+
+    visualizer = ConversationVisualizer(conversation_stats=conversation_stats)
+
+    subtitle = visualizer._format_metrics_subtitle()
+    assert subtitle is not None
+    assert "Tokens (last request)" in subtitle
+    assert "input 150" in subtitle
+    assert "reasoning 20" in subtitle
+    assert "0.1000" in subtitle
+
+    metrics.add_cost(0.05)
+    metrics.add_token_usage(
+        prompt_tokens=200,
+        completion_tokens=75,
+        cache_read_tokens=25,
+        cache_write_tokens=0,
+        reasoning_tokens=0,
+        context_window=8000,
+        response_id="second",
+    )
+
+    subtitle_second = visualizer._format_metrics_subtitle()
+    assert subtitle_second is not None
+    assert "input 200" in subtitle_second
+    assert "output 75" in subtitle_second
+    assert "0.1500" in subtitle_second
+
+
 def test_event_base_fallback_visualize():
     """Test that Event provides fallback visualization."""
     from openhands.sdk.event.base import Event