diff --git a/routstr/auth.py b/routstr/auth.py index 2c1f0207..2385f09d 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -20,7 +20,6 @@ from .payment.cost_calculation import ( MaxCostData, calculate_cost, ) -from .payment.usage import NormalizedUsage from .wallet import credit_balance, deserialize_token_from_string logger = get_logger(__name__) @@ -680,16 +679,14 @@ async def adjust_payment_for_tokens( response_data: dict, session: AsyncSession, deducted_max_cost: int, - usage: NormalizedUsage | None = None, ) -> dict: """ Adjusts the payment based on token usage in the response. This is called after the initial payment and the upstream request is complete. Returns cost data to be included in the response. - ``usage`` carries the upstream provider's normalized token usage (its - ``normalize_usage`` hook); when omitted, the response's usage object is - normalized with the default union parser. + The response's usage object is normalized with the default union parser in + ``calculate_cost``. """ billing_key = await get_billing_key(key, session) model = response_data.get("model", "unknown") @@ -772,7 +769,7 @@ async def adjust_payment_for_tokens( extra={"error": str(e), "fee_msats": fee_msats}, ) - match await calculate_cost(response_data, deducted_max_cost, usage=usage): + match await calculate_cost(response_data, deducted_max_cost): case MaxCostData() as cost: logger.debug( "Using max cost data (no token adjustment)", diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index 774afd02..27b04235 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -5,7 +5,7 @@ from pydantic.v1 import BaseModel from ..core import get_logger from ..core.settings import settings from .price import sats_usd_price -from .usage import NormalizedUsage, normalize_usage, parse_token_count +from .usage import normalize_usage, parse_token_count __all__ = [ "CostData", @@ -44,20 +44,18 @@ class CostDataError(BaseModel): async def calculate_cost( response_data: dict, max_cost: int, - usage: NormalizedUsage | None = None, ) -> CostData | MaxCostData | CostDataError: """Calculate the cost of an API request based on token usage. Args: response_data: Response data containing usage information max_cost: Maximum cost in millisats - usage: Pre-normalized usage from the upstream provider's - ``normalize_usage`` hook. When omitted, the response's usage - object is normalized with the default union parser. This function - holds no vendor-dialect knowledge of its own. Returns: Cost data or error information + + The response's usage object is normalized with the default union parser; + this function holds no vendor-dialect knowledge of its own. """ logger.debug( "Starting cost calculation", @@ -68,8 +66,7 @@ async def calculate_cost( }, ) - if usage is None: - usage = normalize_usage(response_data.get("usage")) + usage = normalize_usage(response_data.get("usage")) if usage is None: logger.warning( diff --git a/routstr/payment/usage.py b/routstr/payment/usage.py index 647fe6b5..02c90055 100644 --- a/routstr/payment/usage.py +++ b/routstr/payment/usage.py @@ -27,8 +27,8 @@ What decides whether cached tokens must be subtracted out of the input count is ``normalize_usage`` maps all of them onto one canonical ``NormalizedUsage`` shape so billing code needs no vendor knowledge. The known dialects' field -names do not collide, so a single union parser is safe; vendors whose fields -would genuinely conflict override ``BaseUpstreamProvider.normalize_usage``. +names do not collide, so a single union parser is safe; a vendor whose fields +would genuinely conflict needs a dedicated branch here. """ from pydantic.v1 import BaseModel diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 10bcbd7f..1503cc67 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -39,7 +39,6 @@ from ..payment.models import ( list_models, ) from ..payment.price import sats_usd_price -from ..payment.usage import NormalizedUsage, normalize_usage from ..wallet import recieve_token, send_token from . import messages_dispatch from .cache_breakpoints import ( @@ -107,17 +106,6 @@ class BaseUpstreamProvider: self._models_cache = [] self._models_by_id = {} - def normalize_usage(self, usage_data: object) -> NormalizedUsage | None: - """Map this provider's usage dialect onto the canonical shape. - - The default union parser covers the known dialects (OpenAI, Anthropic, - DeepSeek), whose field names do not collide. Override in a subclass - when a vendor's usage fields would conflict with another dialect or - need bespoke interpretation; billing code consumes only the canonical - result and holds no vendor knowledge. - """ - return normalize_usage(usage_data) - def get_litellm_provider_prefix(self) -> str: """Resolve the litellm provider prefix for this provider instance. @@ -931,9 +919,6 @@ class BaseUpstreamProvider: adjustment_input, session, max_cost_for_model, - usage=self.normalize_usage( - adjustment_input.get("usage") - ), ) usage_finalized = True except Exception as e: @@ -1076,7 +1061,6 @@ class BaseUpstreamProvider: response_json, session, deducted_max_cost, - usage=self.normalize_usage(response_json.get("usage")), ) await session.refresh(key) @@ -1331,9 +1315,6 @@ class BaseUpstreamProvider: adjustment_input, session, max_cost_for_model, - usage=self.normalize_usage( - adjustment_input.get("usage") - ), ) usage_finalized = True except Exception as e: @@ -1501,7 +1482,6 @@ class BaseUpstreamProvider: response_json, session, deducted_max_cost, - usage=self.normalize_usage(response_json.get("usage")), ) await session.refresh(key) @@ -1700,7 +1680,6 @@ class BaseUpstreamProvider: fallback, new_session, max_cost_for_model, - usage=self.normalize_usage(fallback.get("usage")), ) usage_finalized = True return f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() @@ -1852,9 +1831,6 @@ class BaseUpstreamProvider: combined_data, new_session, max_cost_for_model, - usage=self.normalize_usage( - combined_data.get("usage") - ), ) self.inject_cost_metadata( @@ -1926,7 +1902,6 @@ class BaseUpstreamProvider: response_json, session, deducted_max_cost, - usage=self.normalize_usage(response_json.get("usage")), ) self.inject_cost_metadata(response_json, cost_data, key) @@ -2032,7 +2007,6 @@ class BaseUpstreamProvider: response_json, session, max_cost_for_model, - usage=self.normalize_usage(response_json.get("usage")), ) self.inject_cost_metadata(response_json, cost_data, key) @@ -2174,7 +2148,6 @@ class BaseUpstreamProvider: fallback, new_session, max_cost_for_model, - usage=self.normalize_usage(fallback.get("usage")), ) usage_finalized = True return ( @@ -2248,9 +2221,6 @@ class BaseUpstreamProvider: combined_data, new_session, max_cost_for_model, - usage=self.normalize_usage( - combined_data.get("usage") - ), ) self.inject_cost_metadata( combined_data, cost_data, fresh_key @@ -3112,7 +3082,6 @@ class BaseUpstreamProvider: match await calculate_cost( response_data, max_cost_for_model, - usage=self.normalize_usage(response_data.get("usage")), ): case MaxCostData() as cost: logger.debug( diff --git a/tests/unit/test_usage_normalization.py b/tests/unit/test_usage_normalization.py index f819c2cb..50c6c645 100644 --- a/tests/unit/test_usage_normalization.py +++ b/tests/unit/test_usage_normalization.py @@ -3,39 +3,19 @@ Specifies the seam that keeps vendor usage dialects out of generic billing code: a canonical ``NormalizedUsage`` shape produced by ``routstr.payment.usage.normalize_usage`` (union parser for the known, -non-colliding dialects) and exposed as an overridable -``BaseUpstreamProvider.normalize_usage`` hook for vendors whose fields would -genuinely conflict. ``calculate_cost`` accepts a pre-normalized usage and then -needs no vendor knowledge of its own. +non-colliding dialects). ``calculate_cost`` normalizes the response's usage +object with this parser and needs no vendor knowledge of its own. """ import os -from unittest.mock import patch - -import pytest os.environ.setdefault("UPSTREAM_BASE_URL", "http://test") os.environ.setdefault("UPSTREAM_API_KEY", "test") os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to") -from routstr.core.settings import settings -from routstr.payment.cost_calculation import CostData, calculate_cost +import pytest + from routstr.payment.usage import NormalizedUsage, normalize_usage -from routstr.upstream import BaseUpstreamProvider, GenericUpstreamProvider - - -@pytest.fixture(autouse=True) -def mock_fixed_pricing(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(settings, "fixed_pricing", True) - monkeypatch.setattr(settings, "fixed_per_1k_input_tokens", 0.001) - monkeypatch.setattr(settings, "fixed_per_1k_output_tokens", 0.001) - - -@pytest.fixture(autouse=True) -def patch_sats_usd_price() -> None: # type: ignore[misc] - with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-5): - yield - # ============================================================================ # The union parser: one canonical shape for all known dialects @@ -158,86 +138,3 @@ def test_normalize_usage_never_negative() -> None: assert result is not None assert result.input_tokens == 0 assert result.cache_read_tokens == 150 - - -# ============================================================================ -# The provider hook: default delegates to the union parser, subclasses -# override for vendors whose fields genuinely conflict -# ============================================================================ - - -class ArtificialDumbnessProvider(BaseUpstreamProvider): - """Fictional vendor whose usage dialect collides with nothing we know.""" - - provider_type = "artificial-dumbness" - - def normalize_usage(self, usage_data: object) -> NormalizedUsage | None: - if not isinstance(usage_data, dict): - return None - return NormalizedUsage( - input_tokens=usage_data.get("dumb_tokens_in", 0), - output_tokens=usage_data.get("dumb_tokens_out", 0), - cache_read_tokens=usage_data.get("dumb_tokens_reused", 0), - cache_write_tokens=0, - ) - - -def test_base_provider_hook_delegates_to_union_parser() -> None: - """Without an override, the provider hook equals the module parser, so - DeepSeek-dialect upstreams work through GenericUpstreamProvider with no - subclass.""" - provider = GenericUpstreamProvider(base_url="http://upstream.example") - usage = { - "prompt_tokens": 10000, - "completion_tokens": 500, - "prompt_cache_hit_tokens": 9000, - "prompt_cache_miss_tokens": 1000, - } - assert provider.normalize_usage(usage) == normalize_usage(usage) - - -@pytest.mark.asyncio -async def test_provider_override_is_honored_by_calculate_cost() -> None: - """The escape hatch: a vendor-specific override feeds calculate_cost - through the `usage` parameter, and generic billing needs no knowledge of - the vendor's field names.""" - provider = ArtificialDumbnessProvider(base_url="http://ad.example", api_key="k") - response = { - "model": "dumb-1", - "usage": { - "dumb_tokens_in": 700, - "dumb_tokens_out": 60, - "dumb_tokens_reused": 4300, - }, - } - - result = await calculate_cost( - response, - max_cost=100000, - usage=provider.normalize_usage(response["usage"]), - ) - - assert isinstance(result, CostData) - assert result.input_tokens == 700 - assert result.output_tokens == 60 - assert result.cache_read_input_tokens == 4300 - - -@pytest.mark.asyncio -async def test_explicit_usage_param_wins_over_response_extraction() -> None: - """When a normalized usage is passed, calculate_cost must not re-derive - token counts from the raw response.""" - response = { - "model": "dumb-1", - "usage": {"prompt_tokens": 999999, "completion_tokens": 999999}, - } - - result = await calculate_cost( - response, - max_cost=100000, - usage=NormalizedUsage(input_tokens=10, output_tokens=5), - ) - - assert isinstance(result, CostData) - assert result.input_tokens == 10 - assert result.output_tokens == 5