diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 10347bc7c6..f0f355b489 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -3,14 +3,27 @@ import os import traceback from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Union, + cast, +) from packaging.version import Version import litellm from litellm._logging import verbose_logger from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS -from litellm.litellm_core_utils.core_helpers import safe_deep_copy +from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, + reconstruct_model_name, +) from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info from litellm.llms.custom_httpx.http_handler import _get_httpx_client from litellm.secret_managers.main import str_to_bool @@ -437,12 +450,17 @@ class LangFuseLogger: ) ) + custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) + model_name = reconstruct_model_name( + kwargs.get("model", ""), custom_llm_provider, metadata + ) + trace.generation( CreateGeneration( name=metadata.get("generation_name", "litellm-completion"), startTime=start_time, endTime=end_time, - model=kwargs["model"], + model=model_name, modelParameters=optional_params, prompt=input, completion=output, @@ -543,7 +561,9 @@ class LangFuseLogger: # as we want to fall back to litellm_call_id instead for better traceability. # Note: Users can still explicitly set a UUID trace_id via metadata["trace_id"] (highest priority) if trace_id is None and standard_logging_object is not None: - standard_trace_id = cast(Optional[str], standard_logging_object.get("trace_id")) + standard_trace_id = cast( + Optional[str], standard_logging_object.get("trace_id") + ) # Only use standard_logging_object.trace_id if it's not a UUID # UUIDs are 36 characters with hyphens in format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx # We check for this specific pattern to avoid rejecting valid trace_ids that happen to have hyphens @@ -575,7 +595,9 @@ class LangFuseLogger: mask_output = clean_metadata.pop("mask_output", False) # Look for masking function in the dedicated location first (set by scrub_sensitive_keys_in_metadata) # Fall back to metadata for backwards compatibility - masking_function = litellm_params.get("_langfuse_masking_function") or clean_metadata.pop("langfuse_masking_function", None) + masking_function = litellm_params.get( + "_langfuse_masking_function" + ) or clean_metadata.pop("langfuse_masking_function", None) # Apply custom masking function if provided if masking_function is not None and callable(masking_function): @@ -776,12 +798,17 @@ class LangFuseLogger: if system_fingerprint is not None: optional_params["system_fingerprint"] = system_fingerprint + custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) + model_name = reconstruct_model_name( + kwargs.get("model", ""), custom_llm_provider, metadata + ) + generation_params = { "name": generation_name, "id": clean_metadata.pop("generation_id", generation_id), "start_time": start_time, "end_time": end_time, - "model": kwargs["model"], + "model": model_name, "model_parameters": optional_params, "input": input if not mask_input else "redacted-by-litellm", "output": output if not mask_output else "redacted-by-litellm", @@ -918,7 +945,9 @@ class LangFuseLogger: return Version(self.langfuse_sdk_version) >= Version("2.7.3") @staticmethod - def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any: + def _apply_masking_function( + data: Any, masking_function: Callable[[Any], Any] + ) -> Any: """ Apply a masking function to data, handling different data types. diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 9378ca71f5..dadb36f3fd 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -38,18 +38,18 @@ def safe_divide_seconds( def safe_divide( - numerator: Union[int, float], - denominator: Union[int, float], - default: Union[int, float] = 0 + numerator: Union[int, float], + denominator: Union[int, float], + default: Union[int, float] = 0, ) -> Union[int, float]: """ Safely divide two numbers, returning a default value if denominator is zero. - + Args: numerator: The number to divide denominator: The number to divide by default: Value to return if denominator is zero (defaults to 0) - + Returns: The result of numerator/denominator, or default if denominator is zero """ @@ -153,7 +153,8 @@ def get_metadata_variable_name_from_kwargs( - LiteLLM is now moving to using `litellm_metadata` for our metadata """ return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - + + def get_litellm_metadata_from_kwargs(kwargs: dict): """ Helper to get litellm metadata from all litellm request kwargs @@ -176,6 +177,25 @@ def get_litellm_metadata_from_kwargs(kwargs: dict): return {} +def reconstruct_model_name( + model_name: str, + custom_llm_provider: Optional[str], + metadata: dict, +) -> str: + """Reconstruct full model name with provider prefix for logging.""" + # Check if deployment model name from router metadata is available (has original prefix) + deployment_model_name = metadata.get("deployment") + if deployment_model_name and "/" in deployment_model_name: + # Use the deployment model name which preserves the original provider prefix + return deployment_model_name + elif custom_llm_provider and model_name and "/" not in model_name: + # Only add prefix for Bedrock (not for direct Anthropic API) + # This ensures Bedrock models get the prefix while direct Anthropic models don't + if custom_llm_provider == "bedrock": + return f"{custom_llm_provider}/{model_name}" + return model_name + + # Helper functions used for OTEL logging def _get_parent_otel_span_from_kwargs( kwargs: Optional[dict] = None, @@ -246,8 +266,8 @@ def safe_deep_copy(data): Safe Deep Copy The LiteLLM request may contain objects that cannot be pickled/deep-copied - (e.g., tracing spans, locks, clients). - + (e.g., tracing spans, locks, clients). + This helper deep-copies each top-level key independently; on failure keeps original ref """ @@ -306,23 +326,23 @@ def safe_deep_copy(data): def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: """ Recursively filter out Exception objects and callable objects from dicts/lists. - + This is a defensive utility to prevent deepcopy failures when exception objects are accidentally stored in parameter dictionaries (e.g., optional_params). Also filters callable objects (functions) to prevent JSON serialization errors. Exceptions and callables should not be stored in params - this function removes them. - + Args: data: The data structure to filter (dict, list, or any other type) max_depth: Maximum recursion depth to prevent infinite loops - + Returns: Filtered data structure with Exception and callable objects removed, or None if the entire input was an Exception or callable """ if max_depth <= 0: return data - + # Skip exception objects if isinstance(data, Exception): return None @@ -333,7 +353,7 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: obj_type_name = type(data).__name__ if obj_type_name in ["Logging", "LiteLLMLoggingObj"]: return None - + if isinstance(data, dict): result: dict[str, Any] = {} for k, v in data.items(): @@ -352,7 +372,9 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: result_list: list[Any] = [] for item in data: # Skip exception and callable items - if isinstance(item, Exception) or (callable(item) and not isinstance(item, type)): + if isinstance(item, Exception) or ( + callable(item) and not isinstance(item, type) + ): continue try: filtered = filter_exceptions_from_params(item, max_depth - 1) @@ -366,37 +388,35 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: return data -def filter_internal_params(data: dict, additional_internal_params: Optional[set] = None) -> dict: +def filter_internal_params( + data: dict, additional_internal_params: Optional[set] = None +) -> dict: """ Filter out LiteLLM internal parameters that shouldn't be sent to provider APIs. - + This removes internal/MCP-related parameters that are used by LiteLLM internally but should not be included in API requests to providers. - + Args: data: Dictionary of parameters to filter additional_internal_params: Optional set of additional internal parameter names to filter - + Returns: Filtered dictionary with internal parameters removed """ if not isinstance(data, dict): return data - + # Known internal parameters that should never be sent to provider APIs internal_params = { "skip_mcp_handler", "mcp_handler_context", "_skip_mcp_handler", } - + # Add any additional internal params if provided if additional_internal_params: internal_params.update(additional_internal_params) - + # Filter out internal parameters - return { - k: v - for k, v in data.items() - if k not in internal_params - } \ No newline at end of file + return {k: v for k, v in data.items() if k not in internal_params} diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c0090d1c3e..5c3c56e412 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -59,6 +59,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger +from litellm.litellm_core_utils.core_helpers import reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -332,9 +333,9 @@ class Logging(LiteLLMLoggingBaseClass): self.litellm_trace_id: str = litellm_trace_id or str(uuid.uuid4()) 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 @@ -719,9 +720,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 @@ -789,9 +790,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 ######################################################### @@ -803,9 +804,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 @@ -865,9 +866,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 @@ -896,10 +897,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", ""), @@ -910,34 +911,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: @@ -1238,13 +1239,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 @@ -1423,9 +1424,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: @@ -1451,9 +1452,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 @@ -1603,16 +1604,16 @@ class Logging(LiteLLMLoggingBaseClass): result=logging_result ) - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=logging_result, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - 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=logging_result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) def _transform_usage_objects(self, result): @@ -1667,9 +1668,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 @@ -1706,21 +1707,21 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, ) elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=result, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - 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=result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) 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 @@ -1870,23 +1871,23 @@ 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"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=complete_streaming_response, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - 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=complete_streaming_response, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) callbacks = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_success_callbacks, @@ -2214,10 +2215,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( @@ -2256,10 +2257,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"] @@ -2402,9 +2403,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: @@ -2415,10 +2416,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( @@ -2431,16 +2432,16 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = None ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=complete_streaming_response, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - 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=complete_streaming_response, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) callbacks = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_async_success_callbacks, @@ -2676,18 +2677,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 @@ -3301,7 +3302,9 @@ class Logging(LiteLLMLoggingBaseClass): # Deep copy result and add usage result_copy = result.model_copy(deep=True) - result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) + result_copy.usage = ( + usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) + ) return result_copy @@ -3629,9 +3632,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 endpoint=arize_config.endpoint, ) - 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) @@ -3642,7 +3645,6 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(_arize_otel_logger) return _arize_otel_logger # type: ignore elif logging_integration == "arize_phoenix": - from litellm.integrations.opentelemetry import ( OpenTelemetry, OpenTelemetryConfig, @@ -3658,13 +3660,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) @@ -3672,19 +3674,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 ( @@ -3816,9 +3818,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) @@ -4589,10 +4591,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 @@ -4898,25 +4900,6 @@ def _extract_response_obj_and_hidden_params( return response_obj, hidden_params -def _reconstruct_model_name( - model_name: str, - custom_llm_provider: Optional[str], - metadata: dict, -) -> str: - """Reconstruct full model name with provider prefix for logging.""" - # Check if deployment model name from router metadata is available (has original prefix) - deployment_model_name = metadata.get("deployment") - if deployment_model_name and "/" in deployment_model_name: - # Use the deployment model name which preserves the original provider prefix - return deployment_model_name - elif custom_llm_provider and model_name and "/" not in model_name: - # Only add prefix for Bedrock (not for direct Anthropic API) - # This ensures Bedrock models get the prefix while direct Anthropic models don't - if custom_llm_provider == "bedrock": - return f"{custom_llm_provider}/{model_name}" - return model_name - - def get_standard_logging_object_payload( kwargs: Optional[dict], init_response_obj: Union[Any, BaseModel, dict], @@ -5049,7 +5032,7 @@ def get_standard_logging_object_payload( # This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0" # are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) - model_name = _reconstruct_model_name( + model_name = reconstruct_model_name( kwargs.get("model", "") or "", custom_llm_provider, metadata ) @@ -5205,9 +5188,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/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 1861c44b69..1c457d7bf4 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,7 +11,10 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB, REDACTED_BY_LITELM_STRING -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, + reconstruct_model_name, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.utils import PrismaClient, hash_token @@ -100,9 +103,9 @@ def _get_spend_logs_metadata( clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata - clean_metadata["vector_store_request_metadata"] = ( - _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) - ) + clean_metadata[ + "vector_store_request_metadata" + ] = _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) clean_metadata["guardrail_information"] = guardrail_information clean_metadata["usage_object"] = usage_object clean_metadata["model_map_information"] = model_map_information @@ -393,6 +396,9 @@ def get_logging_payload( # noqa: PLR0915 # Extract agent_id for A2A requests (set directly on model_call_details) agent_id: Optional[str] = kwargs.get("agent_id") + custom_llm_provider = kwargs.get("custom_llm_provider") + raw_model = cast(str, kwargs.get("model") or "") + model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: SpendLogsPayload = SpendLogsPayload( @@ -403,7 +409,7 @@ def get_logging_payload( # noqa: PLR0915 startTime=_ensure_datetime_utc(start_time), endTime=_ensure_datetime_utc(end_time), completionStartTime=_ensure_datetime_utc(completion_start_time), - model=kwargs.get("model", "") or "", + model=model_name, user=metadata.get("user_api_key_user_id", "") or "", team_id=metadata.get("user_api_key_team_id", "") or "", organization_id=metadata.get("user_api_key_org_id") or "", @@ -449,7 +455,7 @@ def get_logging_payload( # noqa: PLR0915 # Explicitly clear large intermediate objects to reduce memory pressure del response_obj_dict, usage, clean_metadata, additional_usage_values - + return payload except Exception as e: verbose_proxy_logger.exception( diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json index 26b712c1cf..bd2f06b502 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json @@ -54,7 +54,7 @@ "id": "time-14-13-16-469836_chatcmpl-3803a9e9-aa68-4493-94d9-247f354830d6", "endTime": "2025-05-26T14:13:16.795438-07:00", "completionStartTime": "2025-05-26T14:13:16.795438-07:00", - "model": "anthropic.claude-3-5-sonnet-20240620-v1:0", + "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", "modelParameters": { "aws_region": "us-east-1" }, diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 89cc11c40d..cd9c401143 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -1,173 +1,45 @@ -import json -import os -import sys -from unittest.mock import MagicMock, patch +"""Tests for litellm_core_utils.core_helpers module.""" -import pytest - -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path - -from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, - safe_divide, - safe_deep_copy -) +from litellm.litellm_core_utils.core_helpers import reconstruct_model_name -def test_get_litellm_metadata_from_kwargs(): - kwargs = { - "litellm_params": { - "litellm_metadata": {}, - "metadata": {"user_api_key": "1234567890"}, - }, - } - assert get_litellm_metadata_from_kwargs(kwargs) == {"user_api_key": "1234567890"} +def test_reconstruct_model_name_prefers_deployment_value(): + """Ensure deployment metadata wins when reconstructing the model name.""" + metadata = {"deployment": "vertex_ai/gemini-1.5-flash"} -def test_add_missing_spend_metadata_to_litellm_metadata(): - litellm_metadata = {"test_key": "test_value"} - metadata = {"user_api_key_hash_value": "1234567890"} - kwargs = { - "litellm_params": { - "litellm_metadata": litellm_metadata, - "metadata": metadata, - }, - } - assert get_litellm_metadata_from_kwargs(kwargs) == { - "test_key": "test_value", - "user_api_key_hash_value": "1234567890", - } - - -def test_preserve_upstream_non_openai_attributes(): - from litellm.litellm_core_utils.core_helpers import ( - preserve_upstream_non_openai_attributes, - ) - from litellm.types.utils import ModelResponseStream - - model_response = ModelResponseStream( - id="123", - object="text_completion", - created=1715811200, - model="gpt-3.5-turbo", + result = reconstruct_model_name( + model_name="gemini-1.5-flash", + custom_llm_provider="vertex_ai", + metadata=metadata, ) - setattr(model_response, "test_key", "test_value") - preserve_upstream_non_openai_attributes( - model_response=ModelResponseStream(), - original_chunk=model_response, + assert result == "vertex_ai/gemini-1.5-flash" + + +def test_reconstruct_model_name_adds_bedrock_prefix_when_missing(): + """Bedrock model names without prefixes should gain the provider prefix.""" + + metadata = {} + + result = reconstruct_model_name( + model_name="us.anthropic.claude-3-sonnet", + custom_llm_provider="bedrock", + metadata=metadata, ) - assert model_response.test_key == "test_value" + assert result == "bedrock/us.anthropic.claude-3-sonnet" -def test_safe_divide_basic(): - """Test basic safe division functionality""" - # Normal division - result = safe_divide(10, 2) - assert result == 5.0, f"Expected 5.0, got {result}" - - # Division with float - result = safe_divide(7.5, 2.5) - assert result == 3.0, f"Expected 3.0, got {result}" - - # Division by zero with default - result = safe_divide(10, 0) - assert result == 0, f"Expected 0, got {result}" - - # Division by zero with custom default - result = safe_divide(10, 0, default=1) - assert result == 1, f"Expected 1, got {result}" - - # Division by zero with custom default as float - result = safe_divide(10, 0, default=0.5) - assert result == 0.5, f"Expected 0.5, got {result}" +def test_reconstruct_model_name_returns_original_for_other_providers(): + """Non-Bedrock providers should not prepend anything.""" + metadata = {} -def test_safe_divide_edge_cases(): - """Test edge cases for safe division""" - # Zero numerator - result = safe_divide(0, 5) - assert result == 0.0, f"Expected 0.0, got {result}" - - # Negative numbers - result = safe_divide(-10, 2) - assert result == -5.0, f"Expected -5.0, got {result}" - - # Negative denominator - result = safe_divide(10, -2) - assert result == -5.0, f"Expected -5.0, got {result}" - - # Both negative - result = safe_divide(-10, -2) - assert result == 5.0, f"Expected 5.0, got {result}" - - # Float division - result = safe_divide(1, 3) - assert abs(result - 0.3333333333333333) < 1e-10, f"Expected ~0.333..., got {result}" + result = reconstruct_model_name( + model_name="claude-3-sonnet", + custom_llm_provider="anthropic", + metadata=metadata, + ) - -def test_safe_divide_weight_scenario(): - """Test safe division in the context of weight calculations""" - # Simulate weight calculation scenario - weights = [3, 7, 0, 2] - total_weight = sum(weights) # 12 - - # Normal case - normalized_weights = [safe_divide(w, total_weight) for w in weights] - expected = [0.25, 7/12, 0.0, 1/6] - - for i, (actual, exp) in enumerate(zip(normalized_weights, expected)): - assert abs(actual - exp) < 1e-10, f"Weight {i}: Expected {exp}, got {actual}" - - # Zero total weight scenario (division by zero) - zero_weights = [0, 0, 0] - zero_total = sum(zero_weights) # 0 - - # Should return default values (0) for all weights - normalized_zero_weights = [safe_divide(w, zero_total) for w in zero_weights] - expected_zero = [0, 0, 0] - - assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}" - - -def test_safe_deep_copy_with_non_pickleables_and_span(): - """ - Verify safe_deep_copy: - - does not crash when non-pickleables are present, - - preserves structure/keys, - - deep-copies JSON-y payloads (e.g., messages), - - keeps non-pickleables by reference, - - redacts OTEL span in the copy and restores it in the original. - """ - import threading - rlock = threading.RLock() - data = { - "metadata": {"litellm_parent_otel_span": rlock, "x": 1}, - "messages": [{"role": "user", "content": "hi"}], - "optional_params": {"handle": rlock}, - "ok": True, - } - - copied = safe_deep_copy(data) - - # Structure preserved - assert set(copied.keys()) == set(data.keys()) - - # Messages are deep-copied (new object, same content) - assert copied["messages"] is not data["messages"] - assert copied["messages"][0] == data["messages"][0] - - # Non-pickleable subtree kept by reference (no crash) - assert copied["optional_params"] is data["optional_params"] - assert copied["optional_params"]["handle"] is rlock - - # OTEL span: redacted in the copy, restored in original - assert copied["metadata"]["litellm_parent_otel_span"] == "placeholder" - assert data["metadata"]["litellm_parent_otel_span"] is rlock - - # Other simple fields unchanged - assert copied["ok"] is True - assert copied["metadata"]["x"] == 1 + assert result == "claude-3-sonnet"