mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 02:34:29 +00:00
103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
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()
|