From 34d65c3a212fe74eba441c350dc38611c065fb2c Mon Sep 17 00:00:00 2001
From: Samir Mlika <samir.mlika@askui.com>
Date: Wed, 25 Mar 2026 16:22:39 +0100
Subject: [PATCH] fix usage tracker

---
 src/askui/callbacks/usage_tracking_callback.py  | 8 ++++----
 src/askui/model_providers/askui_vlm_provider.py | 5 +++++
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/askui/callbacks/usage_tracking_callback.py b/src/askui/callbacks/usage_tracking_callback.py
index a04f02d4..00881f9d 100644
--- a/src/askui/callbacks/usage_tracking_callback.py
+++ b/src/askui/callbacks/usage_tracking_callback.py
@@ -65,10 +65,6 @@ def __init__(
         self._pricing = pricing
         self._summary = UsageSummary()
 
-    @override
-    def on_conversation_start(self, conversation: Conversation) -> None:
-        self._summary = UsageSummary()
-
     @override
     def on_step_end(
         self,
@@ -77,6 +73,10 @@ def on_step_end(
         result: SpeakerResult,
     ) -> None:
         if result.usage:
+            self._reporter.add_message(
+                role="Usage Tracker",
+                content=f"Usage for step {step_index}: {str(result.usage)}",
+            )
             self._accumulate(result.usage)
 
     @override
diff --git a/src/askui/model_providers/askui_vlm_provider.py b/src/askui/model_providers/askui_vlm_provider.py
index d149deff..8451419c 100644
--- a/src/askui/model_providers/askui_vlm_provider.py
+++ b/src/askui/model_providers/askui_vlm_provider.py
@@ -17,6 +17,7 @@
 )
 from askui.models.shared.prompts import SystemPrompt
 from askui.models.shared.tools import ToolCollection
+from askui.utils.model_pricing import ModelPricing
 
 _DEFAULT_MODEL_ID = "claude-sonnet-4-6"
 
@@ -69,6 +70,10 @@ def __init__(
     def model_id(self) -> str:
         return self._model_id_value
 
+    @property
+    def pricing(self) -> ModelPricing | None:
+        return ModelPricing.for_model(self._model_id_value)
+
     @cached_property
     def _messages_api(self) -> AnthropicMessagesApi:
         """Lazily initialise the AnthropicMessagesApi on first use."""