mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-07 04:24:12 +00:00
Merge pull request #22390 from Harshit28j/litellm_langfuse-session-trace-fix
Fix Langfuse failure path kwargs inconsistency
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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():
|
||||
"""
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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():
|
||||
"""
|
||||
|
||||
@@ -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():
|
||||
"""
|
||||
|
||||
@@ -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/<model_id>"
|
||||
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="")
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user