From 828f51d42c7818d5929b6b3dcb856d5e64e6f181 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 17 Sep 2026 02:00:57 +0000 Subject: [PATCH] fix(langfuse): put the Langfuse trace link back into Slack alerts The proxy registers LangfusePromptManagement for callbacks: ["langfuse"], so the alert helper never saw the literal "langfuse" string and returned before looking up the trace id, and the prompt management logger never stored the trace id it got back from log_event_on_langfuse. Recognize LangFuseLogger instances in the callback list, record the returned trace id in the shared service trace id cache, and skip the link when no trace id arrives Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/SlackAlerting/utils.py | 12 ++-- litellm/integrations/langfuse/langfuse.py | 6 +- .../langfuse/langfuse_prompt_management.py | 15 +++- litellm/litellm_core_utils/litellm_logging.py | 18 +---- .../service_trace_id_cache.py | 20 ++++++ litellm/types/integrations/langfuse.py | 5 ++ .../test_slack_alerting_utils.py | 42 ++++++++++- .../test_langfuse_prompt_management.py | 69 ++++++++++++++----- 8 files changed, 141 insertions(+), 46 deletions(-) create mode 100644 litellm/litellm_core_utils/specialty_caches/service_trace_id_cache.py diff --git a/litellm/integrations/SlackAlerting/utils.py b/litellm/integrations/SlackAlerting/utils.py index 77361860327..090213e1c90 100644 --- a/litellm/integrations/SlackAlerting/utils.py +++ b/litellm/integrations/SlackAlerting/utils.py @@ -66,11 +66,11 @@ async def _add_langfuse_trace_id_to_alert( -> trace_id -> litellm_call_id """ - if "langfuse" not in litellm.logging_callback_manager._get_all_callbacks(): + from litellm.integrations.langfuse.langfuse import LangFuseLogger + + callbacks: Final = litellm.logging_callback_manager._get_all_callbacks() + if not any(callback == "langfuse" or isinstance(callback, LangFuseLogger) for callback in callbacks): return None - ######################################################### - # Only run if langfuse is added as a callback - ######################################################### if request_data is not None and request_data.get("litellm_logging_obj", None) is not None: trace_id: str | None = None @@ -81,8 +81,8 @@ async def _add_langfuse_trace_id_to_alert( if trace_id is not None: break await asyncio.sleep(3) # wait 3s before retrying for trace id - ######################################################### - from litellm.integrations.langfuse.langfuse import LangFuseLogger + if trace_id is None: + return None langfuse_object: Final = litellm_logging_obj._get_callback_object(service_name="langfuse") if isinstance(langfuse_object, LangFuseLogger): diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 21cc74a7e93..2311575c023 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -494,7 +494,7 @@ class LangFuseLogger: user_id: str | None = None, level: str = "DEFAULT", status_message: str | None = None, - ) -> dict: + ) -> LangfuseLoggedEvent: """ Logs a success or error event on Langfuse """ @@ -560,10 +560,10 @@ class LangFuseLogger: verbose_logger.debug("Langfuse Layer Logging - final response object: %s", response_obj) verbose_logger.info("Langfuse Layer Logging - logging success") - return {"trace_id": trace_id, "generation_id": generation_id} + return LangfuseLoggedEvent(trace_id=trace_id, generation_id=generation_id) except Exception as e: verbose_logger.exception("Langfuse Layer Error(): Exception occured - %s", e) - return {"trace_id": None, "generation_id": None} + return LangfuseLoggedEvent(trace_id=None, generation_id=None) def _get_langfuse_input_output_content( self, diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 1cf8bbb6f38..d8b3689eb0a 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prompt_management_base import PromptManagementClient from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.types.integrations.langfuse import LangfuseLoggedEvent from litellm.types.llms.openai import AllMessageValues, ChatCompletionSystemMessage from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload @@ -16,6 +17,7 @@ from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPa from ...litellm_core_utils.specialty_caches.dynamic_logging_cache import ( DynamicLoggingCache, ) +from ...litellm_core_utils.specialty_caches.service_trace_id_cache import in_memory_trace_id_cache from ..prompt_management_base import PromptManagementBase from .langfuse import ( LangFuseLogger, @@ -132,6 +134,13 @@ def langfuse_client_init( return client +def _remember_trace_id(litellm_call_id: object, logged: LangfuseLoggedEvent) -> None: + trace_id: Final = logged["trace_id"] + if not isinstance(litellm_call_id, str) or trace_id is None: + return + in_memory_trace_id_cache.set_cache(litellm_call_id=litellm_call_id, service_name="langfuse", trace_id=trace_id) + + class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogger): def __init__( self, @@ -321,13 +330,14 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge standard_callback_dynamic_params=standard_callback_dynamic_params, in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, ) - langfuse_logger_to_use.log_event_on_langfuse( + logged: Final = langfuse_logger_to_use.log_event_on_langfuse( kwargs=kwargs, response_obj=response_obj, start_time=start_time, end_time=end_time, user_id=kwargs.get("user", None), ) + _remember_trace_id(litellm_call_id=kwargs.get("litellm_call_id"), logged=logged) except Exception as e: from litellm._logging import verbose_logger @@ -349,7 +359,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge status_message = str(kwargs.get("exception", "Unknown error")) if standard_logging_object is not None: status_message = standard_logging_object.get("error_str", None) or status_message - langfuse_logger_to_use.log_event_on_langfuse( + logged: Final = langfuse_logger_to_use.log_event_on_langfuse( start_time=start_time, end_time=end_time, response_obj=None, @@ -358,6 +368,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge level="ERROR", kwargs=kwargs, ) + _remember_trace_id(litellm_call_id=kwargs.get("litellm_call_id"), logged=logged) except Exception as e: from litellm._logging import verbose_logger diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 40621a2f68d..bbbce6513c4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -36,7 +36,7 @@ from litellm._logging import ( ) from litellm._uuid import uuid from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final -from litellm.caching.caching import DualCache, InMemoryCache +from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, @@ -203,6 +203,7 @@ from .initialize_dynamic_callback_params import ( initialize_standard_callback_dynamic_params as _initialize_standard_callback_dynamic_params, ) from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache +from .specialty_caches.service_trace_id_cache import in_memory_trace_id_cache if TYPE_CHECKING: from mcp.types import EmbeddedResource, ImageContent, TextContent @@ -329,21 +330,6 @@ last_fetched_at_keys: Final = None #### -class ServiceTraceIDCache: - def __init__(self) -> None: - self.cache = InMemoryCache() - - def get_cache(self, litellm_call_id: str, service_name: str) -> str | None: - key_name: Final = f"{service_name}:{litellm_call_id}" - response: Final = self.cache.get_cache(key=key_name) - return response - - def set_cache(self, litellm_call_id: str, service_name: str, trace_id: str) -> None: - key_name: Final = f"{service_name}:{litellm_call_id}" - self.cache.set_cache(key=key_name, value=trace_id) - - -in_memory_trace_id_cache: Final = ServiceTraceIDCache() in_memory_dynamic_logger_cache: Final = DynamicLoggingCache() # Cached lazy import for PrometheusLogger diff --git a/litellm/litellm_core_utils/specialty_caches/service_trace_id_cache.py b/litellm/litellm_core_utils/specialty_caches/service_trace_id_cache.py new file mode 100644 index 00000000000..f1f60d3e7b8 --- /dev/null +++ b/litellm/litellm_core_utils/specialty_caches/service_trace_id_cache.py @@ -0,0 +1,20 @@ +from typing import Final + +from ...caching import InMemoryCache + + +class ServiceTraceIDCache: + def __init__(self) -> None: + self.cache = InMemoryCache() + + def get_cache(self, litellm_call_id: str, service_name: str) -> str | None: + key_name: Final = f"{service_name}:{litellm_call_id}" + response: Final = self.cache.get_cache(key=key_name) + return response + + def set_cache(self, litellm_call_id: str, service_name: str, trace_id: str) -> None: + key_name: Final = f"{service_name}:{litellm_call_id}" + self.cache.set_cache(key=key_name, value=trace_id) + + +in_memory_trace_id_cache: Final = ServiceTraceIDCache() diff --git a/litellm/types/integrations/langfuse.py b/litellm/types/integrations/langfuse.py index 6742aefea39..fe070a3dd18 100644 --- a/litellm/types/integrations/langfuse.py +++ b/litellm/types/integrations/langfuse.py @@ -14,3 +14,8 @@ class LangfuseUsageDetails(TypedDict): total: int | None cache_creation_input_tokens: int | None cache_read_input_tokens: int | None + + +class LangfuseLoggedEvent(TypedDict): + trace_id: ReadOnly[str | None] + generation_id: ReadOnly[str | None] 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 1d3ae27adc6..2f02ddaf19c 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_utils.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -57,3 +57,43 @@ async def test_langfuse_trace_url_skips_non_langfuse_callback(monkeypatch): logging_obj._get_callback_object.return_value = object() assert await _add_langfuse_trace_id_to_alert({"litellm_logging_obj": logging_obj}) is None + + +@pytest.mark.asyncio +async def test_langfuse_trace_url_when_callback_registered_as_logger_instance(monkeypatch): + from litellm.integrations.langfuse.langfuse import LangFuseLogger + + logger = LangFuseLogger( + langfuse_public_key="pk-slack-instance", + langfuse_secret="sk-slack-instance", + langfuse_host="http://127.0.0.1:1", + ) + monkeypatch.setattr(litellm, "success_callback", [logger]) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + logging_obj = MagicMock() + logging_obj._get_trace_id.return_value = "trace-from-instance" + logging_obj._get_callback_object.return_value = logger + + result = await _add_langfuse_trace_id_to_alert({"litellm_logging_obj": logging_obj}) + + assert result == "http://127.0.0.1:1/trace/trace-from-instance" + + +@pytest.mark.asyncio +async def test_langfuse_trace_url_absent_when_trace_id_never_arrives(monkeypatch): + from litellm.integrations.langfuse.langfuse import LangFuseLogger + + monkeypatch.setattr(litellm, "success_callback", ["langfuse"]) + monkeypatch.setattr("litellm.integrations.SlackAlerting.utils.asyncio.sleep", AsyncMock()) + logging_obj = MagicMock() + logging_obj._get_trace_id.return_value = None + logging_obj._get_callback_object.return_value = LangFuseLogger( + langfuse_public_key="pk-slack-none", + langfuse_secret="sk-slack-none", + langfuse_host="http://127.0.0.1:1", + ) + + assert await _add_langfuse_trace_id_to_alert({"litellm_logging_obj": logging_obj}) is None diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py b/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py index ef8fa34b3f5..b1997b654ac 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py @@ -1,5 +1,4 @@ from datetime import datetime, timezone -from types import MappingProxyType from typing import Final from unittest.mock import MagicMock, patch @@ -21,9 +20,7 @@ class TestLangfusePromptManagement: # This also prevents test-ordering issues when earlier tests remove sys.modules["langfuse"]. self._mock_langfuse = MagicMock() self._mock_langfuse.version.__version__ = "3.0.0" - self._langfuse_patcher = patch.dict( - "sys.modules", {"langfuse": self._mock_langfuse} - ) + self._langfuse_patcher = patch.dict("sys.modules", {"langfuse": self._mock_langfuse}) self._langfuse_patcher.start() def teardown_method(self): @@ -35,9 +32,7 @@ class TestLangfusePromptManagement: patch.object( langfuse_prompt_management, "should_run_prompt_management" ) as mock_should_run_prompt_management, - patch.object( - langfuse_prompt_management, "_get_prompt_from_id" - ) as mock_get_prompt_from_id, + patch.object(langfuse_prompt_management, "_get_prompt_from_id") as mock_get_prompt_from_id, ): mock_should_run_prompt_management.return_value = True langfuse_prompt_management.get_chat_completion_prompt( @@ -55,9 +50,7 @@ class TestLangfusePromptManagement: def test_log_failure_event_runs_async_logger(self): langfuse_prompt_management = LangfusePromptManagement() - with patch( - "litellm.integrations.langfuse.langfuse_prompt_management.run_async_function" - ) as mock_run_async: + with patch("litellm.integrations.langfuse.langfuse_prompt_management.run_async_function") as mock_run_async: kwargs = {"standard_callback_dynamic_params": {}} start_time, end_time = 1, 2 @@ -69,10 +62,7 @@ class TestLangfusePromptManagement: ) mock_run_async.assert_called_once() - assert ( - mock_run_async.call_args[0][0] - == langfuse_prompt_management.async_log_failure_event - ) + assert mock_run_async.call_args[0][0] == langfuse_prompt_management.async_log_failure_event def test_langfuse_client_init_passes_dedicated_httpx_client(self): import httpx @@ -91,13 +81,14 @@ class TestLangfusePromptManagement: "litellm.integrations.langfuse.langfuse_prompt_management.LangFuseLogger._get_langfuse_flush_interval", return_value=1, ), - patch("litellm.integrations.langfuse.langfuse_sdk.Langfuse", mock_langfuse_class), # test-quality-ok: the ctor must be intercepted where acquire_langfuse_client resolves it; a real client spawns export threads + patch( + "litellm.integrations.langfuse.langfuse_sdk.Langfuse", mock_langfuse_class + ), # test-quality-ok: the ctor must be intercepted where acquire_langfuse_client resolves it; a real client spawns export threads patch( "litellm.llms.custom_httpx.http_handler.get_ssl_configuration", return_value=False, ) as mock_get_ssl, ): - langfuse_client_init( langfuse_public_key="pk-1234", langfuse_secret="sk-1234", @@ -118,7 +109,9 @@ class TestLangfusePromptManagement: class _RecordingLangfuseForEnv: last_environment: str | None = None - def __init__(self, *, environment: str | None = None, **parameters: object) -> None: # kwargs-ok: records only environment out of whatever langfuse_client_init forwards + def __init__( + self, *, environment: str | None = None, **parameters: object + ) -> None: # kwargs-ok: records only environment out of whatever langfuse_client_init forwards type(self).last_environment = environment @@ -132,7 +125,9 @@ def test_langfuse_client_init_resolves_deployment_environment(monkeypatch, env_v monkeypatch.setenv("LANGFUSE_HOST", "https://test.langfuse.com") monkeypatch.setenv("LANGFUSE_TRACING_ENVIRONMENT", env_value) monkeypatch.setattr(_RecordingLangfuseForEnv, "last_environment", None) - with patch("litellm.integrations.langfuse.langfuse_sdk.Langfuse", _RecordingLangfuseForEnv): # test-quality-ok: the ctor must be intercepted where acquire_langfuse_client resolves it; a real client spawns export threads + with patch( + "litellm.integrations.langfuse.langfuse_sdk.Langfuse", _RecordingLangfuseForEnv + ): # test-quality-ok: the ctor must be intercepted where acquire_langfuse_client resolves it; a real client spawns export threads langfuse_client_init.cache_clear() langfuse_client_init() langfuse_client_init.cache_clear() @@ -200,3 +195,41 @@ def test_langfuse_client_init_mock_mode_makes_no_network_calls(monkeypatch): LangfuseResourceManager._instances.pop("pk-pm-mock-egress", None) assert received == [], f"LANGFUSE_MOCK still sent spans to the configured host: {received}" + + +@pytest.mark.asyncio +async def test_async_log_failure_event_records_trace_id_for_alerting(monkeypatch): + from langfuse._client.resource_manager import LangfuseResourceManager + + from litellm.integrations.langfuse.langfuse_sdk import resolve_trace_id + from litellm.litellm_core_utils.specialty_caches.service_trace_id_cache import in_memory_trace_id_cache + + monkeypatch.setenv("LANGFUSE_MOCK", "true") + monkeypatch.setenv("LANGFUSE_HOST", "http://127.0.0.1:1") + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-pm-trace-cache") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-pm-trace-cache") + LangfuseResourceManager._instances.pop("pk-pm-trace-cache", None) + langfuse_client_init.cache_clear() + call_id: Final = "call-trace-cache-1" + now: Final = datetime.now(timezone.utc) + kwargs: Final = { + "litellm_call_id": call_id, + "model": "gpt-5.4", + "messages": [{"role": "user", "content": "hi"}], + "litellm_params": {"metadata": {"trace_id": "alert-trace-1"}}, + "optional_params": {}, + "standard_callback_dynamic_params": {}, + "exception": RuntimeError("provider down"), + } + + try: + await LangfusePromptManagement().async_log_failure_event( + kwargs=kwargs, response_obj=None, start_time=now, end_time=now + ) + finally: + langfuse_client_init.cache_clear() + LangfuseResourceManager._instances.pop("pk-pm-trace-cache", None) + + assert in_memory_trace_id_cache.get_cache(litellm_call_id=call_id, service_name="langfuse") == resolve_trace_id( + "alert-trace-1" + )