diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 15c035ceec8..c73f0b22b4b 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 282000a4456..00000000000 --- 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 cbba35942bc..00000000000 --- 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 00000000000..e0fd610e678 --- /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 cbb016ba810..0e150e064c7 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 5929934fb37..f4aa28d98cc 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,