From 8abf2d8e34d3cc37f67ada4feb7e087d6bb93d30 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 14 Mar 2026 12:30:52 -0700 Subject: [PATCH] fix: Fixes https://github.com/BerriAI/litellm/issues/23185 (#23647) --- litellm/litellm_core_utils/litellm_logging.py | 298 +++++++++--------- litellm/llms/custom_httpx/llm_http_handler.py | 5 +- .../test_litellm_logging.py | 32 ++ .../custom_httpx/test_llm_http_handler.py | 74 +++++ tests/test_litellm/test_cost_calculator.py | 73 ++++- tests/test_litellm/test_router.py | 75 +++++ 6 files changed, 392 insertions(+), 165 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e22d057bb6..dbce1f9f70 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -352,9 +352,9 @@ class Logging(LiteLLMLoggingBaseClass): ) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[ - Any - ] = [] # for generating complete stream response + self.sync_streaming_chunks: List[Any] = ( + [] + ) # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks @@ -746,9 +746,9 @@ class Logging(LiteLLMLoggingBaseClass): prompt_spec=prompt_spec, dynamic_callback_params=dynamic_callback_params, ): - self.model_call_details[ - "prompt_integration" - ] = logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + logger.__class__.__name__ + ) return logger except Exception: # If check fails, continue to next logger @@ -816,9 +816,9 @@ class Logging(LiteLLMLoggingBaseClass): if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( non_default_params ): - self.model_call_details[ - "prompt_integration" - ] = anthropic_cache_control_logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + anthropic_cache_control_logger.__class__.__name__ + ) return anthropic_cache_control_logger ######################################################### @@ -830,9 +830,9 @@ class Logging(LiteLLMLoggingBaseClass): internal_usage_cache=None, llm_router=None, ) - self.model_call_details[ - "prompt_integration" - ] = vector_store_custom_logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + vector_store_custom_logger.__class__.__name__ + ) # Add to global callbacks so post-call hooks are invoked if ( vector_store_custom_logger @@ -892,9 +892,9 @@ class Logging(LiteLLMLoggingBaseClass): model ): # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model - self.model_call_details["litellm_params"][ - "api_base" - ] = self._get_masked_api_base(additional_args.get("api_base", "")) + self.model_call_details["litellm_params"]["api_base"] = ( + self._get_masked_api_base(additional_args.get("api_base", "")) + ) def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 # Log the exact input to the LLM API @@ -923,10 +923,10 @@ class Logging(LiteLLMLoggingBaseClass): try: # [Non-blocking Extra Debug Information in metadata] if turn_off_message_logging is True: - _metadata[ - "raw_request" - ] = "redacted by litellm. \ + _metadata["raw_request"] = ( + "redacted by litellm. \ 'litellm.turn_off_message_logging=True'" + ) else: curl_command = self._get_request_curl_command( api_base=additional_args.get("api_base", ""), @@ -937,34 +937,34 @@ class Logging(LiteLLMLoggingBaseClass): _metadata["raw_request"] = str(curl_command) # split up, so it's easier to parse in the UI - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, + self.model_call_details["raw_request_typed_dict"] = ( + RawRequestTypedDict( + raw_request_api_base=str( + additional_args.get("api_base") or "" + ), + raw_request_body=self._get_raw_request_body( + additional_args.get("complete_input_dict", {}) + ), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, + ) ) except Exception as e: - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - error=str(e), + self.model_call_details["raw_request_typed_dict"] = ( + RawRequestTypedDict( + error=str(e), + ) ) - _metadata[ - "raw_request" - ] = "Unable to Log \ + _metadata["raw_request"] = ( + "Unable to Log \ raw request: {}".format( - str(e) + str(e) + ) ) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: @@ -1265,13 +1265,13 @@ class Logging(LiteLLMLoggingBaseClass): for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[ - MCPPostCallResponseObject - ] = await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, + response: Optional[MCPPostCallResponseObject] = ( + await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, + ) ) ###################################################################### # if any of the callbacks modify the response, use the modified response @@ -1466,9 +1466,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info + self.model_call_details["response_cost_failure_debug_information"] = ( + debug_info + ) return None try: @@ -1494,9 +1494,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info + self.model_call_details["response_cost_failure_debug_information"] = ( + debug_info + ) return None @@ -1652,9 +1652,9 @@ class Logging(LiteLLMLoggingBaseClass): result=logging_result ) - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload(logging_result, start_time, end_time) + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload(logging_result, start_time, end_time) + ) if ( standard_logging_payload := self.model_call_details.get( @@ -1732,9 +1732,9 @@ class Logging(LiteLLMLoggingBaseClass): end_time = datetime.datetime.now() if self.completion_start_time is None: self.completion_start_time = end_time - self.model_call_details[ - "completion_start_time" - ] = self.completion_start_time + self.model_call_details["completion_start_time"] = ( + self.completion_start_time + ) self.model_call_details["log_event_type"] = "successful_api_call" self.model_call_details["end_time"] = end_time @@ -1771,10 +1771,10 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, ) elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - result, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + result, start_time, end_time + ) ) if ( standard_logging_payload := self.model_call_details.get( @@ -1783,9 +1783,9 @@ class Logging(LiteLLMLoggingBaseClass): ) is not None: emit_standard_logging_payload(standard_logging_payload) elif standard_logging_object is not None: - self.model_call_details[ - "standard_logging_object" - ] = standard_logging_object + self.model_call_details["standard_logging_object"] = ( + standard_logging_object + ) else: self.model_call_details["response_cost"] = None @@ -1943,17 +1943,17 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( "Logging Details LiteLLM-Success Call streaming complete" ) - self.model_call_details[ - "complete_streaming_response" - ] = complete_streaming_response - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator(result=complete_streaming_response) + self.model_call_details["complete_streaming_response"] = ( + complete_streaming_response + ) + self.model_call_details["response_cost"] = ( + self._response_cost_calculator(result=complete_streaming_response) + ) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) ) if ( standard_logging_payload := self.model_call_details.get( @@ -2287,10 +2287,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} + self.model_call_details["complete_response"] = ( + self.model_call_details.get( + "complete_streaming_response", {} + ) ) result = self.model_call_details["complete_response"] openMeterLogger.log_success_event( @@ -2314,10 +2314,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} + self.model_call_details["complete_response"] = ( + self.model_call_details.get( + "complete_streaming_response", {} + ) ) result = self.model_call_details["complete_response"] @@ -2456,9 +2456,9 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") - self.model_call_details[ - "async_complete_streaming_response" - ] = complete_streaming_response + self.model_call_details["async_complete_streaming_response"] = ( + complete_streaming_response + ) try: if self.model_call_details.get("cache_hit", False) is True: @@ -2469,10 +2469,10 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator( - result=complete_streaming_response + self.model_call_details["response_cost"] = ( + self._response_cost_calculator( + result=complete_streaming_response + ) ) verbose_logger.debug( @@ -2485,10 +2485,10 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = None ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) ) # print standard logging payload @@ -2515,9 +2515,9 @@ class Logging(LiteLLMLoggingBaseClass): # _success_handler_helper_fn if self.model_call_details.get("standard_logging_object") is None: ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload(result, start_time, end_time) + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload(result, start_time, end_time) + ) # print standard logging payload if ( @@ -2760,18 +2760,18 @@ class Logging(LiteLLMLoggingBaseClass): ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=str(exception), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, + self.model_call_details["standard_logging_object"] = ( + get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=str(exception), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) ) return start_time, end_time @@ -3735,9 +3735,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 service_name=arize_config.project_name, ) - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" + ) for callback in _in_memory_loggers: if ( isinstance(callback, ArizeLogger) @@ -3763,13 +3763,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" + ) else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={arize_phoenix_config.project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"openinference.project.name={arize_phoenix_config.project_name}" + ) # Set Phoenix project name from environment variable phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) @@ -3777,19 +3777,19 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"{existing_attrs},openinference.project.name={phoenix_project_name}" + ) else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={phoenix_project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"openinference.project.name={phoenix_project_name}" + ) # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = arize_phoenix_config.otlp_auth_headers + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + arize_phoenix_config.otlp_auth_headers + ) for callback in _in_memory_loggers: if ( @@ -3965,9 +3965,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 exporter="otlp_http", endpoint="https://langtrace.ai/api/trace", ) - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + f"api_key={os.getenv('LANGTRACE_API_KEY')}" + ) for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetry) @@ -4446,15 +4446,17 @@ def use_custom_pricing_for_model(litellm_params: Optional[dict]) -> bool: if litellm_params.get(key) is not None: return True - # Check model_info - metadata: dict = litellm_params.get("metadata", {}) or {} - model_info: dict = metadata.get("model_info", {}) or {} + # Check model_info from metadata or litellm_metadata (generic_api_call routes + # like /responses and /messages store model_info under litellm_metadata) + for metadata_key in ("metadata", "litellm_metadata"): + metadata: dict = litellm_params.get(metadata_key, {}) or {} + model_info: dict = metadata.get("model_info", {}) or {} - if model_info: - matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys() - for key in matching_keys: - if model_info.get(key) is not None: - return True + if model_info: + matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys() + for key in matching_keys: + if model_info.get(key) is not None: + return True return False @@ -4881,10 +4883,10 @@ class StandardLoggingPayloadSetup: for key in StandardLoggingHiddenParams.__annotations__.keys(): if key in hidden_params: if key == "additional_headers": - clean_hidden_params[ - "additional_headers" - ] = StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] + clean_hidden_params["additional_headers"] = ( + StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] + ) ) else: clean_hidden_params[key] = hidden_params[key] # type: ignore @@ -5523,9 +5525,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): ): for k, v in metadata["user_api_key_metadata"].items(): if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[ - k - ] = "scrubbed_by_litellm_for_sensitive_keys" + cleaned_user_api_key_metadata[k] = ( + "scrubbed_by_litellm_for_sensitive_keys" + ) else: cleaned_user_api_key_metadata[k] = v diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a8d649064a..4394343c8e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -22,9 +22,7 @@ import litellm.litellm_core_utils import litellm.types import litellm.types.utils from litellm._logging import verbose_logger -from litellm.anthropic_beta_headers_manager import ( - update_headers_with_filtered_beta, -) +from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -1889,6 +1887,7 @@ class BaseLLMHTTPHandler: optional_params=dict(anthropic_messages_optional_request_params), litellm_params={ "metadata": kwargs.get("metadata", {}), + "litellm_metadata": kwargs.get("litellm_metadata", {}), "preset_cache_key": None, "stream_response": {}, **anthropic_messages_optional_request_params, diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f4aeb27b31..34162bba3f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -148,6 +148,38 @@ def test_use_custom_pricing_for_model(): assert use_custom_pricing_for_model(litellm_params) == True +def test_use_custom_pricing_for_model_via_litellm_metadata(): + """Pricing in litellm_metadata.model_info must be detected. + + Generic API call routes (/messages, /responses) store model_info + under litellm_metadata, not metadata. Regression test for #23185. + """ + from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model + + litellm_params = { + "litellm_metadata": { + "model_info": { + "id": "claude-sonnet-4-custom", + "input_cost_per_token": 0.0003, + "output_cost_per_token": 0.0015, + }, + }, + } + assert use_custom_pricing_for_model(litellm_params) is True + + +def test_use_custom_pricing_not_detected_litellm_metadata_no_pricing(): + """Should return False when litellm_metadata.model_info has no pricing keys.""" + from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model + + litellm_params = { + "litellm_metadata": { + "model_info": {"id": "some-id", "db_model": False}, + }, + } + assert use_custom_pricing_for_model(litellm_params) is False + + def test_logging_prevent_double_logging(logging_obj): """ When using a bridge, log only once from the underlying bridge call. diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 17b4243da1..c418650aa8 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -154,6 +154,80 @@ async def test_async_anthropic_messages_handler_extra_headers(): assert captured_headers["X-Auth-Token"] == "token123" +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_passes_litellm_metadata(): + """Ensure litellm_metadata from kwargs is included in litellm_params + passed to update_environment_variables. + + Routes like /messages store model_info under kwargs['litellm_metadata']. + The handler must forward this into litellm_params so that + use_custom_pricing_for_model can detect custom pricing. Regression test for #23185. + """ + handler = BaseLLMHTTPHandler() + + mock_config = Mock() + mock_config.validate_anthropic_messages_environment = Mock( + return_value=({"x-api-key": "test-key"}, "https://api.anthropic.com") + ) + mock_config.transform_anthropic_messages_request = Mock( + return_value={"model": "claude-sonnet-4-20250514", "messages": []} + ) + + mock_client = AsyncMock() + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hello!"}], + "model": "claude-sonnet-4-20250514", + "stop_reason": "end_turn", + } + mock_client.post = AsyncMock(return_value=mock_response) + + mock_logging_obj = Mock() + mock_logging_obj.update_environment_variables = Mock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.stream = False + + custom_model_info = { + "id": "claude-sonnet-4-custom-pricing", + "input_cost_per_token": 0.0003, + "output_cost_per_token": 0.0015, + } + kwargs = { + "litellm_metadata": { + "model_info": custom_model_info, + "deployment": "anthropic/claude-sonnet-4-20250514", + }, + } + + try: + await handler.async_anthropic_messages_handler( + model="claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_provider_config=mock_config, + anthropic_messages_optional_request_params={}, + custom_llm_provider="anthropic", + litellm_params=GenericLiteLLMParams(), + logging_obj=mock_logging_obj, + client=mock_client, + kwargs=kwargs, + ) + except Exception: + pass + + mock_logging_obj.update_environment_variables.assert_called_once() + call_kwargs = mock_logging_obj.update_environment_variables.call_args + litellm_params_arg = call_kwargs.kwargs.get( + "litellm_params", call_kwargs[1].get("litellm_params", {}) + ) if call_kwargs.kwargs else call_kwargs[1].get("litellm_params", {}) + + assert "litellm_metadata" in litellm_params_arg + assert litellm_params_arg["litellm_metadata"]["model_info"] == custom_model_info + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_header_priority(): """ diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 1204b11934..463f6952d2 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -332,6 +332,62 @@ def test_custom_pricing_with_router_model_id(): assert model_info["cache_read_input_token_cost"] == 0.0000006 +def test_custom_pricing_cost_calc_uses_router_model_id_from_litellm_metadata(): + """When custom pricing is in litellm_metadata.model_info, + use_custom_pricing_for_model should return True and + _select_model_name_for_cost_calc should use router_model_id. + + This tests the full chain that was broken for /messages and /responses + endpoints. Regression test for #23185. + """ + from litellm.cost_calculator import _select_model_name_for_cost_calc + from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model + + custom_model_id = "claude-sonnet-4-custom-pricing-test" + custom_pricing_info = { + "input_cost_per_token": 0.0003, + "output_cost_per_token": 0.0015, + "max_tokens": 8192, + "litellm_provider": "anthropic", + } + litellm.register_model(model_cost={custom_model_id: custom_pricing_info}) + + litellm_params = { + "litellm_metadata": { + "model_info": { + "id": custom_model_id, + "input_cost_per_token": 0.0003, + "output_cost_per_token": 0.0015, + }, + }, + } + + custom_pricing = use_custom_pricing_for_model(litellm_params) + assert custom_pricing is True + + # _select_model_name_for_cost_calc appends provider prefix to the + # selected router_model_id, so the result is "anthropic/" + selected_model = _select_model_name_for_cost_calc( + model="anthropic/claude-sonnet-4-20250514", + completion_response=None, + custom_pricing=custom_pricing, + custom_llm_provider="anthropic", + router_model_id=custom_model_id, + ) + assert selected_model is not None + assert custom_model_id in selected_model + + # Without custom_pricing, the router_model_id is NOT selected + selected_model_no_custom = _select_model_name_for_cost_calc( + model="anthropic/claude-sonnet-4-20250514", + completion_response=None, + custom_pricing=False, + custom_llm_provider="anthropic", + router_model_id=custom_model_id, + ) + assert custom_model_id not in (selected_model_no_custom or "") + + def test_azure_realtime_cost_calculator(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -365,11 +421,7 @@ def test_azure_audio_output_cost_calculation(): Audio tokens should be charged at output_cost_per_audio_token rate, not at the text token rate (output_cost_per_token). """ - from litellm.types.utils import ( - Choices, - CompletionTokensDetailsWrapper, - Message, - ) + from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -471,11 +523,7 @@ def test_default_image_cost_calculator(monkeypatch): def test_cost_calculator_with_cache_creation(): from litellm import completion_cost - from litellm.types.utils import ( - Choices, - Message, - Usage, - ) + from litellm.types.utils import Choices, Message, Usage litellm_model_response = ModelResponse( id="chatcmpl-cc5638bc-fdfe-48e4-8884-57c8f4fb7c63", @@ -896,10 +944,7 @@ def test_azure_ai_cache_cost_calculation(): applied correctly. """ from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token - from litellm.types.utils import ( - PromptTokensDetailsWrapper, - Usage, - ) + from litellm.types.utils import PromptTokensDetailsWrapper, Usage os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8c37214e19..83387ad8b3 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2735,6 +2735,81 @@ def test_credential_name_not_injected_when_absent(): assert kwargs["metadata"]["tags"] == ["A.101"] +def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): + """For generic_api_call, model_info with pricing must go to litellm_metadata. + + Routes like /messages and /responses use generic_api_call which stores + model_info under litellm_metadata. Regression test for #23185. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "claude-sonnet-4", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-20250514", + "api_key": "fake-key", + }, + "model_info": { + "id": "custom-pricing-id", + "input_cost_per_token": 0.0003, + "output_cost_per_token": 0.0015, + }, + }, + ], + ) + + kwargs: dict = {} + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name="generic_api_call" + ) + + assert "litellm_metadata" in kwargs + model_info = kwargs["litellm_metadata"]["model_info"] + assert model_info["id"] == "custom-pricing-id" + assert model_info["input_cost_per_token"] == 0.0003 + assert model_info["output_cost_per_token"] == 0.0015 + + +def test_update_kwargs_with_deployment_model_info_in_metadata(): + """For acompletion (function_name=None), model_info goes to metadata. + + /chat/completions uses acompletion which stores model_info under metadata. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "claude-sonnet-4", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-20250514", + "api_key": "fake-key", + }, + "model_info": { + "id": "custom-pricing-id", + "input_cost_per_token": 0.0003, + "output_cost_per_token": 0.0015, + }, + }, + ], + ) + + kwargs: dict = {} + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name=None + ) + + assert "metadata" in kwargs + model_info = kwargs["metadata"]["model_info"] + assert model_info["id"] == "custom-pricing-id" + assert model_info["input_cost_per_token"] == 0.0003 + assert model_info["output_cost_per_token"] == 0.0015 + + def test_combine_fallback_usage(): """Test that _combine_fallback_usage merges partial and fallback usage.""" from litellm.router import Router