From 62b39775a1b67c2586a35874cb23e999c56d0cb8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 1 Oct 2025 14:09:01 -0700 Subject: [PATCH] fix: add changes from main --- litellm/integrations/opentelemetry.py | 83 +++++++++-------- litellm/proxy/auth/route_checks.py | 12 ++- .../pass_through_endpoints.py | 31 +++++++ litellm/proxy/proxy_server.py | 75 ++++++++++----- .../proxy/auth/test_route_checks.py | 91 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 58 ++++++++++++ 6 files changed, 285 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/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 56218b3345e..39f11e64bb7 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -62,12 +62,22 @@ class RouteChecks: for allowed_route in valid_token.allowed_routes ): for allowed_route in valid_token.allowed_routes: - if allowed_route in LiteLLMRoutes._member_names_: + if allowed_route in LiteLLMRoutes._member_names_: if RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value, ): return True + + ################################################ + # For llm_api_routes, also check registered pass-through endpoints + ################################################ + if allowed_route == "llm_api_routes": + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): + return True # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 34c79718452..274225ee2e3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1111,6 +1111,37 @@ class InitPassThroughEndpointHelpers: "Removed pass-through route from registry: %s", key ) + @staticmethod + def is_registered_pass_through_route(route: str) -> bool: + """ + Check if route is a registered pass-through endpoint from DB + + Uses the in-memory registry to avoid additional DB queries + Optimized for minimal latency + + Args: + route: The route to check + + Returns: + bool: True if route is a registered pass-through endpoint, False otherwise + """ + # Fast path: check if any registered route key contains this path + # Keys are in format: "{endpoint_id}:exact:{path}" or "{endpoint_id}:subpath:{path}" + # Extract unique paths from keys for quick checking + for key in _registered_pass_through_routes.keys(): + parts = key.split(":", 2) # Split into [endpoint_id, type, path] + if len(parts) == 3: + route_type = parts[1] + registered_path = parts[2] + + if route_type == "exact" and route == registered_path: + return True + elif route_type == "subpath": + if route == registered_path or route.startswith(registered_path + "/"): + return True + + return False + async def initialize_pass_through_endpoints( pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f031b9ffbdf..1cda72cada3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -252,7 +252,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, @@ -299,7 +301,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, ) @@ -2599,6 +2603,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 @@ -2606,35 +2635,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/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index ac09917e4cd..4e448f995c9 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -228,3 +228,94 @@ def test_virtual_key_allowed_routes_with_no_member_names_only_explicit(): ) assert "Virtual key is not allowed to call this route" in str(exc_info.value) + + +def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): + """ + Test that virtual keys with llm_api_routes permission can access registered pass-through endpoints. + + This tests the scenario where a pass-through endpoint is registered from the DB + (e.g., /azure-assistant) and a virtual key with llm_api_routes permission should be able to access + both the exact path and subpaths (e.g., /azure-assistant/openai/assistants). + """ + from unittest.mock import patch + + # Mock the registered pass-through routes + mock_registered_routes = { + "test-uuid-1:exact:/azure-assistant": { + "endpoint_id": "test-uuid-1", + "path": "/azure-assistant", + "type": "exact", + }, + "test-uuid-2:subpath:/custom-endpoint": { + "endpoint_id": "test-uuid-2", + "path": "/custom-endpoint", + "type": "subpath", + }, + } + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ): + # Create a virtual key with llm_api_routes permission + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + # Test exact match for registered pass-through endpoint + result1 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/azure-assistant", + valid_token=valid_token, + ) + assert result1 is True + + # Test subpath for registered pass-through endpoint with subpath type + result2 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom-endpoint/openai/assistants", + valid_token=valid_token, + ) + assert result2 is True + + # Test exact match for subpath type + result3 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom-endpoint", + valid_token=valid_token, + ) + assert result3 is True + + +def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): + """ + Test that virtual keys without llm_api_routes permission cannot access registered pass-through endpoints. + """ + from unittest.mock import patch + + # Mock the registered pass-through routes + mock_registered_routes = { + "test-uuid-1:exact:/azure-assistant": { + "endpoint_id": "test-uuid-1", + "path": "/azure-assistant", + "type": "exact", + }, + } + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ): + # Create a virtual key without llm_api_routes permission + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["info_routes"], + ) + + # Test that access is denied + with pytest.raises(Exception) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route="/azure-assistant", + valid_token=valid_token, + ) + + assert "Virtual key is not allowed to call this route" in str(exc_info.value) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d2b516a55d2..a90b3a94465 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1886,3 +1886,61 @@ async def test_add_router_settings_shallow_merge_behavior(): assert merged_settings["nested_setting"] == expected_nested assert merged_settings["top_level"] == "db_top" + +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()