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