chore: change return type of on_conversation_end from UsageTrackingCallback from dict to new `` type

Author: philipph-askui

src/askui/models/shared/usage_tracking_callback.py

@@ -5,18 +5,48 @@
 from typing import TYPE_CHECKING
 
 from opentelemetry import trace
+from pydantic import BaseModel
 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
+from askui.reporting import NULL_REPORTER
 
 if TYPE_CHECKING:
     from askui.models.shared.conversation import Conversation
+    from askui.reporting import Reporter
     from askui.speaker.speaker import SpeakerResult
     from askui.utils.model_pricing import ModelPricing
 
 
+class UsageSummary(BaseModel):
+    """Accumulated token usage and optional cost breakdown for a conversation.
+
+    Args:
+        input_tokens (int | None): Total input tokens sent to the API.
+        output_tokens (int | None): Total output tokens generated.
+        cache_creation_input_tokens (int | None): Tokens used for cache creation.
+        cache_read_input_tokens (int | None): Tokens read from cache.
+        input_cost (float | None): Computed input cost in `currency`.
+        output_cost (float | None): Computed output cost in `currency`.
+        total_cost (float | None): Sum of `input_cost` and `output_cost`.
+        currency (str | None): ISO 4217 currency code (e.g. ``"USD"``).
+        input_cost_per_million_tokens (float | None): Rate used to compute `input_cost`.
+        output_cost_per_million_tokens (float|None): Rate used to compute `output_cost`.
+    """
+
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    cache_creation_input_tokens: int | None = None
+    cache_read_input_tokens: int | None = None
+    input_cost: float | None = None
+    output_cost: float | None = None
+    total_cost: float | None = None
+    currency: str | None = None
+    input_cost_per_million_tokens: float | None = None
+    output_cost_per_million_tokens: float | None = None
+
+
 class UsageTrackingCallback(ConversationCallback):
     """Tracks token usage per step and reports a summary at conversation end.
 
@@ -51,27 +81,40 @@ def on_step_end(
 
     @override
     def on_conversation_end(self, conversation: Conversation) -> None:
-        usage_dict = self._accumulated_usage.model_dump()
+        input_cost: float | None = None
+        output_cost: float | None = None
+        total_cost: float | None = None
+        currency: str | None = None
+        input_cost_per_million_tokens: float | None = None
+        output_cost_per_million_tokens: float | None = None
         if self._pricing is not None:
             input_tokens = self._accumulated_usage.input_tokens or 0
             output_tokens = self._accumulated_usage.output_tokens or 0
             input_cost = (
-                input_tokens * self._pricing.input_cost_per_million_tokens / 1e7
+                input_tokens * self._pricing.input_cost_per_million_tokens / 1e6
             )
             output_cost = (
-                output_tokens * self._pricing.output_cost_per_million_tokens / 1e7
+                output_tokens * self._pricing.output_cost_per_million_tokens / 1e6
             )
-            usage_dict["input_cost"] = input_cost
-            usage_dict["output_cost"] = output_cost
-            usage_dict["total_cost"] = input_cost + output_cost
-            usage_dict["currency"] = self._pricing.currency
-            usage_dict["input_cost_per_million_tokens"] = (
-                self._pricing.input_cost_per_million_tokens
-            )
-            usage_dict["output_cost_per_million_tokens"] = (
+            total_cost = input_cost + output_cost
+            currency = self._pricing.currency
+            input_cost_per_million_tokens = self._pricing.input_cost_per_million_tokens
+            output_cost_per_million_tokens = (
                 self._pricing.output_cost_per_million_tokens
             )
-        self._reporter.add_usage_summary(usage_dict)
+        summary = UsageSummary(
+            input_tokens=self._accumulated_usage.input_tokens,
+            output_tokens=self._accumulated_usage.output_tokens,
+            cache_creation_input_tokens=self._accumulated_usage.cache_creation_input_tokens,
+            cache_read_input_tokens=self._accumulated_usage.cache_read_input_tokens,
+            input_cost=input_cost,
+            output_cost=output_cost,
+            total_cost=total_cost,
+            currency=currency,
+            input_cost_per_million_tokens=input_cost_per_million_tokens,
+            output_cost_per_million_tokens=output_cost_per_million_tokens,
+        )
+        self._reporter.add_usage_summary(summary)
 
     @property
     def accumulated_usage(self) -> UsageParam:

src/askui/reporting.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import base64
 import io
 import json
@@ -9,14 +11,18 @@
 from importlib.metadata import distributions
 from io import BytesIO
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from jinja2 import Template
-from PIL import Image
 from typing_extensions import TypedDict, override
 
 from askui.utils.annotated_image import AnnotatedImage
 
+if TYPE_CHECKING:
+    from PIL import Image
+
+    from askui.models.shared.usage_tracking_callback import UsageSummary
+
 
 def normalize_to_pil_images(
     image: Image.Image | list[Image.Image] | AnnotatedImage | None,
@@ -80,15 +86,14 @@ def add_message(
         raise NotImplementedError
 
     @abstractmethod
-    def add_usage_summary(self, usage: dict[str, int | None]) -> None:
+    def add_usage_summary(self, usage: UsageSummary) -> None:
         """Add usage statistics summary to the report.
 
-        Called at the end of an act() execution with accumulated token usage.
+        Called at the end of an ``act()`` execution with accumulated token
+        usage and optional cost breakdown.
 
         Args:
-            usage (dict[str, int | None]): Accumulated usage statistics containing:
-                - input_tokens: Total input tokens sent to API
-                - output_tokens: Total output tokens generated
+            usage (UsageSummary): Accumulated usage statistics.
         """
         raise NotImplementedError
 
@@ -134,7 +139,7 @@ def add_message(
         pass
 
     @override
-    def add_usage_summary(self, usage: dict[str, int | None]) -> None:
+    def add_usage_summary(self, usage: UsageSummary) -> None:
         pass
 
     @override
@@ -177,7 +182,7 @@ def add_message(
             reporter.add_message(role, content, image)
 
     @override
-    def add_usage_summary(self, usage: dict[str, int | None]) -> None:
+    def add_usage_summary(self, usage: UsageSummary) -> None:
         """Add usage summary to all reporters."""
         for reporter in self._reporters:
             reporter.add_usage_summary(usage)
@@ -215,7 +220,7 @@ def __init__(self, report_dir: str = "reports") -> None:
         self.report_dir = Path(report_dir)
         self.messages: list[dict[str, Any]] = []
         self.system_info = self._collect_system_info()
-        self.usage_summary: dict[str, int | None] | None = None
+        self.usage_summary: UsageSummary | None = None
         self.cache_original_usage: dict[str, int | None] | None = None
         self._start_time: datetime | None = None
 
@@ -264,7 +269,7 @@ def add_message(
         self.messages.append(message)
 
     @override
-    def add_usage_summary(self, usage: dict[str, int | None]) -> None:
+    def add_usage_summary(self, usage: UsageSummary) -> None:
         """Store usage summary for inclusion in the report."""
         self.usage_summary = usage
 
@@ -790,14 +795,14 @@ def generate(self) -> None:
                             </tr>
                             {% endif %}
                             {% if usage_summary is not none %}
-                                {% if usage_summary.get('input_tokens') is not none %}
+                                {% if usage_summary.input_tokens is not none %}
                                 <tr>
                                     <th>Input Tokens</th>
                                     <td>
-                                        {{ "{:,}".format(usage_summary.get('input_tokens')) }}
+                                        {{ "{:,}".format(usage_summary.input_tokens) }}
                                         {% if cache_original_usage and cache_original_usage.get('input_tokens') %}
                                             {% set original = cache_original_usage.get('input_tokens') %}
-                                            {% set current = usage_summary.get('input_tokens') %}
+                                            {% set current = usage_summary.input_tokens %}
                                             {% set saved = original - current %}
                                             {% if saved > 0 and original > 0 %}
                                                 {% set savings_pct = (saved / original * 100) %}
@@ -807,14 +812,14 @@ def generate(self) -> None:
                                     </td>
                                 </tr>
                                 {% endif %}
-                                {% if usage_summary.get('output_tokens') is not none %}
+                                {% if usage_summary.output_tokens is not none %}
                                 <tr>
                                     <th>Output Tokens</th>
                                     <td>
-                                        {{ "{:,}".format(usage_summary.get('output_tokens')) }}
+                                        {{ "{:,}".format(usage_summary.output_tokens) }}
                                         {% if cache_original_usage and cache_original_usage.get('output_tokens') %}
                                             {% set original = cache_original_usage.get('output_tokens') %}
-                                            {% set current = usage_summary.get('output_tokens') %}
+                                            {% set current = usage_summary.output_tokens %}
                                             {% set saved = original - current %}
                                             {% if saved > 0 and original > 0 %}
                                                 {% set savings_pct = (saved / original * 100) %}
@@ -824,14 +829,14 @@ def generate(self) -> None:
                                     </td>
                                 </tr>
                                 {% endif %}
-                                {% if usage_summary.get('total_cost') is not none %}
+                                {% if usage_summary.total_cost is not none %}
                                 <tr>
                                     <th>Estimated Cost <span style="font-weight:normal;color:var(--text-muted);">(actual cost may differ)</span></th>
                                     <td>
-                                        {{ "%.2f"|format(usage_summary.get('total_cost')) }} {{ usage_summary.get('currency', 'USD') }}
+                                        {{ "%.2f"|format(usage_summary.total_cost) }} {{ usage_summary.currency or 'USD' }}
                                         <span style="color: var(--text-muted); margin-left: 8px; font-size: 0.85em;">
-                                            (Input: ${{ "%.2f"|format(usage_summary.get('input_cost_per_million_tokens', 0)) }}/1M tokens,
-                                             Output: ${{ "%.2f"|format(usage_summary.get('output_cost_per_million_tokens', 0)) }}/1M tokens)
+                                            (Input: ${{ "%.2f"|format(usage_summary.input_cost_per_million_tokens or 0) }}/1M tokens,
+                                             Output: ${{ "%.2f"|format(usage_summary.output_cost_per_million_tokens or 0) }}/1M tokens)
                                         </span>
                                     </td>
                                 </tr>
@@ -992,7 +997,7 @@ def add_message(
                     )
 
     @override
-    def add_usage_summary(self, usage: dict[str, int | None]) -> None:
+    def add_usage_summary(self, usage: UsageSummary) -> None:
         """No-op for AllureReporter - usage is not tracked."""
 
     @override

tests/unit/model_providers/test_model_pricing.py

@@ -1,12 +1,14 @@
 """Unit tests for model pricing resolution and cost calculation."""
 
-from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
 
 from askui.models.shared.agent_message_param import UsageParam
-from askui.models.shared.usage_tracking_callback import UsageTrackingCallback
+from askui.models.shared.usage_tracking_callback import (
+    UsageSummary,
+    UsageTrackingCallback,
+)
 from askui.utils.model_pricing import ModelPricing
 
 
@@ -58,7 +60,7 @@ def test_override_costs_unknown_model(self) -> None:
         assert pricing.input_cost_per_million_tokens == 1.0
 
 
-def _get_usage_dict(reporter_mock: MagicMock) -> dict[str, Any]:
+def _get_usage_summary(reporter_mock: MagicMock) -> UsageSummary:
     return reporter_mock.add_usage_summary.call_args[0][0]  # type: ignore[no-any-return]
 
 
@@ -82,13 +84,13 @@ def test_cost_included_when_pricing_set(self) -> None:
         )
         callback.on_conversation_end(MagicMock())
 
-        usage_dict = _get_usage_dict(reporter)
-        assert usage_dict["total_cost"] == pytest.approx(4.5)
-        assert usage_dict["input_cost"] == pytest.approx(3.0)
-        assert usage_dict["output_cost"] == pytest.approx(1.5)
-        assert usage_dict["currency"] == "USD"
-        assert usage_dict["input_cost_per_million_tokens"] == 3.0
-        assert usage_dict["output_cost_per_million_tokens"] == 15.0
+        summary = _get_usage_summary(reporter)
+        assert summary.total_cost == pytest.approx(4.5)
+        assert summary.input_cost == pytest.approx(3.0)
+        assert summary.output_cost == pytest.approx(1.5)
+        assert summary.currency == "USD"
+        assert summary.input_cost_per_million_tokens == 3.0
+        assert summary.output_cost_per_million_tokens == 15.0
 
     def test_no_cost_when_pricing_none(self) -> None:
         callback, reporter = self._make_callback(pricing=None)
@@ -98,9 +100,9 @@ def test_no_cost_when_pricing_none(self) -> None:
         )
         callback.on_conversation_end(MagicMock())
 
-        usage_dict = _get_usage_dict(reporter)
-        assert "total_cost" not in usage_dict
-        assert "currency" not in usage_dict
+        summary = _get_usage_summary(reporter)
+        assert summary.total_cost is None
+        assert summary.currency is None
 
     def test_zero_tokens_produce_zero_cost(self) -> None:
         pricing = ModelPricing(
@@ -114,8 +116,8 @@ def test_zero_tokens_produce_zero_cost(self) -> None:
         )
         callback.on_conversation_end(MagicMock())
 
-        usage_dict = _get_usage_dict(reporter)
-        assert usage_dict["total_cost"] == 0.0
+        summary = _get_usage_summary(reporter)
+        assert summary.total_cost == 0.0
 
     def test_none_tokens_treated_as_zero(self) -> None:
         pricing = ModelPricing(
@@ -126,8 +128,8 @@ def test_none_tokens_treated_as_zero(self) -> None:
         callback._accumulated_usage = UsageParam()
         callback.on_conversation_end(MagicMock())
 
-        usage_dict = _get_usage_dict(reporter)
-        assert usage_dict["total_cost"] == 0.0
+        summary = _get_usage_summary(reporter)
+        assert summary.total_cost == 0.0
 
     def test_cost_calculation_accuracy(self) -> None:
         pricing = ModelPricing(
@@ -141,11 +143,9 @@ def test_cost_calculation_accuracy(self) -> None:
         )
         callback.on_conversation_end(MagicMock())
 
-        usage_dict = _get_usage_dict(reporter)
+        summary = _get_usage_summary(reporter)
         expected_input = 50_000 * 15.0 / 1_000_000
         expected_output = 10_000 * 75.0 / 1_000_000
-        assert usage_dict["input_cost"] == pytest.approx(expected_input)
-        assert usage_dict["output_cost"] == pytest.approx(expected_output)
-        assert usage_dict["total_cost"] == pytest.approx(
-            expected_input + expected_output
-        )
+        assert summary.input_cost == pytest.approx(expected_input)
+        assert summary.output_cost == pytest.approx(expected_output)
+        assert summary.total_cost == pytest.approx(expected_input + expected_output)