Coverage for tools/registry.py: 87%

1 import json

2 from typing import Type

3 from pydantic import BaseModel

4 from qrclaw.logger import get_logger

5

6 logger = get_logger("qrclaw.tools.registry")

7

8# 存所有注册的工具

9# 结构：{ "工具名": {"fn": 函数本身, "model": Pydantic模型, "schema": 给LLM看的描述} }

10 _tools: dict = {}

11

12

13 def register(description: str, args_model: Type[BaseModel], confirm: bool = False):

14 """

15 装饰器：把函数注册成一个工具。

16

17 confirm=True 表示执行前需要用户确认（高风险工具）

18 """

19 def decorator(fn):

20 schema = _build_schema(fn.__name__, description, args_model)

21 _tools[fn.__name__] = {

22 "fn": fn,

23 "model": args_model,

24 "schema": schema,

25 "confirm": confirm,

26 }

27 logger.debug(f"注册工具: {fn.__name__} (需要确认: {confirm})")

28 return fn

29 return decorator

30

31

32 def need_confirm(name: str) -> bool:

33 """判断工具是否需要用户确认"""

34 return _tools.get(name, {}).get("confirm", False)

35

36

37 def _resolve_refs(schema: dict, defs: dict) -> dict:

38 """递归展开 $ref 并清理 title，Gemini 不支持 $ref"""

39 if "$ref" in schema:

40 ref_name = schema["$ref"].split("/")[-1]

41 resolved = _resolve_refs(defs.get(ref_name, {}), defs)

42 other = {k: v for k, v in schema.items() if k != "$ref"}

43 return {**resolved, **other}

44 result = {}

45 for k, v in schema.items():

46 if k in ("$defs", "title"):

47 continue # 去掉 $defs 和所有层级的 title

48 elif isinstance(v, dict):

49 result[k] = _resolve_refs(v, defs)

50 elif isinstance(v, list):

51 result[k] = [_resolve_refs(i, defs) if isinstance(i, dict) else i for i in v]

52 else:

53 result[k] = v

54 return result

55

56

57 def _build_schema(name: str, description: str, args_model: Type[BaseModel]) -> dict:

58 """从 Pydantic 模型生成 OpenAI Tool Schema，兼容 Gemini"""

59 pydantic_schema = args_model.model_json_schema()

60 defs = pydantic_schema.get("$defs", {})

61

62 # 展开 $ref，内联所有引用，同时递归清理 title

63 resolved = _resolve_refs(pydantic_schema, defs)

64

65 properties = dict(resolved.get("properties", {}))

66

67 # 无参数工具：省略 parameters 字段（Gemini 不接受空 properties）

68 if not properties:

69 return {

70 "type": "function",

71 "function": {

72 "name": name,

73 "description": description,

74 },

75 }

76

77 required = pydantic_schema.get("required", [])

78 parameters: dict = {

79 "type": "object",

80 "properties": properties,

81 }

82 # required 为空时省略，避免某些 API 报错

83 if required:

84 parameters["required"] = required

85

86 return {

87 "type": "function",

88 "function": {

89 "name": name,

90 "description": description,

91 "parameters": parameters,

92 },

93 }

94

95

96 def get_schemas() -> list[dict]:

97 """返回所有工具的 schema 列表，发给 LLM 用"""

98 schemas = [item["schema"] for item in _tools.values()]

99 logger.debug(f"获取工具 schemas，共 {len(schemas)} 个工具")

100 return schemas

101

102

103 def execute(name: str, arguments: str) -> str:

104 """

105 执行工具。

106 用 Pydantic 模型校验参数，不合法直接报错，不会传脏数据给工具函数。

107 同时执行安全切面检查。

108

109 Raises:

110 PermissionError: 当权限检查失败时抛出，由上层 agent.py 捕获处理

111 """

112 if name not in _tools:

113 error_msg = f"错误：找不到工具 {name}"

114 logger.error(error_msg)

115 return error_msg

116

117 try:

118 raw_args = json.loads(arguments)

119 logger.debug(f"工具 {name} 原始参数: {raw_args}")

120

121 # 用 Pydantic 校验并解析参数

122 validated = _tools[name]["model"](**raw_args)

123 validated_args = validated.model_dump()

124 logger.debug(f"工具 {name} 校验后参数: {validated_args}")

125

126 # === AOP 安全拦截 (Security Hook) ===

127 try:

128 # 局部导入避免循环引用

129 from qrclaw.agent import get_workspace

130 from qrclaw.security import security_manager

131

132 # 获取当前上下文的工作空间

133 ws = get_workspace()

134

135 # 仅当在 agent 运行上下文中时才检查

136 # 如果是 CLI 直接调试工具或单元测试，可能没有 workspace，此时视为 Full Access

137 if ws:

138 security_manager.check_access(

139 agent_id=ws.agent_id,

140 tool_name=name,

141 args=validated_args,

142 workspace_root=ws.root

143 )

144 except PermissionError:

145 # 权限拒绝：透传给上层 agent.py 处理

146 logger.warning(f"工具 {name} 权限检查失败")

147 raise # 关键：重新抛出异常，而不是返回字符串

148 except ImportError:

149 # 可能是环境问题，忽略

150 pass

151 except Exception as e:

152 # 安全检查本身出错，为了安全起见，选择拦截并报错 (Fail Closed)

153 error_msg = f"系统错误：执行安全检查时发生异常 ({e})"

154 logger.error(error_msg, exc_info=True)

155 return error_msg

156 # ====================================

157

158 # 执行工具

159 result = _tools[name]["fn"](**validated_args)

160 logger.info(f"工具 {name} 执行成功")

161

162 return result

163 except PermissionError:

164 # 权限拒绝：透传给上层 agent.py 处理

165 raise

166 except json.JSONDecodeError as e:

167 error_msg = f"错误：参数不是合法的 JSON: {e}"

168 logger.error(error_msg)

169 return error_msg

170 except Exception as e:

171 error_msg = f"错误：参数校验失败 {e}"

172 logger.error(error_msg, exc_info=True)

173 return error_msg