mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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:
parent
4c29f60aa6
commit
f0ff84112f
4 changed files with 45 additions and 23 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue