2. 为Harness添加工具注册表
任务目标
为AI Agent Harness实现一个线程安全的工具注册表,支持运行时动态注册、注销工具,以及按名称查询工具元数据。工具元数据包含工具名称、描述、参数Schema和可调用函数。
核心实现
# tool_registry.py
import threading
from typing import Any, Callable, Dict, Optional
class ToolMetadata:
"""工具元数据定义"""
def __init__(self, name: str, description: str, parameters: Dict[str, Any], func: Callable):
self.name = name
self.description = description
self.parameters = parameters # JSON Schema格式的参数定义
self.func = func # 实际可调用函数
class ToolRegistry:
"""线程安全的工具注册表"""
def __init__(self):
self._tools: Dict[str, ToolMetadata] = {}
self._lock = threading.RLock()
def register(self, metadata: ToolMetadata) -> bool:
"""注册工具,若名称已存在则返回False"""
with self._lock:
if metadata.name in self._tools:
return False
self._tools[metadata.name] = metadata
return True
def unregister(self, name: str) -> bool:
"""注销工具,工具不存在时返回False"""
with self._lock:
return self._tools.pop(name, None) is not None
def get(self, name: str) -> Optional[ToolMetadata]:
"""按名称查询工具元数据"""
with self._lock:
return self._tools.get(name)
def list_tools(self) -> list:
"""返回所有已注册工具的名称列表"""
with self._lock:
return sorted(self._tools.keys())
def execute(self, name: str, **kwargs) -> Any:
"""调用已注册工具的函数"""
with self._lock:
metadata = self._tools.get(name)
if metadata is None:
raise KeyError(f"工具 '{name}' 未注册")
return metadata.func(**kwargs)
# 使用示例
if __name__ == "__main__":
registry = ToolRegistry()
# 定义两个示例工具
def add(a: float, b: float) -> float:
return a + b
def get_weather(city: str) -> str:
return f"{city}天气:晴,25°C"
# 注册工具
registry.register(ToolMetadata(
name="calculator.add",
description="计算两个数字之和",
parameters={"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}},
func=add
))
registry.register(ToolMetadata(
name="weather.get",
description="查询指定城市天气",
parameters={"type": "object", "properties": {"city": {"type": "string"}}},
func=get_weather
))
# 查询与调用
meta = registry.get("calculator.add")
print(f"工具描述: {meta.description}")
print(f"注册表工具列表: {registry.list_tools()}")
print(f"调用结果: {registry.execute('calculator.add', a=3, b=5)}")
# 注销工具
registry.unregister("weather.get")
print(f"注销后工具列表: {registry.list_tools()}")
运行与验证
- 将上述代码保存为
tool_registry.py,直接运行:python tool_registry.py - 预期输出(与代码逻辑严格一致):
工具描述: 计算两个数字之和
注册表工具列表: ['calculator.add', 'weather.get']
调用结果: 8
注销后工具列表: ['calculator.add']
- 验证通过条件:程序无异常退出,输出与上述预期一致,且
registry.get('weather.get')在注销后返回None。
并发安全验证(可选)
注册表使用RLock保证多线程环境下注册、注销和查询操作的原子性。可用以下方式快速验证并发注册不丢失数据:
import threading
def concurrent_register(registry, prefix, count):
for i in range(count):
registry.register(ToolMetadata(
name=f"{prefix}.{i}",
description="并发测试工具",
parameters={},
func=lambda: None
))
threads = [threading.Thread(target=concurrent_register, args=(registry, f"t{j}", 100)) for j in range(5)]
for t in threads: t.start()
for t in threads: t.join()
print(f"并发注册后工具总数: {len(registry.list_tools())}") # 预期输出: 500
该注册表可直接集成到Harness的Agent工具调用链路中:Agent根据用户意图从注册表查询工具元数据,通过execute方法完成调用。注册表本身不依赖任何第三方库,仅使用Python标准库实现。