Spaces:
Sleeping
Sleeping
File size: 3,228 Bytes
469eae6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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()
|