scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

feat: propagate Langfuse trace_id (#17669)

Browse source
This commit is contained in:
YutaSaito 2025-12-10 05:25:52 +09:00 • committed by GitHub
parent 9253a9c365
commit 80a18f989a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 159 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/my-website/docs/proxy/logging.md
View file

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

11
litellm/integrations/langfuse/langfuse.py
View file

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

2
litellm/integrations/langfuse/langfuse_prompt_management.py
View file

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

44
tests/test_litellm/integrations/langfuse/test_langfuse_prompt_management.py
View file

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

117
tests/test_litellm/integrations/test_langfuse.py
View file

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