Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/tool_registry.py: 29%

1import json 

2from collections.abc import Callable 

3from typing import TYPE_CHECKING, Any, Final 

4 

5from litellm._logging import verbose_logger 

6from litellm.proxy.types_utils.utils import get_instance_fn 

7from litellm.types.mcp_server.tool_registry import MCPTool 

8 

9if TYPE_CHECKING: 9 ↛ 10line 9 didn't jump to line 10 because the condition on line 9 was never true

10 from mcp.types import Tool as MCPToolSDKTool 

11else: 

12 try: 

13 from mcp.types import Tool as MCPToolSDKTool 

14 except ImportError: 

15 MCPToolSDKTool = None 

16 

17 

18class MCPToolRegistry: 

19 """ 

20 A registry for managing MCP tools 

21 """ 

22 

23 def __init__(self): 

24 # Registry to store all registered tools 

25 self.tools: dict[str, MCPTool] = {} 

26 

27 def register_tool( 

28 self, 

29 name: str, 

30 description: str, 

31 input_schema: dict[str, Any], 

32 handler: Callable, 

33 ) -> None: 

34 """ 

35 Register a new tool in the registry 

36 """ 

37 self.tools[name] = MCPTool( 

38 name=name, 

39 description=description, 

40 input_schema=input_schema, 

41 handler=handler, 

42 ) 

43 verbose_logger.debug("Registered tool: %s", name) 

44 

45 def get_tool(self, name: str) -> MCPTool | None: 

46 """ 

47 Get a tool from the registry by name 

48 """ 

49 return self.tools.get(name) 

50 

51 def list_tools(self, tool_prefix: str | None = None) -> list[MCPTool]: 

52 """ 

53 List all registered tools 

54 """ 

55 if tool_prefix: 55 ↛ 57line 55 didn't jump to line 57 because the condition on line 55 was always true

56 return [tool for tool in self.tools.values() if tool.name.startswith(tool_prefix)] 

57 return list(self.tools.values()) 

58 

59 def unregister_tools_with_prefix(self, prefix: str) -> int: 

60 """Remove tools whose registered name starts with ``prefix``. 

61 

62 Used when an OpenAPI-backed MCP server leaves the runtime registry so 

63 stale tool handlers cannot be invoked after eviction. 

64 """ 

65 if not prefix: 

66 return 0 

67 removed = 0 

68 for name in list(self.tools.keys()): 

69 if name.startswith(prefix): 

70 del self.tools[name] 

71 removed += 1 

72 verbose_logger.debug("Unregistered MCP tool %s", name) 

73 return removed 

74 

75 def convert_tools_to_mcp_sdk_tool_type(self, tools: list[MCPTool]) -> list["MCPToolSDKTool"]: 

76 if MCPToolSDKTool is None: 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true

77 raise ImportError("MCP SDK is not installed. Please install it with: pip install 'litellm[proxy]'") 

78 return [ 

79 MCPToolSDKTool( 

80 name=tool.name, 

81 description=tool.description, 

82 inputSchema=tool.input_schema, 

83 ) 

84 for tool in tools 

85 ] 

86 

87 def load_tools_from_config( 

88 self, 

89 mcp_tools_config: dict[str, Any] | None = None, 

90 config_file_path: str | None = None, 

91 ) -> None: 

92 """ 

93 Load and register tools from the proxy config 

94 

95 Args: 

96 mcp_tools_config: The mcp_tools config from the proxy config 

97 config_file_path: Path to the operator's config.yaml. Threaded 

98 through to ``get_instance_fn`` so an ``s3://``/``gcs://`` 

99 ``handler`` declared in the YAML resolves; callers from a 

100 non-YAML path must leave this ``None`` so the runtime gate 

101 fires. 

102 """ 

103 if mcp_tools_config is None: 

104 raise ValueError("mcp_tools_config is required, please set `mcp_tools` in your proxy config") 

105 

106 for tool_config in mcp_tools_config: 

107 if not isinstance(tool_config, dict): 

108 raise ValueError("mcp_tools_config must be a list of dictionaries") 

109 

110 name = tool_config.get("name") 

111 description = tool_config.get("description") 

112 input_schema = tool_config.get("input_schema", {}) 

113 handler_name = tool_config.get("handler") 

114 

115 if not all([name, description, handler_name]): 

116 continue 

117 

118 # Try to resolve the handler 

119 # First check if it's a module path (e.g., "module.submodule.function") 

120 if handler_name is None: 

121 raise ValueError(f"handler is required for tool {name}") 

122 handler = get_instance_fn(handler_name, config_file_path) 

123 

124 if handler is None: 

125 verbose_logger.warning("Warning: Could not find handler %s for tool %s", handler_name, name) 

126 continue 

127 

128 # Register the tool 

129 if name is None: 

130 raise ValueError(f"name is required for tool {name}") 

131 if description is None: 

132 raise ValueError(f"description is required for tool {name}") 

133 

134 self.register_tool( 

135 name=name, 

136 description=description, 

137 input_schema=input_schema, 

138 handler=handler, 

139 ) 

140 verbose_logger.debug("all registered tools: %s", json.dumps(self.tools, indent=4, default=str)) 

141 

142 

143global_mcp_tool_registry: Final = MCPToolRegistry()