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>
This commit is contained in:
yucheng 2026-09-17 02:00:57 +00:00
parent 78bc59b9b7
commit 828f51d42c
8 changed files with 141 additions and 46 deletions

View file

@ -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):

View file

@ -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,

View file

@ -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

View file

@ -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

View file

@ -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()

View file

@ -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]

View file

@ -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

View file

@ -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"
)