From a84ddf40a95142ef3741627de69cd26dbf1cfcb2 Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 7 Sep 2026 15:55:31 +0000 Subject: [PATCH] fix(langfuse): give each call in a session its own trace instead of upserting one trace per session Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/langfuse/langfuse.py | 41 +++++++++--- .../integrations/test_langfuse.py | 63 +++++++++++++++++++ 2 files changed, 94 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 9576eabaa34..4e10c24674a 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -63,6 +63,20 @@ def _object_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, dict) else None +def _resolve_langfuse_trace_id( + metadata_trace_id: object, + standard_trace_id: object, + litellm_trace_id: object, + litellm_call_id: str | None, + session_id: object, +) -> str | None: + """Resolve the Langfuse trace id; a session id never doubles as the trace id.""" + resolved: Final = metadata_trace_id or standard_trace_id or litellm_trace_id or litellm_call_id + if session_id is not None and resolved == session_id: + return litellm_call_id + return str(resolved) if resolved is not None else None + + class _UsageObject(Protocol): """Token-count surface the Langfuse logger reads off a response usage payload.""" @@ -596,19 +610,26 @@ class LangFuseLogger: session_id: Final = clean_metadata.pop("session_id", None) trace_name = cast(str | None, clean_metadata.pop("trace_name", None)) - trace_id = clean_metadata.pop("trace_id", None) - # 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 - if trace_id is None and standard_logging_object is not None: - trace_id = cast(str | None, standard_logging_object.get("trace_id")) - # Fallback to litellm_call_id if no trace_id found - if trace_id is None: - trace_id = kwargs.get("litellm_trace_id") or litellm_call_id + metadata_trace_id: Final = clean_metadata.pop("trace_id", None) + standard_trace_id: Final = ( + cast(str | None, standard_logging_object.get("trace_id")) + if standard_logging_object is not None + else None + ) existing_trace_id: Final = 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 - if existing_trace_id is not None: - trace_id = existing_trace_id + trace_id: Final = ( + existing_trace_id + if existing_trace_id is not None + else _resolve_langfuse_trace_id( + metadata_trace_id=metadata_trace_id, + standard_trace_id=standard_trace_id, + litellm_trace_id=kwargs.get("litellm_trace_id"), + litellm_call_id=litellm_call_id, + session_id=session_id, + ) + ) requested_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) update_trace_keys: Final = ( requested_trace_keys if _as_steering_flag(litellm.langfuse_enable_update_trace_keys) else () diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index d36878e455f..608721998fb 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -902,6 +902,69 @@ class TestLangfuseUsageDetails(unittest.TestCase): # session_id must still be set for session grouping assert self.last_trace_kwargs.get("session_id") == "session-999" + def _drive_session_call(self, litellm_call_id, metadata_trace_id="sess-abc"): + """One _log_langfuse_v2 call shaped like a proxied request that carries a session id.""" + payload = self._build_standard_logging_payload(trace_id="sess-abc") + kwargs = self._build_langfuse_kwargs(payload) + kwargs["litellm_trace_id"] = "sess-abc" + 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": "sess-abc", "trace_id": metadata_trace_id}, + litellm_params={"metadata": {"session_id": "sess-abc", "trace_id": metadata_trace_id}}, + 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=litellm_call_id, + ) + return self.last_trace_kwargs + + def test_log_langfuse_v2_session_header_trace_id_falls_back_to_call_id(self): + """ + When the proxy stamps the session id into every trace_id slot + (metadata, standard_logging_object, litellm_trace_id), the Langfuse + trace id must be the per-call litellm_call_id so each call gets its + own trace inside the shared session. + """ + trace_kwargs = self._drive_session_call(litellm_call_id="call-1") + + assert trace_kwargs.get("id") == "call-1" + assert trace_kwargs.get("session_id") == "sess-abc" + + def test_log_langfuse_v2_session_calls_produce_distinct_trace_ids(self): + """ + Two calls in one session must land as two Langfuse traces. With the + session id used as the trace id, Langfuse upserts both calls onto a + single trace, which is the regression this guards. + """ + first = self._drive_session_call(litellm_call_id="call-1") + second = self._drive_session_call(litellm_call_id="call-2") + + assert first.get("id") == "call-1" + assert second.get("id") == "call-2" + assert first.get("id") != second.get("id") + + def test_log_langfuse_v2_explicit_distinct_trace_id_kept_with_session_id(self): + """ + An explicit metadata.trace_id that differs from the session id is a + real trace id and must be kept, with the session id still applied. + """ + trace_kwargs = self._drive_session_call(litellm_call_id="call-1", metadata_trace_id="trace-xyz") + + assert trace_kwargs.get("id") == "trace-xyz" + assert trace_kwargs.get("session_id") == "sess-abc" + def test_failure_handler_langfuse_kwargs_excludes_original_response(): """