refactor location of proxy

litellm_proxy/_experimental/mcp_server/tool_registry.py

@@ -0,0 +1,103 @@
+import json
+from typing import Any, Callable, Dict, List, Optional
+
+from litellm._logging import verbose_logger
+from litellm_proxy.types_utils.utils import get_instance_fn
+from litellm.types.mcp_server.tool_registry import MCPTool
+
+
+class MCPToolRegistry:
+    """
+    A registry for managing MCP tools
+    """
+
+    def __init__(self):
+        # Registry to store all registered tools
+        self.tools: Dict[str, MCPTool] = {}
+
+    def register_tool(
+        self,
+        name: str,
+        description: str,
+        input_schema: Dict[str, Any],
+        handler: Callable,
+    ) -> None:
+        """
+        Register a new tool in the registry
+        """
+        self.tools[name] = MCPTool(
+            name=name,
+            description=description,
+            input_schema=input_schema,
+            handler=handler,
+        )
+        verbose_logger.debug(f"Registered tool: {name}")
+
+    def get_tool(self, name: str) -> Optional[MCPTool]:
+        """
+        Get a tool from the registry by name
+        """
+        return self.tools.get(name)
+
+    def list_tools(self) -> List[MCPTool]:
+        """
+        List all registered tools
+        """
+        return list(self.tools.values())
+
+    def load_tools_from_config(
+        self, mcp_tools_config: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """
+        Load and register tools from the proxy config
+
+        Args:
+            mcp_tools_config: The mcp_tools config from the proxy config
+        """
+        if mcp_tools_config is None:
+            raise ValueError(
+                "mcp_tools_config is required, please set `mcp_tools` in your proxy config"
+            )
+
+        for tool_config in mcp_tools_config:
+            if not isinstance(tool_config, dict):
+                raise ValueError("mcp_tools_config must be a list of dictionaries")
+
+            name = tool_config.get("name")
+            description = tool_config.get("description")
+            input_schema = tool_config.get("input_schema", {})
+            handler_name = tool_config.get("handler")
+
+            if not all([name, description, handler_name]):
+                continue
+
+            # Try to resolve the handler
+            # First check if it's a module path (e.g., "module.submodule.function")
+            if handler_name is None:
+                raise ValueError(f"handler is required for tool {name}")
+            handler = get_instance_fn(handler_name)
+
+            if handler is None:
+                verbose_logger.warning(
+                    f"Warning: Could not find handler {handler_name} for tool {name}"
+                )
+                continue
+
+            # Register the tool
+            if name is None:
+                raise ValueError(f"name is required for tool {name}")
+            if description is None:
+                raise ValueError(f"description is required for tool {name}")
+
+            self.register_tool(
+                name=name,
+                description=description,
+                input_schema=input_schema,
+                handler=handler,
+            )
+        verbose_logger.debug(
+            "all registered tools: %s", json.dumps(self.tools, indent=4, default=str)
+        )
+
+
+global_mcp_tool_registry = MCPToolRegistry()