From 008ea864a7ba12e49ee713f148116743f5fa9c76 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 11 Aug 2025 22:06:25 -0700 Subject: [PATCH] [Feat] - Add key/team logging for Langfuse OTEL Logger (#13512) * feat - add key/team logging for LF * test_construct_dynamic_otel_headers_with_langfuse_keys * update LangfuseOtelLogger * test_construct_dynamic_otel_headers_with_langfuse_keys * cleanup * OpenTelemetryConfig fixes * fix code qa checks * TestLangfuseOtelIntegration --- .../integrations/langfuse/langfuse_otel.py | 57 ++++++++++++++++--- litellm/litellm_core_utils/litellm_logging.py | 7 ++- litellm/proxy/proxy_config.yaml | 18 +++--- .../integrations/test_langfuse_otel.py | 48 +++++++++++++++- 4 files changed, 105 insertions(+), 25 deletions(-) diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 4072be2a25..8b90e12337 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -1,15 +1,16 @@ import base64 -import os import json # <--- NEW -from typing import TYPE_CHECKING, Any, Union -from urllib.parse import quote +import os +from typing import TYPE_CHECKING, Any, Optional, Union from litellm._logging import verbose_logger from litellm.integrations.arize import _utils +from litellm.integrations.opentelemetry import OpenTelemetry from litellm.types.integrations.langfuse_otel import ( LangfuseOtelConfig, LangfuseSpanAttributes, ) +from litellm.types.utils import StandardCallbackDynamicParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -33,7 +34,11 @@ LANGFUSE_CLOUD_US_ENDPOINT = "https://us.cloud.langfuse.com/api/public/otel" -class LangfuseOtelLogger: +class LangfuseOtelLogger(OpenTelemetry): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @staticmethod def set_langfuse_otel_attributes(span: Span, kwargs, response_obj): """ @@ -174,11 +179,11 @@ class LangfuseOtelLogger: endpoint = LANGFUSE_CLOUD_US_ENDPOINT verbose_logger.debug(f"Using Langfuse US cloud endpoint: {endpoint}") - # Create Basic Auth header - auth_string = f"{public_key}:{secret_key}" - auth_header = base64.b64encode(auth_string.encode()).decode() - # URL encode the entire header value as required by OpenTelemetry specification - otlp_auth_headers = f"Authorization={quote(f'Basic {auth_header}')}" + auth_header = LangfuseOtelLogger._get_langfuse_authorization_header( + public_key=public_key, + secret_key=secret_key + ) + otlp_auth_headers = f"Authorization={auth_header}" # Set standard OTEL environment variables os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint @@ -187,3 +192,37 @@ class LangfuseOtelLogger: return LangfuseOtelConfig( otlp_auth_headers=otlp_auth_headers, protocol="otlp_http" ) + + @staticmethod + def _get_langfuse_authorization_header(public_key: str, secret_key: str) -> str: + """ + Get the authorization header for Langfuse OpenTelemetry. + """ + auth_string = f"{public_key}:{secret_key}" + auth_header = base64.b64encode(auth_string.encode()).decode() + return f'Basic {auth_header}' + + def construct_dynamic_otel_headers( + self, + standard_callback_dynamic_params: StandardCallbackDynamicParams + ) -> Optional[dict]: + """ + Construct dynamic Langfuse headers from standard callback dynamic params + + This is used for team/key based logging. + + Returns: + dict: A dictionary of dynamic Langfuse headers + """ + dynamic_headers = {} + + dynamic_langfuse_public_key = standard_callback_dynamic_params.get("langfuse_public_key") + dynamic_langfuse_secret_key = standard_callback_dynamic_params.get("langfuse_secret_key") + if dynamic_langfuse_public_key and dynamic_langfuse_secret_key: + auth_header = LangfuseOtelLogger._get_langfuse_authorization_header( + public_key=dynamic_langfuse_public_key, + secret_key=dynamic_langfuse_secret_key + ) + dynamic_headers["Authorization"] = auth_header + + return dynamic_headers diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e88074f056..8e4aa43c9c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -131,7 +131,6 @@ from ..integrations.humanloop import HumanloopLogger from ..integrations.lago import LagoLogger from ..integrations.langfuse.langfuse import LangFuseLogger from ..integrations.langfuse.langfuse_handler import LangFuseHandler -from ..integrations.langfuse.langfuse_otel import LangfuseOtelLogger from ..integrations.langfuse.langfuse_prompt_management import LangfusePromptManagement from ..integrations.langsmith import LangsmithLogger from ..integrations.literal_ai import LiteralAILogger @@ -3457,6 +3456,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(langfuse_logger) 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, @@ -3467,15 +3467,16 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 # 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 ( - isinstance(callback, OpenTelemetry) + isinstance(callback, LangfuseOtelLogger) and callback.callback_name == "langfuse_otel" ): return callback # type: ignore - _otel_logger = OpenTelemetry( + _otel_logger = LangfuseOtelLogger( config=otel_config, callback_name="langfuse_otel" ) _in_memory_loggers.append(_otel_logger) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index e1d784bb02..6e9c925257 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,15 +1,11 @@ model_list: - - model_name: gemini/* + - model_name: openai/* litellm_params: - model: gemini/* + model: openai/* + - model_name: anthropic/* + litellm_params: + model: anthropic/* litellm_settings: - callbacks: ["s3_v2"] - s3_callback_params: - s3_bucket_name: litellm-logs # AWS Bucket Name for S3 - s3_region_name: us-west-2 - -general_settings: - cold_storage_custom_logger: s3_v2 - store_prompts_in_cold_storage: true - store_prompts_in_spend_logs: true + callbacks: + - langfuse_otel \ No newline at end of file diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index 20486a4492..41b20c2a23 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/tests/test_litellm/integrations/test_langfuse_otel.py @@ -1,6 +1,6 @@ +import json import os from unittest.mock import MagicMock, patch -import json import pytest @@ -108,7 +108,8 @@ class TestLangfuseOtelIntegration: def test_extract_langfuse_metadata_with_header_enrichment(self, monkeypatch): """_extract_langfuse_metadata should call LangFuseLogger.add_metadata_from_header when available.""" - import sys, types + import sys + import types # Build a stub module + class on-the-fly stub_module = types.ModuleType("litellm.integrations.langfuse.langfuse") @@ -186,6 +187,49 @@ class TestLangfuseOtelIntegration: assert actual == expected, "Mismatch between expected and actual OTEL attribute mapping." + 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.""" + from litellm.types.utils import StandardCallbackDynamicParams + + # Create dynamic params with langfuse keys + dynamic_params = StandardCallbackDynamicParams( + 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): + """Test that construct_dynamic_otel_headers returns empty dict when no langfuse keys are provided.""" + from litellm.types.utils import StandardCallbackDynamicParams + + # 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 == {} + if __name__ == "__main__": pytest.main([__file__]) \ No newline at end of file