108 lines
3.1 KiB
Python
108 lines
3.1 KiB
Python
"""
|
||
工具注册表 - 管理所有可用工具（白名单机制）
|
||
"""
|
||
from typing import Any, Callable, Optional, Dict
|
||
from dataclasses import dataclass, asdict
|
||
from enum import Enum
|
||
|
||
|
||
class SecurityLevel(Enum):
|
||
"""工具安全等级"""
|
||
SAFE = "safe" # 安全操作
|
||
REVIEW = "review" # 需要审核
|
||
DANGER = "danger" # 危险操作
|
||
|
||
|
||
@dataclass
|
||
class ToolMetadata:
|
||
"""工具元数据"""
|
||
name: str
|
||
description: str
|
||
security_level: str
|
||
require_approval: bool = False
|
||
allowed_roles: list = None
|
||
|
||
def dict(self):
|
||
return {
|
||
"name": self.name,
|
||
"description": self.description,
|
||
"security_level": self.security_level,
|
||
"require_approval": self.require_approval
|
||
}
|
||
|
||
|
||
class ToolRegistry:
|
||
"""工具注册表"""
|
||
|
||
def __init__(self):
|
||
self._tools: Dict[str, tuple[Callable, ToolMetadata]] = {}
|
||
self._definitions: Dict[str, dict] = {}
|
||
|
||
def register(
|
||
self,
|
||
name: str,
|
||
func: Callable,
|
||
description: str = "",
|
||
security_level: str = "safe",
|
||
require_approval: bool = False,
|
||
allowed_roles: list = None,
|
||
parameters: dict = None
|
||
):
|
||
"""注册工具到白名单"""
|
||
metadata = ToolMetadata(
|
||
name=name,
|
||
description=description,
|
||
security_level=security_level,
|
||
require_approval=require_approval,
|
||
allowed_roles=allowed_roles or ["user", "admin"]
|
||
)
|
||
|
||
self._tools[name] = (func, metadata)
|
||
|
||
# 生成工具定义（用于 LLM 调用）
|
||
self._definitions[name] = {
|
||
"name": name,
|
||
"description": description,
|
||
"parameters": parameters or {
|
||
"type": "object",
|
||
"properties": {},
|
||
"required": []
|
||
}
|
||
}
|
||
|
||
def get_tool(self, name: str) -> tuple[Callable, ToolMetadata]:
|
||
"""获取工具函数和元数据"""
|
||
if name not in self._tools:
|
||
raise ValueError(f"Tool '{name}' not found in whitelist")
|
||
return self._tools[name]
|
||
|
||
def get_tool_definition(self, name: str) -> Optional[dict]:
|
||
"""获取工具定义（用于 LLM）"""
|
||
return self._definitions.get(name)
|
||
|
||
def list_tools(self) -> list[ToolMetadata]:
|
||
"""列出所有已注册工具"""
|
||
return [meta for _, meta in self._tools.values()]
|
||
|
||
def list_definitions(self) -> list[dict]:
|
||
"""列出所有工具定义（用于LLM）"""
|
||
return list(self._definitions.values())
|
||
|
||
def check_permission(self, tool_name: str, user_role: str) -> bool:
|
||
"""检查用户权限"""
|
||
if tool_name not in self._tools:
|
||
return False
|
||
_, metadata = self._tools[tool_name]
|
||
return user_role in metadata.allowed_roles
|
||
|
||
def need_approval(self, tool_name: str) -> bool:
|
||
"""判断是否需要审批"""
|
||
if tool_name not in self._tools:
|
||
return False
|
||
_, metadata = self._tools[tool_name]
|
||
return metadata.require_approval
|
||
|
||
|
||
# 全局工具注册表
|
||
global_registry = ToolRegistry()
|