From f65a79c528cbfbaa7913b8b744acdfb206585414 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Mon, 9 Mar 2026 17:45:08 +0100 Subject: [PATCH 01/15] feat: add pricing to html reports --- docs/05_bring_your_own_model_provider.md | 43 ++++++ docs/08_reporting.md | 29 ++++ src/askui/agent_base.py | 7 +- src/askui/model_providers/__init__.py | 2 + .../model_providers/anthropic_vlm_provider.py | 24 ++++ .../model_providers/askui_vlm_provider.py | 24 ++++ src/askui/model_providers/vlm_provider.py | 10 ++ .../models/shared/usage_tracking_callback.py | 32 ++++- src/askui/reporting.py | 12 ++ src/askui/utils/model_pricing.py | 57 ++++++++ tests/unit/model_providers/__init__.py | 0 .../model_providers/test_model_pricing.py | 129 ++++++++++++++++++ 12 files changed, 366 insertions(+), 3 deletions(-) create mode 100644 src/askui/utils/model_pricing.py create mode 100644 tests/unit/model_providers/__init__.py create mode 100644 tests/unit/model_providers/test_model_pricing.py diff --git a/docs/05_bring_your_own_model_provider.md b/docs/05_bring_your_own_model_provider.md index 66cb5013..61ff06d3 100644 --- a/docs/05_bring_your_own_model_provider.md +++ b/docs/05_bring_your_own_model_provider.md @@ -135,6 +135,49 @@ class MyImageQAProvider(ImageQAProvider): ``` +### Execution Cost Tracking + +If you want execution cost tracking in your reports, override the `pricing` property on your custom `VlmProvider`: + +```python +from typing import Any +from typing_extensions import override +from askui.model_providers import VlmProvider, ModelPricing +from askui.models.shared.agent_message_param import MessageParam, ThinkingConfigParam, ToolChoiceParam +from askui.models.shared.prompts import SystemPrompt +from askui.models.shared.tools import ToolCollection + + +class MyVlmProvider(VlmProvider): + @property + def model_id(self) -> str: + return "my-model-v1" + + @property + @override + def pricing(self) -> ModelPricing | None: + return ModelPricing( + input_cost_per_million_tokens=1.0, + output_cost_per_million_tokens=5.0, + ) + + @override + def create_message( + self, + messages: list[MessageParam], + tools: ToolCollection | None = None, + max_tokens: int | None = None, + system: SystemPrompt | None = None, + thinking: ThinkingConfigParam | None = None, + tool_choice: ToolChoiceParam | None = None, + temperature: float | None = None, + provider_options: dict[str, Any] | None = None, + ) -> MessageParam: + ... # call your API here +``` + +--- + ## Advanced: Injecting a Custom Client For full control over HTTP settings (timeouts, proxies, retries), you can inject a pre-configured client: diff --git a/docs/08_reporting.md b/docs/08_reporting.md index 10a43f32..ad8fe11a 100644 --- a/docs/08_reporting.md +++ b/docs/08_reporting.md @@ -32,6 +32,35 @@ This generates an HTML file (typically in the current directory) showing: SimpleHtmlReporter(output_dir="./execution_reports", filename="agent_run.html") ``` +### Execution Cost Tracking + +The HTML report automatically shows the estimated API cost when using a `VlmProvider` with pricing information. The built-in Anthropic and AskUI providers include default pricing for supported Claude models. + +To override pricing (for example, if you have a custom pricing agreement): + +```python +from askui import AgentSettings, ComputerAgent +from askui.model_providers import AnthropicVlmProvider +from askui.reporting import SimpleHtmlReporter + +with ComputerAgent( + reporters=[SimpleHtmlReporter()], + settings=AgentSettings( + vlm_provider=AnthropicVlmProvider( + model_id="claude-sonnet-4-6", + input_cost_per_million_tokens=2.5, + output_cost_per_million_tokens=12.0, + ), + ), +) as agent: + agent.act("Open settings") +``` + +The report will display: +- Total estimated cost +- Per-token rates used for the calculation +- Input and output token breakdowns (as before) + ### Custom Reporters Create custom reporters by implementing the `Reporter` interface: diff --git a/src/askui/agent_base.py b/src/askui/agent_base.py index 3f72d84c..5df3e1a2 100644 --- a/src/askui/agent_base.py +++ b/src/askui/agent_base.py @@ -75,7 +75,12 @@ def __init__( # Create conversation with speakers and model providers speakers = Speakers() _callbacks = list(callbacks or []) - _callbacks.append(UsageTrackingCallback(reporter=self._reporter)) + _callbacks.append( + UsageTrackingCallback( + reporter=self._reporter, + pricing=self._vlm_provider.pricing, + ) + ) self._conversation = Conversation( speakers=speakers, vlm_provider=self._vlm_provider, diff --git a/src/askui/model_providers/__init__.py b/src/askui/model_providers/__init__.py index a4f08f0d..add59506 100644 --- a/src/askui/model_providers/__init__.py +++ b/src/askui/model_providers/__init__.py @@ -23,6 +23,7 @@ from askui.model_providers.google_image_qa_provider import GoogleImageQAProvider from askui.model_providers.image_qa_provider import ImageQAProvider from askui.model_providers.vlm_provider import VlmProvider +from askui.utils.model_pricing import ModelPricing __all__ = [ "AnthropicImageQAProvider", @@ -33,5 +34,6 @@ "DetectionProvider", "GoogleImageQAProvider", "ImageQAProvider", + "ModelPricing", "VlmProvider", ] diff --git a/src/askui/model_providers/anthropic_vlm_provider.py b/src/askui/model_providers/anthropic_vlm_provider.py index b0e5d186..ed542340 100644 --- a/src/askui/model_providers/anthropic_vlm_provider.py +++ b/src/askui/model_providers/anthropic_vlm_provider.py @@ -16,6 +16,7 @@ ) from askui.models.shared.prompts import SystemPrompt from askui.models.shared.tools import ToolCollection +from askui.utils.model_pricing import ModelPricing, resolve_default_pricing _DEFAULT_MODEL_ID = "claude-sonnet-4-6" @@ -38,6 +39,11 @@ class AnthropicVlmProvider(VlmProvider): `\"claude-sonnet-4-6\"`. client (Anthropic | None, optional): Pre-configured Anthropic client. If provided, other connection parameters are ignored. + input_cost_per_million_tokens (float | None, optional): Override + cost in USD per 1M input tokens. Both cost params must be set + to override the built-in defaults. + output_cost_per_million_tokens (float | None, optional): Override + cost in USD per 1M output tokens. Example: ```python @@ -60,6 +66,8 @@ def __init__( auth_token: str | None = None, model_id: str | None = None, client: Anthropic | None = None, + input_cost_per_million_tokens: float | None = None, + output_cost_per_million_tokens: float | None = None, ) -> None: self._model_id_value = ( model_id or os.environ.get("VLM_PROVIDER_MODEL_ID") or _DEFAULT_MODEL_ID @@ -72,12 +80,28 @@ def __init__( base_url=base_url, auth_token=auth_token, ) + self._pricing: ModelPricing | None + if ( + input_cost_per_million_tokens is not None + and output_cost_per_million_tokens is not None + ): + self._pricing = ModelPricing( + input_cost_per_million_tokens=input_cost_per_million_tokens, + output_cost_per_million_tokens=output_cost_per_million_tokens, + ) + else: + self._pricing = resolve_default_pricing(self._model_id_value) @property @override def model_id(self) -> str: return self._model_id_value + @property + @override + def pricing(self) -> ModelPricing | None: + return self._pricing + @cached_property def _messages_api(self) -> AnthropicMessagesApi: """Lazily initialise the AnthropicMessagesApi on first use.""" diff --git a/src/askui/model_providers/askui_vlm_provider.py b/src/askui/model_providers/askui_vlm_provider.py index 5dfc9d29..3484a942 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, resolve_default_pricing _DEFAULT_MODEL_ID = "claude-sonnet-4-6" @@ -37,6 +38,11 @@ class AskUIVlmProvider(VlmProvider): `"claude-sonnet-4-6"`. client (Anthropic | None, optional): Pre-configured Anthropic client. If provided, `workspace_id` and `token` are ignored. + input_cost_per_million_tokens (float | None, optional): Override + cost in USD per 1M input tokens. Both cost params must be set + to override the built-in defaults. + output_cost_per_million_tokens (float | None, optional): Override + cost in USD per 1M output tokens. Example: ```python @@ -58,18 +64,36 @@ def __init__( askui_settings: AskUiInferenceApiSettings | None = None, model_id: str | None = None, client: Anthropic | None = None, + input_cost_per_million_tokens: float | None = None, + output_cost_per_million_tokens: float | None = None, ) -> None: self._askui_settings = askui_settings or AskUiInferenceApiSettings() self._model_id_value = ( model_id or os.environ.get("VLM_PROVIDER_MODEL_ID") or _DEFAULT_MODEL_ID ) self._injected_client = client + self._pricing: ModelPricing | None + if ( + input_cost_per_million_tokens is not None + and output_cost_per_million_tokens is not None + ): + self._pricing = ModelPricing( + input_cost_per_million_tokens=input_cost_per_million_tokens, + output_cost_per_million_tokens=output_cost_per_million_tokens, + ) + else: + self._pricing = resolve_default_pricing(self._model_id_value) @property @override def model_id(self) -> str: return self._model_id_value + @property + @override + def pricing(self) -> ModelPricing | None: + return self._pricing + @cached_property def _messages_api(self) -> AnthropicMessagesApi: """Lazily initialise the AnthropicMessagesApi on first use.""" diff --git a/src/askui/model_providers/vlm_provider.py b/src/askui/model_providers/vlm_provider.py index fc1f046c..1e98b972 100644 --- a/src/askui/model_providers/vlm_provider.py +++ b/src/askui/model_providers/vlm_provider.py @@ -10,6 +10,7 @@ ) from askui.models.shared.prompts import SystemPrompt from askui.models.shared.tools import ToolCollection +from askui.utils.model_pricing import ModelPricing class VlmProvider(ABC): @@ -43,6 +44,15 @@ class VlmProvider(ABC): def model_id(self) -> str: """The model identifier used by this provider.""" + @property + def pricing(self) -> ModelPricing | None: + """Pricing information for this provider's model. + + Returns ``None`` if no pricing information is available. + Override in subclasses to provide model-specific pricing. + """ + return None + @abstractmethod def create_message( self, diff --git a/src/askui/models/shared/usage_tracking_callback.py b/src/askui/models/shared/usage_tracking_callback.py index 3bc58206..a32a7df3 100644 --- a/src/askui/models/shared/usage_tracking_callback.py +++ b/src/askui/models/shared/usage_tracking_callback.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from askui.models.shared.conversation import Conversation from askui.speaker.speaker import SpeakerResult + from askui.utils.model_pricing import ModelPricing class UsageTrackingCallback(ConversationCallback): @@ -21,10 +22,17 @@ class UsageTrackingCallback(ConversationCallback): Args: reporter: Reporter to write the final usage summary to. + pricing: Pricing information for cost calculation. If ``None``, + no cost data is included in the usage summary. """ - def __init__(self, reporter: Reporter = NULL_REPORTER) -> None: + def __init__( + self, + reporter: Reporter = NULL_REPORTER, + pricing: ModelPricing | None = None, + ) -> None: self._reporter = reporter + self._pricing = pricing self._accumulated_usage = UsageParam() @override @@ -43,7 +51,27 @@ def on_step_end( @override def on_conversation_end(self, conversation: Conversation) -> None: - self._reporter.add_usage_summary(self._accumulated_usage.model_dump()) + usage_dict = self._accumulated_usage.model_dump() + 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 / 1_000_000 + ) + output_cost = ( + output_tokens * self._pricing.output_cost_per_million_tokens / 1_000_000 + ) + 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"] = ( + self._pricing.output_cost_per_million_tokens + ) + self._reporter.add_usage_summary(usage_dict) @property def accumulated_usage(self) -> UsageParam: diff --git a/src/askui/reporting.py b/src/askui/reporting.py index 4388a2cf..f2ba4242 100644 --- a/src/askui/reporting.py +++ b/src/askui/reporting.py @@ -824,6 +824,18 @@ def generate(self) -> None: {% endif %} + {% if usage_summary.get('total_cost') is not none %} + + Estimated Cost + + {{ "%.2f"|format(usage_summary.get('total_cost')) }} {{ usage_summary.get('currency', 'USD') }} + + (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) + + + + {% endif %} {% endif %} {% if cache_original_usage is not none %} {% if cache_original_usage.get('input_tokens') is not none %} diff --git a/src/askui/utils/model_pricing.py b/src/askui/utils/model_pricing.py new file mode 100644 index 00000000..b0d3ea34 --- /dev/null +++ b/src/askui/utils/model_pricing.py @@ -0,0 +1,57 @@ +"""Pricing information for model API calls.""" + +from pydantic import BaseModel + + +class ModelPricing(BaseModel): + """Cost per 1 million tokens for a model. + + Args: + input_cost_per_million_tokens (float): Cost in USD per 1M input tokens. + output_cost_per_million_tokens (float): Cost in USD per 1M output tokens. + currency (str): ISO 4217 currency code. Defaults to ``"USD"``. + """ + + input_cost_per_million_tokens: float + output_cost_per_million_tokens: float + currency: str = "USD" + + +_DEFAULT_PRICING: dict[str, ModelPricing] = { + "claude-haiku-4-5-20251001": ModelPricing( + input_cost_per_million_tokens=1.0, + output_cost_per_million_tokens=5.0, + ), + "claude-sonnet-4-5-20250929": ModelPricing( + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ), + "claude-opus-4-5-20251101": ModelPricing( + input_cost_per_million_tokens=5.0, + output_cost_per_million_tokens=25.0, + ), + "claude-sonnet-4-6": ModelPricing( + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ), + "claude-opus-4-6": ModelPricing( + input_cost_per_million_tokens=5.0, + output_cost_per_million_tokens=25.0, + ), +} + + +def resolve_default_pricing(model_id: str) -> ModelPricing | None: + """Resolve default pricing for a model ID by prefix matching. + + Tries exact match first, then longest-prefix match. + + Args: + model_id (str): The model identifier. + + Returns: + ModelPricing | None: Default pricing, or ``None`` if no match found. + """ + if model_id in _DEFAULT_PRICING: + return _DEFAULT_PRICING[model_id] + return None diff --git a/tests/unit/model_providers/__init__.py b/tests/unit/model_providers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py new file mode 100644 index 00000000..c8ff5032 --- /dev/null +++ b/tests/unit/model_providers/test_model_pricing.py @@ -0,0 +1,129 @@ +"""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.utils.model_pricing import ModelPricing, resolve_default_pricing + + +class TestResolveDefaultPricing: + def test_prefix_match_sonnet(self) -> None: + pricing = resolve_default_pricing("claude-sonnet-4-6") + assert pricing is not None + assert pricing.input_cost_per_million_tokens == 3.0 + assert pricing.output_cost_per_million_tokens == 15.0 + + def test_prefix_match_opus(self) -> None: + pricing = resolve_default_pricing("claude-opus-4-6") + assert pricing is not None + assert pricing.input_cost_per_million_tokens == 15.0 + assert pricing.output_cost_per_million_tokens == 75.0 + + def test_prefix_match_haiku(self) -> None: + pricing = resolve_default_pricing("claude-haiku-4-5-20251001") + assert pricing is not None + assert pricing.input_cost_per_million_tokens == 0.80 + assert pricing.output_cost_per_million_tokens == 4.0 + + def test_unknown_model_returns_none(self) -> None: + assert resolve_default_pricing("unknown-model-v1") is None + + def test_empty_string_returns_none(self) -> None: + assert resolve_default_pricing("") is None + + +def _get_usage_dict(reporter_mock: MagicMock) -> dict[str, Any]: + return reporter_mock.add_usage_summary.call_args[0][0] # type: ignore[no-any-return] + + +class TestUsageTrackingCallbackCost: + def _make_callback( + self, pricing: ModelPricing | None = None + ) -> tuple[UsageTrackingCallback, MagicMock]: + reporter = MagicMock() + callback = UsageTrackingCallback(reporter=reporter, pricing=pricing) + return callback, reporter + + def test_cost_included_when_pricing_set(self) -> None: + pricing = ModelPricing( + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ) + callback, reporter = self._make_callback(pricing) + callback._accumulated_usage = UsageParam( + input_tokens=1_000_000, + output_tokens=100_000, + ) + 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 + + def test_no_cost_when_pricing_none(self) -> None: + callback, reporter = self._make_callback(pricing=None) + callback._accumulated_usage = UsageParam( + input_tokens=500, + output_tokens=200, + ) + callback.on_conversation_end(MagicMock()) + + usage_dict = _get_usage_dict(reporter) + assert "total_cost" not in usage_dict + assert "currency" not in usage_dict + + def test_zero_tokens_produce_zero_cost(self) -> None: + pricing = ModelPricing( + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ) + callback, reporter = self._make_callback(pricing) + callback._accumulated_usage = UsageParam( + input_tokens=0, + output_tokens=0, + ) + callback.on_conversation_end(MagicMock()) + + usage_dict = _get_usage_dict(reporter) + assert usage_dict["total_cost"] == 0.0 + + def test_none_tokens_treated_as_zero(self) -> None: + pricing = ModelPricing( + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ) + callback, reporter = self._make_callback(pricing) + callback._accumulated_usage = UsageParam() + callback.on_conversation_end(MagicMock()) + + usage_dict = _get_usage_dict(reporter) + assert usage_dict["total_cost"] == 0.0 + + def test_cost_calculation_accuracy(self) -> None: + pricing = ModelPricing( + input_cost_per_million_tokens=15.0, + output_cost_per_million_tokens=75.0, + ) + callback, reporter = self._make_callback(pricing) + callback._accumulated_usage = UsageParam( + input_tokens=50_000, + output_tokens=10_000, + ) + callback.on_conversation_end(MagicMock()) + + usage_dict = _get_usage_dict(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 + ) From 0e6aca6a866befdf8f342ec388cbf913649980bc Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Mon, 9 Mar 2026 18:44:22 +0100 Subject: [PATCH 02/15] feat: add hint to estimated cost in html report that actual cost might differ --- src/askui/reporting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/askui/reporting.py b/src/askui/reporting.py index f2ba4242..758a8557 100644 --- a/src/askui/reporting.py +++ b/src/askui/reporting.py @@ -826,7 +826,7 @@ def generate(self) -> None: {% endif %} {% if usage_summary.get('total_cost') is not none %} - Estimated Cost + Estimated Cost (actual cost may differ) {{ "%.2f"|format(usage_summary.get('total_cost')) }} {{ usage_summary.get('currency', 'USD') }} From 49273a6bfb1826a712b38ddce4ca32fe13bfd113 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Mon, 9 Mar 2026 18:54:25 +0100 Subject: [PATCH 03/15] chore: update docs on pricing --- docs/05_bring_your_own_model_provider.md | 47 +++++++++++------------- 1 file changed, 22 insertions(+), 25 deletions(-) diff --git a/docs/05_bring_your_own_model_provider.md b/docs/05_bring_your_own_model_provider.md index 61ff06d3..a74281c8 100644 --- a/docs/05_bring_your_own_model_provider.md +++ b/docs/05_bring_your_own_model_provider.md @@ -137,43 +137,40 @@ class MyImageQAProvider(ImageQAProvider): ### Execution Cost Tracking -If you want execution cost tracking in your reports, override the `pricing` property on your custom `VlmProvider`: +The built-in providers include default pricing for supported models. You can override the pricing on any provider by passing `input_cost_per_million_tokens` and `output_cost_per_million_tokens`: ```python -from typing import Any -from typing_extensions import override -from askui.model_providers import VlmProvider, ModelPricing -from askui.models.shared.agent_message_param import MessageParam, ThinkingConfigParam, ToolChoiceParam -from askui.models.shared.prompts import SystemPrompt -from askui.models.shared.tools import ToolCollection +from askui import AgentSettings, ComputerAgent +from askui.model_providers import AnthropicVlmProvider +from askui.reporting import SimpleHtmlReporter + +with ComputerAgent( + reporters=[SimpleHtmlReporter()], + settings=AgentSettings( + vlm_provider=AnthropicVlmProvider( + model_id="claude-sonnet-4-6", + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ), + ), +) as agent: + agent.act("Open settings") +``` +If you implement a fully custom `VlmProvider`, override the `pricing` property to enable cost tracking: -class MyVlmProvider(VlmProvider): - @property - def model_id(self) -> str: - return "my-model-v1" +```python +from askui.model_providers import VlmProvider, ModelPricing +class MyVlmProvider(VlmProvider): @property - @override def pricing(self) -> ModelPricing | None: return ModelPricing( input_cost_per_million_tokens=1.0, output_cost_per_million_tokens=5.0, ) - @override - def create_message( - self, - messages: list[MessageParam], - tools: ToolCollection | None = None, - max_tokens: int | None = None, - system: SystemPrompt | None = None, - thinking: ThinkingConfigParam | None = None, - tool_choice: ToolChoiceParam | None = None, - temperature: float | None = None, - provider_options: dict[str, Any] | None = None, - ) -> MessageParam: - ... # call your API here + # ... rest of implementation ``` --- From 5a752cf5b63e3c060ad4478785bf8dd610f9dd21 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Mon, 9 Mar 2026 18:54:36 +0100 Subject: [PATCH 04/15] fix: tests on model pricing --- .../unit/model_providers/test_model_pricing.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py index c8ff5032..798492ce 100644 --- a/tests/unit/model_providers/test_model_pricing.py +++ b/tests/unit/model_providers/test_model_pricing.py @@ -11,23 +11,23 @@ class TestResolveDefaultPricing: - def test_prefix_match_sonnet(self) -> None: + def test_exact_match_sonnet_4_6(self) -> None: pricing = resolve_default_pricing("claude-sonnet-4-6") assert pricing is not None assert pricing.input_cost_per_million_tokens == 3.0 assert pricing.output_cost_per_million_tokens == 15.0 - def test_prefix_match_opus(self) -> None: + def test_exact_match_opus_4_6(self) -> None: pricing = resolve_default_pricing("claude-opus-4-6") assert pricing is not None - assert pricing.input_cost_per_million_tokens == 15.0 - assert pricing.output_cost_per_million_tokens == 75.0 + assert pricing.input_cost_per_million_tokens == 5.0 + assert pricing.output_cost_per_million_tokens == 25.0 - def test_prefix_match_haiku(self) -> None: + def test_exact_match_haiku(self) -> None: pricing = resolve_default_pricing("claude-haiku-4-5-20251001") assert pricing is not None - assert pricing.input_cost_per_million_tokens == 0.80 - assert pricing.output_cost_per_million_tokens == 4.0 + assert pricing.input_cost_per_million_tokens == 1.0 + assert pricing.output_cost_per_million_tokens == 5.0 def test_unknown_model_returns_none(self) -> None: assert resolve_default_pricing("unknown-model-v1") is None @@ -35,6 +35,9 @@ def test_unknown_model_returns_none(self) -> None: def test_empty_string_returns_none(self) -> None: assert resolve_default_pricing("") is None + def test_partial_model_id_returns_none(self) -> None: + assert resolve_default_pricing("claude-sonnet-4") is None + def _get_usage_dict(reporter_mock: MagicMock) -> dict[str, Any]: return reporter_mock.add_usage_summary.call_args[0][0] # type: ignore[no-any-return] From aa2cf3c9f4577b29fe9efe3d3f9a2850d7ca516d Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 06:33:29 +0100 Subject: [PATCH 05/15] chore: cleanup new pricing sections in docs --- docs/05_bring_your_own_model_provider.md | 2 +- docs/08_reporting.md | 20 -------------------- 2 files changed, 1 insertion(+), 21 deletions(-) diff --git a/docs/05_bring_your_own_model_provider.md b/docs/05_bring_your_own_model_provider.md index a74281c8..04f17d48 100644 --- a/docs/05_bring_your_own_model_provider.md +++ b/docs/05_bring_your_own_model_provider.md @@ -137,7 +137,7 @@ class MyImageQAProvider(ImageQAProvider): ### Execution Cost Tracking -The built-in providers include default pricing for supported models. You can override the pricing on any provider by passing `input_cost_per_million_tokens` and `output_cost_per_million_tokens`: +The built-in VLM providers include default pricing for supported models. You can override the pricing on any provider by passing `input_cost_per_million_tokens` and `output_cost_per_million_tokens`: ```python from askui import AgentSettings, ComputerAgent diff --git a/docs/08_reporting.md b/docs/08_reporting.md index ad8fe11a..6041fafd 100644 --- a/docs/08_reporting.md +++ b/docs/08_reporting.md @@ -36,26 +36,6 @@ SimpleHtmlReporter(output_dir="./execution_reports", filename="agent_run.html") The HTML report automatically shows the estimated API cost when using a `VlmProvider` with pricing information. The built-in Anthropic and AskUI providers include default pricing for supported Claude models. -To override pricing (for example, if you have a custom pricing agreement): - -```python -from askui import AgentSettings, ComputerAgent -from askui.model_providers import AnthropicVlmProvider -from askui.reporting import SimpleHtmlReporter - -with ComputerAgent( - reporters=[SimpleHtmlReporter()], - settings=AgentSettings( - vlm_provider=AnthropicVlmProvider( - model_id="claude-sonnet-4-6", - input_cost_per_million_tokens=2.5, - output_cost_per_million_tokens=12.0, - ), - ), -) as agent: - agent.act("Open settings") -``` - The report will display: - Total estimated cost - Per-token rates used for the calculation From 6e1ea2663004ccb6d8523d2000849ef796f950c9 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 06:38:41 +0100 Subject: [PATCH 06/15] chore: move `resolve_default_pricing` to `ModelPricing` --- .../model_providers/anthropic_vlm_provider.py | 18 ++-- .../model_providers/askui_vlm_provider.py | 18 ++-- src/askui/utils/model_pricing.py | 94 ++++++++++++------- .../model_providers/test_model_pricing.py | 35 +++++-- 4 files changed, 97 insertions(+), 68 deletions(-) diff --git a/src/askui/model_providers/anthropic_vlm_provider.py b/src/askui/model_providers/anthropic_vlm_provider.py index ed542340..cc7525ec 100644 --- a/src/askui/model_providers/anthropic_vlm_provider.py +++ b/src/askui/model_providers/anthropic_vlm_provider.py @@ -16,7 +16,7 @@ ) from askui.models.shared.prompts import SystemPrompt from askui.models.shared.tools import ToolCollection -from askui.utils.model_pricing import ModelPricing, resolve_default_pricing +from askui.utils.model_pricing import ModelPricing _DEFAULT_MODEL_ID = "claude-sonnet-4-6" @@ -80,17 +80,11 @@ def __init__( base_url=base_url, auth_token=auth_token, ) - self._pricing: ModelPricing | None - if ( - input_cost_per_million_tokens is not None - and output_cost_per_million_tokens is not None - ): - self._pricing = ModelPricing( - input_cost_per_million_tokens=input_cost_per_million_tokens, - output_cost_per_million_tokens=output_cost_per_million_tokens, - ) - else: - self._pricing = resolve_default_pricing(self._model_id_value) + self._pricing = ModelPricing.for_model( + self._model_id_value, + input_cost_per_million_tokens=input_cost_per_million_tokens, + output_cost_per_million_tokens=output_cost_per_million_tokens, + ) @property @override diff --git a/src/askui/model_providers/askui_vlm_provider.py b/src/askui/model_providers/askui_vlm_provider.py index 3484a942..dcc928d2 100644 --- a/src/askui/model_providers/askui_vlm_provider.py +++ b/src/askui/model_providers/askui_vlm_provider.py @@ -17,7 +17,7 @@ ) from askui.models.shared.prompts import SystemPrompt from askui.models.shared.tools import ToolCollection -from askui.utils.model_pricing import ModelPricing, resolve_default_pricing +from askui.utils.model_pricing import ModelPricing _DEFAULT_MODEL_ID = "claude-sonnet-4-6" @@ -72,17 +72,11 @@ def __init__( model_id or os.environ.get("VLM_PROVIDER_MODEL_ID") or _DEFAULT_MODEL_ID ) self._injected_client = client - self._pricing: ModelPricing | None - if ( - input_cost_per_million_tokens is not None - and output_cost_per_million_tokens is not None - ): - self._pricing = ModelPricing( - input_cost_per_million_tokens=input_cost_per_million_tokens, - output_cost_per_million_tokens=output_cost_per_million_tokens, - ) - else: - self._pricing = resolve_default_pricing(self._model_id_value) + self._pricing = ModelPricing.for_model( + self._model_id_value, + input_cost_per_million_tokens=input_cost_per_million_tokens, + output_cost_per_million_tokens=output_cost_per_million_tokens, + ) @property @override diff --git a/src/askui/utils/model_pricing.py b/src/askui/utils/model_pricing.py index b0d3ea34..568b2c83 100644 --- a/src/askui/utils/model_pricing.py +++ b/src/askui/utils/model_pricing.py @@ -2,6 +2,8 @@ from pydantic import BaseModel +_DEFAULT_PRICING: dict[str, "ModelPricing"] = {} + class ModelPricing(BaseModel): """Cost per 1 million tokens for a model. @@ -16,42 +18,62 @@ class ModelPricing(BaseModel): output_cost_per_million_tokens: float currency: str = "USD" + @classmethod + def for_model( + cls, + model_id: str, + input_cost_per_million_tokens: float | None = None, + output_cost_per_million_tokens: float | None = None, + ) -> "ModelPricing | None": + """Resolve pricing for a model. -_DEFAULT_PRICING: dict[str, ModelPricing] = { - "claude-haiku-4-5-20251001": ModelPricing( - input_cost_per_million_tokens=1.0, - output_cost_per_million_tokens=5.0, - ), - "claude-sonnet-4-5-20250929": ModelPricing( - input_cost_per_million_tokens=3.0, - output_cost_per_million_tokens=15.0, - ), - "claude-opus-4-5-20251101": ModelPricing( - input_cost_per_million_tokens=5.0, - output_cost_per_million_tokens=25.0, - ), - "claude-sonnet-4-6": ModelPricing( - input_cost_per_million_tokens=3.0, - output_cost_per_million_tokens=15.0, - ), - "claude-opus-4-6": ModelPricing( - input_cost_per_million_tokens=5.0, - output_cost_per_million_tokens=25.0, - ), -} - - -def resolve_default_pricing(model_id: str) -> ModelPricing | None: - """Resolve default pricing for a model ID by prefix matching. - - Tries exact match first, then longest-prefix match. + If both cost parameters are provided, creates a ``ModelPricing`` + with those values. Otherwise, looks up built-in defaults by + ``model_id``. - Args: - model_id (str): The model identifier. + Args: + model_id (str): The model identifier. + input_cost_per_million_tokens (float | None, optional): Override + cost in USD per 1M input tokens. + output_cost_per_million_tokens (float | None, optional): Override + cost in USD per 1M output tokens. - Returns: - ModelPricing | None: Default pricing, or ``None`` if no match found. - """ - if model_id in _DEFAULT_PRICING: - return _DEFAULT_PRICING[model_id] - return None + Returns: + ModelPricing | None: Resolved pricing, or ``None`` if no match + and no overrides provided. + """ + if ( + input_cost_per_million_tokens is not None + and output_cost_per_million_tokens is not None + ): + return cls( + input_cost_per_million_tokens=input_cost_per_million_tokens, + output_cost_per_million_tokens=output_cost_per_million_tokens, + ) + return _DEFAULT_PRICING.get(model_id) + + +_DEFAULT_PRICING.update( + { + "claude-haiku-4-5-20251001": ModelPricing( + input_cost_per_million_tokens=1.0, + output_cost_per_million_tokens=5.0, + ), + "claude-sonnet-4-5-20250929": ModelPricing( + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ), + "claude-opus-4-5-20251101": ModelPricing( + input_cost_per_million_tokens=5.0, + output_cost_per_million_tokens=25.0, + ), + "claude-sonnet-4-6": ModelPricing( + input_cost_per_million_tokens=3.0, + output_cost_per_million_tokens=15.0, + ), + "claude-opus-4-6": ModelPricing( + input_cost_per_million_tokens=5.0, + output_cost_per_million_tokens=25.0, + ), + } +) diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py index 798492ce..b572e090 100644 --- a/tests/unit/model_providers/test_model_pricing.py +++ b/tests/unit/model_providers/test_model_pricing.py @@ -7,36 +7,55 @@ from askui.models.shared.agent_message_param import UsageParam from askui.models.shared.usage_tracking_callback import UsageTrackingCallback -from askui.utils.model_pricing import ModelPricing, resolve_default_pricing +from askui.utils.model_pricing import ModelPricing -class TestResolveDefaultPricing: +class TestModelPricingForModel: def test_exact_match_sonnet_4_6(self) -> None: - pricing = resolve_default_pricing("claude-sonnet-4-6") + pricing = ModelPricing.for_model("claude-sonnet-4-6") assert pricing is not None assert pricing.input_cost_per_million_tokens == 3.0 assert pricing.output_cost_per_million_tokens == 15.0 def test_exact_match_opus_4_6(self) -> None: - pricing = resolve_default_pricing("claude-opus-4-6") + pricing = ModelPricing.for_model("claude-opus-4-6") assert pricing is not None assert pricing.input_cost_per_million_tokens == 5.0 assert pricing.output_cost_per_million_tokens == 25.0 def test_exact_match_haiku(self) -> None: - pricing = resolve_default_pricing("claude-haiku-4-5-20251001") + pricing = ModelPricing.for_model("claude-haiku-4-5-20251001") assert pricing is not None assert pricing.input_cost_per_million_tokens == 1.0 assert pricing.output_cost_per_million_tokens == 5.0 def test_unknown_model_returns_none(self) -> None: - assert resolve_default_pricing("unknown-model-v1") is None + assert ModelPricing.for_model("unknown-model-v1") is None def test_empty_string_returns_none(self) -> None: - assert resolve_default_pricing("") is None + assert ModelPricing.for_model("") is None def test_partial_model_id_returns_none(self) -> None: - assert resolve_default_pricing("claude-sonnet-4") is None + assert ModelPricing.for_model("claude-sonnet-4") is None + + def test_override_costs(self) -> None: + pricing = ModelPricing.for_model( + "claude-sonnet-4-6", + input_cost_per_million_tokens=99.0, + output_cost_per_million_tokens=199.0, + ) + assert pricing is not None + assert pricing.input_cost_per_million_tokens == 99.0 + assert pricing.output_cost_per_million_tokens == 199.0 + + def test_override_costs_unknown_model(self) -> None: + pricing = ModelPricing.for_model( + "unknown-model", + input_cost_per_million_tokens=1.0, + output_cost_per_million_tokens=2.0, + ) + assert pricing is not None + assert pricing.input_cost_per_million_tokens == 1.0 def _get_usage_dict(reporter_mock: MagicMock) -> dict[str, Any]: From d90a3ac864338c611574e9dfa643558a5a1e645d Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 06:51:14 +0100 Subject: [PATCH 07/15] fix: minor formatting issue --- src/askui/models/shared/usage_tracking_callback.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/askui/models/shared/usage_tracking_callback.py b/src/askui/models/shared/usage_tracking_callback.py index a32a7df3..17c63966 100644 --- a/src/askui/models/shared/usage_tracking_callback.py +++ b/src/askui/models/shared/usage_tracking_callback.py @@ -56,10 +56,10 @@ def on_conversation_end(self, conversation: Conversation) -> 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 / 1_000_000 + input_tokens * self._pricing.input_cost_per_million_tokens / 1e7 ) output_cost = ( - output_tokens * self._pricing.output_cost_per_million_tokens / 1_000_000 + output_tokens * self._pricing.output_cost_per_million_tokens / 1e7 ) usage_dict["input_cost"] = input_cost usage_dict["output_cost"] = output_cost From cdadaadd05d7bbd15fed8985764c4a0490ab139e Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 07:10:54 +0100 Subject: [PATCH 08/15] chore: change return type of `on_conversation_end` from `UsageTrackingCallback` from dict to new `` type --- .../models/shared/usage_tracking_callback.py | 69 +++++++++++++++---- src/askui/reporting.py | 49 +++++++------ .../model_providers/test_model_pricing.py | 46 ++++++------- 3 files changed, 106 insertions(+), 58 deletions(-) diff --git a/src/askui/models/shared/usage_tracking_callback.py b/src/askui/models/shared/usage_tracking_callback.py index 17c63966..c35df2f0 100644 --- a/src/askui/models/shared/usage_tracking_callback.py +++ b/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: diff --git a/src/askui/reporting.py b/src/askui/reporting.py index 758a8557..ec601286 100644 --- a/src/askui/reporting.py +++ b/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: {% endif %} {% if usage_summary is not none %} - {% if usage_summary.get('input_tokens') is not none %} + {% if usage_summary.input_tokens is not none %} Input Tokens - {{ "{:,}".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: {% endif %} - {% if usage_summary.get('output_tokens') is not none %} + {% if usage_summary.output_tokens is not none %} Output Tokens - {{ "{:,}".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: {% endif %} - {% if usage_summary.get('total_cost') is not none %} + {% if usage_summary.total_cost is not none %} Estimated Cost (actual cost may differ) - {{ "%.2f"|format(usage_summary.get('total_cost')) }} {{ usage_summary.get('currency', 'USD') }} + {{ "%.2f"|format(usage_summary.total_cost) }} {{ usage_summary.currency or 'USD' }} - (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) @@ -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 diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py index b572e090..4ae95527 100644 --- a/tests/unit/model_providers/test_model_pricing.py +++ b/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) From 9b506b875b2b9eb7397cdb7de328066c817f47cc Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 11:42:27 +0100 Subject: [PATCH 09/15] Update src/askui/utils/model_pricing.py Co-authored-by: Dominik Klotz <105296959+programminx-askui@users.noreply.github.com> --- src/askui/utils/model_pricing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/askui/utils/model_pricing.py b/src/askui/utils/model_pricing.py index 568b2c83..2352dc95 100644 --- a/src/askui/utils/model_pricing.py +++ b/src/askui/utils/model_pricing.py @@ -9,8 +9,8 @@ class ModelPricing(BaseModel): """Cost per 1 million tokens for a model. Args: - input_cost_per_million_tokens (float): Cost in USD per 1M input tokens. - output_cost_per_million_tokens (float): Cost in USD per 1M output tokens. + input_cost_per_million_tokens (float): Cost per 1M input tokens in Currency. + output_cost_per_million_tokens (float): Cost per 1M output tokens in Currency. currency (str): ISO 4217 currency code. Defaults to ``"USD"``. """ From bac8638070841808a1ab135c017ce9ba3e4f43fa Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 11:45:37 +0100 Subject: [PATCH 10/15] chore: improve description of currency parameter in docstring --- src/askui/utils/model_pricing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/askui/utils/model_pricing.py b/src/askui/utils/model_pricing.py index 2352dc95..eeba9212 100644 --- a/src/askui/utils/model_pricing.py +++ b/src/askui/utils/model_pricing.py @@ -11,7 +11,7 @@ class ModelPricing(BaseModel): Args: input_cost_per_million_tokens (float): Cost per 1M input tokens in Currency. output_cost_per_million_tokens (float): Cost per 1M output tokens in Currency. - currency (str): ISO 4217 currency code. Defaults to ``"USD"``. + currency (str): descriptor of the currency. Defaults to ``"USD"``. """ input_cost_per_million_tokens: float From ad10a3e61f91e51f8fdb2b548b9c0dcc33402bfd Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 11:46:49 +0100 Subject: [PATCH 11/15] add online comment on model price source --- src/askui/utils/model_pricing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/askui/utils/model_pricing.py b/src/askui/utils/model_pricing.py index eeba9212..03d2e254 100644 --- a/src/askui/utils/model_pricing.py +++ b/src/askui/utils/model_pricing.py @@ -53,6 +53,8 @@ def for_model( return _DEFAULT_PRICING.get(model_id) +# taken from: https://platform.claude.com/docs/en/about-claude/models/overview +# last accessed: March 10, 2026 _DEFAULT_PRICING.update( { "claude-haiku-4-5-20251001": ModelPricing( From 962fe71f0761069145382039aaa0c9707acd59b8 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Tue, 10 Mar 2026 11:49:28 +0100 Subject: [PATCH 12/15] chore: rewrite model_pricing checks as parameterized tests --- .../model_providers/test_model_pricing.py | 165 +++++++++--------- 1 file changed, 80 insertions(+), 85 deletions(-) diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py index 4ae95527..e8381a40 100644 --- a/tests/unit/model_providers/test_model_pricing.py +++ b/tests/unit/model_providers/test_model_pricing.py @@ -13,51 +13,55 @@ class TestModelPricingForModel: - def test_exact_match_sonnet_4_6(self) -> None: - pricing = ModelPricing.for_model("claude-sonnet-4-6") + @pytest.mark.parametrize( + ("model_id", "expected_input", "expected_output"), + [ + ("claude-haiku-4-5-20251001", 1.0, 5.0), + ("claude-sonnet-4-5-20250929", 3.0, 15.0), + ("claude-opus-4-5-20251101", 5.0, 25.0), + ("claude-sonnet-4-6", 3.0, 15.0), + ("claude-opus-4-6", 5.0, 25.0), + ], + ) + def test_known_model_returns_default_pricing( + self, + model_id: str, + expected_input: float, + expected_output: float, + ) -> None: + pricing = ModelPricing.for_model(model_id) assert pricing is not None - assert pricing.input_cost_per_million_tokens == 3.0 - assert pricing.output_cost_per_million_tokens == 15.0 - - def test_exact_match_opus_4_6(self) -> None: - pricing = ModelPricing.for_model("claude-opus-4-6") - assert pricing is not None - assert pricing.input_cost_per_million_tokens == 5.0 - assert pricing.output_cost_per_million_tokens == 25.0 - - def test_exact_match_haiku(self) -> None: - pricing = ModelPricing.for_model("claude-haiku-4-5-20251001") - assert pricing is not None - assert pricing.input_cost_per_million_tokens == 1.0 - assert pricing.output_cost_per_million_tokens == 5.0 - - def test_unknown_model_returns_none(self) -> None: - assert ModelPricing.for_model("unknown-model-v1") is None - - def test_empty_string_returns_none(self) -> None: - assert ModelPricing.for_model("") is None - - def test_partial_model_id_returns_none(self) -> None: - assert ModelPricing.for_model("claude-sonnet-4") is None - - def test_override_costs(self) -> None: - pricing = ModelPricing.for_model( - "claude-sonnet-4-6", - input_cost_per_million_tokens=99.0, - output_cost_per_million_tokens=199.0, - ) - assert pricing is not None - assert pricing.input_cost_per_million_tokens == 99.0 - assert pricing.output_cost_per_million_tokens == 199.0 - - def test_override_costs_unknown_model(self) -> None: + assert pricing.input_cost_per_million_tokens == expected_input + assert pricing.output_cost_per_million_tokens == expected_output + + @pytest.mark.parametrize( + "model_id", + ["unknown-model-v1", "", "claude-sonnet-4"], + ) + def test_unknown_model_returns_none(self, model_id: str) -> None: + assert ModelPricing.for_model(model_id) is None + + @pytest.mark.parametrize( + ("model_id", "input_cost", "output_cost"), + [ + ("claude-sonnet-4-6", 99.0, 199.0), + ("unknown-model", 1.0, 2.0), + ], + ) + def test_override_costs( + self, + model_id: str, + input_cost: float, + output_cost: float, + ) -> None: pricing = ModelPricing.for_model( - "unknown-model", - input_cost_per_million_tokens=1.0, - output_cost_per_million_tokens=2.0, + model_id, + input_cost_per_million_tokens=input_cost, + output_cost_per_million_tokens=output_cost, ) assert pricing is not None - assert pricing.input_cost_per_million_tokens == 1.0 + assert pricing.input_cost_per_million_tokens == input_cost + assert pricing.output_cost_per_million_tokens == output_cost def _get_usage_summary(reporter_mock: MagicMock) -> UsageSummary: @@ -72,25 +76,50 @@ def _make_callback( callback = UsageTrackingCallback(reporter=reporter, pricing=pricing) return callback, reporter - def test_cost_included_when_pricing_set(self) -> None: + @pytest.mark.parametrize( + ( + "input_tokens", + "output_tokens", + "input_rate", + "output_rate", + "expected_input_cost", + "expected_output_cost", + ), + [ + (1_000_000, 100_000, 3.0, 15.0, 3.0, 1.5), + (50_000, 10_000, 15.0, 75.0, 0.75, 0.75), + (0, 0, 3.0, 15.0, 0.0, 0.0), + ], + ) + def test_cost_calculation( + self, + input_tokens: int, + output_tokens: int, + input_rate: float, + output_rate: float, + expected_input_cost: float, + expected_output_cost: float, + ) -> None: pricing = ModelPricing( - input_cost_per_million_tokens=3.0, - output_cost_per_million_tokens=15.0, + input_cost_per_million_tokens=input_rate, + output_cost_per_million_tokens=output_rate, ) callback, reporter = self._make_callback(pricing) callback._accumulated_usage = UsageParam( - input_tokens=1_000_000, - output_tokens=100_000, + input_tokens=input_tokens, + output_tokens=output_tokens, ) callback.on_conversation_end(MagicMock()) 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.input_cost == pytest.approx(expected_input_cost) + assert summary.output_cost == pytest.approx(expected_output_cost) + assert summary.total_cost == pytest.approx( + expected_input_cost + expected_output_cost + ) assert summary.currency == "USD" - assert summary.input_cost_per_million_tokens == 3.0 - assert summary.output_cost_per_million_tokens == 15.0 + assert summary.input_cost_per_million_tokens == input_rate + assert summary.output_cost_per_million_tokens == output_rate def test_no_cost_when_pricing_none(self) -> None: callback, reporter = self._make_callback(pricing=None) @@ -104,21 +133,6 @@ def test_no_cost_when_pricing_none(self) -> None: assert summary.total_cost is None assert summary.currency is None - def test_zero_tokens_produce_zero_cost(self) -> None: - pricing = ModelPricing( - input_cost_per_million_tokens=3.0, - output_cost_per_million_tokens=15.0, - ) - callback, reporter = self._make_callback(pricing) - callback._accumulated_usage = UsageParam( - input_tokens=0, - output_tokens=0, - ) - callback.on_conversation_end(MagicMock()) - - summary = _get_usage_summary(reporter) - assert summary.total_cost == 0.0 - def test_none_tokens_treated_as_zero(self) -> None: pricing = ModelPricing( input_cost_per_million_tokens=3.0, @@ -130,22 +144,3 @@ def test_none_tokens_treated_as_zero(self) -> None: summary = _get_usage_summary(reporter) assert summary.total_cost == 0.0 - - def test_cost_calculation_accuracy(self) -> None: - pricing = ModelPricing( - input_cost_per_million_tokens=15.0, - output_cost_per_million_tokens=75.0, - ) - callback, reporter = self._make_callback(pricing) - callback._accumulated_usage = UsageParam( - input_tokens=50_000, - output_tokens=10_000, - ) - callback.on_conversation_end(MagicMock()) - - summary = _get_usage_summary(reporter) - expected_input = 50_000 * 15.0 / 1_000_000 - expected_output = 10_000 * 75.0 / 1_000_000 - 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) From a871a971e9554ff2c79927b202121b6f7054eaa4 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Wed, 11 Mar 2026 11:01:44 +0100 Subject: [PATCH 13/15] chore: cleanup usage_tracking_callback --- .../models/shared/usage_tracking_callback.py | 94 +++++++++---------- .../model_providers/test_model_pricing.py | 12 +-- 2 files changed, 49 insertions(+), 57 deletions(-) diff --git a/src/askui/models/shared/usage_tracking_callback.py b/src/askui/models/shared/usage_tracking_callback.py index c35df2f0..a600e24c 100644 --- a/src/askui/models/shared/usage_tracking_callback.py +++ b/src/askui/models/shared/usage_tracking_callback.py @@ -8,11 +8,11 @@ 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 if TYPE_CHECKING: + from askui.models.shared.agent_message_param import UsageParam from askui.models.shared.conversation import Conversation from askui.reporting import Reporter from askui.speaker.speaker import SpeakerResult @@ -63,11 +63,11 @@ def __init__( ) -> None: self._reporter = reporter self._pricing = pricing - self._accumulated_usage = UsageParam() + self._summary = UsageSummary() @override def on_conversation_start(self, conversation: Conversation) -> None: - self._accumulated_usage = UsageParam() + self._summary = UsageSummary() @override def on_step_end( @@ -81,60 +81,29 @@ def on_step_end( @override def on_conversation_end(self, conversation: Conversation) -> 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 - 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 / 1e6 - ) - output_cost = ( - output_tokens * self._pricing.output_cost_per_million_tokens / 1e6 - ) - 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 - ) - 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) + self._reporter.add_usage_summary(self._summary) @property - def accumulated_usage(self) -> UsageParam: + def accumulated_usage(self) -> UsageSummary: """Current accumulated usage statistics.""" - return self._accumulated_usage + return self._summary 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 + # Add step tokens to running totals (None counts as 0) + self._summary.input_tokens = (self._summary.input_tokens or 0) + ( + step_usage.input_tokens or 0 + ) + self._summary.output_tokens = (self._summary.output_tokens or 0) + ( + step_usage.output_tokens or 0 + ) + self._summary.cache_creation_input_tokens = ( + self._summary.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 + self._summary.cache_read_input_tokens = ( + self._summary.cache_read_input_tokens or 0 ) + (step_usage.cache_read_input_tokens or 0) + # Record per-step token counts on the current OTel span current_span = trace.get_current_span() current_span.set_attributes( { @@ -146,3 +115,30 @@ def _accumulate(self, step_usage: UsageParam) -> None: "cache_read_input_tokens": (step_usage.cache_read_input_tokens or 0), } ) + + # Update costs from updated totals if pricing values are set + if ( + self._pricing + and self._pricing.input_cost_per_million_tokens + and self._pricing.output_cost_per_million_tokens + ): + input_cost = ( + self._summary.input_tokens + * self._pricing.input_cost_per_million_tokens + / 1e6 + ) + output_cost = ( + self._summary.output_tokens + * self._pricing.output_cost_per_million_tokens + / 1e6 + ) + self._summary.input_cost = input_cost + self._summary.output_cost = output_cost + self._summary.total_cost = input_cost + output_cost + self._summary.currency = self._pricing.currency + self._summary.input_cost_per_million_tokens = ( + self._pricing.input_cost_per_million_tokens + ) + self._summary.output_cost_per_million_tokens = ( + self._pricing.output_cost_per_million_tokens + ) diff --git a/tests/unit/model_providers/test_model_pricing.py b/tests/unit/model_providers/test_model_pricing.py index e8381a40..7fa00c6d 100644 --- a/tests/unit/model_providers/test_model_pricing.py +++ b/tests/unit/model_providers/test_model_pricing.py @@ -105,9 +105,8 @@ def test_cost_calculation( output_cost_per_million_tokens=output_rate, ) callback, reporter = self._make_callback(pricing) - callback._accumulated_usage = UsageParam( - input_tokens=input_tokens, - output_tokens=output_tokens, + callback._accumulate( + UsageParam(input_tokens=input_tokens, output_tokens=output_tokens) ) callback.on_conversation_end(MagicMock()) @@ -123,10 +122,7 @@ def test_cost_calculation( def test_no_cost_when_pricing_none(self) -> None: callback, reporter = self._make_callback(pricing=None) - callback._accumulated_usage = UsageParam( - input_tokens=500, - output_tokens=200, - ) + callback._accumulate(UsageParam(input_tokens=500, output_tokens=200)) callback.on_conversation_end(MagicMock()) summary = _get_usage_summary(reporter) @@ -139,7 +135,7 @@ def test_none_tokens_treated_as_zero(self) -> None: output_cost_per_million_tokens=15.0, ) callback, reporter = self._make_callback(pricing) - callback._accumulated_usage = UsageParam() + callback._accumulate(UsageParam()) callback.on_conversation_end(MagicMock()) summary = _get_usage_summary(reporter) From 47f78b2d3d3c01bca0a510e093301dc666530f63 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Wed, 11 Mar 2026 11:04:34 +0100 Subject: [PATCH 14/15] fix: remove pricing from askui_vlm_provider --- .../model_providers/askui_vlm_provider.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/src/askui/model_providers/askui_vlm_provider.py b/src/askui/model_providers/askui_vlm_provider.py index dcc928d2..d149deff 100644 --- a/src/askui/model_providers/askui_vlm_provider.py +++ b/src/askui/model_providers/askui_vlm_provider.py @@ -17,7 +17,6 @@ ) 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" @@ -38,12 +37,6 @@ class AskUIVlmProvider(VlmProvider): `"claude-sonnet-4-6"`. client (Anthropic | None, optional): Pre-configured Anthropic client. If provided, `workspace_id` and `token` are ignored. - input_cost_per_million_tokens (float | None, optional): Override - cost in USD per 1M input tokens. Both cost params must be set - to override the built-in defaults. - output_cost_per_million_tokens (float | None, optional): Override - cost in USD per 1M output tokens. - Example: ```python from askui import AgentSettings, ComputerAgent @@ -64,30 +57,18 @@ def __init__( askui_settings: AskUiInferenceApiSettings | None = None, model_id: str | None = None, client: Anthropic | None = None, - input_cost_per_million_tokens: float | None = None, - output_cost_per_million_tokens: float | None = None, ) -> None: self._askui_settings = askui_settings or AskUiInferenceApiSettings() self._model_id_value = ( model_id or os.environ.get("VLM_PROVIDER_MODEL_ID") or _DEFAULT_MODEL_ID ) self._injected_client = client - self._pricing = ModelPricing.for_model( - self._model_id_value, - input_cost_per_million_tokens=input_cost_per_million_tokens, - output_cost_per_million_tokens=output_cost_per_million_tokens, - ) @property @override def model_id(self) -> str: return self._model_id_value - @property - @override - def pricing(self) -> ModelPricing | None: - return self._pricing - @cached_property def _messages_api(self) -> AnthropicMessagesApi: """Lazily initialise the AnthropicMessagesApi on first use.""" From a002697de31525fb2702f2f6b4885a1fde4f14c8 Mon Sep 17 00:00:00 2001 From: philipph-askui Date: Thu, 12 Mar 2026 09:23:24 +0100 Subject: [PATCH 15/15] chore: return early if pricing values are not set in usage tacking callback --- .../models/shared/usage_tracking_callback.py | 44 ++++++++++--------- 1 file changed, 23 insertions(+), 21 deletions(-) diff --git a/src/askui/models/shared/usage_tracking_callback.py b/src/askui/models/shared/usage_tracking_callback.py index a600e24c..5e245434 100644 --- a/src/askui/models/shared/usage_tracking_callback.py +++ b/src/askui/models/shared/usage_tracking_callback.py @@ -117,28 +117,30 @@ def _accumulate(self, step_usage: UsageParam) -> None: ) # Update costs from updated totals if pricing values are set - if ( + if not ( self._pricing and self._pricing.input_cost_per_million_tokens and self._pricing.output_cost_per_million_tokens ): - input_cost = ( - self._summary.input_tokens - * self._pricing.input_cost_per_million_tokens - / 1e6 - ) - output_cost = ( - self._summary.output_tokens - * self._pricing.output_cost_per_million_tokens - / 1e6 - ) - self._summary.input_cost = input_cost - self._summary.output_cost = output_cost - self._summary.total_cost = input_cost + output_cost - self._summary.currency = self._pricing.currency - self._summary.input_cost_per_million_tokens = ( - self._pricing.input_cost_per_million_tokens - ) - self._summary.output_cost_per_million_tokens = ( - self._pricing.output_cost_per_million_tokens - ) + return + + input_cost = ( + self._summary.input_tokens + * self._pricing.input_cost_per_million_tokens + / 1e6 + ) + output_cost = ( + self._summary.output_tokens + * self._pricing.output_cost_per_million_tokens + / 1e6 + ) + self._summary.input_cost = input_cost + self._summary.output_cost = output_cost + self._summary.total_cost = input_cost + output_cost + self._summary.currency = self._pricing.currency + self._summary.input_cost_per_million_tokens = ( + self._pricing.input_cost_per_million_tokens + ) + self._summary.output_cost_per_million_tokens = ( + self._pricing.output_cost_per_million_tokens + )