mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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:
parent
78bc59b9b7
commit
828f51d42c
8 changed files with 141 additions and 46 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue