From 6fce7415bd4f57196459d044719ee9f052d53f18 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 28 Feb 2026 17:34:59 +0530 Subject: [PATCH] Fix Langfuse failure path kwargs and add session_id trace tests Fix: The Langfuse failure logging path was passing self.model_call_details (which includes original_response, potentially a coroutine) instead of the clean local kwargs copy. This aligns the failure path with the success path behavior (litellm_logging.py:2956). Reverted the session_id-as-trace_id approach as it causes trace collisions in Langfuse (multiple calls in the same session would overwrite each other). Instead, session_id is correctly used only for Langfuse session grouping via trace_params["session_id"], while each call retains its own unique trace_id. Added 4 tests: - session_id correctly passed as trace session_id (not trace_id) - session_id preserved for ERROR level (failure) logs - explicit trace_id takes priority over session_id - failure path kwargs excludes original_response Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/langfuse/langfuse.py | 5 +- .../integrations/test_langfuse.py | 182 ++++++++++++++++++ 2 files changed, 183 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 70f1161792a..7bf97665fd2 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -604,12 +604,9 @@ class LangFuseLogger: session_id = clean_metadata.pop("session_id", None) trace_name = cast(Optional[str], clean_metadata.pop("trace_name", None)) trace_id = clean_metadata.pop("trace_id", None) - # If session_id is provided, use it as trace_id for consistent trace mapping - if trace_id is None and session_id is not None: - trace_id = session_id # Use standard_logging_object.trace_id if available (when trace_id from metadata is None) # This allows standard trace_id to be used when provided in standard_logging_object - elif trace_id is None and standard_logging_object is not None: + if trace_id is None and standard_logging_object is not None: trace_id = cast( Optional[str], standard_logging_object.get("trace_id") ) diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 7168f5a4332..46ab4fb9f6e 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -455,6 +455,188 @@ class TestLangfuseUsageDetails(unittest.TestCase): assert self.last_trace_kwargs.get("id") == "call-id-xyz" + 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 Langfuse failure logging path passes the local kwargs copy + (which excludes 'original_response') rather than self.model_call_details directly. + This prevents passing coroutines or large response objects to the Langfuse logger. + """ + from litellm.litellm_core_utils.litellm_logging import Logging + + # Create a mock coroutine to simulate original_response + mock_coroutine = MagicMock() + mock_coroutine.__class__.__name__ = "coroutine" + + model_call_details = { + "litellm_call_id": "test-call-id", + "litellm_trace_id": None, + "model": "gpt-4", + "messages": [{"role": "user", "content": "test"}], + "litellm_params": { + "metadata": {"session_id": "test-session"}, + "litellm_session_id": None, + }, + "original_response": mock_coroutine, + "optional_params": {}, + "stream": False, + "call_type": "completion", + "input": [{"role": "user", "content": "test"}], + } + + captured_kwargs = {} + + class MockLangfuseLogger: + def log_event_on_langfuse(self, **log_kwargs): + captured_kwargs.update(log_kwargs) + return {"trace_id": "mock-trace-id", "generation_id": "mock-gen-id"} + + mock_logger = MockLangfuseLogger() + + # Simulate the failure path logic from litellm_logging.py + # This mirrors lines 2937-2957 of the failure_handler + kwargs = {} + for k, v in model_call_details.items(): + if k != "original_response": + kwargs[k] = v + + # Verify the local kwargs does NOT contain original_response + assert "original_response" not in kwargs + # Verify session_id is present in kwargs metadata + assert kwargs["litellm_params"]["metadata"]["session_id"] == "test-session" + + # Call with the local kwargs (as the fix does) + mock_logger.log_event_on_langfuse( + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + response_obj=None, + user_id=kwargs.get("user", None), + status_message="TestError: something failed", + level="ERROR", + kwargs=kwargs, + ) + + # Verify original_response is NOT in the kwargs passed to Langfuse + assert "original_response" not in captured_kwargs.get("kwargs", {}) + # Verify session_id metadata is preserved in the kwargs passed to Langfuse + langfuse_metadata = captured_kwargs["kwargs"]["litellm_params"]["metadata"] + assert langfuse_metadata["session_id"] == "test-session" + + def test_max_langfuse_clients_limit(): """ Test that the max langfuse clients limit is respected when initializing multiple clients