diff --git a/src/agentlab/llm/tracking.py b/src/agentlab/llm/tracking.py index bc8e31715..7238621c2 100644 --- a/src/agentlab/llm/tracking.py +++ b/src/agentlab/llm/tracking.py @@ -320,8 +320,8 @@ def get_effective_cost_from_antrophic_api(self, response) -> float: usage = getattr(response, "usage", {}) new_input_tokens = getattr(usage, "input_tokens", 0) # new input tokens output_tokens = getattr(usage, "output_tokens", 0) - cache_read_tokens = getattr(usage, "cache_input_tokens", 0) - cache_write_tokens = getattr(usage, "cache_creation_input_tokens", 0) + cache_read_tokens = getattr(usage, "cache_read_input_tokens", 0) or 0 + cache_write_tokens = getattr(usage, "cache_creation_input_tokens", 0) or 0 cache_read_cost = self.input_cost * ANTHROPIC_CACHE_PRICING_FACTOR["cache_read_tokens"] cache_write_cost = self.input_cost * ANTHROPIC_CACHE_PRICING_FACTOR["cache_write_tokens"] diff --git a/tests/llm/test_response_api.py b/tests/llm/test_response_api.py index 6174f64de..d9368f36b 100644 --- a/tests/llm/test_response_api.py +++ b/tests/llm/test_response_api.py @@ -168,7 +168,7 @@ def create_mock_anthropic_response( response.usage = MagicMock() response.usage.input_tokens = input_tokens response.usage.output_tokens = output_tokens - response.usage.cache_input_tokens = 0 + response.usage.cache_read_input_tokens = 0 response.usage.cache_creation_input_tokens = 0 return response diff --git a/tests/llm/test_tracking.py b/tests/llm/test_tracking.py index 0809c54a6..328e41863 100644 --- a/tests/llm/test_tracking.py +++ b/tests/llm/test_tracking.py @@ -1,8 +1,17 @@ import os import time from functools import partial +from types import SimpleNamespace +from unittest.mock import patch import pytest +from anthropic.types import Usage as AnthropicUsage +from openai.types.completion_usage import CompletionUsage, PromptTokensDetails +from openai.types.responses.response_usage import ( + InputTokensDetails, + OutputTokensDetails, + ResponseUsage, +) import agentlab.llm.tracking as tracking from agentlab.llm.chat_api import ( @@ -175,3 +184,195 @@ def test_openrouter_chat_model(): answer = chat_model(messages) assert "5" in answer.get("content") assert tracker.stats["cost"] > 0 + + +INPUT_COST = 2e-6 +OUTPUT_COST = 1e-5 + + +def make_pricing_mixin(pricing_api): + mixin = tracking.TrackAPIPricingMixin() + mixin._pricing_api = pricing_api + mixin.input_cost = INPUT_COST + mixin.output_cost = OUTPUT_COST + return mixin + + +def make_chat_completion_usage(prompt_tokens, completion_tokens, cached_tokens): + details = None if cached_tokens is None else PromptTokensDetails(cached_tokens=cached_tokens) + return CompletionUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=details, + ) + + +def make_response_api_usage(input_tokens, output_tokens, cached_tokens): + return ResponseUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + input_tokens_details=InputTokensDetails(cached_tokens=cached_tokens), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + + +def test_anthropic_effective_cost_with_cache_read_and_write(): + mixin = make_pricing_mixin("anthropic") + usage = AnthropicUsage( + input_tokens=100, + output_tokens=50, + cache_read_input_tokens=1000, + cache_creation_input_tokens=200, + ) + response = SimpleNamespace(usage=usage) + + # 100 new input + 50 output + 1000 cache reads at 0.1x + 200 cache writes at 1.25x + expected = 100 * 2e-6 + 50 * 1e-5 + 1000 * 2e-7 + 200 * 2.5e-6 + assert expected == pytest.approx(1.4e-3) + assert mixin.get_effective_cost_from_antrophic_api(response) == pytest.approx(expected) + assert mixin.get_effective_cost(response) == pytest.approx(expected) + + +def test_anthropic_effective_cost_cache_read_is_discounted(): + mixin = make_pricing_mixin("anthropic") + uncached = SimpleNamespace( + usage=AnthropicUsage( + input_tokens=1000, + output_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + ) + ) + cached = SimpleNamespace( + usage=AnthropicUsage( + input_tokens=0, + output_tokens=0, + cache_read_input_tokens=1000, + cache_creation_input_tokens=0, + ) + ) + + assert mixin.get_effective_cost(uncached) == pytest.approx(1000 * INPUT_COST) + assert mixin.get_effective_cost(cached) == pytest.approx( + 1000 * INPUT_COST * tracking.ANTHROPIC_CACHE_PRICING_FACTOR["cache_read_tokens"] + ) + + +def test_anthropic_effective_cost_without_cache_fields(): + mixin = make_pricing_mixin("anthropic") + response = SimpleNamespace(usage=SimpleNamespace(input_tokens=100, output_tokens=50)) + + assert mixin.get_effective_cost(response) == pytest.approx(100 * 2e-6 + 50 * 1e-5) + + +def test_anthropic_effective_cost_with_null_cache_fields(): + mixin = make_pricing_mixin("anthropic") + usage = AnthropicUsage(input_tokens=100, output_tokens=50) + assert usage.cache_read_input_tokens is None + assert usage.cache_creation_input_tokens is None + + assert mixin.get_effective_cost(SimpleNamespace(usage=usage)) == pytest.approx( + 100 * 2e-6 + 50 * 1e-5 + ) + + +def test_openai_chat_completion_effective_cost_with_cached_tokens(): + mixin = make_pricing_mixin("openai") + response = SimpleNamespace( + usage=make_chat_completion_usage( + prompt_tokens=1000, completion_tokens=100, cached_tokens=600 + ) + ) + + # prompt_tokens includes cached tokens: 400 new at 1x, 600 cached at 0.5x + expected = 400 * 2e-6 + 600 * 1e-6 + 100 * 1e-5 + assert expected == pytest.approx(2.4e-3) + assert mixin.get_effective_cost_from_openai_api(response) == pytest.approx(expected) + assert mixin.get_effective_cost(response) == pytest.approx(expected) + + +def test_openai_chat_completion_effective_cost_without_prompt_tokens_details(): + mixin = make_pricing_mixin("openai") + response = SimpleNamespace( + usage=make_chat_completion_usage( + prompt_tokens=1000, completion_tokens=100, cached_tokens=None + ) + ) + + assert mixin.get_effective_cost(response) == pytest.approx(1000 * 2e-6 + 100 * 1e-5) + + +def test_openai_response_api_effective_cost_with_cached_tokens(): + mixin = make_pricing_mixin("openai") + usage = make_response_api_usage(input_tokens=1000, output_tokens=100, cached_tokens=600) + assert not hasattr(usage, "prompt_tokens_details") + response = SimpleNamespace(usage=usage) + + expected = 400 * 2e-6 + 600 * 1e-6 + 100 * 1e-5 + assert mixin.get_effective_cost_from_openai_api(response) == pytest.approx(expected) + assert mixin.get_effective_cost(response) == pytest.approx(expected) + + +def test_openai_effective_cost_without_usage_is_zero(): + mixin = make_pricing_mixin("openai") + + assert mixin.get_effective_cost(SimpleNamespace(usage=None)) == 0.0 + assert mixin.get_effective_cost(SimpleNamespace()) == 0.0 + + +def test_litellm_effective_cost_delegates_to_completion_cost(): + mixin = make_pricing_mixin("litellm") + response = SimpleNamespace(usage=None) + + with patch.object(tracking, "completion_cost", return_value=0.0123) as mock_completion_cost: + assert mixin.get_effective_cost(response) == 0.0123 + mock_completion_cost.assert_called_once_with(response) + + +@pytest.mark.parametrize("pricing_api", ["openrouter", "unknown", None]) +def test_unsupported_provider_effective_cost_is_zero(pricing_api): + mixin = make_pricing_mixin(pricing_api) + response = SimpleNamespace( + usage=make_chat_completion_usage( + prompt_tokens=1000, completion_tokens=100, cached_tokens=600 + ) + ) + + assert mixin.get_effective_cost(response) == 0.0 + + +class FakePricedModel(tracking.TrackAPIPricingMixin): + def __init__(self, pricing_api, response): + self._pricing_api = pricing_api + self.input_cost = INPUT_COST + self.output_cost = OUTPUT_COST + self._response = response + self.reset_stats() + + def _call_api(self, *args, **kwargs): + return self._response + + def _parse_response(self, response): + return "parsed" + + +def test_call_accumulates_cache_aware_effective_cost_in_stats(): + response = SimpleNamespace( + usage=make_chat_completion_usage( + prompt_tokens=1000, completion_tokens=100, cached_tokens=600 + ) + ) + model = FakePricedModel("openai", response) + + with tracking.set_tracker() as tracker: + assert model() == "parsed" + assert model() == "parsed" + + stats = model.stats.stats_dict + assert stats["n_api_calls"] == 2 + assert stats["usage_cached_tokens"] == 1200 + assert stats["effective_cost"] == pytest.approx(2 * 2.4e-3) + # The tracker cost ignores caching and bills all prompt tokens at the full input price. + assert tracker.stats["cost"] == pytest.approx(2 * (1000 * 2e-6 + 100 * 1e-5))