mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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:
parent
edb9d6ef7f
commit
0fbc801bc5
4 changed files with 80 additions and 78 deletions
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue