158 lines
5.3 KiB
Python
158 lines
5.3 KiB
Python
"""
|
|
Tool registry for storing and discovering MCP tools.
|
|
|
|
Manages tool descriptions, provides JSONPath-based discovery,
|
|
and supports sparse mode for context optimization.
|
|
"""
|
|
|
|
import json
|
|
from typing import Dict, Any, List, Optional, Union
|
|
from jsonpath_ng import parse as jsonpath_parse
|
|
from jsonpath_ng.exceptions import JsonPathParserError
|
|
|
|
|
|
class ToolRegistry:
|
|
"""Registry for MCP tools with discovery capabilities."""
|
|
|
|
def __init__(self):
|
|
self.tools: Dict[str, Any] = {}
|
|
self.raw_tool_list: List[Dict[str, Any]] = []
|
|
self._metadata: Dict[str, Any] = {}
|
|
|
|
def update_tools(self, tools: List[Dict[str, Any]]):
|
|
"""
|
|
Update the registry with a list of tools.
|
|
|
|
Args:
|
|
tools: List of tool definitions from MCP server
|
|
"""
|
|
self.raw_tool_list = tools
|
|
self.tools.clear()
|
|
|
|
for tool in tools:
|
|
if "name" in tool:
|
|
self.tools[tool["name"]] = tool
|
|
|
|
def get_tool(self, name: str) -> Optional[Dict[str, Any]]:
|
|
"""Get a tool definition by name."""
|
|
return self.tools.get(name)
|
|
|
|
def get_all_tool_names(self) -> List[str]:
|
|
"""Get all registered tool names."""
|
|
return list(self.tools.keys())
|
|
|
|
def discover(self, jsonpath: str) -> Union[List[Any], Any, None]:
|
|
"""
|
|
Discover tools using JSONPath queries.
|
|
|
|
Args:
|
|
jsonpath: JSONPath expression to query tools
|
|
|
|
Returns:
|
|
Query results or None if no matches
|
|
|
|
Examples:
|
|
$.tools[*].name - Get all tool names
|
|
$.tools[?(@.name=='Bash')] - Get Bash tool details
|
|
$.tools[*].inputSchema - Get all input schemas
|
|
"""
|
|
try:
|
|
expr = jsonpath_parse(jsonpath)
|
|
except (JsonPathParserError, Exception):
|
|
return None
|
|
|
|
# Create a searchable structure
|
|
search_data = {
|
|
"tools": self.raw_tool_list,
|
|
"tool_names": self.get_all_tool_names(),
|
|
"metadata": self._metadata
|
|
}
|
|
|
|
# Execute JSONPath query
|
|
matches = expr.find(search_data)
|
|
|
|
if not matches:
|
|
return None
|
|
elif len(matches) == 1:
|
|
return matches[0].value
|
|
else:
|
|
return [match.value for match in matches]
|
|
|
|
def get_sparse_tools(self) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get minimal tool list for sparse mode.
|
|
|
|
Returns only essential meta-tools for discovery.
|
|
"""
|
|
tool_count = len(self.tools)
|
|
|
|
sparse_tools = [
|
|
{
|
|
"name": "mcp_discover",
|
|
"description": f"Discover available tools using JSONPath. {tool_count} tools available.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"jsonpath": {
|
|
"type": "string",
|
|
"description": "JSONPath expression (e.g., '$.tools[*].name')"
|
|
}
|
|
},
|
|
"required": ["jsonpath"]
|
|
}
|
|
},
|
|
{
|
|
"name": "mcp_call",
|
|
"description": "Execute any MCP tool by constructing a JSON-RPC call.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"method": {
|
|
"type": "string",
|
|
"description": "JSON-RPC method (e.g., 'tools/call')"
|
|
},
|
|
"params": {
|
|
"type": "object",
|
|
"description": "Method parameters"
|
|
}
|
|
},
|
|
"required": ["method", "params"]
|
|
}
|
|
},
|
|
{
|
|
"name": "onboarding",
|
|
"description": "Get or set identity-specific onboarding instructions for AI contexts.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"identity": {
|
|
"type": "string",
|
|
"description": "Identity for onboarding (e.g., 'Claude', project name)"
|
|
},
|
|
"instructions": {
|
|
"type": "string",
|
|
"description": "Optional: Set new instructions. If omitted, retrieves existing."
|
|
},
|
|
"append": {
|
|
"type": "boolean",
|
|
"description": "Append to existing instructions instead of replacing",
|
|
"default": False
|
|
}
|
|
},
|
|
"required": ["identity"]
|
|
}
|
|
}
|
|
]
|
|
|
|
return sparse_tools
|
|
|
|
def set_metadata(self, key: str, value: Any):
|
|
"""Set metadata that can be discovered via JSONPath."""
|
|
self._metadata[key] = value
|
|
|
|
def to_json(self) -> str:
|
|
"""Export registry as JSON for debugging."""
|
|
return json.dumps({
|
|
"tools": self.raw_tool_list,
|
|
"metadata": self._metadata
|
|
}, indent=2) |