resolve review comments

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__)
@@ -713,16 +712,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")
@@ -805,7 +802,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)",

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(

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

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(

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