diff --git a/router_system/tools.py b/router_system/tools.py new file mode 100644 index 0000000..3f72b41 --- /dev/null +++ b/router_system/tools.py @@ -0,0 +1,324 @@ +"""工具调用内核 —— 让 LLM 以 OpenAI function-calling 协议操作工作区文件。 + +组成(对齐《实现方案_v4_模型池与工具智能体.md》D4): +- TOOLS_SPEC:list_dir / read_file / write_file 三个工具的 OpenAI tools 声明 +- WorkspaceTools:被"关押"在根目录内的文件工具(路径越界一律拒绝,Windows pathlib) +- parse_tool_calls:解析 OpenAI 响应里的 tool_calls(arguments 容错为 {}) +- ToolLoop:通用智能体循环。chat_fn 注入(网关传 OpenAI 兼容客户端,测试传假实现), + 本模块只负责循环编排:调用 -> 执行工具 -> 回喂结果 -> 直到模型给出最终答复。 + +工程约束: +- 纯标准库(router_system 零第三方依赖不变) +- 工具结果回喂前截断(防止上下文爆炸),轮数与 token 双上限(金额护栏) +- 事件回调 on_event 逐条产出过程事件(供 SSE 透出"智能体在做什么") +""" +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Awaitable, Callable, Dict, List, Optional + +# 回喂给模型的工具结果/读取内容上限(字符) +MAX_READ_CHARS = 8000 +MAX_LIST_ENTRIES = 200 +MAX_RESULT_CHARS = 8000 +MAX_WRITE_CHARS = 200_000 + +# 默认循环上限 +DEFAULT_MAX_ROUNDS = 8 + +# OpenAI tools 声明(chat/completions 请求的 tools 参数) +TOOLS_SPEC: List[Dict[str, Any]] = [ + { + "type": "function", + "function": { + "name": "list_dir", + "description": "列出工作区内目录的内容(文件与子目录,含大小)。", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string", "description": "工作区内的相对路径,默认根目录"} + }, + }, + }, + }, + { + "type": "function", + "function": { + "name": "read_file", + "description": "读取工作区内一个文本文件的内容(过长自动截断)。", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string", "description": "工作区内的相对路径"} + }, + "required": ["path"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "write_file", + "description": "把文本内容写入(或创建/覆盖)工作区内的一个文件。", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string", "description": "工作区内的相对路径"}, + "content": {"type": "string", "description": "要写入的全文"}, + }, + "required": ["path", "content"], + }, + }, + }, +] + +TOOL_NAMES = {t["function"]["name"] for t in TOOLS_SPEC} + + +class ToolError(Exception): + """工具执行失败(路径越界/不存在/参数非法)。""" + + +class WorkspaceTools: + """被限制在根目录内的文件工具(智能体的"手")。 + + 安全:所有路径先 join 再 resolve,解析结果必须仍位于根目录内 + (根目录自身允许),否则抛 ToolError——防 ../ 越界与绝对路径逃逸。 + """ + + def __init__(self, root: str | Path): + self.root = Path(root).resolve() + self.root.mkdir(parents=True, exist_ok=True) + + # ---------- 路径关押 ---------- + def resolve(self, rel_path: str) -> Path: + rel = (rel_path or "").strip().replace("\\", "/").lstrip("/") + p = (self.root / rel).resolve() + if p != self.root and self.root not in p.parents: + raise ToolError(f"路径越界(不允许访问工作区之外): {rel_path}") + return p + + # ---------- 三个工具 ---------- + def list_dir(self, rel_path: str = "") -> Dict[str, Any]: + d = self.resolve(rel_path) + if not d.exists(): + return {"ok": False, "error": f"目录不存在: {rel_path}"} + if not d.is_dir(): + return {"ok": False, "error": f"不是目录: {rel_path}"} + entries = [] + for child in sorted(d.iterdir(), key=lambda c: (c.is_file(), c.name.lower())): + if child.is_dir(): + entries.append({"name": child.name + "/", "type": "dir"}) + else: + entries.append({ + "name": child.name, "type": "file", + "size": child.stat().st_size, + }) + if len(entries) >= MAX_LIST_ENTRIES: + entries.append({"name": f"…(超过 {MAX_LIST_ENTRIES} 项已截断)", "type": "notice"}) + break + return {"ok": True, "path": rel_path or ".", "entries": entries} + + def read_file(self, rel_path: str) -> Dict[str, Any]: + p = self.resolve(rel_path) + if not p.exists(): + return {"ok": False, "error": f"文件不存在: {rel_path}"} + if not p.is_file(): + return {"ok": False, "error": f"不是文件: {rel_path}"} + try: + text = p.read_text(encoding="utf-8") + except UnicodeDecodeError: + return {"ok": False, "error": f"非文本文件(UTF-8 解码失败): {rel_path}"} + truncated = len(text) > MAX_READ_CHARS + return { + "ok": True, + "path": rel_path, + "content": text[:MAX_READ_CHARS], + "truncated": truncated, + "total_chars": len(text), + } + + def write_file(self, rel_path: str, content: str) -> Dict[str, Any]: + if len(content) > MAX_WRITE_CHARS: + return {"ok": False, "error": f"内容过长(>{MAX_WRITE_CHARS} 字符),拒绝写入"} + p = self.resolve(rel_path) + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(content, encoding="utf-8") + return {"ok": True, "path": rel_path, "bytes_written": len(content.encode("utf-8"))} + + # ---------- 统一执行入口 ---------- + def execute(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]: + """按名字执行工具;任何异常折叠为 {"ok": False, "error": ...}。""" + try: + if name == "list_dir": + return self.list_dir(str(arguments.get("path", ""))) + if name == "read_file": + return self.read_file(str(arguments.get("path", ""))) + if name == "write_file": + return self.write_file( + str(arguments.get("path", "")), str(arguments.get("content", ""))) + return {"ok": False, "error": f"未知工具: {name}"} + except ToolError as e: + return {"ok": False, "error": str(e)} + except OSError as e: + return {"ok": False, "error": f"文件系统错误: {type(e).__name__}: {e}"} + + +def parse_tool_calls(message: Dict[str, Any]) -> List[Dict[str, Any]]: + """从 OpenAI 响应的 message 解析 tool_calls。 + + 返回 [{"id","name","arguments"(dict)}];arguments 非法 JSON 时容错为 {}。 + """ + out: List[Dict[str, Any]] = [] + for tc in message.get("tool_calls") or []: + fn = tc.get("function") or {} + raw = fn.get("arguments") + if isinstance(raw, dict): + args = raw + elif isinstance(raw, str) and raw.strip(): + try: + parsed = json.loads(raw) + args = parsed if isinstance(parsed, dict) else {} + except json.JSONDecodeError: + args = {} + else: + args = {} + out.append({ + "id": tc.get("id") or f"call_{len(out)}", + "name": fn.get("name") or "", + "arguments": args, + }) + return out + + +class ToolLoop: + """通用智能体工具循环(zcode 式:思考 -> 调工具 -> 看结果 -> 再思考)。 + + chat_fn(messages, tools_spec) -> {"content": str|None, + "tool_calls": [ {id,name,arguments}, ... ], + "usage": {"prompt_tokens", "completion_tokens"}} + 由网关注入真实 OpenAI 兼容客户端;测试注入脚本化假实现。 + + on_event(ev) 为可选同步回调,逐条收到过程事件: + {"type":"round","round":n} + {"type":"tool_call","round":n,"name":...,"arguments":...} + {"type":"tool_result","round":n,"name":...,"ok":...,"preview":...} + {"type":"usage","prompt_tokens":...,"completion_tokens":...} + {"type":"final","round":n,"reason":"answer"|"max_rounds"|"token_cap"|"error"} + """ + + def __init__( + self, + tools: WorkspaceTools, + chat_fn: Callable[[List[Dict[str, Any]], List[Dict[str, Any]]], Awaitable[Dict[str, Any]]], + max_rounds: int = DEFAULT_MAX_ROUNDS, + token_cap: int = 0, + on_event: Optional[Callable[[Dict[str, Any]], None]] = None, + result_preview_chars: int = MAX_RESULT_CHARS, + ): + self.tools = tools + self.chat_fn = chat_fn + self.max_rounds = max(1, int(max_rounds)) + self.token_cap = int(token_cap) # 0 = 不限 + self.on_event = on_event + self.result_preview_chars = result_preview_chars + + def _emit(self, ev: Dict[str, Any]) -> None: + if self.on_event is not None: + try: + self.on_event(ev) + except Exception: + pass # 事件回调不允许打断主循环 + + def _total_tokens(self, usage: Dict[str, int]) -> int: + return int(usage.get("prompt_tokens", 0)) + int(usage.get("completion_tokens", 0)) + + async def run(self, task: str, system: str = "") -> Dict[str, Any]: + """执行任务直到模型给出最终答复或触顶。返回最终结果与账目。""" + messages: List[Dict[str, Any]] = [] + if system: + messages.append({"role": "system", "content": system}) + messages.append({"role": "user", "content": task}) + + total_in = 0 + total_out = 0 + last_content = "" + + for round_no in range(1, self.max_rounds + 1): + self._emit({"type": "round", "round": round_no}) + try: + resp = await self.chat_fn(messages, TOOLS_SPEC) + except Exception as e: + self._emit({"type": "final", "round": round_no, "reason": "error", + "error": f"{type(e).__name__}: {e}"}) + return {"response": "", "rounds": round_no, "reason": "error", + "error": f"{type(e).__name__}: {e}", + "prompt_tokens": total_in, "completion_tokens": total_out} + + usage = resp.get("usage") or {} + total_in += int(usage.get("prompt_tokens", 0)) + total_out += int(usage.get("completion_tokens", 0)) + self._emit({"type": "usage", "prompt_tokens": total_in, + "completion_tokens": total_out}) + + # 金额护栏(D6 同源):token 触顶立即停 + if self.token_cap and total_in + total_out > self.token_cap: + self._emit({"type": "final", "round": round_no, "reason": "token_cap"}) + return {"response": last_content, "rounds": round_no, "reason": "token_cap", + "error": f"token 熔断({total_in + total_out}/{self.token_cap})", + "prompt_tokens": total_in, "completion_tokens": total_out} + + calls = resp.get("tool_calls") or [] + if not calls: + self._emit({"type": "final", "round": round_no, "reason": "answer"}) + return {"response": resp.get("content") or "", "rounds": round_no, + "reason": "answer", "error": None, + "prompt_tokens": total_in, "completion_tokens": total_out} + + # 有工具调用:回填 assistant 消息(带原始 tool_calls 结构)+ 逐个执行 + messages.append({ + "role": "assistant", + "content": resp.get("content") or None, + "tool_calls": [ + { + "id": c["id"], + "type": "function", + "function": {"name": c["name"], + "arguments": json.dumps(c["arguments"], ensure_ascii=False)}, + } + for c in calls + ], + }) + for c in calls: + self._emit({"type": "tool_call", "round": round_no, + "name": c["name"], "arguments": c["arguments"]}) + result = self.tools.execute(c["name"], c["arguments"]) + preview = json.dumps(result, ensure_ascii=False) + if len(preview) > self.result_preview_chars: + preview = preview[:self.result_preview_chars] + "…(截断)" + self._emit({"type": "tool_result", "round": round_no, + "name": c["name"], "ok": bool(result.get("ok")), + "preview": preview}) + messages.append({ + "role": "tool", + "tool_call_id": c["id"], + "content": preview, + }) + last_content = resp.get("content") or last_content + + # 轮次耗尽:不再给工具,让模型立即总结(无 tools 的最后一次调用) + self._emit({"type": "final", "round": self.max_rounds, "reason": "max_rounds"}) + try: + messages.append({"role": "user", + "content": "工具轮次已达上限。请基于以上信息立即给出最终答复,不要再调用工具。"}) + resp = await self.chat_fn(messages, []) + usage = resp.get("usage") or {} + total_in += int(usage.get("prompt_tokens", 0)) + total_out += int(usage.get("completion_tokens", 0)) + final_text = resp.get("content") or last_content + except Exception: + final_text = last_content + return {"response": final_text, "rounds": self.max_rounds, "reason": "max_rounds", + "error": "工具轮次达上限,已强制总结", + "prompt_tokens": total_in, "completion_tokens": total_out} diff --git a/tests/test_tools.py b/tests/test_tools.py new file mode 100644 index 0000000..fec1e27 --- /dev/null +++ b/tests/test_tools.py @@ -0,0 +1,176 @@ +"""工具内核测试:路径关押、文件工具往返、ToolLoop 循环编排(假 chat_fn,零外部依赖)。""" +import asyncio +import json + +import pytest + + +def asyncio_run(coro): + """与 test_pipeline.py 相同的事件循环包装(不依赖 pytest-asyncio 插件模式)。""" + return asyncio.run(coro) + +from router_system.tools import ( + ToolError, + ToolLoop, + WorkspaceTools, + parse_tool_calls, +) + + +@pytest.fixture() +def ws(tmp_path): + return WorkspaceTools(tmp_path / "ws") + + +# ---------------- WorkspaceTools ---------------- + +def test_jail_blocks_traversal(ws): + with pytest.raises(ToolError): + ws.resolve("../outside.txt") + with pytest.raises(ToolError): + ws.resolve("../../etc/passwd") + with pytest.raises(ToolError): + ws.resolve("a/../../b.txt") + # 绝对路径逃逸 + with pytest.raises(ToolError): + ws.resolve(str(ws.root.parent / "secret.txt")) + + +def test_write_read_roundtrip(ws): + r = ws.write_file("dir/a.txt", "你好,工作区") + assert r["ok"] is True + got = ws.read_file("dir/a.txt") + assert got["ok"] is True + assert got["content"] == "你好,工作区" + assert got["truncated"] is False + # 越界写入被折叠为 ok=False(不抛出) + bad = ws.execute("write_file", {"path": "../evil.txt", "content": "x"}) + assert bad["ok"] is False + assert not (ws.root.parent / "evil.txt").exists() + + +def test_list_dir(ws): + ws.write_file("b.txt", "x" * 10) + ws.write_file("sub/c.txt", "y") + r = ws.list_dir("") + names = [e["name"] for e in r["entries"]] + assert "b.txt" in names and "sub/" in names + r2 = ws.list_dir("sub") + assert r2["entries"][0]["name"] == "c.txt" + # 不存在的目录 + assert ws.list_dir("nope")["ok"] is False + + +def test_read_truncation(ws): + ws.write_file("big.txt", "字" * 20000) + got = ws.read_file("big.txt") + assert got["truncated"] is True + assert len(got["content"]) == 8000 + + +def test_unknown_tool_and_missing_file(ws): + assert ws.execute("rm_rf", {"path": "."})["ok"] is False + assert ws.execute("read_file", {"path": "ghost.txt"})["ok"] is False + + +# ---------------- parse_tool_calls ---------------- + +def test_parse_tool_calls(): + msg = {"tool_calls": [ + {"id": "c1", "function": {"name": "read_file", "arguments": '{"path": "a.txt"}'}}, + {"id": "c2", "function": {"name": "write_file", "arguments": "{broken json"}}, + ]} + calls = parse_tool_calls(msg) + assert calls[0]["name"] == "read_file" + assert calls[0]["arguments"] == {"path": "a.txt"} + assert calls[1]["arguments"] == {} # 非法 JSON 容错 + assert parse_tool_calls({}) == [] + + +# ---------------- ToolLoop ---------------- + +def _mk_chat(script): + """script: 依次弹出的响应列表。记录每次收到的 messages。""" + calls = [] + + async def chat_fn(messages, tools_spec): + calls.append([dict(m) for m in messages]) + return script.pop(0) + + chat_fn.calls = calls + return chat_fn + + +def test_toolloop_answer_direct(ws): + chat = _mk_chat([{"content": "最终答案", "tool_calls": [], + "usage": {"prompt_tokens": 10, "completion_tokens": 5}}]) + events = [] + loop = ToolLoop(ws, chat, on_event=events.append) + result = asyncio_run(loop.run("任务")) + assert result["response"] == "最终答案" + assert result["reason"] == "answer" + assert result["prompt_tokens"] == 10 and result["completion_tokens"] == 5 + assert events[0]["type"] == "round" + assert events[-1]["type"] == "final" and events[-1]["reason"] == "answer" + + +def test_toolloop_write_then_answer(ws): + """第一轮调 write_file,第二轮给最终答复;工具结果应回喂。""" + chat = _mk_chat([ + {"content": None, + "tool_calls": [{"id": "c1", "name": "write_file", + "arguments": {"path": "hello.txt", "content": "hi"}}], + "usage": {"prompt_tokens": 20, "completion_tokens": 4}}, + {"content": "已写入 hello.txt", "tool_calls": [], + "usage": {"prompt_tokens": 30, "completion_tokens": 6}}, + ]) + events = [] + loop = ToolLoop(ws, chat, on_event=events.append) + result = asyncio_run(loop.run("写个文件")) + assert result["reason"] == "answer" + assert (ws.root / "hello.txt").read_text(encoding="utf-8") == "hi" + kinds = [e["type"] for e in events] + assert kinds.count("tool_call") == 1 and kinds.count("tool_result") == 1 + # 第二轮 messages 应包含 assistant(tool_calls) + tool 结果 + second = chat.calls[1] + assert second[1]["role"] == "assistant" + assert json.loads(second[1]["tool_calls"][0]["function"]["arguments"])["path"] == "hello.txt" + assert second[2]["role"] == "tool" + assert second[2]["tool_call_id"] == "c1" + + +def test_toolloop_max_rounds_forces_summary(ws): + chat = _mk_chat([ + {"content": None, + "tool_calls": [{"id": "c1", "name": "list_dir", "arguments": {}}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1}}, + ] * 3 + [ + {"content": "强制总结", "tool_calls": [], "usage": {"prompt_tokens": 5, "completion_tokens": 2}}, + ]) + loop = ToolLoop(ws, chat, max_rounds=3) + result = asyncio_run(loop.run("任务")) + assert result["reason"] == "max_rounds" + assert result["response"] == "强制总结" + assert result["error"] + + +def test_toolloop_token_cap(ws): + chat = _mk_chat([ + {"content": "x", "tool_calls": [], + "usage": {"prompt_tokens": 100, "completion_tokens": 100}}, + ] * 5) + loop = ToolLoop(ws, chat, token_cap=150) + result = asyncio_run(loop.run("任务")) + assert result["reason"] == "token_cap" + assert "熔断" in result["error"] + assert len(chat.calls) == 1 # 触顶后不再继续调用 + + +def test_toolloop_chat_error(ws): + async def bad_chat(messages, tools_spec): + raise RuntimeError("网络断了") + + loop = ToolLoop(ws, bad_chat) + result = asyncio_run(loop.run("任务")) + assert result["reason"] == "error" + assert "RuntimeError" in result["error"]