Merge 7f93300647 into 3b8ad37296

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

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)
             )

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():
     """

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