diff --git a/litellm/integrations/SlackAlerting/utils.py b/litellm/integrations/SlackAlerting/utils.py index 2e203b5ab98..eb3a7f80f72 100644 --- a/litellm/integrations/SlackAlerting/utils.py +++ b/litellm/integrations/SlackAlerting/utils.py @@ -3,9 +3,11 @@ Utils used for slack alerting """ import asyncio +from collections.abc import Callable from typing import TYPE_CHECKING, Any, Final import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import AlertType from litellm.secret_managers.main import get_secret @@ -68,7 +70,9 @@ async def _add_langfuse_trace_id_to_alert( """ from litellm.integrations.langfuse.langfuse import LangFuseLogger, resolve_langfuse_host - callbacks: Final = litellm.logging_callback_manager._get_all_callbacks() + callbacks: Final[list[CustomLogger | Callable[..., object] | str]] = ( + litellm.logging_callback_manager._get_all_callbacks() + ) if not any(callback == "langfuse" or isinstance(callback, LangFuseLogger) for callback in callbacks): return None @@ -76,7 +80,12 @@ async def _add_langfuse_trace_id_to_alert( return None litellm_logging_obj: Final[Logging] = request_data["litellm_logging_obj"] - host: Final = resolve_langfuse_host(litellm_logging_obj.standard_callback_dynamic_params.get("langfuse_host")) + instance_host: Final = next( + (callback.langfuse_host for callback in callbacks if isinstance(callback, LangFuseLogger)), None + ) + host: Final = resolve_langfuse_host( + litellm_logging_obj.standard_callback_dynamic_params.get("langfuse_host") or instance_host + ) for _ in range(3): if (trace_id := litellm_logging_obj._get_trace_id(service_name="langfuse")) is not None: return f"{host}/trace/{trace_id}" diff --git a/litellm/integrations/langfuse/langfuse_sdk.py b/litellm/integrations/langfuse/langfuse_sdk.py index 51bc22afb0b..f27d8f5e72c 100644 --- a/litellm/integrations/langfuse/langfuse_sdk.py +++ b/litellm/integrations/langfuse/langfuse_sdk.py @@ -604,6 +604,14 @@ def acquire_langfuse_tracing( return created +def flush_langfuse_tracing(timeout_millis: int = 30_000) -> bool: + """Force-flush every export channel this process acquired; ``True`` when all of them succeeded.""" + with _TRACING_LOCK: + channels: Final = tuple(_TRACING.values()) + results: Final = tuple(channel.flush(timeout_millis) for channel in channels) + return all(results) + + def build_langfuse_tracing( *, exporter: SpanExporter, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7f88d7b005..23ee029cd78 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1027,14 +1027,11 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N if shutdown_billing_metrics_recorder is not None: shutdown_billing_metrics_recorder() - # flush remaining langfuse logs - if "langfuse" in litellm.success_callback: + if "litellm.integrations.langfuse.langfuse_sdk" in sys.modules: try: - # flush langfuse logs on shutdow - from litellm.utils import langFuseLogger + from litellm.integrations.langfuse.langfuse_sdk import flush_langfuse_tracing - if langFuseLogger is not None: - langFuseLogger.flush() + flush_langfuse_tracing() except Exception: # [DO NOT BLOCK shutdown events for this] pass diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py index e27f0ac96d6..d7e9a1c97cf 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py @@ -76,9 +76,10 @@ async def test_langfuse_trace_url_when_callback_registered_as_logger_instance(mo monkeypatch.setattr(litellm, "_async_success_callback", []) monkeypatch.setattr(litellm, "_async_failure_callback", []) monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setenv("LANGFUSE_HOST", "http://env-host.invalid") logging_obj = MagicMock() logging_obj._get_trace_id.return_value = "trace-from-instance" - logging_obj.standard_callback_dynamic_params = {"langfuse_host": "http://127.0.0.1:1"} + logging_obj.standard_callback_dynamic_params = {} result = await _add_langfuse_trace_id_to_alert({"litellm_logging_obj": logging_obj}) diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py index 1e591f42cf0..9f365097780 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_sdk.py @@ -35,6 +35,7 @@ from litellm.integrations.langfuse.langfuse_sdk import ( build_langfuse_client, build_langfuse_tracing, configured_sample_rate, + flush_langfuse_tracing, observation_attributes, resolve_observation_id, resolve_trace_id, @@ -502,6 +503,26 @@ def test_changed_credentials_or_settings_get_their_own_channel(override): assert _acquire() is not _acquire(**override) +def test_flush_langfuse_tracing_exports_the_queued_spans_of_every_channel(monkeypatch: pytest.MonkeyPatch): + """The proxy shutdown hook flushes through this, so a span finished just before a + graceful restart must reach the exporter without waiting for the batch interval.""" + exporters: Final[ + list[InMemorySpanExporter] + ] = [] # mutable-ok: collects the exporters the patched builder hands out + + def build_in_memory(*, public_key: str, secret_key: str, base_url: str) -> InMemorySpanExporter: + exporters.append(InMemorySpanExporter()) + return exporters[-1] + + monkeypatch.setattr("litellm.integrations.langfuse.langfuse_sdk._build_span_exporter", build_in_memory) + for public_key in ("pk-flush-test-a", "pk-flush-test-b"): + _acquire(public_key=public_key, mock_mode=False, flush_interval=600.0).tracer.start_span("generation").end() + + assert [len(exporter.get_finished_spans()) for exporter in exporters] == [0, 0] + assert flush_langfuse_tracing() is True + assert [len(exporter.get_finished_spans()) for exporter in exporters] == [1, 1] + + def test_a_changed_sample_rate_rebuilds_the_channel(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("LANGFUSE_SAMPLE_RATE", "0.25") quarter = _acquire(public_key="pk-resample-test") @@ -624,9 +645,29 @@ def test_ssl_exporter_carries_litellm_tls_material(monkeypatch, tmp_path): ("https://lf.internal.example", "/otel/traces", "https://lf.internal.example/otel/traces"), ("https://lf.internal.example/", "/otel/traces", "https://lf.internal.example/otel/traces"), ("https://lf.internal.example", "otel/traces", "https://lf.internal.example/otel/traces"), + ( + "https://lf.internal.example", + "//elsewhere.example/otel", + "https://lf.internal.example/elsewhere.example/otel", + ), + ( + "https://lf.internal.example", + "https://elsewhere.example/otel", + "https://lf.internal.example/https://elsewhere.example/otel", + ), + ], + ids=[ + "default", + "leading-slash", + "both-slashes", + "no-slash", + "scheme-relative-stays-on-host", + "absolute-stays-on-host", ], ) -def test_export_endpoint_never_doubles_the_slash(monkeypatch, base_url, export_path, expected): +def test_export_endpoint_never_doubles_the_slash_or_leaves_the_configured_host( + monkeypatch, base_url, export_path, expected +): if export_path is None: monkeypatch.delenv("LANGFUSE_OTEL_TRACES_EXPORT_PATH", raising=False) else: diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index deb7289d2d1..acfd765e40b 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -130,6 +130,28 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): } +@pytest.mark.asyncio +async def test_proxy_shutdown_flushes_every_langfuse_export_channel(monkeypatch): + """A generation finished just before a graceful restart is still queued in its batch + processor, so shutdown must flush every acquired export channel.""" + from litellm.integrations.langfuse import langfuse_sdk + + flushed = MagicMock(return_value=True) + monkeypatch.setattr(langfuse_sdk, "flush_langfuse_tracing", flushed) + monkeypatch.setattr(ps, "prisma_client", None, raising=False) + monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + await proxy_shutdown_event() + + assert flushed.call_count == 1 + + @pytest.mark.asyncio async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monkeypatch): """