diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 03b6966809c..ded860ddc6e 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -12,6 +12,7 @@ from litellm.integrations.arize import _utils from litellm.integrations.opentelemetry import OpenTelemetry from litellm.types.integrations.arize import ArizeConfig from litellm.types.services import ServiceLoggerPayload +from litellm.types.utils import StandardCallbackDynamicParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -102,3 +103,30 @@ class ArizeLogger(OpenTelemetry): ): """Arize is used mainly for LLM I/O tracing, sending Proxy Server Request adds bloat to arize logs""" pass + + + @staticmethod + def construct_dynamic_arize_headers( + standard_callback_dynamic_params: StandardCallbackDynamicParams + ): + """ + Construct dynamic Arize headers from standard callback dynamic params + + Returns: + dict: A dictionary of dynamic Arize headers + """ + dynamic_headers = {} + if standard_callback_dynamic_params.get("arize_space_key"): + dynamic_headers["space_key"] = standard_callback_dynamic_params.get( + "arize_space_key" + ) + if standard_callback_dynamic_params.get("arize_api_key"): + dynamic_headers["api_key"] = standard_callback_dynamic_params.get( + "arize_api_key" + ) + + if standard_callback_dynamic_params.get("arize_space_id"): + dynamic_headers["arize-space-id"] = standard_callback_dynamic_params.get( + "arize_space_id" + ) + return dynamic_headers diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index c51447c1169..e42d00315c8 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -429,6 +429,8 @@ class OpenTelemetry(CustomLogger): """ from opentelemetry import trace + from litellm.integrations.arize.arize import ArizeLogger + standard_callback_dynamic_params: Optional[ StandardCallbackDynamicParams ] = kwargs.get("standard_callback_dynamic_params") @@ -439,14 +441,7 @@ class OpenTelemetry(CustomLogger): dynamic_headers = {} # Handle Arize headers - if standard_callback_dynamic_params.get("arize_space_key"): - dynamic_headers["space_key"] = standard_callback_dynamic_params.get( - "arize_space_key" - ) - if standard_callback_dynamic_params.get("arize_api_key"): - dynamic_headers["api_key"] = standard_callback_dynamic_params.get( - "arize_api_key" - ) + dynamic_headers = ArizeLogger.construct_dynamic_arize_headers(standard_callback_dynamic_params=standard_callback_dynamic_params) # Only create a span processor if we have headers to use if len(dynamic_headers) > 0: @@ -579,7 +574,9 @@ class OpenTelemetry(CustomLogger): ) return elif self.callback_name == "langfuse_otel": - from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger + from litellm.integrations.langfuse.langfuse_otel import ( + LangfuseOtelLogger, + ) LangfuseOtelLogger.set_langfuse_otel_attributes( span, kwargs, response_obj diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e49e592ac6d..f7ab4e06eb6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -141,20 +141,49 @@ def convert_key_logging_metadata_to_callback( return team_callback_settings_obj +class KeyAndTeamLoggingSettings: + """ + Helper class to get the dynamic logging settings for the key and team + """ + @staticmethod + def get_key_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): + if user_api_key_dict.metadata is not None and "logging" in user_api_key_dict.metadata: + return user_api_key_dict.metadata["logging"] + return None + + @staticmethod + def get_team_dynamic_logging_settings(user_api_key_dict: UserAPIKeyAuth): + if user_api_key_dict.team_metadata is not None and "logging" in user_api_key_dict.team_metadata: + return user_api_key_dict.team_metadata["logging"] + return None def _get_dynamic_logging_metadata( user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig ) -> Optional[TeamCallbackMetadata]: callback_settings_obj: Optional[TeamCallbackMetadata] = None - if ( - user_api_key_dict.metadata is not None - and "logging" in user_api_key_dict.metadata - ): - for item in user_api_key_dict.metadata["logging"]: + key_dynamic_logging_settings: Optional[dict] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict) + team_dynamic_logging_settings: Optional[dict] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict) + ######################################################################################### + # Key-based callbacks + ######################################################################################### + if key_dynamic_logging_settings is not None: + for item in key_dynamic_logging_settings: callback_settings_obj = convert_key_logging_metadata_to_callback( data=AddTeamCallback(**item), team_callback_settings_obj=callback_settings_obj, ) + ######################################################################################### + # Team-based callbacks + ######################################################################################### + elif team_dynamic_logging_settings is not None: + for item in team_dynamic_logging_settings: + callback_settings_obj = convert_key_logging_metadata_to_callback( + data=AddTeamCallback(**item), + team_callback_settings_obj=callback_settings_obj, + ) + ######################################################################################### + # Deprecated format - maintained for backwards compatibility + ######################################################################################### elif ( user_api_key_dict.team_metadata is not None and "callback_settings" in user_api_key_dict.team_metadata @@ -174,6 +203,9 @@ def _get_dynamic_logging_metadata( verbose_proxy_logger.debug( "Team callback settings activated: %s", callback_settings_obj ) + ######################################################################################### + # Enter here when configured on the config.yaml file. + ######################################################################################### elif user_api_key_dict.team_id is not None: callback_settings_obj = ( LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config( diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/test_litellm/integrations/arize/test_arize_utils.py index 4286398aca0..83df42c73a1 100644 --- a/tests/test_litellm/integrations/arize/test_arize_utils.py +++ b/tests/test_litellm/integrations/arize/test_arize_utils.py @@ -232,3 +232,42 @@ async def test_arize_dynamic_params(): test_arize_logger.standard_callback_dynamic_params.get("arize_space_key") == "test_space_key_dynamic" ) + + +def test_construct_dynamic_arize_headers(): + """ + Test the construct_dynamic_arize_headers method with various input scenarios. + Ensures that dynamic Arize headers are properly constructed from callback parameters. + """ + + # Test with all parameters present + dynamic_params_full = { + "arize_space_key": "test_space_key", + "arize_api_key": "test_api_key", + "arize_space_id": "test_space_id" + } + + headers = ArizeLogger.construct_dynamic_arize_headers(dynamic_params_full) + expected_headers = { + "space_key": "test_space_key", + "api_key": "test_api_key", + "arize-space-id": "test_space_id" + } + assert headers == expected_headers + + # Test with only space_id + dynamic_params_space_id_only = { + "arize_space_id": "test_space_id" + } + + headers = ArizeLogger.construct_dynamic_arize_headers(dynamic_params_space_id_only) + expected_headers = { + "arize-space-id": "test_space_id" + } + assert headers == expected_headers + + # Test with empty parameters dict + dynamic_params_empty = {} + + headers = ArizeLogger.construct_dynamic_arize_headers(dynamic_params_empty) + assert headers == {} diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 4b0b1a13a4f..626d82e95cc 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -8,9 +8,11 @@ from unittest.mock import MagicMock, patch import pytest from fastapi import Request -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( + KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup, + _get_dynamic_logging_metadata, _get_enforced_params, add_litellm_data_to_request, check_if_token_is_service_account, @@ -216,3 +218,123 @@ async def test_add_litellm_data_to_request_audio_transcription_multipart(): "jobID:214590dsff09fds", "taskName:run_page_classification", ] + + +def test_key_dynamic_logging_settings(): + """ + Test KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings method with arize and langfuse callbacks + """ + # Test with arize logging + key_with_arize = UserAPIKeyAuth( + api_key="test-key", + metadata={ + "logging": [ + {"callback_name": "arize", "callback_type": "success"} + ] + }, + team_metadata={} + ) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_arize) + assert result == [{"callback_name": "arize", "callback_type": "success"}] + + # Test with langfuse logging + key_with_langfuse = UserAPIKeyAuth( + api_key="test-key", + metadata={ + "logging": [ + {"callback_name": "langfuse", "callback_type": "success"} + ] + }, + team_metadata={} + ) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse) + assert result == [{"callback_name": "langfuse", "callback_type": "success"}] + + # Test with no logging metadata + key_without_logging = UserAPIKeyAuth( + api_key="test-key", + metadata={}, + team_metadata={} + ) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_without_logging) + assert result is None + + +def test_team_dynamic_logging_settings(): + """ + Test KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings method with arize and langfuse callbacks + """ + # Test with arize team logging + key_with_team_arize = UserAPIKeyAuth( + api_key="test-key", + metadata={}, + team_metadata={ + "logging": [ + {"callback_name": "arize", "callback_type": "failure"} + ] + } + ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_arize) + assert result == [{"callback_name": "arize", "callback_type": "failure"}] + + # Test with langfuse team logging + key_with_team_langfuse = UserAPIKeyAuth( + api_key="test-key", + metadata={}, + team_metadata={ + "logging": [ + {"callback_name": "langfuse", "callback_type": "success"} + ] + } + ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_langfuse) + assert result == [{"callback_name": "langfuse", "callback_type": "success"}] + + # Test with no team logging metadata + key_without_team_logging = UserAPIKeyAuth( + api_key="test-key", + metadata={}, + team_metadata={} + ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_without_team_logging) + assert result is None + + +def test_get_dynamic_logging_metadata_with_arize_team_logging(): + """ + Test _get_dynamic_logging_metadata function with arize team logging and dynamic parameters + """ + # Setup user with arize team logging including callback_vars + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + metadata={}, + team_metadata={ + "logging": [ + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": { + "arize_api_key": "test_arize_api_key", + "arize_space_id": "test_arize_space_id" + } + } + ] + } + ) + + # Mock proxy_config (not used in this test path since we have team dynamic logging) + mock_proxy_config = MagicMock() + + # Call the function + result = _get_dynamic_logging_metadata( + user_api_key_dict=user_api_key_dict, + proxy_config=mock_proxy_config + ) + + # Verify the result + assert result is not None + assert isinstance(result, TeamCallbackMetadata) + assert result.success_callback == ["arize"] + assert result.callback_vars is not None + assert result.callback_vars["arize_api_key"] == "test_arize_api_key" + assert result.callback_vars["arize_space_id"] == "test_arize_space_id"