refactor(aibox): 工具基类 Tool/ToolRegistry 抽离至 tools 包
延续 common 包的拆分模式,形成 common <- tools <- aibox 单向依赖, 为后续工具子类逐步迁移铺路;aibox.py 瘦身约 230 行,工具子类与注册逻辑零改动。
This commit is contained in:
1 parent
91b30040c7
commit
14db64f7a2
5 files changed
+285
-233
No files matched your search
@@ -0,0 +1,141 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.registry - 工具注册器
|
||||
|
||||
抽离自 aibox.py 的 ToolRegistry 类。仅依赖 common 与标准库,
|
||||
不反向依赖 aibox 任何符号(skill_manager 为前向引用,运行期注入)。
|
||||
"""
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import Logger, i18n
|
||||
|
||||
from .base import Tool
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
"""工具注册器 - 支持技能热重载"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self._tools: Dict[str, Tool] = {}
|
||||
self._builtin_tools: Dict[str, Tool] = {} # 内置工具单独存储
|
||||
self.logger = logger or Logger("ToolRegistry")
|
||||
self.current_output_queue = None # 用于流式输出的队列
|
||||
self.skill_manager: Optional['SkillManager'] = None # 技能管理器引用
|
||||
|
||||
def register(self, tool: Tool, is_builtin: bool = False):
|
||||
"""注册工具,可以标记是否为内置工具"""
|
||||
self._tools[tool.get_name()] = tool
|
||||
if is_builtin:
|
||||
self._builtin_tools[tool.get_name()] = tool
|
||||
self.logger.info(f"工具已注册: {tool.get_name()} {'(内置)' if is_builtin else ''}")
|
||||
|
||||
def register_many(self, tools: List[Tool], is_builtin: bool = False):
|
||||
"""批量注册工具"""
|
||||
for tool in tools:
|
||||
self.register(tool, is_builtin)
|
||||
|
||||
def get_tool(self, name: str) -> Optional[Tool]:
|
||||
"""获取工具"""
|
||||
return self._tools.get(name)
|
||||
|
||||
def get_all_tools(self) -> List[Tool]:
|
||||
"""获取所有工具"""
|
||||
return list(self._tools.values())
|
||||
|
||||
def get_tools_schemas(self) -> List[Dict]:
|
||||
"""获取所有工具的OpenAI函数调用格式"""
|
||||
return [tool.to_function_call_format() for tool in self._tools.values()]
|
||||
|
||||
def set_output_queue(self, queue: queue.Queue):
|
||||
"""设置输出队列"""
|
||||
self.current_output_queue = queue
|
||||
# 同时设置所有工具的队列
|
||||
for tool in self._tools.values():
|
||||
tool.output_queue = queue
|
||||
|
||||
def execute_tool(self, name: str, **kwargs) -> str:
|
||||
"""执行工具"""
|
||||
tool = self.get_tool(name)
|
||||
if not tool:
|
||||
return f"❌ 错误:找不到工具 '{name}'"
|
||||
|
||||
self.logger.debug(f"工具调用: {name}({kwargs})")
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 检查是否有队列可用,并且工具支持流式输出
|
||||
if self.current_output_queue and hasattr(tool, 'execute_streaming'):
|
||||
# 设置工具的流式输出属性
|
||||
tool.streaming_output = True
|
||||
tool.output_queue = self.current_output_queue
|
||||
|
||||
# 创建队列来收集流式输出
|
||||
output_queue = queue.Queue()
|
||||
|
||||
# 执行流式输出
|
||||
thread = threading.Thread(
|
||||
target=tool.execute_streaming,
|
||||
args=(output_queue,),
|
||||
kwargs=kwargs
|
||||
)
|
||||
thread.daemon = True
|
||||
thread.start()
|
||||
|
||||
# 从队列收集输出并发送到当前队列
|
||||
final_result = None
|
||||
while thread.is_alive() or not output_queue.empty():
|
||||
try:
|
||||
msg_type, content = output_queue.get(timeout=0.1)
|
||||
if msg_type == "line":
|
||||
# 发送到当前队列(会通过WebSocket发送)
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(("line", content))
|
||||
elif msg_type == "complete":
|
||||
final_result = content
|
||||
# 也发送完整结果
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(("complete", content))
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
thread.join(timeout=5)
|
||||
if thread.is_alive():
|
||||
return "❌ 工具执行超时"
|
||||
|
||||
result = final_result or ""
|
||||
tool.streaming_output = False
|
||||
tool.output_queue = None
|
||||
else:
|
||||
# 普通执行
|
||||
result = tool.execute(**kwargs)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
self.logger.debug(f"工具结果 ({elapsed:.2f}s): {result[:100]}...")
|
||||
return result
|
||||
except TypeError as e:
|
||||
return f"❌ 工具执行错误: 参数错误 - {str(e)}"
|
||||
except Exception as e:
|
||||
return f"❌ 工具执行错误: {str(e)}"
|
||||
|
||||
def reload_skills(self, skills_dir: str = None) -> int:
|
||||
"""
|
||||
重新加载所有扩展技能(热重载)。
|
||||
委托给 SkillManager 重新扫描目录并执行 --help 收集简介。
|
||||
"""
|
||||
self.logger.info("开始重新加载扩展技能...")
|
||||
if self.skill_manager is not None:
|
||||
count = self.skill_manager.reload()
|
||||
self.logger.info(i18n.get('skills_reloaded', count))
|
||||
return count
|
||||
self.logger.warning("skill_manager 未初始化,无法重载技能")
|
||||
return 0
|
||||
|
||||
def get_skills_count(self) -> Tuple[int, int]:
|
||||
"""获取内置工具和扩展技能的数量"""
|
||||
builtin_count = len(self._builtin_tools)
|
||||
skill_count = len(self._tools) - builtin_count
|
||||
return builtin_count, skill_count
|
||||
Reference in new issue
Block a user