131 lines
5.3 KiB
Python
131 lines
5.3 KiB
Python
"""Tool 基类: 子类只需声明 name/description/parameters 和 execute。"""
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
# docker 沙箱把 user_root(`<workspace>/users/<uid>`)bind 到容器内 `/workspace`
|
|
# (agent_builder _build_system_prompt / executor_docker)。host-side 文件工具
|
|
# (send_email / wechat_push)在宿主进程跑,agent 给的可能是容器绝对路径
|
|
# `/workspace/<wd>/x`,需翻回宿主 `user_root/<wd>/x`。
|
|
_CONTAINER_ROOT = "/workspace"
|
|
|
|
|
|
class FileOutOfBounds(Exception):
|
|
"""agent 给的文件路径解析后落在 user_root 之外(防 `../` 读到别人/系统文件)。"""
|
|
|
|
|
|
def compact_tool_output(
|
|
text: str,
|
|
*,
|
|
max_chars: int = 8_000,
|
|
head_chars: int = 4_000,
|
|
tail_chars: int = 2_000,
|
|
) -> str:
|
|
"""压缩长工具输出,保留头尾和截断说明。"""
|
|
if len(text) <= max_chars:
|
|
return text
|
|
head_chars = max(0, min(head_chars, max_chars))
|
|
tail_chars = max(0, min(tail_chars, max_chars - head_chars))
|
|
removed = len(text) - head_chars - tail_chars
|
|
return (
|
|
text[:head_chars]
|
|
+ f"\n[... truncated, {removed} chars omitted ...]\n"
|
|
+ (text[-tail_chars:] if tail_chars else "")
|
|
)
|
|
|
|
|
|
_TIMEOUT_HINT = (
|
|
"若是批量循环(检索 / PDF 抽取 / 下载等):① 据上面已完成的部分**续跑剩下的**,别整批重跑;"
|
|
"② 每完成一项就落盘 / 打印进度,超时也不丢;③ 大批量 / 大下载改用 background=true 后台跑,"
|
|
"再用 check_process 取结果;④ 单次别塞太多项。"
|
|
)
|
|
|
|
|
|
def format_timeout_result(stdout: str, stderr: str, timeout_s: int) -> str:
|
|
"""run_python / shell 超时结果:带上超时前已捕获的部分输出 + 续跑提示。
|
|
|
|
进程被 kill 前跑出的 stdout/stderr 本已捕获(subprocess 标准行为:TimeoutExpired
|
|
带 .stdout/.stderr、docker 路径 kill 后 communicate() 续读),旧实现直接丢弃只回一句
|
|
超时 → 模型看不到「跑到哪了」、整批重来(工具失败面板 #8:批量检索 / PDF 抽取 / 下载
|
|
超时反复重跑,累计最多)。这里把部分输出一并返回,让模型据此续跑未完成的。"""
|
|
parts = []
|
|
if stdout and stdout.strip():
|
|
parts.append(f"[stdout]\n{stdout.rstrip()}")
|
|
if stderr and stderr.strip():
|
|
parts.append(f"[stderr]\n{stderr.rstrip()}")
|
|
parts.append(
|
|
f"[Error] command timed out after {timeout_s}s(进程已被杀,上方为超时前的部分输出)。"
|
|
)
|
|
parts.append(_TIMEOUT_HINT)
|
|
return "\n".join(parts)
|
|
|
|
|
|
class Tool(ABC):
|
|
name: str = ""
|
|
description: str = ""
|
|
parameters: dict = {}
|
|
|
|
def __init__(self, base_dir: Optional[Path] = None, user_root: Optional[Path] = None) -> None:
|
|
self.base_dir: Path = Path(base_dir) if base_dir else Path.cwd()
|
|
# tool 输出渲染路径用:user_root 内的 path 渲成相对 POSIX 串,user_root 外
|
|
# (用户 --working-dir 指向外部目录)保持绝对。None → 全部按绝对渲染。
|
|
# 目的:不让 tool result 文本里出现 user_id / 部署绝对路径,SPA 截图分享更安全;
|
|
# 顺便让 web SPA 的 artifact chip 抽取(限定 <wd>/ 前缀)更稳。
|
|
self.user_root: Optional[Path] = Path(user_root) if user_root else None
|
|
|
|
@abstractmethod
|
|
def execute(self, **kwargs) -> str:
|
|
...
|
|
|
|
@property
|
|
def schema(self) -> dict:
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": self.name,
|
|
"description": self.description,
|
|
"parameters": self.parameters,
|
|
},
|
|
}
|
|
|
|
def _resolve(self, path: str) -> Path:
|
|
p = Path(path)
|
|
return p if p.is_absolute() else (self.base_dir / p)
|
|
|
|
def _resolve_user_file(self, raw: str) -> Path:
|
|
"""把 agent 给的路径解析成**宿主**绝对 Path,并强制落在 user_root 内。
|
|
|
|
host-side 文件工具(send_email / wechat_push)专用。处理三种入参:
|
|
- 容器绝对路径 `/workspace/<wd>/x` → 翻回宿主 `user_root/<wd>/x`(docker bind);
|
|
- 相对路径 `x` / `<wd>/x` → 拼到 `base_dir`(应为该 task 的宿主工作目录);
|
|
- 宿主绝对路径 → 原样。
|
|
解析后越界(不在 user_root 下)抛 `FileOutOfBounds(raw)`。不校验存在性,
|
|
调用方按需 `.is_file()`。host 模式 agent 不会产出 `/workspace/` 前缀,翻译分支天然不触发。
|
|
"""
|
|
raw = (raw or "").strip()
|
|
if self.user_root is not None and (
|
|
raw == _CONTAINER_ROOT or raw.startswith(_CONTAINER_ROOT + "/")
|
|
):
|
|
rest = raw[len(_CONTAINER_ROOT):].lstrip("/")
|
|
p = self.user_root / rest
|
|
else:
|
|
p = self._resolve(raw)
|
|
p = p.resolve()
|
|
if self.user_root is not None:
|
|
try:
|
|
p.relative_to(self.user_root.resolve())
|
|
except ValueError:
|
|
raise FileOutOfBounds(raw)
|
|
return p
|
|
|
|
def _display(self, p: Path) -> str:
|
|
"""对外渲染路径:在 user_root 内 → POSIX 相对串;否则原绝对。"""
|
|
if self.user_root is not None:
|
|
try:
|
|
return p.resolve().relative_to(self.user_root.resolve()).as_posix()
|
|
except (ValueError, OSError):
|
|
pass
|
|
return str(p)
|