diff --git a/litellm/proxy/pass_through_endpoints/response_body_usage.py b/litellm/proxy/pass_through_endpoints/response_body_usage.py new file mode 100644 index 00000000000..ca1e970ec95 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/response_body_usage.py @@ -0,0 +1,87 @@ +"""Token usage reported in the body of a generic pass-through response. + +A pass-through target LiteLLM has no provider handler for still, most of the +time, returns a ``usage`` object next to its answer, in one of the two shapes +the industry settled on: OpenAI's ``prompt_tokens``/``completion_tokens`` or +Anthropic's ``input_tokens``/``output_tokens``. Without a handler those counts +were dropped, so every spend row and every trace for such a route showed 0 +tokens even though the body said otherwise. +""" + +from collections.abc import Mapping +from typing import Final + +from pydantic import ConfigDict, TypeAdapter, ValidationError + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.types.utils import Usage + +_PROMPT_TOKEN_KEYS: Final = ("prompt_tokens", "input_tokens") +_COMPLETION_TOKEN_KEYS: Final = ("completion_tokens", "output_tokens") +_TOTAL_TOKEN_KEYS: Final = ("total_tokens",) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + +def _json_object(value: object) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +def _token_count(usage: Mapping[str, object], keys: tuple[str, ...]) -> int | None: + """The first usable count under ``keys``; ``None`` when the body reports none. + + ``bool`` is an ``int`` subclass, so it is ruled out explicitly: a target that + answers ``"input_tokens": true`` is not reporting one token. + """ + for key in keys: + if isinstance(value := usage.get(key), int) and not isinstance(value, bool) and value >= 0: + return value + return None + + +def parse_response_body_usage(response_body: object) -> Usage | None: + """Read the token counts a pass-through target reported in its JSON body. + + Returns ``None`` when the body is not a JSON object, carries no ``usage`` + object, or none of the known keys holds a non-negative integer. A missing + side counts as 0 (a decisions API, for instance, legitimately reports + ``output_tokens: 0``). An explicit ``total_tokens`` is kept as sent; + otherwise the total is the sum of both sides. + """ + body: Final = _json_object(response_body) + if body is None: + return None + usage: Final = _json_object(body.get("usage")) + if usage is None: + return None + prompt_tokens: Final = _token_count(usage, _PROMPT_TOKEN_KEYS) + completion_tokens: Final = _token_count(usage, _COMPLETION_TOKEN_KEYS) + if prompt_tokens is None and completion_tokens is None: + return None + reported_total: Final = _token_count(usage, _TOTAL_TOKEN_KEYS) + prompt: Final = prompt_tokens or 0 + completion: Final = completion_tokens or 0 + return Usage( + prompt_tokens=prompt, + completion_tokens=completion, + total_tokens=reported_total if reported_total is not None else prompt + completion, + ) + + +def apply_response_body_usage(logging_obj: LiteLLMLoggingObj, response_body: object) -> Usage | None: + """Record the body-reported counts on the request's logging object. + + Runs after the upstream header contract has had its say: a + ``combined_usage_object`` already on the logging object (an upstream that + reported ``x-litellm-total-tokens``) is left untouched, since a target that + reports its own totals owns them. + """ + if logging_obj.model_call_details.get("combined_usage_object") is not None: + return None + usage: Final = parse_response_body_usage(response_body) + if usage is None: + return None + logging_obj.model_call_details["combined_usage_object"] = usage # rebind-ok: shared per-request logging record + return usage diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 9c6c3111209..a7c4b93df57 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -47,6 +47,7 @@ from .llm_provider_handlers.transcribe_passthrough_logging_handler import ( from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) +from .response_body_usage import apply_response_body_usage from .upstream_usage_headers import has_upstream_reported_usage cohere_passthrough_logging_handler: Final = CoherePassthroughLoggingHandler() @@ -501,6 +502,12 @@ class PassThroughEndpointLogging: ] kwargs = normalized_llm_passthrough_logging_payload["kwargs"] if standard_logging_response_object is None: + # No provider handler read this body, so the only token counts + # we will ever get are the ones the target wrote in it. + apply_response_body_usage( + logging_obj=logging_obj, + response_body=response_body, # pyright: ignore[reportUnknownArgumentType] # untyped dict from the relayed body + ) standard_logging_response_object = StandardPassThroughResponseObject( response=_safe_response_text(httpx_response) ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index db0a74d360a..2fb408d189c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -6726,6 +6726,7 @@ async def _run_upstream_reporting_passthrough( status_code: int = 200, cost_per_request: float | None = None, request_headers: Mapping[str, str] | None = None, + upstream_body: bytes = b'{"answer": "ok"}', ) -> tuple[list[object], MagicMock]: """Drive a generic pass-through against an upstream that reports its own cost/usage. Returns (recorded standard logging payloads, proxy logging mock).""" @@ -6735,7 +6736,7 @@ async def _run_upstream_reporting_passthrough( _FakeUpstreamTransport( status_code=status_code, headers={"content-type": "application/json", **upstream_headers}, - stream=_RecordingUpstreamByteStream((b'{"answer": "ok"}',)), + stream=_RecordingUpstreamByteStream((upstream_body,)), ), timeout=None, ) @@ -6814,6 +6815,50 @@ async def test_passthrough_records_zero_when_upstream_reports_nothing(): assert payloads[0]["total_tokens"] == 0 +@pytest.mark.asyncio +async def test_passthrough_records_tokens_reported_in_the_response_body(): + """ + A generic pass-through target (no provider handler, no x-litellm-* headers) + that writes a ``usage`` object in its JSON body, here a decisions API with + the Anthropic-style ``input_tokens``/``output_tokens`` keys. Before this, + the counts were dropped and the spend row showed 0 tokens. + """ + payloads, _ = await _run_upstream_reporting_passthrough( + {}, + upstream_body=b'{"answers": {"is_spam": {"type": "noul", "noul": 0.93}}, "usage": {"input_tokens": 1874, "output_tokens": 0}}', + ) + + assert len(payloads) == 1 + assert payloads[0]["prompt_tokens"] == 1874 + assert payloads[0]["completion_tokens"] == 0 + assert payloads[0]["total_tokens"] == 1874 + + +@pytest.mark.asyncio +async def test_passthrough_records_openai_shaped_body_usage(): + payloads, _ = await _run_upstream_reporting_passthrough( + {}, + upstream_body=b'{"choices": [], "usage": {"prompt_tokens": 40, "completion_tokens": 12, "total_tokens": 52}}', + ) + + assert len(payloads) == 1 + assert payloads[0]["prompt_tokens"] == 40 + assert payloads[0]["completion_tokens"] == 12 + assert payloads[0]["total_tokens"] == 52 + + +@pytest.mark.asyncio +async def test_passthrough_upstream_total_tokens_header_wins_over_body_usage(): + payloads, _ = await _run_upstream_reporting_passthrough( + {"x-litellm-response-cost": "0.000415", "x-litellm-total-tokens": "930"}, + upstream_body=b'{"answer": "ok", "usage": {"input_tokens": 1874, "output_tokens": 0}}', + ) + + assert len(payloads) == 1 + assert payloads[0]["response_cost"] == 0.000415 + assert payloads[0]["total_tokens"] == 930 + + @pytest.mark.asyncio async def test_passthrough_records_upstream_reported_cost_on_error_response(): """ diff --git a/tests/unit/proxy/pass_through_endpoints/test_response_body_usage.py b/tests/unit/proxy/pass_through_endpoints/test_response_body_usage.py new file mode 100644 index 00000000000..cf517a0732c --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/test_response_body_usage.py @@ -0,0 +1,96 @@ +from datetime import datetime, timezone + +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy.pass_through_endpoints.response_body_usage import ( + apply_response_body_usage, + parse_response_body_usage, +) +from litellm.types.utils import Usage + + +def _logging_obj() -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="decider", + messages=[], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime(2025, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + litellm_call_id="body-usage-1", + function_id="body-usage", + ) + + +def test_parse_reads_openai_shape(): + usage = parse_response_body_usage({"usage": {"prompt_tokens": 1874, "completion_tokens": 12, "total_tokens": 1886}}) + + assert usage == Usage(prompt_tokens=1874, completion_tokens=12, total_tokens=1886) + + +def test_parse_reads_anthropic_shape_and_sums_the_total(): + usage = parse_response_body_usage({"answers": {}, "usage": {"input_tokens": 1874, "output_tokens": 0}}) + + assert usage == Usage(prompt_tokens=1874, completion_tokens=0, total_tokens=1874) + + +def test_parse_counts_a_missing_side_as_zero(): + usage = parse_response_body_usage({"usage": {"input_tokens": 40}}) + + assert usage == Usage(prompt_tokens=40, completion_tokens=0, total_tokens=40) + + +def test_parse_keeps_an_explicit_total_over_the_sum(): + usage = parse_response_body_usage({"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 20}}) + + assert usage == Usage(prompt_tokens=10, completion_tokens=5, total_tokens=20) + + +@pytest.mark.parametrize( + "response_body", + [ + None, + "not json", + [{"usage": {"input_tokens": 3}}], + {"answer": "ok"}, + {"usage": None}, + {"usage": "1874"}, + {"usage": {}}, + {"usage": {"input_tokens": "1874", "output_tokens": "0"}}, + {"usage": {"input_tokens": -1}}, + {"usage": {"input_tokens": True}}, + {"usage": {"input_tokens": 12.5}}, + ], +) +def test_parse_returns_none_when_the_body_reports_no_usable_count(response_body: object): + assert parse_response_body_usage(response_body) is None + + +def test_parse_skips_an_unusable_value_but_keeps_the_usable_side(): + usage = parse_response_body_usage({"usage": {"input_tokens": "oops", "output_tokens": 7}}) + + assert usage == Usage(prompt_tokens=0, completion_tokens=7, total_tokens=7) + + +def test_apply_records_the_body_usage_on_the_logging_object(): + logging_obj = _logging_obj() + + applied = apply_response_body_usage(logging_obj, {"usage": {"input_tokens": 1874, "output_tokens": 0}}) + + assert applied == Usage(prompt_tokens=1874, completion_tokens=0, total_tokens=1874) + assert logging_obj.model_call_details["combined_usage_object"] == applied + + +def test_apply_leaves_an_upstream_reported_total_untouched(): + logging_obj = _logging_obj() + logging_obj.model_call_details["combined_usage_object"] = Usage(total_tokens=930) + + assert apply_response_body_usage(logging_obj, {"usage": {"input_tokens": 1874, "output_tokens": 0}}) is None + assert logging_obj.model_call_details["combined_usage_object"] == Usage(total_tokens=930) + + +def test_apply_writes_nothing_without_usage_in_the_body(): + logging_obj = _logging_obj() + + assert apply_response_body_usage(logging_obj, {"answer": "ok"}) is None + assert "combined_usage_object" not in logging_obj.model_call_details