144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
"""Tool system core module"""
|
|
from dataclasses import dataclass, field
|
|
from typing import Callable, Any, Dict, List, Optional
|
|
from enum import IntEnum
|
|
|
|
|
|
class CommandPermission(IntEnum):
|
|
"""Command permission level - higher value means more permissions"""
|
|
READ_ONLY = 1
|
|
WRITE = 2
|
|
EXECUTE = 3
|
|
ADMIN = 4
|
|
|
|
|
|
@dataclass
|
|
class ToolContext:
|
|
"""Tool execution context - explicitly passed to tool execution"""
|
|
workspace: Optional[str] = None
|
|
user_id: Optional[int] = None
|
|
username: Optional[str] = None
|
|
extra: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class ToolDefinition:
|
|
"""Tool definition"""
|
|
name: str
|
|
description: str
|
|
parameters: Dict[str, Any]
|
|
handler: Callable
|
|
category: str = "general"
|
|
required_permission: CommandPermission = CommandPermission.READ_ONLY
|
|
|
|
def to_openai_format(self) -> Dict[str, Any]:
|
|
"""Convert to OpenAI format"""
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"parameters": self.parameters
|
|
}
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ToolResult:
|
|
"""Tool execution result"""
|
|
success: bool
|
|
data: Any = None
|
|
error: Optional[str] = None
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert to dictionary"""
|
|
return {"success": self.success, "data": self.data, "error": self.error}
|
|
|
|
@classmethod
|
|
def ok(cls, data: Any) -> "ToolResult":
|
|
"""Create success result"""
|
|
return cls(success=True, data=data)
|
|
|
|
@classmethod
|
|
def fail(cls, error: str) -> "ToolResult":
|
|
"""Create failure result"""
|
|
return cls(success=False, error=error)
|
|
|
|
|
|
class ToolRegistry:
|
|
"""Tool registry (singleton pattern)"""
|
|
_instance: Optional["ToolRegistry"] = None
|
|
_tools: Dict[str, ToolDefinition] = {}
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
cls._instance._tools = {}
|
|
return cls._instance
|
|
|
|
def register(self, tool: ToolDefinition) -> None:
|
|
"""Register tool"""
|
|
self._tools[tool.name] = tool
|
|
|
|
def get(self, name: str) -> Optional[ToolDefinition]:
|
|
"""Get tool definition"""
|
|
return self._tools.get(name)
|
|
|
|
def list_all(self) -> List[Dict[str, Any]]:
|
|
"""List all tools"""
|
|
return [t.to_openai_format() for t in self._tools.values()]
|
|
|
|
def list_by_category(self, category: str) -> List[Dict[str, Any]]:
|
|
"""List tools by category"""
|
|
return [
|
|
t.to_openai_format()
|
|
for t in self._tools.values()
|
|
if t.category == category
|
|
]
|
|
|
|
def execute(self, name: str, arguments: dict, context: ToolContext = None) -> Dict[str, Any]:
|
|
"""Execute tool with optional context and automatic permission check"""
|
|
tool = self.get(name)
|
|
if not tool:
|
|
return {"success": False, "error": f"Tool '{name}' not found"}
|
|
|
|
# Automatic permission check (transparent to tool function)
|
|
if context is not None and context.user_id is not None:
|
|
user_level = context.extra.get("user_permission_level", CommandPermission.READ_ONLY)
|
|
if user_level < tool.required_permission:
|
|
return {
|
|
"success": False,
|
|
"error": f"Permission denied: requires {tool.required_permission.name}, you have {user_level.name}"
|
|
}
|
|
|
|
try:
|
|
# Pass context to handler if it accepts it
|
|
if context is not None:
|
|
result = tool.handler(arguments, context=context)
|
|
else:
|
|
result = tool.handler(arguments)
|
|
if isinstance(result, ToolResult):
|
|
return result.to_dict()
|
|
return result
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e)}
|
|
|
|
def clear(self) -> None:
|
|
"""Clear all tools"""
|
|
self._tools.clear()
|
|
|
|
def remove(self, name: str) -> bool:
|
|
"""Remove tool"""
|
|
if name in self._tools:
|
|
del self._tools[name]
|
|
return True
|
|
return False
|
|
|
|
def tool_count(self) -> int:
|
|
"""Tool count"""
|
|
return len(self._tools)
|
|
|
|
|
|
# Global registry instance
|
|
registry = ToolRegistry()
|