diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index eaee27b1341..bf8d1d8e5fd 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -487,7 +487,12 @@ def _run_teardowns(state: _LangfuseLifecycleState, clients: tuple[Langfuse, ...] def _evict_if_stale_locked( - *, public_key: object, secret_key: object, base_url: object + *, + public_key: object, + secret_key: object, + base_url: object, + mock_mode: bool | None = None, + sample_rate: float | None = None, ) -> LangfuseResourceManager | None: """Assumes ``LangfuseResourceManager._lock`` is held; returns the still-valid bundle, evicting a stale one.""" if not public_key: @@ -495,7 +500,16 @@ def _evict_if_stale_locked( cached: Final = LangfuseResourceManager._instances.get(public_key) # pyright: ignore[reportPrivateUsage] # registry has no public accessor if cached is None: return None - if getattr(cached, "secret_key", None) == secret_key and getattr(cached, "base_url", None) == base_url: + same_exporter_kind: Final = mock_mode is None or ( + isinstance(getattr(cached, "span_exporter", None), DiscardingSpanExporter) == mock_mode + ) + same_sample_rate: Final = sample_rate is None or getattr(cached, "sample_rate", None) == sample_rate + if ( + getattr(cached, "secret_key", None) == secret_key + and getattr(cached, "base_url", None) == base_url + and same_exporter_kind + and same_sample_rate + ): return cached LangfuseResourceManager._instances.pop(public_key, None) # pyright: ignore[reportPrivateUsage] # registry has no public accessor return None @@ -611,6 +625,8 @@ def acquire_langfuse_client( public_key=public_key, secret_key=parameters.get("secret_key"), base_url=parameters.get("base_url"), + mock_mode=mock_mode, + sample_rate=sample_rate, ) client: Final = Langfuse( **parameters, # pyright: ignore[reportArgumentType] # kwargs-ok: dict mirrors the typed ctor, values resolved by the callers diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index cd5a9d57291..8bc8ed75b10 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -1230,3 +1230,38 @@ def test_second_client_on_the_same_key_does_not_build_another_provider(): assert len(_litellm_built_providers) == providers_after_first finally: LangfuseResourceManager._instances.pop(pk, None) + + +def test_leaving_mock_mode_on_the_same_key_stops_using_the_discarding_exporter(): + from litellm.integrations.langfuse.langfuse_sdk import DiscardingSpanExporter + + pk = "pk-mock-to-live-test" + LangfuseResourceManager._instances.pop(pk, None) + parameters = {"public_key": pk, "secret_key": "sk-live", "base_url": "http://127.0.0.1:1"} + try: + mocked = acquire_langfuse_client(parameters=parameters, environment=None, release=None, mock_mode=True) + assert isinstance(mocked._resources.span_exporter, DiscardingSpanExporter) + + live = acquire_langfuse_client(parameters=parameters, environment=None, release=None, mock_mode=False) + assert live._resources is not mocked._resources + assert not isinstance(live._resources.span_exporter, DiscardingSpanExporter) + assert LangfuseResourceManager._instances.get(pk) is live._resources + finally: + LangfuseResourceManager._instances.pop(pk, None) + + +def test_a_changed_sample_rate_on_the_same_key_rebuilds_the_bundle(monkeypatch: pytest.MonkeyPatch): + pk = "pk-resample-test" + LangfuseResourceManager._instances.pop(pk, None) + parameters = {"public_key": pk, "secret_key": "sk-resample", "base_url": "http://127.0.0.1:1"} + try: + monkeypatch.setenv("LANGFUSE_SAMPLE_RATE", "0.25") + quarter = acquire_langfuse_client(parameters=parameters, environment=None, release=None, mock_mode=True) + monkeypatch.setenv("LANGFUSE_SAMPLE_RATE", "1") + full = acquire_langfuse_client(parameters=parameters, environment=None, release=None, mock_mode=True) + + assert quarter._resources.sample_rate == 0.25 + assert full._resources is not quarter._resources + assert full._resources.sample_rate == 1.0 + finally: + LangfuseResourceManager._instances.pop(pk, None)