File size: 3,194 Bytes
49cb9f5 |
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 |
import functools
import inspect
import typing as ty
import gradio as gr
# Define a registry to store tools and their metadata
class MCPToolRegistry:
def __init__(self):
# Stores tool data: {'api_name': {'name': str, 'tool_func': func, 'ui_builder': func}}
self.tools = {}
def define(self, name: str, api_name: str):
"""Decorator to define an MCP tool."""
def decorator(tool_func: ty.Callable):
if api_name in self.tools:
raise ValueError(f"Tool with api_name '{api_name}' already defined.")
# Store the tool function and metadata
self.tools[api_name] = {
'name': name,
'tool_func': tool_func,
'ui_builder': None, # Will be filled by @build_ui_control
}
# Gradio's MCP server needs the function itself to be decorated
# with an attribute carrying the MCP metadata.
# This is a Gradio-specific requirement for mcp_server=True
# The structure below is based on Gradio's internal handling.
setattr(tool_func, "__mcp_tool__", {"name": name, "api_name": api_name})
# Return the original function so it can be called if needed
return tool_func
return decorator
def build_ui_control(self, api_name: str):
"""Decorator to associate a UI builder function with a tool."""
def decorator(ui_builder_func: ty.Callable[..., ty.Union[gr.components.Component, tuple[gr.components.Component, ...]]]):
if api_name not in self.tools:
raise ValueError(f"Tool with api_name '{api_name}' not defined. Define it using @mcp_tool.define first.")
# Store the UI builder function
self.tools[api_name]['ui_builder'] = ui_builder_func
# Return the original UI builder function
return ui_builder_func
return decorator
def get_tools_list(self) -> list[tuple[str, str]]:
"""Returns a list of (tool_name, api_name) for all defined tools."""
return [(data['name'], api_name) for api_name, data in self.tools.items()]
def get_tool_info(self, api_name: str) -> ty.Optional[dict]:
"""Returns the full info dict for a given api_name."""
return self.tools.get(api_name)
def get_tool_ui_builder(self, api_name: str) -> ty.Optional[ty.Callable]:
"""Returns the UI builder function for a given api_name."""
info = self.get_tool_info(api_name)
return info['ui_builder'] if info else None
def get_tool_function(self, api_name: str) -> ty.Optional[ty.Callable]:
"""Returns the tool function for a given api_name."""
info = self.get_tool_info(api_name)
return info['tool_func'] if info else None
# Instantiate the registry. This instance will be used by decorators
mcp_tool = MCPToolRegistry()
# Note: Gradio's mcp_server=True inspects functions decorated with the
# __mcp_tool__ attribute. Our @mcp_tool.define decorator handles this.
# The registry (`mcp_tool` instance) is primarily for the Gradio UI
# part to list tools, get docstrings, and build dynamic interfaces. |