diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index fc20a12e..cba9ce6e 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -1,4 +1,5 @@ import math +from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal from pydantic.v1 import BaseModel @@ -157,11 +158,18 @@ async def calculate_cost( }, ) try: - input_usd = _coerce_usd( - usage_data.get("cost_details", {}).get("input_cost", 0) + cost_details = usage_data.get("cost_details", {}) + if not isinstance(cost_details, dict): + cost_details = {} + input_usd = _first_usd( + cost_details, + "input_cost", + "upstream_inference_prompt_cost", ) - output_usd = _coerce_usd( - usage_data.get("cost_details", {}).get("output_cost", 0) + output_usd = _first_usd( + cost_details, + "output_cost", + "upstream_inference_completions_cost", ) return _calculate_from_usd_cost( usd_cost, @@ -256,6 +264,15 @@ def _coerce_usd(value: object) -> float: return 0.0 +def _first_usd(source: dict, *fields: str) -> float: + """Return the first positive USD value among equivalent provider fields.""" + for field in fields: + value = _coerce_usd(source.get(field)) + if value > 0: + return value + return 0.0 + + def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float: """Resolve USD cost with clear priority order. @@ -263,7 +280,11 @@ def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float: """ cost_details = usage_data.get("cost_details") if isinstance(cost_details, dict): - cost = _coerce_usd(cost_details.get("total_cost")) + cost = _first_usd( + cost_details, + "total_cost", + "upstream_inference_cost", + ) if cost > 0: return cost @@ -359,19 +380,32 @@ def _calculate_from_usd_cost( ) -> CostData: """Calculate cost from USD figures, deriving input/output split from tokens.""" provider_fee = _resolve_provider_fee(response_data.get("model", "")) - usd_cost = usd_cost * provider_fee - input_usd = input_usd * provider_fee - output_usd = output_usd * provider_fee - sats_per_usd = 1.0 / sats_usd_price() - cost_in_sats = usd_cost * sats_per_usd - cost_in_msats = math.ceil(cost_in_sats * 1000) + fee_decimal = Decimal(str(provider_fee)) + usd_cost_decimal = Decimal(str(usd_cost)) * fee_decimal + input_usd_decimal = Decimal(str(input_usd)) * fee_decimal + output_usd_decimal = Decimal(str(output_usd)) * fee_decimal + sats_usd_decimal = Decimal(str(sats_usd_price())) - if input_usd > 0 or output_usd > 0: + usd_cost = float(usd_cost_decimal) + input_usd = float(input_usd_decimal) + output_usd = float(output_usd_decimal) + cost_in_sats = float(usd_cost_decimal / sats_usd_decimal) + cost_in_msats = int( + (usd_cost_decimal * Decimal(1000) / sats_usd_decimal).to_integral_value( + rounding=ROUND_CEILING + ) + ) + + if input_usd_decimal > 0 or output_usd_decimal > 0: # The total is the authoritative billed amount. Allocating that integer # total proportionally avoids losing sub-millisatoshi remainders when # input and output components are each truncated independently. - component_usd = input_usd + output_usd - input_msats = math.floor(cost_in_msats * input_usd / component_usd) + component_usd = input_usd_decimal + output_usd_decimal + input_msats = int( + ( + Decimal(cost_in_msats) * input_usd_decimal / component_usd + ).to_integral_value(rounding=ROUND_FLOOR) + ) output_msats = cost_in_msats - input_msats else: # Providers often report only a total USD cost. Derive the visible diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index f8ea2d4a..e2318ded 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -1748,14 +1748,19 @@ class BaseUpstreamProvider: total_cost = max( total_cost, _coerce_usd(cd.get("total_cost")), + _coerce_usd(cd.get("upstream_inference_cost")), ) input_cost = max( input_cost, _coerce_usd(cd.get("input_cost")), + _coerce_usd(cd.get("upstream_inference_prompt_cost")), ) output_cost = max( output_cost, _coerce_usd(cd.get("output_cost")), + _coerce_usd( + cd.get("upstream_inference_completions_cost") + ), ) for field in ("total_cost", "cost"): total_cost = max( diff --git a/routstr/upstream/messages_dispatch.py b/routstr/upstream/messages_dispatch.py index 3d689922..affcaf94 100644 --- a/routstr/upstream/messages_dispatch.py +++ b/routstr/upstream/messages_dispatch.py @@ -317,6 +317,27 @@ def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent: total_cost += _coerce_float(usage.get("total_cost")) input_cost += _coerce_float(usage.get("input_cost")) output_cost += _coerce_float(usage.get("output_cost")) + cost_details = usage.get("cost_details") + if isinstance(cost_details, dict): + total_cost = max( + total_cost, + _coerce_float(cost_details.get("total_cost")), + _coerce_float(cost_details.get("upstream_inference_cost")), + ) + input_cost = max( + input_cost, + _coerce_float(cost_details.get("input_cost")), + _coerce_float( + cost_details.get("upstream_inference_prompt_cost") + ), + ) + output_cost = max( + output_cost, + _coerce_float(cost_details.get("output_cost")), + _coerce_float( + cost_details.get("upstream_inference_completions_cost") + ), + ) msg_for_meta = event.get("message") if isinstance(msg_for_meta, dict): @@ -340,14 +361,21 @@ def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent: total_cost = max( total_cost, _coerce_float(root_cost_details.get("total_cost")), + _coerce_float(root_cost_details.get("upstream_inference_cost")), ) input_cost = max( input_cost, _coerce_float(root_cost_details.get("input_cost")), + _coerce_float( + root_cost_details.get("upstream_inference_prompt_cost") + ), ) output_cost = max( output_cost, _coerce_float(root_cost_details.get("output_cost")), + _coerce_float( + root_cost_details.get("upstream_inference_completions_cost") + ), ) event_type = str(event.get("type") or "") diff --git a/tests/unit/test_cost_calculation_caching.py b/tests/unit/test_cost_calculation_caching.py index f9f53f58..dff31d26 100644 --- a/tests/unit/test_cost_calculation_caching.py +++ b/tests/unit/test_cost_calculation_caching.py @@ -532,6 +532,35 @@ async def test_total_only_usd_cost_uses_model_prices_for_component_split( assert result.output_msats == 6671 +@pytest.mark.asyncio +async def test_upstream_inference_cost_details_set_nonzero_components() -> None: + """OpenAI-compatible upstream inference aliases retain their exact split.""" + response = { + "model": "z-ai/glm-5.2-20260616", + "usage": { + "prompt_tokens": 211, + "completion_tokens": 500, + "total_tokens": 711, + "cost": 0.00242155, + "cost_details": { + "upstream_inference_cost": 0.00242155, + "upstream_inference_prompt_cost": 0.00022155, + "upstream_inference_completions_cost": 0.0022, + }, + }, + } + + result = await calculate_cost(response, max_cost=100000) + + assert isinstance(result, CostData) + assert result.input_tokens == 211 + assert result.output_tokens == 500 + assert result.input_msats == 4431 + assert result.output_msats == 44000 + assert result.total_msats == 48431 + assert result.input_msats + result.output_msats == result.total_msats + + # ============================================================================ # Test 13: Missing Usage Block # ============================================================================ diff --git a/tests/unit/test_messages_dispatch_cost_accumulation.py b/tests/unit/test_messages_dispatch_cost_accumulation.py index 3bb22bd8..27a7bf7d 100644 --- a/tests/unit/test_messages_dispatch_cost_accumulation.py +++ b/tests/unit/test_messages_dispatch_cost_accumulation.py @@ -161,6 +161,31 @@ def test_cost_details_extracted_from_event_root() -> None: assert result.output_cost == 0.005 +@pytest.mark.unit +def test_upstream_inference_cost_details_extracted_from_usage() -> None: + """Provider inference aliases inside usage preserve the USD split.""" + event = { + "usage": { + "input_tokens": 211, + "output_tokens": 500, + "cost": 0.00242155, + "cost_details": { + "upstream_inference_cost": 0.00242155, + "upstream_inference_prompt_cost": 0.00022155, + "upstream_inference_completions_cost": 0.0022, + }, + } + } + + result = annotate_event(event, None) + + assert result.input_tokens == 211 + assert result.output_tokens == 500 + assert result.total_cost == 0.00242155 + assert result.input_cost == 0.00022155 + assert result.output_cost == 0.0022 + + # ============================================================================ # Test 7: No Duplicated Dict Lookups # ============================================================================