From c8c73e6fa5f78dff1d8f02f33c91d87d1586918c Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Mon, 5 Jan 2026 10:41:24 +0900 Subject: [PATCH 1/2] fix: MCP handling in unified guardrail --- litellm/llms/mcp/__init__.py | 7 ++ .../mcp/guardrail_translation/__init__.py | 14 +++ .../llms/mcp/guardrail_translation/handler.py | 89 +++++++++++++++++++ .../mcp_server/mcp_server_manager.py | 9 +- .../unified_guardrail/unified_guardrail.py | 17 ++-- litellm/proxy/utils.py | 72 +++++++-------- .../mcp/guardrail_translation/test_handler.py | 76 ++++++++++++++++ .../test_unified_guardrail.py | 82 +++++++++++++++++ 8 files changed, 322 insertions(+), 44 deletions(-) create mode 100644 litellm/llms/mcp/__init__.py create mode 100644 litellm/llms/mcp/guardrail_translation/__init__.py create mode 100644 litellm/llms/mcp/guardrail_translation/handler.py create mode 100644 tests/test_litellm/llms/mcp/guardrail_translation/test_handler.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py diff --git a/litellm/llms/mcp/__init__.py b/litellm/llms/mcp/__init__.py new file mode 100644 index 0000000000..282000a445 --- /dev/null +++ b/litellm/llms/mcp/__init__.py @@ -0,0 +1,7 @@ +"""Guardrail-only MCP helpers within litellm.llms. + +This package lives under ``litellm.llms`` purely so unified guardrail discovery +can find MCP-specific translations. It does not expose a traditional provider. +""" + +__all__ = ["guardrail_translation"] diff --git a/litellm/llms/mcp/guardrail_translation/__init__.py b/litellm/llms/mcp/guardrail_translation/__init__.py new file mode 100644 index 0000000000..cbba35942b --- /dev/null +++ b/litellm/llms/mcp/guardrail_translation/__init__.py @@ -0,0 +1,14 @@ +"""Guardrail translation handler wiring for MCP tool calls.""" + +from litellm.llms.mcp.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) +from litellm.types.utils import CallTypes + +# These mappings live under litellm.llms so the unified guardrail discovery can +# treat MCP tooling like any other provider endpoint. +guardrail_translation_mappings = { + CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, +} + +__all__ = ["guardrail_translation_mappings", "MCPGuardrailTranslationHandler"] diff --git a/litellm/llms/mcp/guardrail_translation/handler.py b/litellm/llms/mcp/guardrail_translation/handler.py new file mode 100644 index 0000000000..8d6d236b88 --- /dev/null +++ b/litellm/llms/mcp/guardrail_translation/handler.py @@ -0,0 +1,89 @@ +""" +MCP Guardrail Handler for Unified Guardrails. + +This handler works with the synthetic "messages" payload generated by +`ProxyLogging._convert_mcp_to_llm_format`, which always produces a single user +message whose `content` string encodes the MCP tool name and arguments. The +handler simply feeds that text through the configured guardrail and writes the +result back onto the message. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from mcp.types import CallToolResult + + +class MCPGuardrailTranslationHandler(BaseTranslation): + """Guardrail translation handler for MCP tool calls.""" + + async def process_input_messages( + self, + data: Dict[str, Any], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + ) -> Dict[str, Any]: + messages = data.get("messages") + if not isinstance(messages, list) or not messages: + verbose_proxy_logger.debug("MCP Guardrail: No messages to process") + return data + + first_message = messages[0] + content: Optional[str] = None + if isinstance(first_message, dict): + content = first_message.get("content") + else: + content = getattr(first_message, "content", None) + + if not isinstance(content, str): + verbose_proxy_logger.debug( + "MCP Guardrail: Message content missing or not a string", + ) + return data + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=[content]), + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + guardrailed_texts = ( + guardrailed_inputs.get("texts", []) if guardrailed_inputs else [] + ) + + if guardrailed_texts: + new_content = guardrailed_texts[0] + if isinstance(first_message, dict): + first_message["content"] = new_content + else: + setattr(first_message, "content", new_content) + + verbose_proxy_logger.debug( + "MCP Guardrail: Updated content for tool %s", + data.get("mcp_tool_name"), + ) + else: + verbose_proxy_logger.debug( + "MCP Guardrail: Guardrail returned no text updates for tool %s", + data.get("mcp_tool_name"), + ) + + return data + + async def process_output_response( + self, + response: "CallToolResult", + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + user_api_key_dict: Optional[Any] = None, + ) -> Any: + # Not implemented: MCP guardrail translation never calls this path today. + verbose_proxy_logger.debug( + "MCP Guardrail: Output processing not implemented for MCP tools", + ) + return response diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 15c41ecbd5..d193fc27fb 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -30,6 +30,7 @@ from pydantic import AnyUrl import litellm from litellm._logging import verbose_logger +from litellm.types.utils import CallTypes from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.experimental_mcp_client.client import MCPClient from litellm.llms.custom_httpx.http_handler import get_async_httpx_client @@ -1676,11 +1677,11 @@ class MCPServerManager: ) try: - # Use standard pre_call_hook with call_type="mcp_call" + # Use standard pre_call_hook modified_data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_auth, # type: ignore data=synthetic_llm_data, - call_type="mcp_call", # type: ignore + call_type=CallTypes.call_mcp_tool.value, ) if modified_data: # Convert response back to MCP format and apply modifications @@ -1737,7 +1738,7 @@ class MCPServerManager: proxy_logging_obj.during_call_hook( user_api_key_dict=user_api_key_auth, data=synthetic_llm_data, - call_type="mcp_call", # type: ignore + call_type=CallTypes.call_mcp_tool.value, ) ) @@ -1893,7 +1894,7 @@ class MCPServerManager: ######################################################### # Pre MCP Tool Call Hook # Allow validation and modification of tool calls before execution - # Using standard pre_call_hook with call_type="mcp_call" + # Using standard pre_call_hook ######################################################### if proxy_logging_obj: await self.pre_call_tool_check( diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index a1bbf36ac0..f66341fde5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -28,7 +28,6 @@ class UnifiedLLMGuardrails(CustomLogger): self, **kwargs, ): - # store kwargs as optional_params self.optional_params = kwargs @@ -63,6 +62,9 @@ class UnifiedLLMGuardrails(CustomLogger): return data event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call + if call_type == CallTypes.call_mcp_tool.value: + event_type = GuardrailEventHooks.pre_mcp_call + if ( guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True @@ -114,6 +116,9 @@ class UnifiedLLMGuardrails(CustomLogger): return data event_type: GuardrailEventHooks = GuardrailEventHooks.during_call + if call_type == CallTypes.call_mcp_tool.value: + event_type = GuardrailEventHooks.during_mcp_call + if ( guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True @@ -128,7 +133,10 @@ class UnifiedLLMGuardrails(CustomLogger): endpoint_guardrail_translation_mappings = ( load_guardrail_translation_mappings() ) - if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: + if ( + call_type is not None + and CallTypes(call_type) not in endpoint_guardrail_translation_mappings + ): return data endpoint_translation = endpoint_guardrail_translation_mappings[ @@ -180,8 +188,8 @@ class UnifiedLLMGuardrails(CustomLogger): call_type: Optional[CallTypesLiteral] = None if user_api_key_dict.request_route is not None: call_types = get_call_types_for_route(user_api_key_dict.request_route) - if call_types is not None and len(call_types) > 0: # type: ignore - call_type = call_types[0] # type: ignore + if call_types is not None and len(call_types) > 0: # type: ignore + call_type = call_types[0] # type: ignore if call_type is None: call_type = _infer_call_type(call_type=None, completion_response=response) # type: ignore @@ -330,7 +338,6 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: - verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d595db4a2e..d1a78534da 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -151,25 +151,25 @@ def _get_email_logger_class(): """ Determine which email logger class to use based on environment variables. Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback) - + Returns: The email logger class to use, or None if BaseEmailLogger is not available """ if BaseEmailLogger is None: return None - + # Check for SendGrid API key if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"): return SendGridEmailLogger - + # Check for Resend API key if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"): return ResendEmailLogger - + # Check for SMTP configuration if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"): return SMTPEmailLogger - + # Fallback to BaseEmailLogger (though it won't actually send emails) return BaseEmailLogger @@ -452,7 +452,6 @@ class ProxyLogging: litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # type: ignore for callback in litellm.callbacks: if isinstance(callback, str): - callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore cast(_custom_logger_compatible_callbacks_literal, callback), internal_usage_cache=self.internal_usage_cache.dual_cache, @@ -965,7 +964,7 @@ class ProxyLogging: # Determine the event type based on call type event_type = GuardrailEventHooks.pre_call - if call_type == "mcp_call": + if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.pre_mcp_call # Check if the guardrail should run for this request @@ -1038,7 +1037,6 @@ class ProxyLogging: data.pop("prompt_id", None) if custom_logger and prompt_spec is not None: - ( model, messages, @@ -1261,7 +1259,7 @@ class ProxyLogging: from litellm.types.guardrails import GuardrailEventHooks event_type = GuardrailEventHooks.during_call - if call_type == "mcp_call": + if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.during_mcp_call if ( @@ -1270,7 +1268,7 @@ class ProxyLogging: ): continue # Convert user_api_key_dict to proper format for async_moderation_hook - if call_type == "mcp_call": + if call_type == CallTypes.call_mcp_tool.value: user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( user_api_key_dict ) @@ -1288,7 +1286,6 @@ class ProxyLogging: call_type=call_type, ) else: - guardrail_task = callback.async_moderation_hook( data=data, user_api_key_dict=user_api_key_auth_dict, # type: ignore @@ -1337,7 +1334,7 @@ class ProxyLogging: if self.alerting is None: # do nothing if alerting is not switched on return - + if "slack" in self.alerting: await self.slack_alerting_instance.budget_alerts( type=type, @@ -1548,7 +1545,10 @@ class ProxyLogging: traceback_str=traceback_str, ) # If callback returned an HTTPException, use it (first one wins) - if isinstance(hook_result, HTTPException) and transformed_exception is None: + if ( + isinstance(hook_result, HTTPException) + and transformed_exception is None + ): transformed_exception = hook_result except HTTPException as e: # If callback raised an HTTPException, use it (first one wins) @@ -1849,7 +1849,6 @@ class ProxyLogging: current_response = response for callback in litellm.callbacks: - _callback: Optional[CustomLogger] = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -3568,11 +3567,13 @@ class ProxyUpdateSpend: ) # Atomically read and remove logs to process (protected by lock) async with prisma_client._spend_log_transactions_lock: - logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] + logs_to_process = prisma_client.spend_log_transactions[ + :MAX_LOGS_PER_INTERVAL + ] # Remove the logs we're about to process - prisma_client.spend_log_transactions = ( - prisma_client.spend_log_transactions[len(logs_to_process):] - ) + prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ + len(logs_to_process) : + ] start_time = time.time() try: for i in range(n_retry_times + 1): @@ -3675,9 +3676,7 @@ async def update_spend( # noqa: PLR0915 # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - verbose_proxy_logger.debug( - "Spend Logs transactions: {}".format(queue_size) - ) + verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size)) # Process spend log transactions when called directly. # This keeps backwards compatibility with the old behavior. @@ -3699,19 +3698,19 @@ async def update_spend_logs_job( ): """ Job to process spend_log_transactions queue. - + This job is triggered based on queue size rather than time. Processes spend log transactions when the queue reaches a threshold. """ n_retry_times = 3 - + # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - + if queue_size == 0: return - + await ProxyUpdateSpend.update_spend_logs( n_retry_times=n_retry_times, prisma_client=prisma_client, @@ -3728,7 +3727,7 @@ async def _monitor_spend_logs_queue( """ Background task that monitors the spend_log_transactions queue size and triggers processing when the threshold is reached. - + Args: prisma_client: Prisma client instance db_writer_client: Optional HTTP handler for external spend logs endpoint @@ -3738,23 +3737,23 @@ async def _monitor_spend_logs_queue( SPEND_LOG_QUEUE_POLL_INTERVAL, SPEND_LOG_QUEUE_SIZE_THRESHOLD, ) - + threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL max_backoff = 30.0 # Maximum backoff interval in seconds backoff_multiplier = 1.5 # Exponential backoff multiplier current_interval = base_interval - + verbose_proxy_logger.info( f"Starting spend logs queue monitor (threshold: {threshold}, poll_interval: {base_interval}s)" ) - + while True: try: # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - + if queue_size > 0: if queue_size >= threshold: verbose_proxy_logger.debug( @@ -3767,8 +3766,10 @@ async def _monitor_spend_logs_queue( f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff" ) # Exponential backoff when below threshold but still processing - current_interval = min(current_interval * backoff_multiplier, max_backoff) - + current_interval = min( + current_interval * backoff_multiplier, max_backoff + ) + await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, @@ -3776,8 +3777,10 @@ async def _monitor_spend_logs_queue( ) else: # Exponential backoff when no logs to process - current_interval = min(current_interval * backoff_multiplier, max_backoff) - + current_interval = min( + current_interval * backoff_multiplier, max_backoff + ) + await asyncio.sleep(current_interval) except Exception as e: verbose_proxy_logger.error( @@ -3788,7 +3791,6 @@ async def _monitor_spend_logs_queue( await asyncio.sleep(current_interval) - def _raise_failed_update_spend_exception( e: Exception, start_time: float, proxy_logging_obj: ProxyLogging ): diff --git a/tests/test_litellm/llms/mcp/guardrail_translation/test_handler.py b/tests/test_litellm/llms/mcp/guardrail_translation/test_handler.py new file mode 100644 index 0000000000..cbb016ba81 --- /dev/null +++ b/tests/test_litellm/llms/mcp/guardrail_translation/test_handler.py @@ -0,0 +1,76 @@ +"""Tests for the MCP guardrail translation handler.""" + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.mcp.guardrail_translation.handler import MCPGuardrailTranslationHandler + + +class MockGuardrail(CustomGuardrail): + """Simple guardrail mock that records invocations.""" + + def __init__(self, return_texts=None): + super().__init__(guardrail_name="mock-mcp-guardrail") + self.return_texts = return_texts + self.call_count = 0 + self.last_inputs = None + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.call_count += 1 + self.last_inputs = inputs + + if self.return_texts is not None: + return {"texts": self.return_texts} + + texts = inputs.get("texts", []) + return {"texts": [f"{text} [SAFE]" for text in texts]} + + +@pytest.mark.asyncio +async def test_process_input_messages_updates_content(): + """Handler should update the synthetic message content when guardrail modifies text.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MockGuardrail() + + original_content = "Tool: weather\nArguments: {'city': 'tokyo'}" + data = { + "messages": [{"role": "user", "content": original_content}], + "mcp_tool_name": "weather", + } + + result = await handler.process_input_messages(data, guardrail) + + assert result["messages"][0]["content"].endswith("[SAFE]") + assert guardrail.last_inputs == {"texts": [original_content]} + assert guardrail.call_count == 1 + + +@pytest.mark.asyncio +async def test_process_input_messages_skips_when_no_messages(): + """Handler should skip guardrail invocation if messages array is missing or empty.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MockGuardrail() + + data = {"mcp_tool_name": "noop"} + result = await handler.process_input_messages(data, guardrail) + + assert result == data + assert guardrail.call_count == 0 + + +@pytest.mark.asyncio +async def test_process_input_messages_handles_empty_guardrail_result(): + """Handler should leave content untouched when guardrail returns no text updates.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MockGuardrail(return_texts=[]) + + original_content = "Tool: calendar\nArguments: {'date': '2024-12-25'}" + data = { + "messages": [{"role": "user", "content": original_content}], + "mcp_tool_name": "calendar", + } + + result = await handler.process_input_messages(data, guardrail) + + assert result["messages"][0]["content"] == original_content + assert guardrail.call_count == 1 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py new file mode 100644 index 0000000000..5929934fb3 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -0,0 +1,82 @@ +"""Tests for unified guardrail.""" + +import pytest + +from litellm.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.mcp.guardrail_translation.handler import MCPGuardrailTranslationHandler +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes + + +class RecordingGuardrail(CustomGuardrail): + """Records the event types it is asked to run for.""" + + def __init__(self): + super().__init__(guardrail_name="recording-guardrail") + self.event_history = [] + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + self.event_history.append(event_type) + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + return {"texts": inputs.get("texts", [])} + + +@pytest.fixture(autouse=True) +def _inject_mcp_handler_mapping(): + """Inject MCP handler mapping so the unified guardrail can run inside tests.""" + unified_module.endpoint_guardrail_translation_mappings = { + CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, + } + yield + unified_module.endpoint_guardrail_translation_mappings = None + + +@pytest.mark.asyncio +async def test_pre_call_hook_uses_mcp_event_type(): + """pre_call hook should swap to GuardrailEventHooks.pre_mcp_call for MCP calls.""" + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + cache = DualCache() + + data = { + "guardrail_to_apply": guardrail, + "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}], + "model": "mcp-tool-call", + } + + await handler.async_pre_call_hook( + user_api_key_dict=None, + cache=cache, + data=data, + call_type=CallTypes.call_mcp_tool.value, + ) + + assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call] + + +@pytest.mark.asyncio +async def test_moderation_hook_uses_mcp_event_type(): + """moderation hook should request GuardrailEventHooks.during_mcp_call for MCP calls.""" + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + + data = { + "guardrail_to_apply": guardrail, + "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}], + "model": "mcp-tool-call", + } + + await handler.async_moderation_hook( + data=data, + user_api_key_dict=None, + call_type=CallTypes.call_mcp_tool.value, + ) + + assert guardrail.event_history == [GuardrailEventHooks.during_mcp_call] From 0509fc79da241ce9e565cf2520889bd3058c4ab7 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Mon, 5 Jan 2026 12:05:16 +0900 Subject: [PATCH 2/2] fix: move code from litellm/llms to the mcp_server dir --- litellm/llms/__init__.py | 16 ++++++++++++++++ litellm/llms/mcp/__init__.py | 7 ------- .../llms/mcp/guardrail_translation/__init__.py | 14 -------------- .../mcp_server/guardrail_translation/__init__.py | 16 ++++++++++++++++ .../mcp_server}/guardrail_translation/handler.py | 0 .../test_mcp_guardrail_handler.py} | 4 +++- .../unified_guardrails/test_unified_guardrail.py | 4 +++- 7 files changed, 38 insertions(+), 23 deletions(-) delete mode 100644 litellm/llms/mcp/__init__.py delete mode 100644 litellm/llms/mcp/guardrail_translation/__init__.py create mode 100644 litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py rename litellm/{llms/mcp => proxy/_experimental/mcp_server}/guardrail_translation/handler.py (100%) rename tests/test_litellm/{llms/mcp/guardrail_translation/test_handler.py => proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py} (95%) diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 15c035ceec..c73f0b22b4 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -45,6 +45,7 @@ def get_cost_for_web_search_request( return 0.0 elif custom_llm_provider == "xai": from .xai.cost_calculator import cost_per_web_search_request + return cost_per_web_search_request(usage=usage, model_info=model_info) else: return None @@ -110,6 +111,21 @@ def discover_guardrail_translation_mappings() -> ( verbose_logger.error(f"Error processing {module_path}: {e}") continue + try: + from litellm.proxy._experimental.mcp_server.guardrail_translation import ( + guardrail_translation_mappings as mcp_guardrail_translation_mappings, + ) + + discovered_mappings.update(mcp_guardrail_translation_mappings) + verbose_logger.debug( + "Loaded MCP guardrail translation mappings: %s", + list(mcp_guardrail_translation_mappings.keys()), + ) + except ImportError: + verbose_logger.debug( + "MCP guardrail translation mappings not available; skipping" + ) + verbose_logger.debug( f"Discovered {len(discovered_mappings)} guardrail translation mappings: {list(discovered_mappings.keys())}" ) diff --git a/litellm/llms/mcp/__init__.py b/litellm/llms/mcp/__init__.py deleted file mode 100644 index 282000a445..0000000000 --- a/litellm/llms/mcp/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Guardrail-only MCP helpers within litellm.llms. - -This package lives under ``litellm.llms`` purely so unified guardrail discovery -can find MCP-specific translations. It does not expose a traditional provider. -""" - -__all__ = ["guardrail_translation"] diff --git a/litellm/llms/mcp/guardrail_translation/__init__.py b/litellm/llms/mcp/guardrail_translation/__init__.py deleted file mode 100644 index cbba35942b..0000000000 --- a/litellm/llms/mcp/guardrail_translation/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Guardrail translation handler wiring for MCP tool calls.""" - -from litellm.llms.mcp.guardrail_translation.handler import ( - MCPGuardrailTranslationHandler, -) -from litellm.types.utils import CallTypes - -# These mappings live under litellm.llms so the unified guardrail discovery can -# treat MCP tooling like any other provider endpoint. -guardrail_translation_mappings = { - CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, -} - -__all__ = ["guardrail_translation_mappings", "MCPGuardrailTranslationHandler"] diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py new file mode 100644 index 0000000000..e0fd610e67 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py @@ -0,0 +1,16 @@ +"""Guardrail translation mapping for MCP tool calls.""" + +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) +from litellm.types.utils import CallTypes + +# This mapping lives alongside the MCP server implementation because MCP +# integrations are managed by the proxy subsystem, not litellm.llms providers. +# Unified guardrails import this module explicitly to register the handler. + +guardrail_translation_mappings = { + CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, +} + +__all__ = ["guardrail_translation_mappings", "MCPGuardrailTranslationHandler"] diff --git a/litellm/llms/mcp/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py similarity index 100% rename from litellm/llms/mcp/guardrail_translation/handler.py rename to litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py diff --git a/tests/test_litellm/llms/mcp/guardrail_translation/test_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py similarity index 95% rename from tests/test_litellm/llms/mcp/guardrail_translation/test_handler.py rename to tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py index cbb016ba81..0e150e064c 100644 --- a/tests/test_litellm/llms/mcp/guardrail_translation/test_handler.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -3,7 +3,9 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.llms.mcp.guardrail_translation.handler import MCPGuardrailTranslationHandler +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) class MockGuardrail(CustomGuardrail): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 5929934fb3..f4aa28d98c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -4,7 +4,9 @@ import pytest from litellm.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.llms.mcp.guardrail_translation.handler import MCPGuardrailTranslationHandler +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails,