From 14785e4cdedfe60fccf1d0d240e8f0052d4fe025 Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Wed, 17 Jun 2026 23:40:16 +0200 Subject: [PATCH] deepseek cache calculation fix --- routstr/payment/cost_calculation.py | 4 +- routstr/upstream/base.py | 65 ++++++++++++-------- tests/unit/test_cache_pricing.py | 6 +- tests/unit/test_streaming_sse_providers.py | 71 ++++++++++++++++++++++ 4 files changed, 117 insertions(+), 29 deletions(-) diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index 27b04235..a635a0d4 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -418,9 +418,11 @@ def _calculate_from_tokens( }, ) + visible_input_msats = int(calc_input_msats + calc_cache_read_msats) + return CostData( base_msats=0, - input_msats=int(calc_input_msats), + input_msats=visible_input_msats, output_msats=int(calc_output_msats), total_msats=token_based_cost, total_usd=total_usd, diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 1503cc67..05cea0a5 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -39,6 +39,7 @@ from ..payment.models import ( list_models, ) from ..payment.price import sats_usd_price +from ..payment.usage import normalize_usage from ..wallet import recieve_token, send_token from . import messages_dispatch from .cache_breakpoints import ( @@ -157,48 +158,56 @@ class BaseUpstreamProvider: @staticmethod def _fold_cache_into_input_tokens(usage: object) -> None: - """Fold cache token counts into ``input_tokens`` / ``prompt_tokens``. + """Fold additive cache token counts into Anthropic ``input_tokens``. Cost calculation has already used the per-bucket counts to bill the - request correctly; what the client sees in the visible token total - should be a single rolled-up prompt count *including* the cache - portion. The standalone ``cache_read_input_tokens`` / + request correctly; what the client sees in Anthropic-shaped visible + token totals should be a single rolled-up input count *including* the + cache portion. OpenAI-compatible ``prompt_tokens`` is already inclusive + (DeepSeek, OpenAI, OpenRouter, litellm), so adding cache fields there + would double-count. + + The standalone ``cache_read_input_tokens`` / ``cache_creation_input_tokens`` fields are left in place for clients that want the breakdown. - - For Anthropic-shaped responses (``input_tokens`` present), the cache - fields are forced to ``0`` when the upstream omitted them, so the - client always sees a consistent shape. """ if not isinstance(usage, dict): return - # Normalise missing cache fields to 0 on Anthropic-shaped usage so - # downstream consumers can rely on them being present. - if "input_tokens" in usage: - usage.setdefault("cache_read_input_tokens", 0) - usage.setdefault("cache_creation_input_tokens", 0) + # ``prompt_tokens`` is an inclusive OpenAI-compatible grand total. + # Fold only native Anthropic-style usage where ``input_tokens`` excludes + # cache reads/writes. + if "input_tokens" not in usage or "prompt_tokens" in usage: + return + + usage.setdefault("cache_read_input_tokens", 0) + usage.setdefault("cache_creation_input_tokens", 0) try: cache_read = int(usage.get("cache_read_input_tokens") or 0) cache_creation = int(usage.get("cache_creation_input_tokens") or 0) + input_tokens = int(usage.get("input_tokens") or 0) except (TypeError, ValueError): return extra = cache_read + cache_creation - if extra <= 0: + if extra > 0: + usage["input_tokens"] = input_tokens + extra + + @staticmethod + def _add_normalized_usage_fields(response_json: object) -> None: + """Preserve raw usage while adding canonical fields for billing/display.""" + if not isinstance(response_json, dict): return - if "input_tokens" in usage: - try: - usage["input_tokens"] = int(usage.get("input_tokens") or 0) + extra - except (TypeError, ValueError): - pass - if "prompt_tokens" in usage: - try: - usage["prompt_tokens"] = ( - int(usage.get("prompt_tokens") or 0) + extra - ) - except (TypeError, ValueError): - pass + usage = response_json.get("usage") + if not isinstance(usage, dict): + return + normalized = normalize_usage(usage) + if normalized is None: + return + usage.setdefault("input_tokens", normalized.input_tokens) + usage.setdefault("output_tokens", normalized.output_tokens) + usage.setdefault("cache_read_input_tokens", normalized.cache_read_tokens) + usage.setdefault("cache_creation_input_tokens", normalized.cache_write_tokens) def _apply_provider_field(self, response_json: object) -> None: """Stamp the routstr ``provider`` field onto an upstream response payload. @@ -858,6 +867,7 @@ class BaseUpstreamProvider: k: v for k, v in obj.items() if k != "choices" } usage_chunk_data["choices"] = [] + self._add_normalized_usage_fields(usage_chunk_data) # Forward the content now, without usage, so token # usage is reported exactly once (in the trailer). forward = {k: v for k, v in obj.items() if k != "usage"} @@ -869,6 +879,7 @@ class BaseUpstreamProvider: ) return usage_chunk_data = obj + self._add_normalized_usage_fields(usage_chunk_data) return yield prefix + b"data: " + json.dumps(obj).encode() + b"\n\n" else: @@ -906,6 +917,8 @@ class BaseUpstreamProvider: if fresh_key: cost_data: dict try: + if usage_chunk_data is not None: + self._add_normalized_usage_fields(usage_chunk_data) adjustment_input = ( usage_chunk_data if usage_chunk_data is not None diff --git a/tests/unit/test_cache_pricing.py b/tests/unit/test_cache_pricing.py index 14eb7965..80e027d8 100644 --- a/tests/unit/test_cache_pricing.py +++ b/tests/unit/test_cache_pricing.py @@ -181,8 +181,10 @@ async def test_deepseek_cache_hits_billed_at_cache_rate(model_pricing: Mock) -> 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 - assert result.input_msats == 1000 + # 1000 input @ 1 msat + 9000 cache reads @ 0.1 msat + 500 output @ 2 msat. + # Client-visible input_msats includes cache-read input cost for display, + # while cache_read_msats keeps the detailed breakdown. + assert result.input_msats == 1900 assert result.cache_read_msats == 900 assert result.output_msats == 1000 assert result.total_msats == 2900 diff --git a/tests/unit/test_streaming_sse_providers.py b/tests/unit/test_streaming_sse_providers.py index 77d02903..f86d981c 100644 --- a/tests/unit/test_streaming_sse_providers.py +++ b/tests/unit/test_streaming_sse_providers.py @@ -20,6 +20,7 @@ comment ever reaches the client. That invariant is exactly what the buggy import json from collections.abc import AsyncGenerator +from typing import cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -107,6 +108,76 @@ def _assert_clean(out: list[bytes]) -> list[dict]: return objs +@pytest.mark.asyncio +async def test_deepseek_usage_chunk_is_normalized_before_billing() -> None: + """DeepSeek stream trailers keep raw fields and add canonical billing fields.""" + usage = { + "prompt_tokens": 10000, + "completion_tokens": 500, + "total_tokens": 10500, + "prompt_cache_hit_tokens": 9000, + "prompt_cache_miss_tokens": 1000, + } + chunks = [ + b'data: {"id":"ds","model":"deepseek-chat","choices":[{"delta":{"content":"ok"}}]}\n\n', + b"data: " + + json.dumps( + { + "id": "ds", + "model": "deepseek-chat", + "choices": [], + "usage": usage, + } + ).encode() + + b"\n\n", + b"data: [DONE]\n\n", + ] + + await _drive(chunks) + + adjust_mock = cast(AsyncMock, base.adjust_payment_for_tokens) + adjustment_input = adjust_mock.call_args.args[1] + billed_usage = adjustment_input["usage"] + assert billed_usage["prompt_tokens"] == 10000 + assert billed_usage["prompt_cache_hit_tokens"] == 9000 + assert billed_usage["prompt_cache_miss_tokens"] == 1000 + assert billed_usage["input_tokens"] == 1000 + assert billed_usage["output_tokens"] == 500 + assert billed_usage["cache_read_input_tokens"] == 9000 + assert billed_usage["cache_creation_input_tokens"] == 0 + + +@pytest.mark.asyncio +async def test_fold_cache_tokens_does_not_double_count_inclusive_prompt_tokens() -> None: + """Visible usage mutation must not inflate OpenAI-compatible prompt totals.""" + usage = { + "prompt_tokens": 10000, + "completion_tokens": 500, + "cache_read_input_tokens": 9000, + "prompt_cache_hit_tokens": 9000, + } + + BaseUpstreamProvider._fold_cache_into_input_tokens(usage) + + assert usage["prompt_tokens"] == 10000 + assert usage["cache_read_input_tokens"] == 9000 + + +@pytest.mark.asyncio +async def test_fold_cache_tokens_still_rolls_up_anthropic_input_tokens() -> None: + """Anthropic native input_tokens excludes cache and still needs rollup.""" + usage = { + "input_tokens": 1000, + "output_tokens": 500, + "cache_read_input_tokens": 9000, + "cache_creation_input_tokens": 200, + } + + BaseUpstreamProvider._fold_cache_into_input_tokens(usage) + + assert usage["input_tokens"] == 10200 + + @pytest.mark.asyncio async def test_openai_style_plain_stream() -> None: """OpenAI / Groq / Fireworks / xAI / Perplexity: plain data + [DONE]."""