Spaces:
Sleeping
Sleeping
File size: 4,217 Bytes
f2b5c2a | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """
MCP (Model Context Protocol) Server for AI Chatbot.
This module initializes and manages MCP tools that the AI can use
to interact with the task management system.
"""
import logging
from typing import Dict, Any, Callable, List
from pydantic import BaseModel
logger = logging.getLogger(__name__)
class MCPTool(BaseModel):
"""Base model for MCP tool definition."""
name: str
description: str
parameters: Dict[str, Any]
function: Any # Will be the actual callable
class MCPServer:
"""
MCP Server managing tool registration and execution.
This server acts as a registry for tools that the AI can invoke.
All tools must be registered before they can be used.
"""
def __init__(self, name: str = "todo-tools", version: str = "1.0.0"):
self.name = name
self.version = version
self.tools: Dict[str, MCPTool] = {}
logger.info(f"MCP Server initialized: {name} v{version}")
def register_tool(
self,
name: str,
description: str,
parameters: Dict[str, Any],
function: Callable
):
"""
Register a new tool with the MCP server.
Args:
name: Tool name (must be unique)
description: Human-readable description of what the tool does
parameters: JSON schema describing the tool's parameters
function: The actual function to execute when tool is called
"""
if name in self.tools:
logger.warning(f"Tool '{name}' already registered, overwriting")
tool = MCPTool(
name=name,
description=description,
parameters=parameters,
function=function
)
self.tools[name] = tool
logger.info(f"Registered tool: {name}")
def get_tool(self, name: str) -> MCPTool:
"""Get a registered tool by name."""
if name not in self.tools:
raise ValueError(f"Tool '{name}' not found in MCP server")
return self.tools[name]
def list_tools(self) -> List[Dict[str, Any]]:
"""
List all registered tools in Cohere-compatible format.
Returns:
List of tool definitions for Cohere API v2
"""
tools_list = []
for tool in self.tools.values():
# Convert parameter definitions to JSON Schema format
properties = {}
required = []
for param_name, param_def in tool.parameters.items():
properties[param_name] = {
"type": param_def.get("type", "string"),
"description": param_def.get("description", "")
}
if param_def.get("required", False):
required.append(param_name)
tools_list.append({
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": {
"type": "object",
"properties": properties,
"required": required
}
}
})
return tools_list
async def execute_tool(
self,
tool_name: str,
parameters: Dict[str, Any],
db: Any = None
) -> Dict[str, Any]:
"""
Execute a registered tool with given parameters.
Args:
tool_name: Name of the tool to execute
parameters: Parameters to pass to the tool
db: Database session to pass to the tool
Returns:
Tool execution result as dictionary
"""
tool = self.get_tool(tool_name)
try:
logger.info(f"Executing tool: {tool_name} with params: {parameters}")
# Pass database session to tool
result = await tool.function(**parameters, db=db)
logger.info(f"Tool {tool_name} executed successfully")
return result
except Exception as e:
logger.error(f"Tool {tool_name} execution failed: {str(e)}")
raise
# Global MCP server instance
mcp_server = MCPServer()
|