diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index c15789cf903..f52323a89f5 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -902,6 +902,13 @@ class LangFuseLogger: public=_trace_public_flag(trace_params.get("public")), attributes=_generation_attributes(generation_params, propagated=propagated_trace_attributes), ) + if existing_trace_id is not None and ("input" in update_trace_keys or "output" in update_trace_keys): + # with a real parent the generation is not the trace root, so trace-level + # I/O has to be stamped explicitly; v2 updated the trace object directly + generation.set_trace_io( # pyright: ignore[reportDeprecated] # the SDK keeps it exactly for this legacy trace-level contract + 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, + ) 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. diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index 63a1e821604..b5a73cd9d81 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -18,6 +18,7 @@ from opentelemetry.context import Context from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult +from opentelemetry.sdk.trace.sampling import TraceIdRatioBased __all__ = ( "AS_ROOT_ATTRIBUTE", @@ -171,8 +172,14 @@ def build_isolated_tracer_provider(*, environment: str | None, release: str | No and langfuse spans to every other litellm destination. The resource is rebuilt here because langfuse only applies ``environment`` - and ``release`` when it constructs the provider itself. + and ``release`` when it constructs the provider itself, and the sampler is + rebuilt for the same reason: ``LANGFUSE_SAMPLE_RATE`` is otherwise silently + ignored and every trace exports. """ + raw_sample_rate: Final = os.environ.get("LANGFUSE_SAMPLE_RATE") + sample_rate: Final = float(raw_sample_rate) if raw_sample_rate is not None else 1.0 + if not 0.0 <= sample_rate <= 1.0: + raise ValueError(f"Sample rate must be between 0.0 and 1.0, got {sample_rate}") attributes: Final = MappingProxyType( { key: value @@ -180,7 +187,10 @@ def build_isolated_tracer_provider(*, environment: str | None, release: str | No if value is not None } ) - provider: Final = TracerProvider(resource=Resource.create(dict(attributes))) + provider: Final = TracerProvider( + resource=Resource.create(dict(attributes)), + sampler=TraceIdRatioBased(sample_rate) if sample_rate < 1 else None, + ) _litellm_built_providers.add(provider) return provider diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index 80b05512f8d..1902c4a55e5 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -346,6 +346,70 @@ def test_isolated_provider_carries_environment_and_release(): assert attributes["langfuse.release"] == "v9" +def test_langfuse_sample_rate_drops_spans_on_the_isolated_provider(monkeypatch): + """The SDK only installs its sampler on providers it builds itself; v2 sampled via the same env var.""" + monkeypatch.setenv("LANGFUSE_SAMPLE_RATE", "0") + dropped_exporter = InMemorySpanExporter() + dropping_provider = build_isolated_tracer_provider(environment=None, release=None) + dropping_provider.add_span_processor(SimpleSpanProcessor(dropped_exporter)) + dropping_provider.get_tracer("test").start_span("dropped").end() + assert not dropped_exporter.get_finished_spans() + + monkeypatch.delenv("LANGFUSE_SAMPLE_RATE") + kept_exporter = InMemorySpanExporter() + keeping_provider = build_isolated_tracer_provider(environment=None, release=None) + keeping_provider.add_span_processor(SimpleSpanProcessor(kept_exporter)) + keeping_provider.get_tracer("test").start_span("kept").end() + assert [span.name for span in kept_exporter.get_finished_spans()] == ["kept"] + + +def test_invalid_sample_rate_fails_at_construction_like_the_sdk(monkeypatch): + monkeypatch.setenv("LANGFUSE_SAMPLE_RATE", "1.5") + with pytest.raises(ValueError, match=r"between 0\.0 and 1\.0"): + build_isolated_tracer_provider(environment=None, release=None) + + +def test_environment_override_lands_per_span_despite_shared_resources(): + """The SDK registry is keyed on public key alone, so a second client for the + same key adopts the first client's provider; the observation wrapper stamps + each span with its own client's environment, which the server prefers over + the resource-level value.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + first = Langfuse( + public_key=PUBLIC_KEY, + secret_key="sk-original", + host="http://127.0.0.1:1", + environment="prod", + tracer_provider=provider, + span_exporter=exporter, + ) + second = Langfuse( + public_key=PUBLIC_KEY, + secret_key="sk-original", + host="http://127.0.0.1:1", + environment="staging", + ) + assert second._resources is first._resources + + for client, environment in ((first, "prod"), (second, "staging")): + context, claim_trace_root = open_trace_context(client=client, trace_id="a" * 32, parent_observation_id=None) + start_generation( + client=client, + context=context, + name=f"generation-{environment}", + start_time=CALL_START, + claim_trace_root=claim_trace_root, + attributes={}, + ).end() + first.flush() + + spans = {span.name: span for span in exporter.get_finished_spans()} + assert spans["generation-prod"].attributes["langfuse.environment"] == "prod" + assert spans["generation-staging"].attributes["langfuse.environment"] == "staging" + + def test_client_does_not_take_over_the_process_tracer_provider(): # the global provider can only be set once per process, so assert it is left # alone rather than assuming this test is the one that installed it diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 221866de729..46485d7145e 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1544,6 +1544,42 @@ def test_update_trace_keys_input_and_output_are_gated_too(): assert "input" in on and "output" in on +def test_update_trace_keys_input_output_reach_the_trace_even_under_a_parent(): + """With a real parent the generation is not the trace root, so trace-level + I/O must be stamped explicitly; v2 updated the trace object directly.""" + rig = _steering_logger() + + with patch.object(litellm, "langfuse_enable_update_trace_keys", True): + _, _, span = _emit( + rig, + metadata={ + "existing_trace_id": "trace-1", + "parent_observation_id": "b" * 16, + "update_trace_keys": ["input", "output"], + }, + ) + + assert "the-input" in str(span.attributes["langfuse.trace.input"]) + assert "the-output" in str(span.attributes["langfuse.trace.output"]) + + +def test_trace_io_is_not_stamped_when_update_trace_keys_does_not_ask(): + rig = _steering_logger() + + with patch.object(litellm, "langfuse_enable_update_trace_keys", True): + _, _, span = _emit( + rig, + metadata={ + "existing_trace_id": "trace-1", + "parent_observation_id": "b" * 16, + "update_trace_keys": ["trace_release"], + }, + ) + + assert "langfuse.trace.input" not in (span.attributes or {}) + assert "langfuse.trace.output" not in (span.attributes or {}) + + def test_update_trace_keys_from_the_request_body_list_applies_when_enabled(): rig = _steering_logger()