mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-16 22:23:46 +00:00
Merge pull request #18573 from BerriAI/litellm_fix_normalize-model-name
fix: unify model names to provider-defined names
This commit is contained in:
YutaSaito
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
return {k: v for k, v in data.items() if k not in internal_params}
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
+1
-1
@@ -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"
|
||||
},
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user