2026-04-16 17:14:06 +08:00

62 lines
2.0 KiB
Python

"""
Tool registry implementation with permission management and metadata.
"""
import asyncio
import time
from typing import Dict, Any, Callable, Optional, List
from dataclasses import dataclass, asdict
import logging
logger = logging.getLogger(__name__)
@dataclass
class ToolMetadata:
"""Metadata for a tool including description, parameters, and permissions."""
name: str
description: str
parameters: Dict[str, Dict[str, Any]]
permissions: List[str]
examples: List[Dict[str, Any]]
security_notes: str = ""
timeout: int = 300
max_retries: int = 3
class ToolRegistry:
"""Registry for managing available tools with metadata and permissions."""
def __init__(self):
self.tools: Dict[str, Dict[str, Any]] = {}
self.logger = logger
def register_tool(self, name: str, func: Callable, metadata: ToolMetadata):
"""Register a tool with its metadata."""
self.tools[name] = {
'function': func,
'metadata': asdict(metadata)
}
self.logger.info(f"Registered tool: {name}")
def get_tool(self, name: str) -> Optional[Dict[str, Any]]:
"""Get tool information by name."""
return self.tools.get(name)
def list_tools(self) -> List[str]:
"""List all registered tool names."""
return list(self.tools.keys())
def get_tool_metadata(self, name: str) -> Optional[Dict[str, Any]]:
"""Get tool metadata by name."""
tool = self.get_tool(name)
return tool['metadata'] if tool else None
def has_permission(self, tool_name: str, user_permissions: List[str]) -> bool:
"""Check if user has required permissions for a tool."""
metadata = self.get_tool_metadata(tool_name)
if not metadata:
return False
required_perms = metadata.get('permissions', [])
if not required_perms:
return True
return any(perm in user_permissions for perm in required_perms)