Merge pull request #22390 from Harshit28j/litellm_langfuse-session-trace-fix

Fix Langfuse failure path kwargs inconsistency
This commit is contained in:
yuneng-jiang
2026-03-14 12:49:04 -07:00
committed by GitHub
11 changed files with 883 additions and 171 deletions
+1 -1
View File
@@ -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,
)
+152 -150
View File
@@ -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():
"""
+59 -14
View File
@@ -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="")
+75
View File
@@ -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