diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 9d6ddd0f1e..6ac337d99a 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -610,7 +610,7 @@ class LangFuseLogger: trace_id = cast(Optional[str], standard_logging_object.get("trace_id")) # Fallback to litellm_call_id if no trace_id found if trace_id is None: - trace_id = litellm_call_id + trace_id = kwargs.get("litellm_trace_id") or litellm_call_id existing_trace_id = clean_metadata.pop("existing_trace_id", None) # If existing_trace_id is provided, use it as the trace_id to return # This allows continuing an existing trace while still returning the correct trace_id diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index f0e1e30b68..03a93cd988 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -338,14 +338,17 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None), ) - if standard_logging_object is None: - return + status_message = str(kwargs.get("exception", "Unknown error")) + if standard_logging_object is not None: + status_message = standard_logging_object.get( + "error_str", None + ) or status_message langfuse_logger_to_use.log_event_on_langfuse( start_time=start_time, end_time=end_time, response_obj=None, user_id=kwargs.get("user", None), - status_message=standard_logging_object["error_str"], + status_message=status_message, level="ERROR", kwargs=kwargs, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f567363df1..27bdaee116 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 @@ -406,7 +406,7 @@ class Logging(LiteLLMLoggingBaseClass): self.passthrough_guardrails_config: Optional[Dict[str, Any]] = None self.model_call_details: Dict[str, Any] = { - "litellm_trace_id": litellm_trace_id, + "litellm_trace_id": self.litellm_trace_id, "litellm_call_id": litellm_call_id, "input": _input, "litellm_params": litellm_params, @@ -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 @@ -2950,7 +2950,7 @@ class Logging(LiteLLMLoggingBaseClass): user_id=kwargs.get("user", None), status_message=str(exception), level="ERROR", - kwargs=self.model_call_details, + kwargs=kwargs, ) if _response is not None and isinstance(_response, dict): _trace_id = _response.get("trace_id", None) @@ -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 ( @@ -3974,9 +3974,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) @@ -4461,15 +4461,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 @@ -4896,10 +4898,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 @@ -5538,9 +5540,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/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 49ff94d0a9..43cfd93019 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -115,6 +115,22 @@ class _ProxyDBLogger(CustomLogger): "custom_llm_provider" ) or request_data.get("custom_llm_provider", "") + # Propagate standard_logging_object and litellm_trace_id from the + # Logging instance so that _get_session_id_for_spend_log uses the same + # trace_id that Langfuse received (via async_failure_handler). + # Without this, the DB session_id would be a random UUID that doesn't + # match the Langfuse trace_id, making failed requests unsearchable. + _litellm_logging_obj = request_data.get("litellm_logging_obj") + if _litellm_logging_obj is not None: + if not request_data.get("standard_logging_object"): + request_data["standard_logging_object"] = getattr( + _litellm_logging_obj, "model_call_details", {} + ).get("standard_logging_object") + if request_data.get("litellm_trace_id") is None: + request_data["litellm_trace_id"] = getattr( + _litellm_logging_obj, "litellm_trace_id", None + ) + await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key_dict.api_key, response_cost=0.0, diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 10d3323a25..b402870921 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -467,6 +467,410 @@ class TestLangfuseUsageDetails(unittest.TestCase): assert self.last_trace_kwargs.get("id") == "call-id-xyz" + def test_log_langfuse_v2_uses_litellm_trace_id_fallback_over_call_id(self): + """ + When standard_logging_object has no trace_id, but kwargs contains + litellm_trace_id (the same ID the DB stores as Session ID), Langfuse + should use litellm_trace_id — NOT litellm_call_id. This ensures the + trace_id in Langfuse matches the Session ID shown in LiteLLM logs. + """ + payload = self._build_standard_logging_payload() # no trace_id + kwargs = self._build_langfuse_kwargs(payload) + kwargs["litellm_trace_id"] = "trace-id-from-kwargs" + self.last_trace_kwargs = {} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kwargs: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata={}, + litellm_params={"metadata": {}}, + output=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="ERROR", + litellm_call_id="call-id-xyz", + ) + + # litellm_trace_id should be preferred over litellm_call_id + assert self.last_trace_kwargs.get("id") == "trace-id-from-kwargs" + + def test_log_langfuse_v2_uses_litellm_trace_id_when_standard_logging_object_none(self): + """ + When standard_logging_object is None (failure case where + get_standard_logging_object_payload threw), litellm_trace_id from kwargs + should be used as the Langfuse trace_id. This matches the DB Session ID. + """ + kwargs = { + "standard_logging_object": None, + "model": "gpt-4", + "call_type": "completion", + "cache_hit": False, + "messages": [], + "litellm_trace_id": "trace-id-failure", + } + self.last_trace_kwargs = {} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kwargs: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata={}, + litellm_params={"metadata": {}}, + output=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="ERROR", + litellm_call_id="call-id-different", + ) + + # Must use litellm_trace_id, not litellm_call_id + assert self.last_trace_kwargs.get("id") == "trace-id-failure" + + def test_log_langfuse_v2_session_id_passed_as_trace_session_id(self): + """ + Test that metadata.session_id is correctly passed as trace_params["session_id"] + for Langfuse session grouping, and does NOT override trace_id. + Each LLM call should get its own unique trace_id while sharing the session_id. + """ + payload = self._build_standard_logging_payload(trace_id="std-trace-123") + kwargs = self._build_langfuse_kwargs(payload) + self.last_trace_kwargs = {} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kwargs: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata={"session_id": "my-session-abc"}, + litellm_params={"metadata": {"session_id": "my-session-abc"}}, + output=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="INFO", + litellm_call_id="call-id-456", + ) + + # session_id should be set for Langfuse session grouping + assert self.last_trace_kwargs.get("session_id") == "my-session-abc" + # trace_id should remain the standard trace_id, NOT the session_id + assert self.last_trace_kwargs.get("id") == "std-trace-123" + + def test_log_langfuse_v2_session_id_preserved_for_error_level(self): + """ + Test that session_id is correctly passed in trace_params even when + the log level is ERROR (failure case). This verifies the fix for + failed requests losing session_id mapping in Langfuse. + """ + payload = self._build_standard_logging_payload(trace_id="std-trace-err") + kwargs = self._build_langfuse_kwargs(payload) + self.last_trace_kwargs = {} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kwargs: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata={"session_id": "error-session-xyz"}, + litellm_params={"metadata": {"session_id": "error-session-xyz"}}, + output="BadRequestError: model not found", + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + kwargs=kwargs, + optional_params={}, + input={"messages": [{"role": "user", "content": "test"}]}, + response_obj=None, + level="ERROR", + litellm_call_id="call-id-err-789", + ) + + # session_id must be preserved even for ERROR level logs + assert self.last_trace_kwargs.get("session_id") == "error-session-xyz" + # trace_id should be the standard trace_id, not the session_id + assert self.last_trace_kwargs.get("id") == "std-trace-err" + # status_message should be set for error traces + assert self.last_trace_kwargs.get("status_message") is not None + + def test_log_langfuse_v2_explicit_trace_id_takes_priority_over_session_id(self): + """ + Test that when both trace_id and session_id are provided in metadata, + trace_id takes priority as the trace identifier. + """ + payload = self._build_standard_logging_payload() + kwargs = self._build_langfuse_kwargs(payload) + self.last_trace_kwargs = {} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kwargs: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata={ + "session_id": "session-999", + "trace_id": "explicit-trace-id-777", + }, + litellm_params={ + "metadata": { + "session_id": "session-999", + "trace_id": "explicit-trace-id-777", + } + }, + output=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="DEFAULT", + litellm_call_id="call-id-aaa", + ) + + # Explicit trace_id must take priority + assert self.last_trace_kwargs.get("id") == "explicit-trace-id-777" + # session_id must still be set for session grouping + assert self.last_trace_kwargs.get("session_id") == "session-999" + + +def test_failure_handler_langfuse_kwargs_excludes_original_response(): + """ + Test that the actual Logging.failure_handler() passes kwargs without + 'original_response' to the Langfuse logger. Exercises the real code path + rather than simulating the filtering logic. + """ + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + + # Create a Logging instance + logging_obj = Logging( + model="gpt-4", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.utcnow(), + litellm_call_id="test-call-id-failure", + function_id="test-function-id", + ) + + # Set up model_call_details with original_response (simulates a coroutine) + mock_coroutine = MagicMock() + logging_obj.model_call_details["original_response"] = mock_coroutine + logging_obj.model_call_details["litellm_params"] = { + "metadata": {"session_id": "test-session-failure"}, + "litellm_session_id": None, + } + logging_obj.model_call_details["optional_params"] = {} + + # Capture what gets passed to log_event_on_langfuse + captured_kwargs = {} + mock_langfuse_logger = MagicMock() + + def capture_log_event(**log_kwargs): + captured_kwargs.update(log_kwargs) + return {"trace_id": "mock-trace-id", "generation_id": "mock-gen-id"} + + mock_langfuse_logger.log_event_on_langfuse.side_effect = capture_log_event + + # Set "langfuse" as a failure callback so the failure_handler processes it + original_failure_callback = litellm.failure_callback + litellm.failure_callback = ["langfuse"] + + try: + # Mock LangFuseHandler to return our capturing mock logger + with patch( + "litellm.litellm_core_utils.litellm_logging.LangFuseHandler" + ) as mock_handler_class: + mock_handler_class.get_langfuse_logger_for_request.return_value = ( + mock_langfuse_logger + ) + + # Call the actual failure_handler + test_exception = Exception("TestError: model not found") + logging_obj.failure_handler( + exception=test_exception, + traceback_exception="Traceback: test", + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + ) + + # Verify log_event_on_langfuse was actually called + assert mock_langfuse_logger.log_event_on_langfuse.called, ( + "log_event_on_langfuse was not called" + ) + + # Verify original_response is NOT in the kwargs passed to Langfuse + langfuse_kwargs = captured_kwargs.get("kwargs", {}) + assert "original_response" not in langfuse_kwargs, ( + "original_response should be excluded from kwargs passed to Langfuse" + ) + + # Verify session_id metadata is preserved in the kwargs + langfuse_metadata = langfuse_kwargs.get("litellm_params", {}).get( + "metadata", {} + ) + assert langfuse_metadata.get("session_id") == "test-session-failure", ( + "session_id should be preserved in kwargs passed to Langfuse" + ) + + # Verify level is ERROR + assert captured_kwargs.get("level") == "ERROR" + finally: + litellm.failure_callback = original_failure_callback + + +@pytest.mark.asyncio +async def test_async_log_failure_event_logs_to_langfuse(): + """ + Test that LangfusePromptManagement.async_log_failure_event() calls + log_event_on_langfuse with level=ERROR even when standard_logging_object + is present. This is the code path the proxy uses for failed LLM calls. + """ + from litellm.integrations.langfuse.langfuse_prompt_management import ( + LangfusePromptManagement, + ) + + mock_langfuse_module = MagicMock() + mock_langfuse_module.version.__version__ = "3.0.0" + + with patch.dict( + "os.environ", + { + "LANGFUSE_SECRET_KEY": "test-secret", + "LANGFUSE_PUBLIC_KEY": "test-public", + "LANGFUSE_HOST": "https://test.langfuse.com", + }, + ), patch.dict("sys.modules", {"langfuse": mock_langfuse_module}): + prompt_mgmt = LangfusePromptManagement() + + # Mock the langfuse logger returned by get_langfuse_logger_for_request + mock_logger = MagicMock() + mock_logger.log_event_on_langfuse.return_value = { + "trace_id": "mock-trace", + "generation_id": "mock-gen", + } + + with patch( + "litellm.integrations.langfuse.langfuse_prompt_management.LangFuseHandler" + ) as mock_handler: + mock_handler.get_langfuse_logger_for_request.return_value = mock_logger + + kwargs = { + "litellm_params": { + "metadata": {"session_id": "test-session-fail"}, + }, + "litellm_call_id": "call-fail-123", + "user": "test-user", + "exception": Exception("API error: model not found"), + "standard_logging_object": { + "error_str": "API error: model not found", + "trace_id": "std-trace-fail", + "metadata": {}, + }, + } + + await prompt_mgmt.async_log_failure_event( + kwargs=kwargs, + response_obj=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + ) + + # Verify log_event_on_langfuse was called + assert mock_logger.log_event_on_langfuse.called, ( + "log_event_on_langfuse was not called for failure event" + ) + call_kwargs = mock_logger.log_event_on_langfuse.call_args[1] + assert call_kwargs["level"] == "ERROR" + assert call_kwargs["status_message"] == "API error: model not found" + assert call_kwargs["response_obj"] is None + + +@pytest.mark.asyncio +async def test_async_log_failure_event_works_without_standard_logging_object(): + """ + Test that async_log_failure_event() still logs to Langfuse even when + standard_logging_object is None (e.g. when get_standard_logging_object_payload + threw an exception). This is the critical fix — before, it silently returned. + """ + from litellm.integrations.langfuse.langfuse_prompt_management import ( + LangfusePromptManagement, + ) + + mock_langfuse_module = MagicMock() + mock_langfuse_module.version.__version__ = "3.0.0" + + with patch.dict( + "os.environ", + { + "LANGFUSE_SECRET_KEY": "test-secret", + "LANGFUSE_PUBLIC_KEY": "test-public", + "LANGFUSE_HOST": "https://test.langfuse.com", + }, + ), patch.dict("sys.modules", {"langfuse": mock_langfuse_module}): + prompt_mgmt = LangfusePromptManagement() + + mock_logger = MagicMock() + mock_logger.log_event_on_langfuse.return_value = { + "trace_id": "mock-trace", + "generation_id": "mock-gen", + } + + with patch( + "litellm.integrations.langfuse.langfuse_prompt_management.LangFuseHandler" + ) as mock_handler: + mock_handler.get_langfuse_logger_for_request.return_value = mock_logger + + kwargs = { + "litellm_params": { + "metadata": {"session_id": "test-session-no-slo"}, + }, + "litellm_call_id": "call-no-slo-456", + "user": "test-user", + "exception": Exception("InternalServerError: something broke"), + "standard_logging_object": None, # This is the key — it's None + } + + await prompt_mgmt.async_log_failure_event( + kwargs=kwargs, + response_obj=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + ) + + # CRITICAL: log_event_on_langfuse MUST still be called + assert mock_logger.log_event_on_langfuse.called, ( + "log_event_on_langfuse was NOT called when standard_logging_object " + "is None — failure trace would be silently dropped" + ) + call_kwargs = mock_logger.log_event_on_langfuse.call_args[1] + assert call_kwargs["level"] == "ERROR" + # Falls back to exception from kwargs + assert "InternalServerError" in call_kwargs["status_message"] + def test_max_langfuse_clients_limit(): """ 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/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index c46b8df5ef..d269a9531f 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -169,6 +169,68 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj(): + """ + When an LLM call fails, the proxy calls post_call_failure_hook with + request_data that doesn't contain standard_logging_object. But the + litellm_logging_obj (set by function_setup) is in request_data and + holds the standard_logging_object with the correct trace_id. + + The failure hook should propagate this so the DB spend log's session_id + matches the Langfuse trace_id. + """ + logger = _ProxyDBLogger() + + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + team_id="test_team_id", + ) + + # Simulate a litellm_logging_obj with model_call_details containing + # the standard_logging_object (as set by _failure_handler_helper_fn) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "trace-id-from-logging-obj" + mock_logging_obj.model_call_details = { + "standard_logging_object": { + "trace_id": "trace-id-from-logging-obj", + "error_str": "InternalServerError", + } + } + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + "litellm_logging_obj": mock_logging_obj, + # Note: no "standard_logging_object" and no "litellm_trace_id" + } + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database: + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider error"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_kwargs = mock_update_database.call_args[1]["kwargs"] + + # standard_logging_object should have been propagated from logging obj + assert call_kwargs.get("standard_logging_object") is not None + assert ( + call_kwargs["standard_logging_object"]["trace_id"] + == "trace-id-from-logging-obj" + ) + # litellm_trace_id should also be propagated as a fallback + assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" + + @pytest.mark.asyncio async def test_enrich_failure_metadata_with_team_alias(): """ 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 b563a2e3c5..f4a7c4f499 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2741,6 +2741,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