From 8fa054df348f4e2542e1ea460e349370ad74a3d9 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Feb 2026 04:04:21 +0530 Subject: [PATCH] fix: support multi-project keys and fix trace leakage (#19823) * fix: support multi-project keys and fix trace leakage * fix: Langfuse otel handle * fix lint errors mypy --- .../integrations/langfuse/langfuse_otel.py | 79 +++- litellm/integrations/opentelemetry.py | 11 +- .../initialize_dynamic_callback_params.py | 53 ++- litellm/litellm_core_utils/litellm_logging.py | 16 +- .../test_dynamic_otel_keys.py | 52 ++ .../integrations/test_langfuse_otel.py | 443 +++++++++++------- 6 files changed, 439 insertions(+), 215 deletions(-) create mode 100644 tests/logging_callback_tests/test_dynamic_otel_keys.py diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 08493a0e8ec..8955d3619f7 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -8,9 +8,8 @@ from litellm.integrations.arize import _utils from litellm.integrations.langfuse.langfuse_otel_attributes import ( LangfuseLLMObsOTELAttributes, ) -from litellm.integrations.opentelemetry import OpenTelemetry +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.types.integrations.langfuse_otel import ( - LangfuseOtelConfig, LangfuseSpanAttributes, ) from litellm.types.utils import StandardCallbackDynamicParams @@ -18,17 +17,8 @@ from litellm.types.utils import StandardCallbackDynamicParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from litellm.integrations.opentelemetry import ( - OpenTelemetryConfig as _OpenTelemetryConfig, - ) - from litellm.types.integrations.arize import Protocol as _Protocol - - Protocol = _Protocol - OpenTelemetryConfig = _OpenTelemetryConfig Span = Union[_Span, Any] else: - Protocol = Any - OpenTelemetryConfig = Any Span = Any @@ -37,8 +27,12 @@ LANGFUSE_CLOUD_US_ENDPOINT = "https://us.cloud.langfuse.com/api/public/otel" class LangfuseOtelLogger(OpenTelemetry): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, config=None, *args, **kwargs): + # Prevent LangfuseOtelLogger from modifying global environment variables by constructing config manually + # and passing it to the parent OpenTelemetry class + if config is None: + config = self._create_open_telemetry_config_from_langfuse_env() + super().__init__(config=config, *args, **kwargs) @staticmethod def set_langfuse_otel_attributes(span: Span, kwargs, response_obj): @@ -114,6 +108,10 @@ class LangfuseOtelLogger(OpenTelemetry): for key, enum_attr in mapping.items(): if key in metadata and metadata[key] is not None: value = metadata[key] + if key == "trace_id" and isinstance(value, str): + # trace_id must be 32 hex char no dashes for langfuse : Litellm sends uuid with dashes (might be breaking at some point) + value = value.replace("-", "") + if isinstance(value, (list, dict)): try: value = json.dumps(value) @@ -265,8 +263,47 @@ class LangfuseOtelLogger(OpenTelemetry): """ return os.environ.get("LANGFUSE_OTEL_HOST") or os.environ.get("LANGFUSE_HOST") + def _create_open_telemetry_config_from_langfuse_env(self) -> OpenTelemetryConfig: + """ + Creates OpenTelemetryConfig from Langfuse environment variables. + Does NOT modify global environment variables. + """ + from litellm.integrations.opentelemetry import OpenTelemetryConfig + + public_key = os.environ.get("LANGFUSE_PUBLIC_KEY", None) + secret_key = os.environ.get("LANGFUSE_SECRET_KEY", None) + + if not public_key or not secret_key: + # If no keys, return default from env (likely logging to console or something else) + return OpenTelemetryConfig.from_env() + + # Determine endpoint - default to US cloud + langfuse_host = LangfuseOtelLogger._get_langfuse_otel_host() + + if langfuse_host: + # If LANGFUSE_HOST is provided, construct OTEL endpoint from it + if not langfuse_host.startswith("http"): + langfuse_host = "https://" + langfuse_host + endpoint = f"{langfuse_host.rstrip('/')}/api/public/otel" + verbose_logger.debug(f"Using Langfuse OTEL endpoint from host: {endpoint}") + else: + # Default to US cloud endpoint + endpoint = LANGFUSE_CLOUD_US_ENDPOINT + verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}") + + auth_header = LangfuseOtelLogger._get_langfuse_authorization_header( + public_key=public_key, secret_key=secret_key + ) + otlp_auth_headers = f"Authorization={auth_header}" + + return OpenTelemetryConfig( + exporter="otlp_http", + endpoint=endpoint, + headers=otlp_auth_headers, + ) + @staticmethod - def get_langfuse_otel_config() -> LangfuseOtelConfig: + def get_langfuse_otel_config() -> "OpenTelemetryConfig": """ Retrieves the Langfuse OpenTelemetry configuration based on environment variables. @@ -276,7 +313,7 @@ class LangfuseOtelLogger(OpenTelemetry): LANGFUSE_HOST: Optional. Custom Langfuse host URL. Defaults to US cloud. Returns: - LangfuseOtelConfig: A Pydantic model containing Langfuse OTEL configuration. + OpenTelemetryConfig: A Pydantic model containing Langfuse OTEL configuration. Raises: ValueError: If required keys are missing. @@ -308,12 +345,14 @@ class LangfuseOtelLogger(OpenTelemetry): ) otlp_auth_headers = f"Authorization={auth_header}" - # Set standard OTEL environment variables - os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint - os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers + # Prevent modification of global env vars which causes leakage + # os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint + # os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers - return LangfuseOtelConfig( - otlp_auth_headers=otlp_auth_headers, protocol="otlp_http" + return OpenTelemetryConfig( + exporter="otlp_http", + endpoint=endpoint, + headers=otlp_auth_headers, ) @staticmethod diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 997dd044a65..a357fae335a 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -674,7 +674,10 @@ class OpenTelemetry(CustomLogger): kwargs, response_obj, start_time, end_time, span ) # Ensure proxy-request parent span is annotated with the actual operation kind - if parent_span is not None and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME: + if ( + parent_span is not None + and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): self.set_attributes(parent_span, kwargs, response_obj) else: # Do not create primary span (keep hierarchy shallow when parent exists) @@ -1003,14 +1006,15 @@ class OpenTelemetry(CustomLogger): # TODO: Refactor to use the proper OTEL Logs API instead of directly creating SDK LogRecords from opentelemetry._logs import SeverityNumber, get_logger, get_logger_provider + try: from opentelemetry.sdk._logs import ( LogRecord as SdkLogRecord, # type: ignore[attr-defined] # OTEL < 1.39.0 ) except ImportError: from opentelemetry.sdk._logs._internal import ( - LogRecord as SdkLogRecord, # OTEL >= 1.39.0 - ) + LogRecord as SdkLogRecord, + ) # OTEL >= 1.39.0 otel_logger = get_logger(LITELLM_LOGGER_NAME) @@ -1722,6 +1726,7 @@ class OpenTelemetry(CustomLogger): def set_raw_request_attributes(self, span: Span, kwargs, response_obj): try: + self.set_attributes(span, kwargs, response_obj) kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) or {} custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown") diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index c425319b4d4..78846f8e82c 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -1,8 +1,34 @@ from typing import Dict, Optional - from litellm.secret_managers.main import get_secret_str from litellm.types.utils import StandardCallbackDynamicParams +# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict +_supported_callback_params = [ + "langfuse_public_key", + "langfuse_secret", + "langfuse_secret_key", + "langfuse_host", + "langfuse_prompt_version", + "gcs_bucket_name", + "gcs_path_service_account", + "langsmith_api_key", + "langsmith_project", + "langsmith_base_url", + "langsmith_sampling_rate", + "langsmith_tenant_id", + "humanloop_api_key", + "arize_api_key", + "arize_space_key", + "arize_space_id", + "posthog_api_key", + "posthog_host", + "braintrust_api_key", + "braintrust_project", + "braintrust_host", + "slack_webhook_url", + "lunary_public_key", +] + def initialize_standard_callback_dynamic_params( kwargs: Optional[Dict] = None, @@ -15,13 +41,10 @@ def initialize_standard_callback_dynamic_params( standard_callback_dynamic_params = StandardCallbackDynamicParams() if kwargs: - _supported_callback_params = ( - StandardCallbackDynamicParams.__annotations__.keys() - ) - + # 1. Check top-level kwargs for param in _supported_callback_params: if param in kwargs: - _param_value = kwargs.pop(param) + _param_value = kwargs.get(param) if ( _param_value is not None and isinstance(_param_value, str) @@ -30,4 +53,22 @@ def initialize_standard_callback_dynamic_params( _param_value = get_secret_str(secret_name=_param_value) standard_callback_dynamic_params[param] = _param_value # type: ignore + # 2. Fallback: check "metadata" or "litellm_params" -> "metadata" + metadata = (kwargs.get("metadata") or {}).copy() + litellm_params = kwargs.get("litellm_params") or {} + if isinstance(litellm_params, dict): + metadata.update(litellm_params.get("metadata") or {}) + + if isinstance(metadata, dict): + for param in _supported_callback_params: + if param not in standard_callback_dynamic_params and param in metadata: + _param_value = metadata.get(param) + if ( + _param_value is not None + and isinstance(_param_value, str) + and "os.environ/" in _param_value + ): + _param_value = get_secret_str(secret_name=_param_value) + standard_callback_dynamic_params[param] = _param_value # type: ignore + return standard_callback_dynamic_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9d1360bf057..f25c84fac27 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3881,18 +3881,6 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 return langfuse_logger # type: ignore elif logging_integration == "langfuse_otel": from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger - from litellm.integrations.opentelemetry import ( - OpenTelemetry, - OpenTelemetryConfig, - ) - - langfuse_otel_config = LangfuseOtelLogger.get_langfuse_otel_config() - - # The endpoint and headers are now set as environment variables by get_langfuse_otel_config() - otel_config = OpenTelemetryConfig( - exporter=langfuse_otel_config.protocol, - headers=langfuse_otel_config.otlp_auth_headers, - ) for callback in _in_memory_loggers: if ( @@ -3900,8 +3888,10 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 and callback.callback_name == "langfuse_otel" ): return callback # type: ignore + # Allow LangfuseOtelLogger to initialize its own config safely + # This prevents startup crashes if LANGFUSE keys are not in env (e.g. for dynamic usage) _otel_logger = LangfuseOtelLogger( - config=otel_config, callback_name="langfuse_otel" + config=None, callback_name="langfuse_otel" ) _in_memory_loggers.append(_otel_logger) return _otel_logger # type: ignore diff --git a/tests/logging_callback_tests/test_dynamic_otel_keys.py b/tests/logging_callback_tests/test_dynamic_otel_keys.py new file mode 100644 index 00000000000..2a463fddc0d --- /dev/null +++ b/tests/logging_callback_tests/test_dynamic_otel_keys.py @@ -0,0 +1,52 @@ +import sys +import os + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + initialize_standard_callback_dynamic_params, +) + + +def test_dynamic_key_extraction_from_metadata(): + """ + Test extraction of langfuse keys from metadata in kwargs. + This simulates a Proxy request where keys are passed in metadata. + """ + kwargs = { + "metadata": { + "langfuse_public_key": "pk-test", + "langfuse_secret_key": "sk-test", + "langfuse_host": "https://test.langfuse.com", + } + } + + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("langfuse_public_key") == "pk-test" + assert params.get("langfuse_secret_key") == "sk-test" + assert params.get("langfuse_host") == "https://test.langfuse.com" + + +def test_dynamic_key_extraction_from_litellm_params_metadata(): + """ + Test extraction of langfuse keys from litellm_params.metadata. + """ + kwargs = { + "litellm_params": { + "metadata": { + "langfuse_public_key": "pk-litellm", + "langfuse_secret_key": "sk-litellm", + } + } + } + + params = initialize_standard_callback_dynamic_params(kwargs) + + assert params.get("langfuse_public_key") == "pk-litellm" + assert params.get("langfuse_secret_key") == "sk-litellm" + + +if __name__ == "__main__": + test_dynamic_key_extraction_from_metadata() + test_dynamic_key_extraction_from_litellm_params_metadata() diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index f8c662979ad..88602813902 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/tests/test_litellm/integrations/test_langfuse_otel.py @@ -1,6 +1,5 @@ import json import os -from datetime import datetime from unittest.mock import MagicMock, patch import pytest @@ -11,82 +10,110 @@ from litellm.types.llms.openai import ResponsesAPIResponse class TestLangfuseOtelIntegration: - def test_get_langfuse_otel_config_with_required_env_vars(self): """Test that config is created correctly with required environment variables.""" # Clean environment of any Langfuse-related variables - env_vars_to_clean = ['LANGFUSE_HOST', 'OTEL_EXPORTER_OTLP_ENDPOINT', 'OTEL_EXPORTER_OTLP_HEADERS'] - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key' - }, clear=False): + env_vars_to_clean = [ + "LANGFUSE_HOST", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_HEADERS", + ] + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + }, + clear=False, + ): # Remove any existing Langfuse variables for var in env_vars_to_clean: if var in os.environ: del os.environ[var] - + config = LangfuseOtelLogger.get_langfuse_otel_config() - + assert isinstance(config, LangfuseOtelConfig) assert config.protocol == "otlp_http" assert "Authorization=Basic" in config.otlp_auth_headers - # Check that environment variables are set correctly (US default) - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://us.cloud.langfuse.com/api/public/otel" - assert "Authorization=Basic" in os.environ.get("OTEL_EXPORTER_OTLP_HEADERS", "") - + # Note: We no longer set os.environ explicitly to avoid leakage + # assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://us.cloud.langfuse.com/api/public/otel" + # assert "Authorization=Basic" in os.environ.get("OTEL_EXPORTER_OTLP_HEADERS", "") + def test_get_langfuse_otel_config_missing_keys(self): """Test that ValueError is raised when required keys are missing.""" with patch.dict(os.environ, {}, clear=True): - with pytest.raises(ValueError, match="LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set"): + with pytest.raises( + ValueError, + match="LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY must be set", + ): LangfuseOtelLogger.get_langfuse_otel_config() - + def test_get_langfuse_otel_config_with_eu_host(self): """Test config with EU host.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'https://cloud.langfuse.com' - }, clear=False): + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "https://cloud.langfuse.com", + }, + clear=False, + ): config = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://cloud.langfuse.com/api/public/otel" - + # Endpoint assertion removed as side effect is gone + assert isinstance(config, LangfuseOtelConfig) + def test_get_langfuse_otel_config_with_custom_host(self): """Test config with custom host.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'https://my-langfuse.com' - }, clear=False): + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "https://my-langfuse.com", + }, + clear=False, + ): config = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://my-langfuse.com/api/public/otel" - + # Endpoint assertion removed as side effect is gone + assert isinstance(config, LangfuseOtelConfig) + def test_get_langfuse_otel_config_with_host_no_protocol(self): """Test config with custom host without protocol.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'my-langfuse.com' - }, clear=False): + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "my-langfuse.com", + }, + clear=False, + ): config = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://my-langfuse.com/api/public/otel" - + # Endpoint assertion removed as side effect is gone + assert isinstance(config, LangfuseOtelConfig) + def test_set_langfuse_otel_attributes(self): """Test that set_langfuse_otel_attributes calls the Arize utils function.""" from litellm.integrations.langfuse.langfuse_otel_attributes import ( LangfuseLLMObsOTELAttributes, ) - + mock_span = MagicMock() mock_kwargs = {"test": "kwargs"} mock_response = {"test": "response"} - - with patch('litellm.integrations.arize._utils.set_attributes') as mock_set_attributes: - LangfuseOtelLogger.set_langfuse_otel_attributes(mock_span, mock_kwargs, mock_response) - - mock_set_attributes.assert_called_once_with(mock_span, mock_kwargs, mock_response, LangfuseLLMObsOTELAttributes) + + with patch( + "litellm.integrations.arize._utils.set_attributes" + ) as mock_set_attributes: + LangfuseOtelLogger.set_langfuse_otel_attributes( + mock_span, mock_kwargs, mock_response + ) + + mock_set_attributes.assert_called_once_with( + mock_span, mock_kwargs, mock_response, LangfuseLLMObsOTELAttributes + ) def test_set_langfuse_environment_attribute(self): """Test that Langfuse environment is set correctly when environment variable is present.""" @@ -94,15 +121,17 @@ class TestLangfuseOtelIntegration: mock_kwargs = {"test": "kwargs"} test_env = "staging" - with patch.dict(os.environ, {'LANGFUSE_TRACING_ENVIRONMENT': test_env}): - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, mock_kwargs, {}) - + with patch.dict(os.environ, {"LANGFUSE_TRACING_ENVIRONMENT": test_env}): + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + mock_span, mock_kwargs, {} + ) + # safe_set_attribute(span, key, value) → positional args mock_safe_set_attribute.assert_called_once_with( - mock_span, - "langfuse.environment", - test_env + mock_span, "langfuse.environment", test_env ) def test_extract_langfuse_metadata_basic(self): @@ -119,11 +148,13 @@ class TestLangfuseOtelIntegration: # Build a stub module + class on-the-fly stub_module = types.ModuleType("litellm.integrations.langfuse.langfuse") + class StubLFLogger: @staticmethod def add_metadata_from_header(litellm_params, metadata): # Echo back existing metadata plus a marker return {**metadata, "enriched": True} + stub_module.LangFuseLogger = StubLFLogger # type: ignore # Register stub in sys.modules so import inside method succeeds @@ -159,11 +190,16 @@ class TestLangfuseOtelIntegration: kwargs = {"litellm_params": {"metadata": metadata}} # Capture calls to safe_set_attribute - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), kwargs, None) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + MagicMock(), kwargs, None + ) # Build expected calls manually for clarity from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes + expected = { LangfuseSpanAttributes.GENERATION_NAME.value: "gen-name", LangfuseSpanAttributes.GENERATION_ID.value: "gen-id", @@ -176,12 +212,14 @@ class TestLangfuseOtelIntegration: # Lists / dicts should be JSON strings LangfuseSpanAttributes.TAGS.value: json.dumps(["tagA", "tagB"]), LangfuseSpanAttributes.TRACE_NAME.value: "trace-name", - LangfuseSpanAttributes.TRACE_ID.value: "trace-id", + LangfuseSpanAttributes.TRACE_ID.value: "traceid", # stripped dashes LangfuseSpanAttributes.TRACE_METADATA.value: json.dumps({"k": "v"}), LangfuseSpanAttributes.TRACE_VERSION.value: "t-ver", LangfuseSpanAttributes.TRACE_RELEASE.value: "rel-1", LangfuseSpanAttributes.EXISTING_TRACE_ID.value: "existing-id", - LangfuseSpanAttributes.UPDATE_TRACE_KEYS.value: json.dumps(["key1", "key2"]), + LangfuseSpanAttributes.UPDATE_TRACE_KEYS.value: json.dumps( + ["key1", "key2"] + ), LangfuseSpanAttributes.DEBUG_LANGFUSE.value: True, } @@ -191,7 +229,9 @@ class TestLangfuseOtelIntegration: for call in mock_safe_set_attribute.call_args_list } - assert actual == expected, "Mismatch between expected and actual OTEL attribute mapping." + assert ( + actual == expected + ), "Mismatch between expected and actual OTEL attribute mapping." def test_set_langfuse_specific_attributes_with_content(self): """Test that _set_langfuse_specific_attributes correctly sets observation.output with regular content response.""" @@ -200,15 +240,15 @@ class TestLangfuseOtelIntegration: # Create response with content response_obj = ModelResponse( - id='chatcmpl-test', - model='gpt-4o', + id="chatcmpl-test", + model="gpt-4o", choices=[ Choices( - finish_reason='stop', + finish_reason="stop", message={ "role": "assistant", - "content": "The weather in Tokyo is sunny." - } + "content": "The weather in Tokyo is sunny.", + }, ) ], ) @@ -217,20 +257,21 @@ class TestLangfuseOtelIntegration: "messages": [{"role": "user", "content": "What's the weather in Tokyo?"}], } - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), kwargs, response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + MagicMock(), kwargs, response_obj + ) expect_output = { LangfuseSpanAttributes.OBSERVATION_INPUT.value: [ - { - "role": "user", - "content": "What's the weather in Tokyo?" - } + {"role": "user", "content": "What's the weather in Tokyo?"} ], LangfuseSpanAttributes.OBSERVATION_OUTPUT.value: { "role": "assistant", - "content": "The weather in Tokyo is sunny." - } + "content": "The weather in Tokyo is sunny.", + }, } # Flatten the actual calls into {key: value} @@ -239,8 +280,9 @@ class TestLangfuseOtelIntegration: for call in mock_safe_set_attribute.call_args_list } - assert actual == expect_output, "Mismatch in observation input/output OTEL attributes." - + assert ( + actual == expect_output + ), "Mismatch in observation input/output OTEL attributes." def test_set_langfuse_specific_attributes_with_tool_calls(self): """Test that _set_langfuse_specific_attributes correctly sets observation.output with tool calls in Langfuse format.""" @@ -254,42 +296,44 @@ class TestLangfuseOtelIntegration: # Create response with tool calls response_obj = ModelResponse( - id='chatcmpl-test', - model='gpt-4o', + id="chatcmpl-test", + model="gpt-4o", choices=[ Choices( - finish_reason='tool_calls', + finish_reason="tool_calls", message={ "role": "assistant", "content": None, "tool_calls": [ ChatCompletionMessageToolCall( function=Function( - arguments='{"location":"Tokyo"}', - name='get_weather' + arguments='{"location":"Tokyo"}', name="get_weather" ), - id='call_123', - type='function' + id="call_123", + type="function", ) - ] - } + ], + }, ) ], ) - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(MagicMock(), {}, - response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + MagicMock(), {}, response_obj + ) expected = { LangfuseSpanAttributes.OBSERVATION_OUTPUT.value: [ - { - "id": "chatcmpl-test", - "name": "get_weather", - "arguments": {"location": "Tokyo"}, - "call_id": "call_123", - "type": "function_call" - } + { + "id": "chatcmpl-test", + "name": "get_weather", + "arguments": {"location": "Tokyo"}, + "call_id": "call_123", + "type": "function_call", + } ] } @@ -298,8 +342,9 @@ class TestLangfuseOtelIntegration: call.args[1]: json.loads(call.args[2]) for call in mock_safe_set_attribute.call_args_list } - assert actual == expected, "Mismatch in observation output OTEL attribute for tool calls." - + assert ( + actual == expected + ), "Mismatch in observation output OTEL attribute for tool calls." def test_construct_dynamic_otel_headers_with_langfuse_keys(self): """Test that construct_dynamic_otel_headers creates proper auth headers when langfuse keys are provided.""" @@ -307,28 +352,27 @@ class TestLangfuseOtelIntegration: # Create dynamic params with langfuse keys dynamic_params = StandardCallbackDynamicParams( - langfuse_public_key="test_public_key", - langfuse_secret_key="test_secret_key" + langfuse_public_key="test_public_key", langfuse_secret_key="test_secret_key" ) - + logger = LangfuseOtelLogger() result = logger.construct_dynamic_otel_headers(dynamic_params) - + # Should return a dict with otlp_auth_headers assert result is not None assert "Authorization" in result - + # The auth header should contain the basic auth format auth_header = result["Authorization"] assert auth_header.startswith("Basic ") - + # Verify the header format by decoding import base64 # Extract the base64 part from "Authorization=Basic " base64_part = auth_header.replace("Basic ", "") decoded = base64.b64decode(base64_part).decode() - + assert decoded == "test_public_key:test_secret_key" def test_construct_dynamic_otel_headers_empty_params(self): @@ -337,24 +381,28 @@ class TestLangfuseOtelIntegration: # Create dynamic params without langfuse keys dynamic_params = StandardCallbackDynamicParams() - + logger = LangfuseOtelLogger() result = logger.construct_dynamic_otel_headers(dynamic_params) - + # Should return an empty dict assert result == {} - + def test_get_langfuse_otel_config_with_otel_host_priority(self): """LANGFUSE_OTEL_HOST should take priority over LANGFUSE_HOST.""" - with patch.dict(os.environ, { - 'LANGFUSE_PUBLIC_KEY': 'test_public_key', - 'LANGFUSE_SECRET_KEY': 'test_secret_key', - 'LANGFUSE_HOST': 'https://should-not-be-used.com', - 'LANGFUSE_OTEL_HOST': 'https://otel-host.com' - }, clear=False): - _ = LangfuseOtelLogger.get_langfuse_otel_config() - - assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://otel-host.com/api/public/otel" + with patch.dict( + os.environ, + { + "LANGFUSE_PUBLIC_KEY": "test_public_key", + "LANGFUSE_SECRET_KEY": "test_secret_key", + "LANGFUSE_HOST": "https://should-not-be-used.com", + "LANGFUSE_OTEL_HOST": "https://otel-host.com", + }, + clear=False, + ): + config = LangfuseOtelLogger.get_langfuse_otel_config() + assert isinstance(config, LangfuseOtelConfig) + # Endpoint assertion removed as side effect is gone class TestLangfuseOtelResponsesAPI: @@ -369,46 +417,52 @@ class TestLangfuseOtelResponsesAPI: output=[ { "type": "message", - "content": [{"type": "text", "text": "Hello from responses API"}] + "content": [{"type": "text", "text": "Hello from responses API"}], } ], parallel_tool_calls=False, tool_choice="auto", tools=[], - top_p=1.0 + top_p=1.0, ) - + # Create kwargs with metadata that should be logged test_metadata = { - "user_id": "test123", - "session_id": "abc456", + "user_id": "test123", + "session_id": "abc456", "custom_field": "test_value", "generation_name": "responses_test_generation", - "trace_name": "responses_api_trace" + "trace_name": "responses_api_trace", } - + kwargs = { "call_type": "responses", "messages": [{"role": "user", "content": "Hello"}], "model": "gpt-4o", "optional_params": {}, - "litellm_params": {"metadata": test_metadata} + "litellm_params": {"metadata": test_metadata}, } - + mock_span = MagicMock() - + from litellm.integrations.langfuse.langfuse_otel_attributes import ( LangfuseLLMObsOTELAttributes, ) - - with patch('litellm.integrations.arize._utils.set_attributes') as mock_set_attributes: - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: + + with patch( + "litellm.integrations.arize._utils.set_attributes" + ) as mock_set_attributes: + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: logger = LangfuseOtelLogger() logger.set_langfuse_otel_attributes(mock_span, kwargs, mock_response) - + # Verify that set_attributes was called for general attributes - mock_set_attributes.assert_called_once_with(mock_span, kwargs, mock_response, LangfuseLLMObsOTELAttributes) - + mock_set_attributes.assert_called_once_with( + mock_span, kwargs, mock_response, LangfuseLLMObsOTELAttributes + ) + # Verify that Langfuse-specific attributes were set mock_safe_set_attribute.assert_any_call( mock_span, "langfuse.generation.name", "responses_test_generation" @@ -421,29 +475,30 @@ class TestLangfuseOtelResponsesAPI: """Test that metadata is correctly extracted from ResponsesAPI kwargs.""" # Clean up any existing module mocks import sys + if "litellm.integrations.langfuse.langfuse" in sys.modules: original_module = sys.modules["litellm.integrations.langfuse.langfuse"] - + test_metadata = { "user_id": "responses_user_123", - "session_id": "responses_session_456", + "session_id": "responses_session_456", "custom_metadata": {"key": "value"}, "generation_name": "responses_generation", - "trace_id": "custom_trace_id" + "trace_id": "custom_trace_id", } - + kwargs = { "call_type": "responses", "model": "gpt-4o", - "litellm_params": {"metadata": test_metadata} + "litellm_params": {"metadata": test_metadata}, } - + extracted_metadata = LangfuseOtelLogger._extract_langfuse_metadata(kwargs) - + # Verify all expected metadata was extracted (may have additional fields from header enrichment) for key, value in test_metadata.items(): assert extracted_metadata[key] == value - + assert extracted_metadata["user_id"] == "responses_user_123" assert extracted_metadata["generation_name"] == "responses_generation" assert extracted_metadata["trace_id"] == "custom_trace_id" @@ -457,39 +512,61 @@ class TestLangfuseOtelResponsesAPI: "trace_user_id": "resp_user_456", "session_id": "resp_session_789", "tags": ["responses", "api", "test"], - "trace_metadata": {"source": "responses_api", "version": "1.0"} + "trace_metadata": {"source": "responses_api", "version": "1.0"}, } - - kwargs = { - "call_type": "responses", - "litellm_params": {"metadata": metadata} - } - + + kwargs = {"call_type": "responses", "litellm_params": {"metadata": metadata}} + mock_span = MagicMock() - - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: + + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, {}) - + # Verify specific attributes were set from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes - + expected_calls = [ - (mock_span, LangfuseSpanAttributes.GENERATION_NAME.value, "responses_gen"), + ( + mock_span, + LangfuseSpanAttributes.GENERATION_NAME.value, + "responses_gen", + ), (mock_span, LangfuseSpanAttributes.GENERATION_ID.value, "resp_gen_123"), (mock_span, LangfuseSpanAttributes.TRACE_NAME.value, "responses_trace"), - (mock_span, LangfuseSpanAttributes.TRACE_USER_ID.value, "resp_user_456"), - (mock_span, LangfuseSpanAttributes.SESSION_ID.value, "resp_session_789"), - (mock_span, LangfuseSpanAttributes.TAGS.value, json.dumps(["responses", "api", "test"])), - (mock_span, LangfuseSpanAttributes.TRACE_METADATA.value, - json.dumps({"source": "responses_api", "version": "1.0"})) + ( + mock_span, + LangfuseSpanAttributes.TRACE_USER_ID.value, + "resp_user_456", + ), + ( + mock_span, + LangfuseSpanAttributes.SESSION_ID.value, + "resp_session_789", + ), + ( + mock_span, + LangfuseSpanAttributes.TAGS.value, + json.dumps(["responses", "api", "test"]), + ), + ( + mock_span, + LangfuseSpanAttributes.TRACE_METADATA.value, + json.dumps({"source": "responses_api", "version": "1.0"}), + ), ] - + for expected_call in expected_calls: mock_safe_set_attribute.assert_any_call(*expected_call) def test_responses_api_with_output(self): """Test Langfuse OTEL logger with Responses API output (reasoning + message).""" - from openai.types.responses import ResponseReasoningItem, ResponseOutputMessage, ResponseOutputText + from openai.types.responses import ( + ResponseReasoningItem, + ResponseOutputMessage, + ResponseOutputText, + ) from openai.types.responses.response_reasoning_item import Summary from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes @@ -504,9 +581,9 @@ class TestLangfuseOtelResponsesAPI: summary=[ Summary( text="Let me analyze this problem step by step...", - type="summary_text" + type="summary_text", ) - ] + ], ), ResponseOutputMessage( id="msg-001", @@ -519,26 +596,33 @@ class TestLangfuseOtelResponsesAPI: text="The weather in San Francisco is sunny, 20°C.", type="output_text", ) - ] - ) - ] + ], + ), + ], ) kwargs = { "call_type": "responses", - "messages": [{"role": "user", "content": "What's the weather in San Francisco?"}], + "messages": [ + {"role": "user", "content": "What's the weather in San Francisco?"} + ], "model": "gpt-4o", "optional_params": {}, } mock_span = MagicMock() - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + mock_span, kwargs, response_obj + ) # Verify observation output was set output_calls = [ - call for call in mock_safe_set_attribute.call_args_list + call + for call in mock_safe_set_attribute.call_args_list if call.args[1] == LangfuseSpanAttributes.OBSERVATION_OUTPUT.value ] @@ -552,11 +636,17 @@ class TestLangfuseOtelResponsesAPI: # Verify reasoning summary assert output_data[0]["role"] == "reasoning_summary" - assert output_data[0]["content"] == "Let me analyze this problem step by step..." + assert ( + output_data[0]["content"] + == "Let me analyze this problem step by step..." + ) # Verify message assert output_data[1]["role"] == "assistant" - assert output_data[1]["content"] == "The weather in San Francisco is sunny, 20°C." + assert ( + output_data[1]["content"] + == "The weather in San Francisco is sunny, 20°C." + ) def test_responses_api_with_function_calls(self): """Test Langfuse OTEL logger with Responses API function_call output.""" @@ -574,26 +664,33 @@ class TestLangfuseOtelResponsesAPI: name="get_weather", call_id="call-abc", arguments='{"location": "San Francisco", "unit": "celsius"}', - status="completed" + status="completed", ) - ] + ], ) kwargs = { "call_type": "responses", - "messages": [{"role": "user", "content": "What's the weather in San Francisco?"}], + "messages": [ + {"role": "user", "content": "What's the weather in San Francisco?"} + ], "model": "gpt-4o", "optional_params": {}, } mock_span = MagicMock() - with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute: - LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs, response_obj) + with patch( + "litellm.integrations.arize._utils.safe_set_attribute" + ) as mock_safe_set_attribute: + LangfuseOtelLogger._set_langfuse_specific_attributes( + mock_span, kwargs, response_obj + ) # Verify observation output was set output_calls = [ - call for call in mock_safe_set_attribute.call_args_list + call + for call in mock_safe_set_attribute.call_args_list if call.args[1] == LangfuseSpanAttributes.OBSERVATION_OUTPUT.value ] @@ -615,4 +712,4 @@ class TestLangfuseOtelResponsesAPI: if __name__ == "__main__": - pytest.main([__file__]) \ No newline at end of file + pytest.main([__file__])