diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index cf36963b7e1..a99651cb4a4 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -371,6 +371,8 @@ export LANGFUSE_PUBLIC_KEY="pk_kk" export LANGFUSE_SECRET_KEY="sk_ss" # Optional, defaults to https://cloud.langfuse.com export LANGFUSE_HOST="https://xxx.langfuse.com" +# Optional - When True, forwards LiteLLM's logging trace_id to Langfuse +LANGFUSE_PROPAGATE_TRACE_ID=True ``` **Step 4**: Start the proxy, make a test request diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 12eb00efa9e..11c6108ecc2 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -70,6 +70,7 @@ class LangFuseLogger: self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval( flush_interval ) + self.langfuse_propagate_trace_id = str_to_bool(os.getenv("LANGFUSE_PROPAGATE_TRACE_ID", "False")) is True http_client = _get_httpx_client() self.langfuse_client = http_client.client @@ -536,7 +537,15 @@ class LangFuseLogger: session_id = clean_metadata.pop("session_id", None) trace_name = cast(Optional[str], clean_metadata.pop("trace_name", None)) - trace_id = clean_metadata.pop("trace_id", litellm_call_id) + trace_id = clean_metadata.pop("trace_id", None) + if ( + trace_id is None + and self.langfuse_propagate_trace_id is True + and standard_logging_object is not None + ): + trace_id = cast(Optional[str], standard_logging_object.get("trace_id")) + if trace_id is None: + trace_id = litellm_call_id existing_trace_id = clean_metadata.pop("existing_trace_id", None) update_trace_keys = cast(list, clean_metadata.pop("update_trace_keys", [])) debug = clean_metadata.pop("debug_langfuse", None) diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 58698ef35a5..a9a1937da30 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -12,6 +12,7 @@ from typing_extensions import TypeAlias 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.secret_managers.main import str_to_bool from litellm.types.llms.openai import AllMessageValues, ChatCompletionSystemMessage from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload @@ -124,6 +125,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge langfuse_host=langfuse_host, flush_interval=flush_interval, ) + self.langfuse_propagate_trace_id = str_to_bool(os.getenv("LANGFUSE_PROPAGATE_TRACE_ID", "False")) is True @property def integration_name(self): 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 7a5b8bfb1cb..d4ac22d37bd 100644 --- a/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py @@ -15,18 +15,40 @@ class TestLangfusePromptManagement: langfuse_prompt_management, "_get_prompt_from_id" ) as mock_get_prompt_from_id: mock_should_run_prompt_management.return_value = True - chat_completion_prompt = ( - langfuse_prompt_management.get_chat_completion_prompt( - model="langfuse/langfuse-model", - messages=[{"role": "user", "content": "Hello, how are you?"}], - non_default_params={}, - prompt_id="test-chat-prompt", - prompt_variables={}, - dynamic_callback_params={}, - prompt_version=4, - ) + langfuse_prompt_management.get_chat_completion_prompt( + model="langfuse/langfuse-model", + messages=[{"role": "user", "content": "Hello, how are you?"}], + non_default_params={}, + prompt_id="test-chat-prompt", + prompt_variables={}, + dynamic_callback_params={}, + prompt_version=4, ) mock_get_prompt_from_id.assert_called_once() - print(mock_get_prompt_from_id.call_args.kwargs) assert mock_get_prompt_from_id.call_args.kwargs["prompt_version"] == 4 + + def test_trace_id_propagation_flag_from_env(self): + with patch.dict( + os.environ, + { + "LANGFUSE_SECRET_KEY": "secret", + "LANGFUSE_PUBLIC_KEY": "public", + "LANGFUSE_PROPAGATE_TRACE_ID": "True", + }, + clear=True, + ): + pm = LangfusePromptManagement() + assert pm.langfuse_propagate_trace_id is True + + with patch.dict( + os.environ, + { + "LANGFUSE_SECRET_KEY": "secret", + "LANGFUSE_PUBLIC_KEY": "public", + "LANGFUSE_PROPAGATE_TRACE_ID": "False", + }, + clear=True, + ): + pm2 = LangfusePromptManagement() + assert pm2.langfuse_propagate_trace_id is False diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 34470a4a80f..a04ee28b418 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1,8 +1,7 @@ -import asyncio import datetime -import json import os import sys +import types import unittest from typing import Optional from unittest.mock import MagicMock, patch @@ -48,7 +47,13 @@ class TestLangfuseUsageDetails(unittest.TestCase): # Setup the trace and generation chain self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation - self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace + self.last_trace_kwargs = {} + + def _trace_side_effect(*args, **kwargs): + self.last_trace_kwargs = kwargs + return self.mock_langfuse_trace + + self.mock_langfuse_client.trace.side_effect = _trace_side_effect # Mock the langfuse module that's imported locally in methods self.langfuse_module_patcher = patch.dict( @@ -108,8 +113,6 @@ class TestLangfuseUsageDetails(unittest.TestCase): ) # Bind the method to the instance - import types - self.logger.log_event_on_langfuse = types.MethodType( log_event_on_langfuse, self.logger ) @@ -343,6 +346,110 @@ class TestLangfuseUsageDetails(unittest.TestCase): mock_add_prompt_params.assert_called_once() + def _build_standard_logging_payload(self, trace_id: Optional[str] = None): + payload = { + "id": "payload-id", + "call_type": "completion", + "response_cost": 0.0, + "status": "success", + "total_tokens": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "startTime": 0.0, + "endTime": 0.0, + "completionStartTime": 0.0, + "model": "gpt-4", + "model_id": "model-123", + "model_group": "openai", + "api_base": "https://api.openai.com", + "metadata": { + "user_api_key_end_user_id": None, + "prompt_management_metadata": None, + "session_id": None, + "trace_name": None, + "trace_version": None, + "headers": None, + "endpoint": None, + "caching_groups": None, + "previous_models": None, + }, + "hidden_params": {}, + "request_tags": [], + "messages": [], + "response": {"id": "resp"}, + "model_parameters": {}, + "guardrail_information": None, + "standard_built_in_tools_params": None, + } + if trace_id is not None: + payload["trace_id"] = trace_id + return payload + + def _build_langfuse_kwargs(self, standard_logging_payload): + return { + "standard_logging_object": standard_logging_payload, + "model": standard_logging_payload["model"], + "call_type": standard_logging_payload["call_type"], + "cache_hit": False, + "messages": [], + } + + def test_log_langfuse_v2_propagates_standard_trace_id_when_enabled(self): + self.logger.langfuse_propagate_trace_id = True + payload = self._build_standard_logging_payload(trace_id="std-trace-id") + kwargs = self._build_langfuse_kwargs(payload) + self.last_trace_kwargs = {} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kwargs: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata={}, + litellm_params={"metadata": {}}, + output=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="INFO", + litellm_call_id="call-id-xyz", + ) + + assert self.last_trace_kwargs.get("id") == "std-trace-id" + + def test_log_langfuse_v2_defaults_to_call_id_when_propagation_disabled(self): + self.logger.langfuse_propagate_trace_id = False + payload = self._build_standard_logging_payload(trace_id="std-trace-id") + kwargs = self._build_langfuse_kwargs(payload) + self.last_trace_kwargs = {} + + with patch( + "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", + side_effect=lambda generation_params, **kwargs: generation_params, + create=True, + ): + self.logger._log_langfuse_v2( + user_id="user-1", + metadata={}, + litellm_params={"metadata": {}}, + output=None, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + kwargs=kwargs, + optional_params={}, + input=None, + response_obj=None, + level="INFO", + litellm_call_id="call-id-xyz", + ) + + assert self.last_trace_kwargs.get("id") == "call-id-xyz" + def test_max_langfuse_clients_limit(): """