From 068fb3572ffa7c857aa629e22b3ca2552ff14d0c Mon Sep 17 00:00:00 2001 From: Jeroen Ubbink Date: Thu, 11 Jun 2026 21:17:11 +0200 Subject: [PATCH] refactor: drop unused session parameter from calculate_cost The session was needed when model pricing lived in the DB (73d3613) and has been dead since pricing moved to the in-memory model map (0da08fb), yet every caller was still obliged to supply one. get_x_cashu_cost even opened a DB session per x-cashu request solely to feed it. Co-Authored-By: Claude Fable 5 --- routstr/auth.py | 2 +- routstr/payment/cost_calculation.py | 2 - routstr/upstream/base.py | 84 ++++++++++----------- tests/unit/test_cache_pricing.py | 8 +- tests/unit/test_cost_calculation_caching.py | 80 +++++++++----------- tests/unit/test_usage_normalization.py | 4 +- 6 files changed, 84 insertions(+), 96 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index a5e1d396..2c1f0207 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -772,7 +772,7 @@ async def adjust_payment_for_tokens( extra={"error": str(e), "fee_msats": fee_msats}, ) - match await calculate_cost(response_data, deducted_max_cost, session, usage=usage): + match await calculate_cost(response_data, deducted_max_cost, usage=usage): 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 3dd52854..774afd02 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -3,7 +3,6 @@ import math from pydantic.v1 import BaseModel from ..core import get_logger -from ..core.db import AsyncSession from ..core.settings import settings from .price import sats_usd_price from .usage import NormalizedUsage, normalize_usage, parse_token_count @@ -45,7 +44,6 @@ class CostDataError(BaseModel): async def calculate_cost( response_data: dict, max_cost: int, - session: AsyncSession, usage: NormalizedUsage | None = None, ) -> CostData | MaxCostData | CostDataError: """Calculate the cost of an API request based on token usage. diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 56cd80b8..6c2d11db 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -3071,49 +3071,47 @@ class BaseUpstreamProvider: extra={"model": model, "has_usage": "usage" in response_data}, ) - async with create_session() as session: - match await calculate_cost( - response_data, - max_cost_for_model, - session, - usage=self.normalize_usage(response_data.get("usage")), - ): - case MaxCostData() as cost: - logger.debug( - "Using max cost pricing", - extra={"model": model, "max_cost_msats": cost.total_msats}, - ) - return cost - case CostData() as cost: - logger.debug( - "Using token-based pricing", - extra={ - "model": model, - "total_cost_msats": cost.total_msats, - "input_msats": cost.input_msats, - "output_msats": cost.output_msats, - }, - ) - return cost - case CostDataError() as error: - logger.error( - "Cost calculation error", - extra={ - "model": model, - "error_message": error.message, - "error_code": error.code, - }, - ) - raise HTTPException( - status_code=400, - detail={ - "error": { - "message": error.message, - "type": "invalid_request_error", - "code": error.code, - } - }, - ) + match await calculate_cost( + response_data, + max_cost_for_model, + usage=self.normalize_usage(response_data.get("usage")), + ): + case MaxCostData() as cost: + logger.debug( + "Using max cost pricing", + extra={"model": model, "max_cost_msats": cost.total_msats}, + ) + return cost + case CostData() as cost: + logger.debug( + "Using token-based pricing", + extra={ + "model": model, + "total_cost_msats": cost.total_msats, + "input_msats": cost.input_msats, + "output_msats": cost.output_msats, + }, + ) + return cost + case CostDataError() as error: + logger.error( + "Cost calculation error", + extra={ + "model": model, + "error_message": error.message, + "error_code": error.code, + }, + ) + raise HTTPException( + status_code=400, + detail={ + "error": { + "message": error.message, + "type": "invalid_request_error", + "code": error.code, + } + }, + ) return None async def send_refund( diff --git a/tests/unit/test_cache_pricing.py b/tests/unit/test_cache_pricing.py index 9d73d231..14eb7965 100644 --- a/tests/unit/test_cache_pricing.py +++ b/tests/unit/test_cache_pricing.py @@ -13,7 +13,7 @@ Specifies two things: """ import os -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import Mock, patch import litellm import pytest @@ -178,7 +178,7 @@ async def test_deepseek_cache_hits_billed_at_cache_rate(model_pricing: Mock) -> } with patch("routstr.proxy.get_model_instance", return_value=model_pricing): - result = await calculate_cost(response, max_cost=100000, session=AsyncMock()) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) # 1000 input @ 1 msat + 9000 cache reads @ 0.1 msat + 500 output @ 2 msat @@ -203,7 +203,7 @@ async def test_anthropic_cache_write_billed_at_write_rate(model_pricing: Mock) - } with patch("routstr.proxy.get_model_instance", return_value=model_pricing): - result = await calculate_cost(response, max_cost=100000, session=AsyncMock()) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) # 300 @ 1 + 500 @ 0.1 + 2000 @ 1.25 + 100 @ 2 @@ -234,7 +234,7 @@ async def test_missing_cache_rate_falls_back_to_input_rate( } with patch("routstr.proxy.get_model_instance", return_value=model): - result = await calculate_cost(response, max_cost=100000, session=AsyncMock()) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) # 1000 @ 1 + 9000 @ 1 (fallback) + 500 @ 2 diff --git a/tests/unit/test_cost_calculation_caching.py b/tests/unit/test_cost_calculation_caching.py index 8b8e09ea..2eb7a7c0 100644 --- a/tests/unit/test_cost_calculation_caching.py +++ b/tests/unit/test_cost_calculation_caching.py @@ -5,7 +5,7 @@ edge cases, and billing accuracy. """ import os -from unittest.mock import AsyncMock, patch +from unittest.mock import patch import pytest @@ -17,12 +17,6 @@ from routstr.core.settings import settings from routstr.payment.cost_calculation import CostData, MaxCostData, calculate_cost -@pytest.fixture -def mock_session() -> AsyncMock: - """Mock AsyncSession for cost calculation tests.""" - return AsyncMock() - - @pytest.fixture(autouse=True) def mock_fixed_pricing(monkeypatch: pytest.MonkeyPatch) -> None: """Mock settings and price to use fixed pricing.""" @@ -42,7 +36,7 @@ def patch_sats_usd_price() -> None: # type: ignore[misc] # Test 1: OpenAI Cache Format # ============================================================================ @pytest.mark.asyncio -async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None: +async def test_openai_cache_subtraction() -> None: """OpenAI includes cached_tokens in prompt_tokens, subtract them.""" response = { "model": "gpt-4", @@ -54,7 +48,7 @@ async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None: } } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 1000 # 2000 - 1000 @@ -66,7 +60,7 @@ async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None: # Test 2: Anthropic Cache Format # ============================================================================ @pytest.mark.asyncio -async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_anthropic_cache_additive(mock_fixed_pricing: None) -> None: """Anthropic cache tokens are separate (additive) from input_tokens.""" response = { "model": "claude-3-5-sonnet", @@ -77,7 +71,7 @@ async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pric "cache_read_input_tokens": 0, } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 500 @@ -90,7 +84,7 @@ async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pric # Test 3: Invalid Cache (Edge Case) # ============================================================================ @pytest.mark.asyncio -async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_cache_read_exceeds_prompt_tokens(mock_fixed_pricing: None) -> None: """Handle buggy upstream reporting cached > prompt_tokens.""" response = { "model": "gpt-4", @@ -102,7 +96,7 @@ async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fi } } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) # Should not go negative assert isinstance(result, CostData) @@ -115,7 +109,7 @@ async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fi # Test 4: Malformed Token Values # ============================================================================ @pytest.mark.asyncio -async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_malformed_cache_tokens_coerce_to_zero(mock_fixed_pricing: None) -> None: """Handle non-numeric cache token values.""" response = { "model": "gpt-4", @@ -128,7 +122,7 @@ async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mo } } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) # Both should coerce to 0 assert isinstance(result, CostData) @@ -140,7 +134,7 @@ async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mo # Test 5: Anthropic Cache Not Subtracted # ============================================================================ @pytest.mark.asyncio -async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_anthropic_cache_not_subtracted(mock_fixed_pricing: None) -> None: """Anthropic cache fields should NOT be subtracted from input_tokens.""" response = { "model": "claude-3-5-sonnet", @@ -150,7 +144,7 @@ async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixe "cache_read_input_tokens": 200, # ← Additive, don't subtract } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) # Anthropic: input_tokens stays as-is assert isinstance(result, CostData) @@ -162,7 +156,7 @@ async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixe # Test 6: Only Cache Read, No Regular Input # ============================================================================ @pytest.mark.asyncio -async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_only_cache_read_tokens(mock_fixed_pricing: None) -> None: """Handle response with only cache read tokens.""" response = { "model": "gpt-4", @@ -174,7 +168,7 @@ async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricin } } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 0 # max(0, 0 - 1000) @@ -186,7 +180,7 @@ async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricin # Test 7: Only Cache Creation # ============================================================================ @pytest.mark.asyncio -async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_only_cache_creation_tokens(mock_fixed_pricing: None) -> None: """Handle response with only cache creation tokens (Anthropic).""" response = { "model": "claude-3-5-sonnet", @@ -197,7 +191,7 @@ async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pr "cache_read_input_tokens": 0, } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 500 @@ -210,7 +204,7 @@ async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pr # Test 8: Both Cache Read and Creation # ============================================================================ @pytest.mark.asyncio -async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_both_cache_read_and_creation(mock_fixed_pricing: None) -> None: """Handle response with both cache read and creation.""" response = { "model": "claude-3-5-sonnet", @@ -221,7 +215,7 @@ async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_ "cache_read_input_tokens": 500, } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 300 @@ -234,7 +228,7 @@ async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_ # Test 9: Token Field Fallback # ============================================================================ @pytest.mark.asyncio -async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_token_field_fallback_order(mock_fixed_pricing: None) -> None: """Verify fallback order for token extraction.""" # When prompt_tokens is not present, fall back to input_tokens response = { @@ -244,7 +238,7 @@ async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pr "completion_tokens": 50, } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 250 @@ -255,7 +249,7 @@ async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pr # Test 10: Float Token Values # ============================================================================ @pytest.mark.asyncio -async def test_float_token_values_coerced_to_int(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_float_token_values_coerced_to_int(mock_fixed_pricing: None) -> None: """Handle float token values by converting to int.""" response = { "model": "gpt-4", @@ -265,7 +259,7 @@ async def test_float_token_values_coerced_to_int(mock_session: AsyncMock, mock_f "cache_read_input_tokens": 25.9, # Float } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 100 # Floored @@ -277,7 +271,7 @@ async def test_float_token_values_coerced_to_int(mock_session: AsyncMock, mock_f # Test 11: Boolean Cache Tokens # ============================================================================ @pytest.mark.asyncio -async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_boolean_cache_tokens_coerced_to_zero(mock_fixed_pricing: None) -> None: """Handle boolean cache token values by coercing to zero.""" response = { "model": "gpt-4", @@ -287,7 +281,7 @@ async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, moc "cache_read_input_tokens": True, # Boolean } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.cache_read_input_tokens == 0 # Boolean coerced to 0 @@ -298,7 +292,7 @@ async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, moc # Test 12: Zero Cache Tokens # ============================================================================ @pytest.mark.asyncio -async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_zero_cache_tokens(mock_fixed_pricing: None) -> None: """Handle explicit zero cache tokens.""" response = { "model": "gpt-4", @@ -310,7 +304,7 @@ async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: No } } } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.cache_read_input_tokens == 0 @@ -326,7 +320,7 @@ async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: No # them as regular input is a large overcharge. # ============================================================================ @pytest.mark.asyncio -async def test_deepseek_cache_hit_tokens_extracted(mock_session: AsyncMock) -> None: +async def test_deepseek_cache_hit_tokens_extracted() -> None: """DeepSeek cache hits are extracted and removed from regular input. Payload shape verbatim from the DeepSeek API reference (usage object). @@ -341,7 +335,7 @@ async def test_deepseek_cache_hit_tokens_extracted(mock_session: AsyncMock) -> N "prompt_cache_miss_tokens": 1000, }, } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 1000 # only the cache misses @@ -350,7 +344,7 @@ async def test_deepseek_cache_hit_tokens_extracted(mock_session: AsyncMock) -> N @pytest.mark.asyncio -async def test_deepseek_all_tokens_cached(mock_session: AsyncMock) -> None: +async def test_deepseek_all_tokens_cached() -> None: """A fully cached DeepSeek prompt bills zero regular input tokens.""" response = { "model": "deepseek-chat", @@ -361,7 +355,7 @@ async def test_deepseek_all_tokens_cached(mock_session: AsyncMock) -> None: "prompt_cache_miss_tokens": 0, }, } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 0 @@ -369,7 +363,7 @@ async def test_deepseek_all_tokens_cached(mock_session: AsyncMock) -> None: @pytest.mark.asyncio -async def test_dialect_precedence_never_double_subtracts(mock_session: AsyncMock) -> None: +async def test_dialect_precedence_never_double_subtracts() -> None: """If a vendor emits both OpenAI-style and DeepSeek-style cache fields for the same cached tokens, they are counted once, not subtracted twice.""" response = { @@ -382,7 +376,7 @@ async def test_dialect_precedence_never_double_subtracts(mock_session: AsyncMock "prompt_cache_miss_tokens": 1000, }, } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 1000 # 10000 - 9000, applied exactly once @@ -390,7 +384,7 @@ async def test_dialect_precedence_never_double_subtracts(mock_session: AsyncMock @pytest.mark.asyncio -async def test_deepseek_malformed_hit_tokens_coerce_to_zero(mock_session: AsyncMock) -> None: +async def test_deepseek_malformed_hit_tokens_coerce_to_zero() -> None: """Malformed DeepSeek cache fields degrade to billing all input at full rate instead of crashing or going negative.""" response = { @@ -402,7 +396,7 @@ async def test_deepseek_malformed_hit_tokens_coerce_to_zero(mock_session: AsyncM "prompt_cache_miss_tokens": -5, }, } - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, CostData) assert result.input_tokens == 1000 @@ -413,10 +407,10 @@ async def test_deepseek_malformed_hit_tokens_coerce_to_zero(mock_session: AsyncM # Test 13: Missing Usage Block # ============================================================================ @pytest.mark.asyncio -async def test_missing_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_missing_usage_block(mock_fixed_pricing: None) -> None: """When usage is missing, return MaxCostData with zero tokens.""" response = {"model": "gpt-4", "choices": [{"message": {"content": "test"}}]} - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, MaxCostData) assert result.input_tokens == 0 @@ -428,10 +422,10 @@ async def test_missing_usage_block(mock_session: AsyncMock, mock_fixed_pricing: # Test 14: Null Usage Block # ============================================================================ @pytest.mark.asyncio -async def test_null_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None: +async def test_null_usage_block(mock_fixed_pricing: None) -> None: """When usage is null, return MaxCostData with zero tokens.""" response = {"model": "gpt-4", "usage": None} - result = await calculate_cost(response, max_cost=100000, session=mock_session) + result = await calculate_cost(response, max_cost=100000) assert isinstance(result, MaxCostData) assert result.input_tokens == 0 diff --git a/tests/unit/test_usage_normalization.py b/tests/unit/test_usage_normalization.py index 8401c355..0412d55e 100644 --- a/tests/unit/test_usage_normalization.py +++ b/tests/unit/test_usage_normalization.py @@ -10,7 +10,7 @@ needs no vendor knowledge of its own. """ import os -from unittest.mock import AsyncMock, patch +from unittest.mock import patch import pytest @@ -175,7 +175,6 @@ async def test_provider_override_is_honored_by_calculate_cost() -> None: result = await calculate_cost( response, max_cost=100000, - session=AsyncMock(), usage=provider.normalize_usage(response["usage"]), ) @@ -197,7 +196,6 @@ async def test_explicit_usage_param_wins_over_response_extraction() -> None: result = await calculate_cost( response, max_cost=100000, - session=AsyncMock(), usage=NormalizedUsage(input_tokens=10, output_tokens=5), )