fix: keep cost normalization separate from provider callbacks

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-11-09 15:10:22 +00:00
parent 4c29f60aa6
commit f0ff84112f
4 changed files with 45 additions and 23 deletions

View file

@ -2978,10 +2978,10 @@ class Logging(LiteLLMLoggingBaseClass):
_update_completion_start_time = update_completion_start_time
def normalize_endpoint_logging_result(self, result: object) -> object:
def normalize_endpoint_logging_result(self, result: object, *, for_billing: bool = False) -> object:
call_type: Final[object] = self.call_type
if call_type == CallTypes.anthropic_messages.value:
return self._anthropic_messages_logged_response(result=result)
return self._anthropic_messages_logged_response(result=result, for_billing=for_billing)
if call_type in (CallTypes.generate_content.value, CallTypes.agenerate_content.value):
return self._handle_non_streaming_google_genai_generate_content_response_logging(result=result)
if call_type in (CallTypes.asend_message.value, CallTypes.send_message.value):
@ -5324,7 +5324,7 @@ class Logging(LiteLLMLoggingBaseClass):
return result
return {**body, "usage": {**reported, "speed": requested}}
def _anthropic_messages_logged_response(self, result: object) -> ModelResponse:
def _anthropic_messages_logged_response(self, result: object, *, for_billing: bool = False) -> ModelResponse:
"""
The ModelResponse a /v1/messages spend_logs row is built from.
@ -5332,18 +5332,21 @@ class Logging(LiteLLMLoggingBaseClass):
caller was served is minted locally rather than issued upstream, so it is absent from the
response the row would otherwise be keyed on and has to be carried over here.
"""
logged: Final = self._handle_anthropic_messages_response_logging(result=result)
logged: Final = self._handle_anthropic_messages_response_logging(result=result, for_billing=for_billing)
streamed_message_id: Final = self.streamed_anthropic_message_id
if streamed_message_id is None:
return logged
return logged.model_copy(update={"id": streamed_message_id})
def _handle_anthropic_messages_response_logging(self, result: object) -> ModelResponse:
def _handle_anthropic_messages_response_logging(
self, result: object, *, for_billing: bool = False
) -> ModelResponse:
"""
Handles logging for Anthropic messages responses.
Args:
result: The response object from the model call
for_billing: Convert the served body without replaying provider callbacks
Returns:
The the response object from the model call
@ -5375,10 +5378,14 @@ class Logging(LiteLLMLoggingBaseClass):
)
provider_response_id: Final = _provider_response_id(raw_body)
bridged: Final = self.bridged_model_response
if bridged is not None and isinstance(result, dict):
return bridged if provider_response_id is None else bridged.model_copy(update={"id": provider_response_id})
httpx_response: Final = self.model_call_details.get("httpx_response", None)
if not self.stream and httpx_response and isinstance(httpx_response, httpx.Response):
if (
bridged is not None
and isinstance(result, dict)
and (for_billing or self.stream or not isinstance(httpx_response, httpx.Response))
):
return bridged if provider_response_id is None else bridged.model_copy(update={"id": provider_response_id})
if not for_billing and not self.stream and httpx_response and isinstance(httpx_response, httpx.Response):
result = litellm.AnthropicConfig().transform_response(
raw_response=httpx_response,
model_response=litellm.ModelResponse(id=provider_response_id),
@ -5392,22 +5399,32 @@ class Logging(LiteLLMLoggingBaseClass):
json_mode=False,
litellm_params={},
)
if bridged is not None and served_body is not None:
return (
bridged if provider_response_id is None else bridged.model_copy(update={"id": provider_response_id})
)
elif result is None:
verbose_logger.warning("LiteLLM: the anthropic_messages stream assembled no response, logging an empty one")
return litellm.ModelResponse(model=self.model)
else:
billing_body: Final = (
_SERVED_ANTHROPIC_BODY.validate_json(httpx_response.content)
if for_billing and not self.stream and isinstance(httpx_response, httpx.Response)
else served_body
)
anthropic_body: Final = (
served_body
if served_body is not None
billing_body
if billing_body is not None
else _SERVED_ANTHROPIC_BODY.validate_python(
result.model_dump() if isinstance(result, BaseModel) else result
)
)
result = litellm.AnthropicConfig().transform_parsed_response(
completion_response={"stop_sequence": None, **anthropic_body},
raw_response=httpx.Response(
status_code=200,
headers={},
raw_response=(
httpx_response
if not self.stream and isinstance(httpx_response, httpx.Response)
else httpx.Response(status_code=200, headers={})
),
model_response=litellm.ModelResponse(id=provider_response_id),
json_mode=None,

View file

@ -114,16 +114,17 @@ class MessagesUsageCostStamper(UsageCostStamper):
body: Final = self.response if billing_body is None else billing_body
if body.get("type") != "message" or not isinstance(body.get("usage"), Mapping):
return False
priced: Final = logging_obj.normalize_endpoint_logging_result(body)
priced: Final = logging_obj.normalize_endpoint_logging_result(body, for_billing=True)
cost: Final = validated_usage_cost(logging_obj.response_cost_calculator(result=priced, cache_hit=cache_hit))
if cost is not None:
self.stamp(cost)
if cost is None:
return False
self.stamp(cost)
if isinstance(priced, ModelResponse):
stamp_usage_cost(priced, cost)
logging_obj.record_bridged_model_response(priced)
except Exception: # noqa: BLE001 # optional billing metadata must not fail a completed response or stream
return False
return cost is not None
return True
def stamp_usage_cost(response: object, cost: float | None) -> None:

View file

@ -13,6 +13,7 @@ from typing import Final
from unittest.mock import MagicMock, patch
import pytest
import httpx
import litellm
import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod
@ -46,15 +47,16 @@ def test_metadata_apply_preserves_model_hidden_params(
assert response.model_dump() == before
@pytest.mark.parametrize("shape", ["chat", "responses", "messages"])
@pytest.mark.parametrize("shape", ["chat", "responses", "messages", "messages_httpx"])
@pytest.mark.parametrize("input_tokens,output_tokens", [(3, 4), (0, 0)])
def test_nonstream_usage_cost_matches_the_bill(shape: str, input_tokens: int, output_tokens: int) -> None:
start: Final = datetime.datetime(2026, 1, 1)
messages_response: Final = shape.startswith("messages")
logging_obj: Final = Logging(
model="gpt-6.1-sol",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="anthropic_messages" if shape == "messages" else "completion",
call_type="anthropic_messages" if messages_response else "completion",
start_time=start,
litellm_call_id="usage-cost-" + shape,
function_id="usage-cost",
@ -77,7 +79,7 @@ def test_nonstream_usage_cost_matches_the_bill(shape: str, input_tokens: int, ou
"stop_sequence": None,
"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
}
if shape == "messages"
if messages_response
else ResponsesAPIResponse(
id="resp_cost",
created_at=0,
@ -94,6 +96,9 @@ def test_nonstream_usage_cost_matches_the_bill(shape: str, input_tokens: int, ou
model="gpt-6.1-sol", usage=Usage(prompt_tokens=input_tokens, completion_tokens=output_tokens)
)
)
if shape == "messages_httpx" and isinstance(response, dict):
logging_obj.model_call_details["httpx_response"] = httpx.Response(200, json=response)
response["usage"] = {"input_tokens": 0, "output_tokens": 0}
with patch.object(litellm, "response_cost_calculator", wraps=litellm.response_cost_calculator) as calculator:
update_response_metadata(response, logging_obj, "gpt-6.1-sol", {}, start, start)
logging_obj.prepare_success_logging(
@ -111,7 +116,7 @@ def test_nonstream_usage_cost_matches_the_bill(shape: str, input_tokens: int, ou
assert cost == pytest.approx(input_tokens * 0.01 + output_tokens * 0.02)
priced: Final = logging_obj.normalize_endpoint_logging_result(response)
assert logging_obj.response_cost_calculator(result=priced) == pytest.approx(cost)
hidden_response: Final = logging_obj.bridged_model_response if shape == "messages" else response
hidden_response: Final = logging_obj.bridged_model_response if messages_response else response
assert hidden_response is not None
assert hidden_response._hidden_params["response_cost"] == pytest.approx(cost)
assert logging_obj.model_call_details["response_cost"] == pytest.approx(cost)

View file

@ -4418,8 +4418,7 @@ async def test_native_compaction_wire_roundtrip(
if messages_api:
assert response["stop_reason"] == "compaction"
if not chat_wire:
assert isinstance(response["usage"]["cost"], float)
assert response["usage"] == {**usage, "cost": response["usage"]["cost"]}
assert {key: value for key, value in response["usage"].items() if key != "cost"} == usage
if not gateway:
replay: Final = {"role": "assistant", "content": blocks} if messages_api else message
await call(messages=[replay, {"role": "user", "content": "Continue"}], **params)