From 388761f52d6b933448f14882ae08d99e3cba2ea7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 1 Oct 2025 15:33:22 -0700 Subject: [PATCH] [Fix] LiteLLM UI - Ensure OTEL settings are saved in DB after set on UI (#15118) * fix: fix _add_callback_from_db_to_in_memory_litellm_callbacks * test_add_callback_from_db_to_in_memory_litellm_callbacks * fix otel * fix: fix _add_callback_from_db_to_in_memory_litellm_callbacks --- litellm/integrations/opentelemetry.py | 83 ++++++++++--------- litellm/proxy/proxy_server.py | 76 +++++++++++------ tests/test_litellm/proxy/test_proxy_server.py | 59 +++++++++++++ 3 files changed, 153 insertions(+), 65 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index e6f265ded58..39047dbfea4 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -645,7 +645,7 @@ class OpenTelemetry(CustomLogger): if not self.config.enable_events: return - from opentelemetry._logs import get_logger, LogRecord + from opentelemetry._logs import LogRecord, get_logger otel_logger = get_logger(LITELLM_LOGGER_NAME) parent_ctx = span.get_span_context() @@ -1115,51 +1115,56 @@ class OpenTelemetry(CustomLogger): span.set_attribute(key, primitive_value) def set_raw_request_attributes(self, span: 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") + try: + kwargs.get("optional_params", {}) + litellm_params = kwargs.get("litellm_params", {}) or {} + custom_llm_provider = litellm_params.get("custom_llm_provider", "Unknown") - _raw_response = kwargs.get("original_response") - _additional_args = kwargs.get("additional_args", {}) or {} - complete_input_dict = _additional_args.get("complete_input_dict") - ############################################# - ########## LLM Request Attributes ########### - ############################################# + _raw_response = kwargs.get("original_response") + _additional_args = kwargs.get("additional_args", {}) or {} + complete_input_dict = _additional_args.get("complete_input_dict") + ############################################# + ########## LLM Request Attributes ########### + ############################################# - # OTEL Attributes for the RAW Request to https://docs.anthropic.com/en/api/messages - if complete_input_dict and isinstance(complete_input_dict, dict): - for param, val in complete_input_dict.items(): - self.safe_set_attribute( - span=span, key=f"llm.{custom_llm_provider}.{param}", value=val - ) + # OTEL Attributes for the RAW Request to https://docs.anthropic.com/en/api/messages + if complete_input_dict and isinstance(complete_input_dict, dict): + for param, val in complete_input_dict.items(): + self.safe_set_attribute( + span=span, key=f"llm.{custom_llm_provider}.{param}", value=val + ) - ############################################# - ########## LLM Response Attributes ########## - ############################################# - if _raw_response and isinstance(_raw_response, str): - # cast sr -> dict - import json + ############################################# + ########## LLM Response Attributes ########## + ############################################# + if _raw_response and isinstance(_raw_response, str): + # cast sr -> dict + import json + + try: + _raw_response = json.loads(_raw_response) + for param, val in _raw_response.items(): + self.safe_set_attribute( + span=span, + key=f"llm.{custom_llm_provider}.{param}", + value=val, + ) + except json.JSONDecodeError: + verbose_logger.debug( + "litellm.integrations.opentelemetry.py::set_raw_request_attributes() - raw_response not json string - {}".format( + _raw_response + ) + ) - try: - _raw_response = json.loads(_raw_response) - for param, val in _raw_response.items(): self.safe_set_attribute( span=span, - key=f"llm.{custom_llm_provider}.{param}", - value=val, + key=f"llm.{custom_llm_provider}.stringified_raw_response", + value=_raw_response, ) - except json.JSONDecodeError: - verbose_logger.debug( - "litellm.integrations.opentelemetry.py::set_raw_request_attributes() - raw_response not json string - {}".format( - _raw_response - ) - ) - - self.safe_set_attribute( - span=span, - key=f"llm.{custom_llm_provider}.stringified_raw_response", - value=_raw_response, - ) + except Exception as e: + verbose_logger.exception( + "OpenTelemetry logging error in set_raw_request_attributes %s", str(e) + ) def _to_ns(self, dt): return int(dt.timestamp() * 1e9) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b1a269b31cd..f9cd1003a90 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -155,7 +155,6 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( router as mcp_discoverable_endpoints_router, ) - from litellm.proxy._experimental.mcp_server.rest_endpoints import ( router as mcp_rest_endpoints_router, ) @@ -254,7 +253,9 @@ from litellm.proxy.management_endpoints.customer_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_update, +) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -301,7 +302,9 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + set_files_config, +) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -2608,6 +2611,31 @@ class ProxyConfig: proxy_logging_obj=proxy_logging_obj, ) + def _add_callback_from_db_to_in_memory_litellm_callbacks( + self, + callback: str, + event_types: List[Literal["success", "failure"]], + existing_callbacks: list, + ) -> None: + """ + Helper method to add a single callback to litellm for specified event types. + + Args: + callback: The callback name to add + event_types: List of event types (e.g., ["success"], ["failure"], or ["success", "failure"]) + existing_callbacks: The existing callback list to check against + """ + if callback in litellm._known_custom_logger_compatible_callbacks: + for event_type in event_types: + _add_custom_logger_callback_to_specific_event(callback, event_type) + elif callback not in existing_callbacks: + if event_types == ["success"]: + litellm.logging_callback_manager.add_litellm_success_callback(callback) + elif event_types == ["failure"]: + litellm.logging_callback_manager.add_litellm_failure_callback(callback) + else: # Both success and failure + litellm.logging_callback_manager.add_litellm_callback(callback) + def _add_callbacks_from_db_config(self, config_data: dict) -> None: """ Adds callbacks from DB config to litellm @@ -2615,35 +2643,31 @@ class ProxyConfig: litellm_settings = config_data.get("litellm_settings", {}) or {} success_callbacks = litellm_settings.get("success_callback", None) failure_callbacks = litellm_settings.get("failure_callback", None) + callbacks = litellm_settings.get("callbacks", None) if success_callbacks is not None and isinstance(success_callbacks, list): for success_callback in success_callbacks: - if ( - success_callback - in litellm._known_custom_logger_compatible_callbacks - ): - _add_custom_logger_callback_to_specific_event( - success_callback, "success" - ) - elif success_callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback( - success_callback - ) + self._add_callback_from_db_to_in_memory_litellm_callbacks( + callback=success_callback, + event_types=["success"], + existing_callbacks=litellm.success_callback, + ) - # Add failure callbacks from DB to litellm if failure_callbacks is not None and isinstance(failure_callbacks, list): for failure_callback in failure_callbacks: - if ( - failure_callback - in litellm._known_custom_logger_compatible_callbacks - ): - _add_custom_logger_callback_to_specific_event( - failure_callback, "failure" - ) - elif failure_callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback( - failure_callback - ) + self._add_callback_from_db_to_in_memory_litellm_callbacks( + callback=failure_callback, + event_types=["failure"], + existing_callbacks=litellm.failure_callback, + ) + + if callbacks is not None and isinstance(callbacks, list): + for callback in callbacks: + self._add_callback_from_db_to_in_memory_litellm_callbacks( + callback=callback, + event_types=["success", "failure"], + existing_callbacks=litellm.callbacks, + ) def _encrypt_env_variables( self, environment_variables: dict, new_encryption_key: Optional[str] = None diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 1cbe6420f6e..63436899fcc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1971,3 +1971,62 @@ async def test_model_info_v1_oci_secrets_not_leaked(): assert "aa:bb:cc:dd:ee:ff:11:22:33:44:55:66:77:88:99:00" not in result_str assert "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str assert "/path/to/oci_api_key.pem" not in result_str + + +def test_add_callback_from_db_to_in_memory_litellm_callbacks(): + """ + Test that _add_callback_from_db_to_in_memory_litellm_callbacks correctly adds callbacks + for success, failure, and combined event types. + """ + from unittest.mock import MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Mock the callback manager + mock_callback_manager = MagicMock() + + with patch("litellm.proxy.proxy_server.litellm") as mock_litellm: + # Set up mock litellm attributes + mock_litellm._known_custom_logger_compatible_callbacks = [] + mock_litellm.logging_callback_manager = mock_callback_manager + + # Test Case 1: Add success callback + mock_success_callbacks = [] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="prometheus", + event_types=["success"], + existing_callbacks=mock_success_callbacks, + ) + mock_callback_manager.add_litellm_success_callback.assert_called_once_with("prometheus") + mock_callback_manager.reset_mock() + + # Test Case 2: Add failure callback + mock_failure_callbacks = [] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="langfuse", + event_types=["failure"], + existing_callbacks=mock_failure_callbacks, + ) + mock_callback_manager.add_litellm_failure_callback.assert_called_once_with("langfuse") + mock_callback_manager.reset_mock() + + # Test Case 3: Add callback for both success and failure + mock_callbacks = [] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="s3", + event_types=["success", "failure"], + existing_callbacks=mock_callbacks, + ) + mock_callback_manager.add_litellm_callback.assert_called_once_with("s3") + mock_callback_manager.reset_mock() + + # Test Case 4: Don't add callback if it already exists + existing_callbacks_with_item = ["prometheus"] + proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( + callback="prometheus", + event_types=["success"], + existing_callbacks=existing_callbacks_with_item, + ) + mock_callback_manager.add_litellm_success_callback.assert_not_called()