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()}")

运行与验证

  1. 将上述代码保存为 tool_registry.py,直接运行:python tool_registry.py
  2. 预期输出(与代码逻辑严格一致):
工具描述: 计算两个数字之和
注册表工具列表: ['calculator.add', 'weather.get']
调用结果: 8
注销后工具列表: ['calculator.add']
  1. 验证通过条件:程序无异常退出,输出与上述预期一致,且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标准库实现。