diff --git a/litellm/__init__.py b/litellm/__init__.py index 056dd532f5f..8961de940a0 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -172,6 +172,7 @@ callbacks: List[ callback_settings: Dict[str, Dict[str, Any]] = {} initialized_langfuse_clients: int = 0 langfuse_default_tags: Optional[List[str]] = None +langfuse_enable_update_trace_keys: bool = False langsmith_batch_size: Optional[int] = None prometheus_initialize_budget_metrics: Optional[bool] = False prometheus_latency_buckets: Optional[List[float]] = None diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 8720f561e14..6d31f22b422 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -568,7 +568,10 @@ class LangFuseLogger: # 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 - update_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ())) + 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 () + ) debug: Final = clean_metadata.pop("debug_langfuse", None) mask_input: Final = _as_steering_flag(clean_metadata.pop("mask_input", False)) mask_output: Final = _as_steering_flag(clean_metadata.pop("mask_output", False)) diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index de04a65c310..6e57a36c5b6 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1,4 +1,5 @@ import datetime +import json import os import sys import types @@ -1332,35 +1333,72 @@ def test_mask_input_from_the_request_body_is_unchanged(mask_input, expect_redact assert (trace_params["input"] == _LANGFUSE_REDACTED) is expect_redacted -def test_update_trace_keys_header_applies_every_key(): +@pytest.mark.parametrize("flag", [True, "true"]) +def test_update_trace_keys_header_applies_every_key_when_enabled(flag): logger = _steering_logger() - trace_params, _ = _emit( - logger, - headers={ - "langfuse_existing_trace_id": "trace-1", - "langfuse_update_trace_keys": "trace_release, trace_tail", - "langfuse_trace_release": "v1.2.3", - "langfuse_trace_tail": "last", - }, - ) + with patch.object(litellm, "langfuse_enable_update_trace_keys", flag): + trace_params, _ = _emit( + logger, + headers={ + "langfuse_existing_trace_id": "trace-1", + "langfuse_update_trace_keys": "trace_release, trace_tail", + "langfuse_trace_release": "v1.2.3", + "langfuse_trace_tail": "last", + }, + ) assert trace_params["release"] == "v1.2.3" assert trace_params["tail"] == "last" -def test_update_trace_keys_from_the_request_body_list_is_unchanged(): +def test_update_trace_keys_is_off_by_default(): + """ + The caller picks the key name, so while the feature is on they can name + user_api_key_auth and have the resolved auth object, including team callback + credentials, serialized onto the trace. It stays inert until an operator opts in. + """ logger = _steering_logger() trace_params, _ = _emit( logger, metadata={ "existing_trace_id": "trace-1", - "update_trace_keys": ["trace_release"], + "update_trace_keys": ["user_api_key_auth", "trace_release"], + "user_api_key_auth": {"team_metadata": {"logging": [{"callback_vars": {"secret": "sk-canary"}}]}}, "trace_release": "v1.2.3", }, ) + assert "user_api_key_auth" not in trace_params + assert "release" not in trace_params + assert "sk-canary" not in json.dumps(trace_params, default=repr) + + +def test_update_trace_keys_input_and_output_are_gated_too(): + logger = _steering_logger() + + off, _ = _emit(logger, metadata={"existing_trace_id": "trace-1", "update_trace_keys": ["input", "output"]}) + with patch.object(litellm, "langfuse_enable_update_trace_keys", True): + on, _ = _emit(logger, metadata={"existing_trace_id": "trace-1", "update_trace_keys": ["input", "output"]}) + + assert "input" not in off and "output" not in off + assert "input" in on and "output" in on + + +def test_update_trace_keys_from_the_request_body_list_applies_when_enabled(): + logger = _steering_logger() + + with patch.object(litellm, "langfuse_enable_update_trace_keys", True): + trace_params, _ = _emit( + logger, + metadata={ + "existing_trace_id": "trace-1", + "update_trace_keys": ["trace_release"], + "trace_release": "v1.2.3", + }, + ) + assert trace_params["release"] == "v1.2.3"