DESKTOP-72TV0V4\caoxiaozhu
b2bc9988a9
- 新增 Go 语言后端服务(server/)，包含用户认证、Agent管理、数据库连接等API - 新增 Python Agent 服务(agent/)，实现Agent核心逻辑和工具集 - 前端从原生HTML迁移到Vue.js框架(web/src/) - 添加 Docker Compose 支持(docker-compose.yml) - 添加项目架构文档(docs/ARCHITECTURE.md) - 添加环境变量示例(.env.example)和本地启动脚本(start-local.ps1) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
"""
|
||
工具注册表 - 管理所有可用工具（白名单机制）
|
||
"""
|
||
from typing import Any, Callable, Optional
|
||
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 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
|