diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 4b55c1eef76..46283d8a499 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -45,14 +45,13 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - from langfuse import Langfuse - from opentelemetry.context import Context + from langfuse import Langfuse, LangfuseGeneration from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache else: - Context = Any DynamicLoggingCache = Any Langfuse = Any + LangfuseGeneration = Any _DENIED_STEERING_KEYS: Final = frozenset({"headers", "endpoint", "caching_groups", "previous_models"}) @@ -984,18 +983,6 @@ class LangFuseLogger: parent_observation_id=resolve_observation_id(parent_observation_id), # pyright: ignore[reportArgumentType] # metadata value, str or None at runtime existing_trace=existing_trace_id is not None, ) - log_provider_specific_information_as_span( - client=self.Langfuse, - context=trace_context, - enrichments=enrichments, - claim_trace_root=claim_trace_root, - ) - self._log_guardrail_information_as_span( - client=self.Langfuse, - context=trace_context, - standard_logging_object=standard_logging_object, - claim_trace_root=claim_trace_root, - ) generation: Final = start_generation( client=self.Langfuse, context=trace_context, @@ -1014,6 +1001,12 @@ class LangFuseLogger: input=trace_params.get("input") if "input" in update_trace_keys else None, output=trace_params.get("output") if "output" in update_trace_keys else None, ) + log_provider_specific_information_as_span( + client=self.Langfuse, parent=generation, enrichments=enrichments + ) + self._log_guardrail_information_as_span( + client=self.Langfuse, parent=generation, standard_logging_object=standard_logging_object + ) generation.end(end_time=to_unix_nanos(end_time)) # log_event_on_langfuse tuple-unpacks this and re-wraps it in the dict callers cache. @@ -1158,9 +1151,8 @@ class LangFuseLogger: def _log_guardrail_information_as_span( self, client: "Langfuse", - context: "Context", + parent: "LangfuseGeneration", standard_logging_object: StandardLoggingPayload | None, - claim_trace_root: bool, ): """ Log guardrail information as a span @@ -1193,10 +1185,9 @@ class LangFuseLogger: span = start_child_span( client=client, - context=context, + parent=parent, name="guardrail", start_time=guardrail_entry.get("start_time", None), - claim_trace_root=claim_trace_root, attributes={ # mutable-ok: langfuse serializes this payload, a proxy is not json-encodable "input": guardrail_entry.get("guardrail_request", None), "output": guardrail_entry.get("guardrail_response", None), @@ -1288,20 +1279,10 @@ def _add_prompt_to_generation_params( def log_provider_specific_information_as_span( *, client: "Langfuse", - context: "Context", + parent: "LangfuseGeneration", enrichments: Mapping[str, Any], - claim_trace_root: bool, ): - """ - Logs provider-specific information as spans. - - Parameters: - trace: The tracing object used to log spans. - enrichments: The litellm-computed fields on the emitted payload. - - Returns: - None - """ + """Logs provider-specific information as spans under the generation.""" _hidden_params: Final[Mapping[str, object] | None] = enrichments.get("hidden_params", None) if _hidden_params is None: @@ -1314,38 +1295,23 @@ def log_provider_specific_information_as_span( for elem in vertex_ai_grounding_metadata: if isinstance(elem, dict): for key, value in elem.items(): - _end_grounding_span( - client=client, context=context, name=key, value=value, claim_trace_root=claim_trace_root - ) + _end_grounding_span(client=client, parent=parent, name=key, value=value) else: - _end_grounding_span( - client=client, - context=context, - name="vertex_ai_grounding_metadata", - value=elem, - claim_trace_root=claim_trace_root, - ) + _end_grounding_span(client=client, parent=parent, name="vertex_ai_grounding_metadata", value=elem) else: _end_grounding_span( - client=client, - context=context, - name="vertex_ai_grounding_metadata", - value=vertex_ai_grounding_metadata, - claim_trace_root=claim_trace_root, + client=client, parent=parent, name="vertex_ai_grounding_metadata", value=vertex_ai_grounding_metadata ) -def _end_grounding_span( - *, client: "Langfuse", context: "Context", name: str, value: object, claim_trace_root: bool -) -> None: +def _end_grounding_span(*, client: "Langfuse", parent: "LangfuseGeneration", name: str, value: object) -> None: from litellm.integrations.langfuse.langfuse_sdk import start_child_span start_child_span( client=client, - context=context, + parent=parent, name=name, start_time=None, - claim_trace_root=claim_trace_root, attributes={"input": value}, # mutable-ok: langfuse serializes this payload ).end() diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index fc6e8a5e880..eaee27b1341 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -168,23 +168,22 @@ def start_generation( def start_child_span( *, client: Langfuse, - context: Context, + parent: LangfuseGeneration, name: str, start_time: datetime | float | None, - claim_trace_root: bool, attributes: Mapping[str, object], ) -> LangfuseSpan: - """Create a sibling observation inside the same trace, keeping its own window. + """Create an observation under the generation, keeping its own time window. - When the shared parent is the fabricated remote span, every observation must - claim trace root itself — the SDK's own remote-parent paths stamp each span — - or it exports with a parent id that is never exported. + The server derives the trace's name and I/O from every observation marked + root, last start time wins, so only the generation may claim root. Nesting + the rest under it keeps a post-call guardrail from rewriting the trace. """ otel_span: Final = client._otel_tracer.start_span( # pyright: ignore[reportPrivateUsage] # only route to a historical start time - name=name, context=context, start_time=to_unix_nanos(start_time) + name=name, + context=otel_trace.set_span_in_context(parent._otel_span), # pyright: ignore[reportPrivateUsage] # the wrapper exposes no public span handle + start_time=to_unix_nanos(start_time), ) - if claim_trace_root: - otel_span.set_attribute(AS_ROOT_ATTRIBUTE, True) return LangfuseSpan(otel_span=otel_span, langfuse_client=client, **attributes) # pyright: ignore[reportArgumentType] # kwargs-ok: callback-built params, v2 accepted the same shapes diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index bf4a380d573..cd5a9d57291 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -110,17 +110,17 @@ def test_guardrail_span_with_float_timestamps_does_not_break_the_generation(clie lf, exporter = client context, claim_root = open_trace_context(client=lf, trace_id="9" * 32, parent_observation_id=None) guardrail_start = 1709294400.0 + generation = start_generation( + client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={} + ) start_child_span( client=lf, - context=context, + parent=generation, name="guardrail", start_time=guardrail_start, - claim_trace_root=claim_root, attributes={}, ).end(end_time=to_unix_nanos(guardrail_start + 2)) - start_generation( - client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={} - ).end(end_time=to_unix_nanos(CALL_END)) + generation.end(end_time=to_unix_nanos(CALL_END)) lf.flush() guardrail = _only_span(exporter, "guardrail") @@ -152,30 +152,35 @@ def test_generation_claims_trace_root_only_without_a_real_parent(client): assert _only_span(exporter, "child-gen").attributes.get(AS_ROOT_ATTRIBUTE) is None -def test_child_span_keeps_its_own_window_and_stays_a_sibling(client): +def test_child_span_keeps_its_own_window_and_only_the_generation_claims_root(client): + """The server takes trace name and I/O from every root observation, latest start wins. + + A post-call guardrail starts after the model call, so if it also claimed root the + trace would show the guardrail's I/O instead of the model's. + """ lf, exporter = client context, claim_root = open_trace_context(client=lf, trace_id="d" * 32, parent_observation_id=None) - guardrail_start = CALL_START + timedelta(seconds=1) + generation = start_generation( + client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={} + ) + guardrail_start = CALL_END + timedelta(seconds=1) start_child_span( client=lf, - context=context, + parent=generation, name="guardrail", start_time=guardrail_start, - claim_trace_root=claim_root, attributes={}, ).end(end_time=to_unix_nanos(guardrail_start + timedelta(seconds=2))) - start_generation( - client=lf, context=context, name="gen", start_time=CALL_START, claim_trace_root=claim_root, attributes={} - ).end(end_time=to_unix_nanos(CALL_END)) + generation.end(end_time=to_unix_nanos(CALL_END)) lf.flush() guardrail = _only_span(exporter, "guardrail") - generation = _only_span(exporter, "gen") + exported_generation = _only_span(exporter, "gen") assert (guardrail.end_time - guardrail.start_time) == 2 * 1_000_000_000 - assert guardrail.context.trace_id == generation.context.trace_id - # the shared remote parent is fabricated and never exported, so both must claim trace root - assert guardrail.attributes.get(AS_ROOT_ATTRIBUTE) is True - assert generation.attributes.get(AS_ROOT_ATTRIBUTE) is True + assert guardrail.context.trace_id == exported_generation.context.trace_id + assert guardrail.parent.span_id == exported_generation.context.span_id + assert AS_ROOT_ATTRIBUTE not in guardrail.attributes + assert exported_generation.attributes.get(AS_ROOT_ATTRIBUTE) is True def test_release_is_carried_on_the_root_observation(client): diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 77a0b9c43ec..a825e3732c4 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -557,11 +557,13 @@ class TestLangfuseUsageDetails(unittest.TestCase): self.logger.Langfuse.flush() return [span for span in self.span_exporter.get_finished_spans() if span.name == name] - def _drive_with_canary(self, extra_metadata=None, hidden_params=None): + def _drive_with_canary(self, extra_metadata=None, hidden_params=None, guardrail_information=None): metadata = {**self._canary_request_metadata(), **(extra_metadata or {})} payload = self._build_standard_logging_payload(trace_id="canary-trace-id") if hidden_params is not None: payload["hidden_params"] = hidden_params + if guardrail_information is not None: + payload["guardrail_information"] = guardrail_information kwargs = {**self._build_langfuse_kwargs(payload), "response_cost": 0.25} self.use_real_langfuse_client() @@ -641,6 +643,36 @@ class TestLangfuseUsageDetails(unittest.TestCase): assert span_inputs == ["ground-a", "ground-b"] assert self.CANARY not in self._emitted_payload_text() + def test_only_the_generation_claims_the_trace_root(self): + """ + Langfuse v4 derives trace name and I/O from every observation marked root, and + the one with the latest start wins. A post-call guardrail starts after the model + call, so it must nest under the generation instead of claiming root itself, or + the trace shows the guardrail's request instead of the model's. + """ + self._drive_with_canary( + hidden_params={"vertex_ai_grounding_metadata": ["ground-a"]}, + guardrail_information=[ + { + "guardrail_name": "pii-post", + "guardrail_mode": "post_call", + "guardrail_request": {"texts": ["post-call scan"]}, + "guardrail_response": {"flagged": False}, + "start_time": 1704110402.0, + "end_time": 1704110403.0, + } + ], + ) + + [generation] = [span for span in self.span_exporter.get_finished_spans() if span.name.startswith("litellm-")] + [guardrail] = self.exported_spans_named("guardrail") + [grounding] = self.exported_spans_named("vertex_ai_grounding_metadata") + assert generation.attributes.get("langfuse.internal.as_root") is True + for child in (guardrail, grounding): + assert child.parent.span_id == generation.context.span_id + assert child.context.trace_id == generation.context.trace_id + assert "langfuse.internal.as_root" not in child.attributes + def test_caller_cannot_spoof_an_allowlisted_identity_field(self): """ Request metadata never reaches the blob, so a caller naming user_api_key_alias