utils/mcp_decorator.py

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.