Merge 6994fbfab5 into c58fc0c996

litellm/llms/anthropic/pass_through/adapters/handler.py

@@ -37,6 +37,7 @@ from litellm.llms.anthropic.pass_through.utils import (
     proxy_sub_call_logging_kwargs,
     proxy_sub_call_trace_id,
 )
+from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
 from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
 from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,
@@ -81,6 +82,18 @@ def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool:
     return False
 
 
+def _record_bridged_response_cost(completion_response: ModelResponse, logging_obj: LiteLLMLoggingObject | None) -> None:
+    """
+    The Anthropic-shaped response can no longer be priced correctly (its usage drops the image token breakdown),
+    and the success handlers that store this cost run on another thread, so the proxy's x-litellm-response-cost
+    header could read the logging object before they did
+    """
+    response_cost: Final = get_hidden_params_dict(completion_response).get("response_cost")
+    if logging_obj is None or not isinstance(response_cost, (int, float)):
+        return
+    logging_obj.model_call_details["response_cost"] = response_cost  # rebind-ok: the record spend logging reads
+
+
 def _proxy_router_fallback() -> "Router | None":
     try:
         from litellm.proxy.proxy_server import llm_router as _proxy_router
@@ -799,9 +812,11 @@ class LiteLLMMessagesToCompletionTransformationHandler:
                 return transformed_stream
             raise ValueError("Failed to transform streaming response")
         else:
-            _record_bridged_model_response(logging_obj, completion_response)
+            bridged_response: Final = cast(ModelResponse, completion_response)
+            _record_bridged_model_response(logging_obj, bridged_response)
+            _record_bridged_response_cost(bridged_response, logging_obj)
             anthropic_response: Final = ANTHROPIC_ADAPTER.translate_completion_output_params(
-                cast(ModelResponse, completion_response),
+                bridged_response,
                 tool_name_mapping=tool_name_mapping,
                 polyfill_result=polyfill_result,
             )
@@ -943,9 +958,11 @@ class LiteLLMMessagesToCompletionTransformationHandler:
                 return transformed_stream
             raise ValueError("Failed to transform streaming response")
         else:
-            _record_bridged_model_response(logging_obj, completion_response)
+            bridged_response: Final = cast(ModelResponse, completion_response)
+            _record_bridged_model_response(logging_obj, bridged_response)
+            _record_bridged_response_cost(bridged_response, logging_obj)
             anthropic_response: Final = ANTHROPIC_ADAPTER.translate_completion_output_params(
-                cast(ModelResponse, completion_response),
+                bridged_response,
                 tool_name_mapping=tool_name_mapping,
                 polyfill_result=polyfill_result,
             )

tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py

@@ -3,11 +3,14 @@ import copy
 import json
 import os
 import uuid
+from collections.abc import Iterator
+from datetime import datetime
 from typing import Any, Dict, Final, List
 
 import httpx
 import respx
 import pytest
+import respx
 from fastapi.testclient import TestClient
 from pydantic import ValidationError
 
@@ -17,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import litellm
 from litellm.anthropic_interface import messages
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@@ -1731,6 +1735,92 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet
     assert response["safeguard_results"] == safeguard_results
 
 
+_IMAGE_MODEL = "gemini-image-cost-probe"
+_IMAGE_INPUT_RATE = 1e-06
+_IMAGE_TEXT_OUTPUT_RATE = 2e-06
+_IMAGE_OUTPUT_RATE = 5e-05
+_IMAGE_PROMPT_TOKENS = 7
+_IMAGE_OUTPUT_TOKENS = 1120
+
+
+class _LoggingBeforeSuccessHandlersRan(Logging):
+    """The proxy's logging object as the response headers see it when the success handlers have not run yet"""
+
+    def success_handler(self, *args, **kwargs):
+        return None
+
+    async def async_success_handler(self, *args, **kwargs):
+        return None
+
+
+@pytest.fixture
+def gemini_image_response(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    monkeypatch.setitem(
+        litellm.model_cost,
+        f"gemini/{_IMAGE_MODEL}",
+        {
+            "input_cost_per_token": _IMAGE_INPUT_RATE,
+            "output_cost_per_token": _IMAGE_TEXT_OUTPUT_RATE,
+            "output_cost_per_image_token": _IMAGE_OUTPUT_RATE,
+            "litellm_provider": "gemini",
+            "mode": "chat",
+        },
+    )
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    litellm.get_model_info.cache_clear()
+    body = {
+        "candidates": [{"content": {"role": "model", "parts": [{"text": "a red square"}]}, "finishReason": "STOP"}],
+        "usageMetadata": {
+            "promptTokenCount": _IMAGE_PROMPT_TOKENS,
+            "candidatesTokenCount": _IMAGE_OUTPUT_TOKENS,
+            "candidatesTokensDetails": [{"modality": "IMAGE", "tokenCount": _IMAGE_OUTPUT_TOKENS}],
+            "totalTokenCount": _IMAGE_PROMPT_TOKENS + _IMAGE_OUTPUT_TOKENS,
+        },
+        "modelVersion": _IMAGE_MODEL,
+    }
+    with respx.mock(assert_all_called=True) as router:
+        router.route(host="generativelanguage.googleapis.com").mock(return_value=httpx.Response(200, json=body))
+        yield
+    litellm.get_model_info.cache_clear()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+async def test_bridged_messages_response_leaves_its_image_priced_cost_on_the_logging_object(
+    gemini_image_response, is_async
+):
+    from litellm.llms.anthropic.pass_through.adapters.handler import (
+        LiteLLMMessagesToCompletionTransformationHandler,
+    )
+
+    logging_obj: Final = _LoggingBeforeSuccessHandlersRan(
+        model=_IMAGE_MODEL,
+        messages=[],
+        stream=False,
+        call_type="anthropic_messages",
+        start_time=datetime(2026, 1, 1),
+        litellm_call_id="call-1",
+        function_id="fn-1",
+    )
+    request: Final = {
+        "max_tokens": 2048,
+        "messages": [{"role": "user", "content": "a red square on a white background"}],
+        "model": f"gemini/{_IMAGE_MODEL}",
+        "api_key": "test-key",
+        "litellm_logging_obj": logging_obj,
+    }
+
+    if is_async:
+        response = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler(**request)
+    else:
+        response = LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler(**request)
+
+    assert response["usage"]["output_tokens"] == _IMAGE_OUTPUT_TOKENS, response
+    assert logging_obj.model_call_details.get("response_cost") == pytest.approx(
+        _IMAGE_PROMPT_TOKENS * _IMAGE_INPUT_RATE + _IMAGE_OUTPUT_TOKENS * _IMAGE_OUTPUT_RATE
+    ), "x-litellm-response-cost reads this before the success handlers run, so it must already hold the image rate"
+
+
 _CURL_TOOL_USE: Final = {
     "type": "tool_use",
     "id": "toolu_curl",