scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

Make the usage.cost call site one testable entry point

Browse source 
The call site was a two-line conditional inside _process_llm_request, so the injecting branch was only reachable through a full request and went uncovered. Collapsing it into _maybe_set_usage_cost leaves one unconditional line at the call site and puts the opt-in decision next to the write, where both branches can be exercised directly.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Ben Langfeld 2026-09-16 21:39:50 -04:00
parent 5e6fbc749d
commit aa617a9f09
No known key found for this signature in database
2 changed files with 24 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm/proxy/common_request_processing.py
View file

@ -2823,8 +2823,7 @@ class ProxyBaseLLMRequestProcessing:
# Same value the x-litellm-response-cost header carries below, so the body and
# the header cannot report different costs for one request.
if self._should_include_cost_in_usage(request):
self._set_usage_cost(response, response_cost_for_headers)
self._maybe_set_usage_cost(request, response, response_cost_for_headers)
# Always return the client-requested model name (not provider-prefixed internal identifiers)
# for OpenAI-compatible responses.
@ -4097,6 +4096,17 @@ class ProxyBaseLLMRequestProcessing:
return cost_from_logging_obj
return ProxyBaseLLMRequestProcessing._completion_cost_or_none(model_response, model_name, service_tier)
@staticmethod
def _maybe_set_usage_cost(request: Request, response: object, cost: float | str | None) -> None:
"""
Record the gateway's cost on ``usage.cost``, if this request asked for it.
Kept as one entry point so the decision and the write are exercised together
rather than only through a full request.
"""
if ProxyBaseLLMRequestProcessing._should_include_cost_in_usage(request):
ProxyBaseLLMRequestProcessing._set_usage_cost(response, cost)
@staticmethod
def _should_include_cost_in_usage(request: Request) -> bool:
"""

12
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -9384,6 +9384,18 @@ class TestIncludeCostInUsage:
response: Final = _response_with_usage(prompt_tokens=11, completion_tokens=5, total_tokens=16)
assert "cost" not in response.model_dump()["usage"]
def test_opted_in_request_gets_the_cost_recorded(self):
"""The decision and the write, exercised together as the request path runs them."""
response: Final = _response_with_usage(prompt_tokens=11, completion_tokens=5, total_tokens=16)
request: Final = _request_with_headers(x_litellm_include_cost_in_usage="true")
ProxyBaseLLMRequestProcessing._maybe_set_usage_cost(request, response, 5.85e-06)
assert response.model_dump()["usage"]["cost"] == 5.85e-06
def test_opted_out_request_is_left_untouched(self):
response: Final = _response_with_usage(prompt_tokens=11, completion_tokens=5, total_tokens=16)
ProxyBaseLLMRequestProcessing._maybe_set_usage_cost(_request_with_headers(), response, 5.85e-06)
assert "cost" not in response.model_dump()["usage"]
def test_a_response_without_usage_is_left_alone(self):
sentinel: Final = SimpleNamespace(usage=None)
ProxyBaseLLMRequestProcessing._set_usage_cost(sentinel, 5.85e-06)