chore: refactor usage tracking to integrate via callback

Author: philipph-askui

src/askui/agent_base.py

@@ -23,6 +23,7 @@
     LocateSettings,
 )
 from askui.models.shared.tools import Tool, ToolCollection
+from askui.models.shared.usage_tracking_callback import UsageTrackingCallback
 from askui.prompts.act_prompts import CACHE_USE_PROMPT, create_default_prompt
 from askui.tools.agent_os import AgentOs
 from askui.tools.android.agent_os import AndroidAgentOs
@@ -73,13 +74,15 @@ def __init__(
 
         # Create conversation with speakers and model providers
         speakers = Speakers()
+        _callbacks = list(callbacks or [])
+        _callbacks.append(UsageTrackingCallback(reporter=self._reporter))
         self._conversation = Conversation(
             speakers=speakers,
             vlm_provider=self._vlm_provider,
             image_qa_provider=self._image_qa_provider,
             detection_provider=self._detection_provider,
             reporter=self._reporter,
-            callbacks=callbacks,
+            callbacks=_callbacks,
         )
 
         # Provider-based tools

src/askui/models/shared/conversation.py

@@ -9,10 +9,7 @@
 from askui.model_providers.detection_provider import DetectionProvider
 from askui.model_providers.image_qa_provider import ImageQAProvider
 from askui.model_providers.vlm_provider import VlmProvider
-from askui.models.shared.agent_message_param import (
-    MessageParam,
-    UsageParam,
-)
+from askui.models.shared.agent_message_param import MessageParam
 from askui.models.shared.settings import ActSettings
 from askui.models.shared.tools import ToolCollection
 from askui.models.shared.truncation_strategies import (
@@ -84,7 +81,6 @@ def __init__(
         # Speakers and current state
         self.speakers = speakers
         self.current_speaker = speakers[speakers.default_speaker]
-        self.accumulated_usage = UsageParam()
 
         # Model providers - accessible by speakers via conversation instance
         self.vlm_provider = vlm_provider
@@ -184,7 +180,6 @@ def _setup_control_loop(
         reporters: list[Reporter] | None = None,
     ) -> None:
         # Reset state
-        self.accumulated_usage = UsageParam()
         self._executed_from_cache = False
         self.speakers.reset_state()
 
@@ -221,9 +216,6 @@ def _conclude_control_loop(self) -> None:
         if self.cache_manager is not None and not self._executed_from_cache:
             self.cache_manager.finish_recording(self.get_messages())
 
-        # Report final usage
-        self._reporter.add_usage_summary(self.accumulated_usage.model_dump())
-
     def _setup_speaker_handoff(self) -> None:
         """Set up speaker handoff infrastructure.
 
@@ -316,10 +308,6 @@ def _execute_step(self) -> bool:
         status_continue = self._handle_result_status(result)
         continue_loop = continue_loop or status_continue
 
-        # 5. Collect Statistics
-        if result.usage:
-            self._accumulate_usage(result.usage)
-
         self._on_step_end(self._step_index, result)
         self._step_index += 1
 
@@ -450,34 +438,3 @@ def get_truncation_strategy(self) -> TruncationStrategy | None:
             Current truncation strategy or None if not initialized
         """
         return self._truncation_strategy
-
-    def _accumulate_usage(self, step_usage: UsageParam) -> None:
-        """Accumulate token usage statistics.
-
-        Args:
-            step_usage: Usage from a single step
-        """
-        self.accumulated_usage.input_tokens = (
-            self.accumulated_usage.input_tokens or 0
-        ) + (step_usage.input_tokens or 0)
-        self.accumulated_usage.output_tokens = (
-            self.accumulated_usage.output_tokens or 0
-        ) + (step_usage.output_tokens or 0)
-        self.accumulated_usage.cache_creation_input_tokens = (
-            self.accumulated_usage.cache_creation_input_tokens or 0
-        ) + (step_usage.cache_creation_input_tokens or 0)
-        self.accumulated_usage.cache_read_input_tokens = (
-            self.accumulated_usage.cache_read_input_tokens or 0
-        ) + (step_usage.cache_read_input_tokens or 0)
-
-        current_span = trace.get_current_span()
-        current_span.set_attributes(
-            {
-                "input_tokens": step_usage.input_tokens or 0,
-                "output_tokens": step_usage.output_tokens or 0,
-                "cache_creation_input_tokens": (
-                    step_usage.cache_creation_input_tokens or 0
-                ),
-                "cache_read_input_tokens": step_usage.cache_read_input_tokens or 0,
-            }
-        )

src/askui/models/shared/usage_tracking_callback.py

@@ -0,0 +1,77 @@
+"""Callback for tracking token usage and reporting usage summaries."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from opentelemetry import trace
+from typing_extensions import override
+
+from askui.models.shared.agent_message_param import UsageParam
+from askui.models.shared.conversation_callback import ConversationCallback
+from askui.reporting import NULL_REPORTER, Reporter
+
+if TYPE_CHECKING:
+    from askui.models.shared.conversation import Conversation
+    from askui.speaker.speaker import SpeakerResult
+
+
+class UsageTrackingCallback(ConversationCallback):
+    """Tracks token usage per step and reports a summary at conversation end.
+
+    Args:
+        reporter: Reporter to write the final usage summary to.
+    """
+
+    def __init__(self, reporter: Reporter = NULL_REPORTER) -> None:
+        self._reporter = reporter
+        self._accumulated_usage = UsageParam()
+
+    @override
+    def on_conversation_start(self, conversation: Conversation) -> None:
+        self._accumulated_usage = UsageParam()
+
+    @override
+    def on_step_end(
+        self,
+        conversation: Conversation,
+        step_index: int,
+        result: SpeakerResult,
+    ) -> None:
+        if result.usage:
+            self._accumulate(result.usage)
+
+    @override
+    def on_conversation_end(self, conversation: Conversation) -> None:
+        self._reporter.add_usage_summary(self._accumulated_usage.model_dump())
+
+    @property
+    def accumulated_usage(self) -> UsageParam:
+        """Current accumulated usage statistics."""
+        return self._accumulated_usage
+
+    def _accumulate(self, step_usage: UsageParam) -> None:
+        self._accumulated_usage.input_tokens = (
+            self._accumulated_usage.input_tokens or 0
+        ) + (step_usage.input_tokens or 0)
+        self._accumulated_usage.output_tokens = (
+            self._accumulated_usage.output_tokens or 0
+        ) + (step_usage.output_tokens or 0)
+        self._accumulated_usage.cache_creation_input_tokens = (
+            self._accumulated_usage.cache_creation_input_tokens or 0
+        ) + (step_usage.cache_creation_input_tokens or 0)
+        self._accumulated_usage.cache_read_input_tokens = (
+            self._accumulated_usage.cache_read_input_tokens or 0
+        ) + (step_usage.cache_read_input_tokens or 0)
+
+        current_span = trace.get_current_span()
+        current_span.set_attributes(
+            {
+                "input_tokens": step_usage.input_tokens or 0,
+                "output_tokens": step_usage.output_tokens or 0,
+                "cache_creation_input_tokens": (
+                    step_usage.cache_creation_input_tokens or 0
+                ),
+                "cache_read_input_tokens": (step_usage.cache_read_input_tokens or 0),
+            }
+        )