fix(langfuse): nest guardrail and grounding spans under the generation

Langfuse v4 derives the trace name and I/O from every observation marked as_root, and the one with the latest start time wins. Guardrail and grounding spans used to claim root next to the generation, so a post_call guardrail could replace the model's request and response on the trace with its own. Only the generation claims root now; the sibling spans become its children

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-17 03:56:20 +00:00
parent edb9d6ef7f
commit 0fbc801bc5
4 changed files with 80 additions and 78 deletions

View file

@ -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()

View file

@ -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

View file

@ -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):

View file

@ -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