Compare commits
141
Commits
62fab2bc2c
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2832566e87 | ||
|
|
7d5d7b9e36 | ||
|
|
4497d80a94 | ||
|
|
04f1375b45 | ||
|
|
d9b2359d25 | ||
|
|
cc72248e26 | ||
|
|
c4154d4db9 | ||
|
|
0ed03ac095 | ||
|
|
2ea3fa14d8 | ||
|
|
dbfb925f65 | ||
|
|
67acae708c | ||
|
|
7e44025671 | ||
|
|
2dc79f0637 | ||
|
|
5894e64287 | ||
|
|
3e860c16e0 | ||
|
|
68a8e49896 | ||
|
|
5cb35945a4 | ||
|
|
f3bb09e6b5 | ||
|
|
6f1f2bd296 | ||
|
|
38be729509 | ||
|
|
d4e627fbb2 | ||
|
|
3e6e171831 | ||
|
|
d1357ae483 | ||
|
|
9db2bc112d | ||
|
|
9064750f33 | ||
|
|
5dc5ec9af2 | ||
|
|
e43aeeaf44 | ||
|
|
7dc6796f3a | ||
|
|
f7fc06c0ff | ||
|
|
04e3357dd0 | ||
|
|
2c6d1e8ce6 | ||
|
|
457bb900fb | ||
|
|
f4fcf175ec | ||
|
|
d513c5f892 | ||
|
|
4cb97318b5 | ||
|
|
31b2a4d6ca | ||
|
|
db54921066 | ||
|
|
9fee9a1478 | ||
|
|
d3acccd146 | ||
|
|
d1f780e381 | ||
|
|
acb6ddbb4c | ||
|
|
cccc946660 | ||
|
|
aa15841d73 | ||
|
|
0246943165 | ||
|
|
d53eba77f0 | ||
|
|
3d7f5ca3cb | ||
|
|
d29bb5504a | ||
|
|
a583f44d5d | ||
|
|
5f3ca9a246 | ||
|
|
77262f8c3a | ||
|
|
84a3f6dd5e | ||
|
|
0e9efe4895 | ||
|
|
4480bce684 | ||
|
|
58725c8a12 | ||
|
|
90f5544673 | ||
|
|
5bcf22f9a7 | ||
|
|
6ef5d11d9d | ||
|
|
eb88a77238 | ||
|
|
a833cd0d5f | ||
|
|
207137c0dd | ||
|
|
b2ecb8e12b | ||
|
|
910ab8aa38 | ||
|
|
8de7f6696c | ||
|
|
1348660cd8 | ||
|
|
177a9ac9a6 | ||
|
|
ce142fe32f | ||
|
|
d03dbc7f91 | ||
|
|
a5af4bb404 | ||
|
|
9e0d39c912 | ||
|
|
ce3fb7fc74 | ||
|
|
2b133f1531 | ||
|
|
40523f6f8a | ||
|
|
8a5162c912 | ||
|
|
116d6b63bd | ||
|
|
0fa32cb5e7 | ||
|
|
63c29680ce | ||
|
|
76afa37c6b | ||
|
|
293d71e86b | ||
|
|
d8b601ab9a | ||
|
|
286e632288 | ||
|
|
37feb09185 | ||
|
|
8a9a061662 | ||
|
|
0be9e5d0d2 | ||
|
|
a80aaabffe | ||
|
|
01126a5786 | ||
|
|
205a60f039 | ||
|
|
13beba09b4 | ||
|
|
8cf3dd2d93 | ||
|
|
9000bb4b04 | ||
|
|
7f21b65428 | ||
|
|
26de527a09 | ||
|
|
a18a7de1ca | ||
|
|
99c1bc1739 | ||
|
|
1cf428f384 | ||
|
|
fc310c9081 | ||
|
|
9179281420 | ||
|
|
ad3a624b4c | ||
|
|
35ee614ca9 | ||
|
|
b73628b2e5 | ||
|
|
9bbbc3fdf4 | ||
|
|
85bf68f505 | ||
|
|
adfa61aa3c | ||
|
|
aa271e395a | ||
|
|
7ec8c9d0e9 | ||
|
|
af0ff56848 | ||
|
|
a2c3a579fd | ||
|
|
cd4d512ecd | ||
|
|
b289e7d17c | ||
|
|
a2ba797410 | ||
|
|
3e773debc9 | ||
|
|
cc2f7bca79 | ||
|
|
fe12dd7b33 | ||
|
|
5af863a2c5 | ||
|
|
85133868f3 | ||
|
|
8cb822b511 | ||
|
|
bfa6d8deb6 | ||
|
|
3fc213e7b9 | ||
|
|
d47ef6efb4 | ||
|
|
791fac0efd | ||
|
|
b9f8abdfeb | ||
|
|
eadba81c26 | ||
|
|
8af8dec1b0 | ||
|
|
823e91a30e | ||
|
|
316758d80e | ||
|
|
52c0ade3c4 | ||
|
|
8c58770280 | ||
|
|
a6e1781185 | ||
|
|
14964b9574 | ||
|
|
2068319c3a | ||
|
|
0d16832612 | ||
|
|
0cfb3a7581 | ||
|
|
e77ef8cf76 | ||
|
|
175fa5d244 | ||
|
|
96d9f96e6d | ||
|
|
f606f7ec56 | ||
|
|
02d95a2bd9 | ||
|
|
82875b3021 | ||
|
|
4d6e1e1fd4 | ||
|
|
8fe2afdfde | ||
|
|
4b7773daec | ||
|
|
6309358262 |
No files matched your search
@@ -22,3 +22,12 @@ waf/certs/
|
||||
# 本地 IDE / 工具
|
||||
.qoder/
|
||||
.vscode/
|
||||
|
||||
# 运行时日志(HailoRT 等)
|
||||
hailort.log
|
||||
|
||||
# 录制产物(可复现:mtc DSL → 编译 → pybot --record)
|
||||
actpred/data_v3
|
||||
|
||||
# pybot 训练/评估数据(pybot_record_* 动作录制 + eval_recorded*.json)
|
||||
data/
|
||||
@@ -69,6 +69,7 @@
|
||||
| [dnfs](dnfs/) | Python(Flask) | 8006 | ~~树状网络文件系统~~(遗留,已由 dnfs2 平替) |
|
||||
| [db2](db2/) | C++20 | 8005 | 通用数据库服务,基于 pymsg 协议(TCP 长度前缀帧 + JSON)的 SQL-over-TCP,参数化查询防注入;内置集中式日志收集(Logs 表 + 自动清理) |
|
||||
| [db](db/) | Python(pymsg) | 8005 | ~~通用数据库服务~~(遗留,已由 db2 平替) |
|
||||
| [memdb](memdb/) | C++20 | 8011 | 内存缓存数据库(libmemdb 模板库 + libmsg RPC),TTL 惰性驱逐,9 路由(ping/set/get/del/update/exists/size/show/expire),线程安全 + 优雅退出;暂无业务调用方,预留热点数据/会话缓存 |
|
||||
| [ai_review](ai_review/) | Python(Flask) | 8007 | AI 内容审核服务,DeepSeek LLM + ChromaDB RAG,支持人工反馈纠偏 |
|
||||
| [captcha2](captcha2/) | C++20 | 8009 | 验证码服务(pymsg/TCP),生成图形验证码(5 分钟过期、一次性),经 ca 的 `GET /captcha/image` 代理给浏览器 |
|
||||
| [captcha](captcha/) | Python(pymsg) | 8009 | ~~验证码服务~~(遗留,已由 captcha2 平替) |
|
||||
@@ -77,6 +78,9 @@
|
||||
| [aibox](aibox/) | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度,由 core 拉起但与主干微服务无调用关系 |
|
||||
| [BotOps](BotOps/) | Python | — | aibox 智能体 TUI 前端(Textual),WebSocket 连接 aibox,支持 Markdown 增量流式渲染、思考模式展示、任务独立侧栏、自动重连 |
|
||||
| [bot4cj](bot4cj/) | 仓颉(Cangjie) | — | 机器人硬件脚本解析器,解析 .mt 矩阵脚本通过 I2C 驱动 PCA9685 舵机 / MPU6050 陀螺仪,面向树莓派 Dashun 蜘蛛机器人平台(独立硬件项目) |
|
||||
| [pybot](pybot/) | Python | — | bot4cj 的 Python 迁移版:.mt 动作脚本解析器(零第三方依赖,与仓颉版逐语义对齐);支持执行 / `--dry-run` / `--record` 录制 actpred episode(含 imu.npy)/ `--loop` 实时预测闭环(HEF 后端自动异步,推理与执行并行) |
|
||||
| [mtc](mtc/) | Python | — | 命令式动作脚本(Python DSL)→ .mt 编译器:非阻塞时间轴 + 10ms 栅格展开,输出可由 pybot/bot4cj 直接执行;零第三方依赖 |
|
||||
| [actpred](actpred/) | Python(PyTorch) | — | 动作预测模型(图像 + 动作历史 → 下一动作),bot4cj/pybot 的「大脑」配套;两阶段训练(.mt 预训练 → episode 微调),ONNX/Hailo HEF 部署(实验) |
|
||||
| [lci](lci/) | C++ | — | 嵌入式帧缓冲图形界面(LCI),基于 `/dev/fb0` 直接操作显存,无需 X11/Wayland,矢量字体、图片渲染、实时数据可视化,面向树莓派等嵌入式环境 |
|
||||
|
||||
### core(核心模块 · 进程托管服务)
|
||||
@@ -122,7 +126,7 @@
|
||||
- 支持文件浏览、视频/音频在线播放(HTTP Range 分段请求,返回 206/416)、一键下载。
|
||||
- 路径遍历防护(realpath 前缀校验),仅允许访问 `--root` 之内。
|
||||
- **dnfs2 代理**:`/get_notes`、`/add_note`、`/delete_note`、`/update_note` 透传 Authorization 转发到 `dnfs2:8006`(解决浏览器跨域)。
|
||||
- **账号旁路**:`/api/check_user`、`/api/reset_password` 等接口直接读写 `gateway/sqlite3.db` 的 Users 表;`/yhdet_callback`、`/api/yhdet_*` 对接 yhdet.top OAuth 登录。
|
||||
- **账号旁路已收口**:原 static 时代的 `/api/check_user`、`/api/reset_password`、yhdet OAuth 等旁路接口已移除/迁至 ca(忘记密码走 ca `/reset/send_code` + `/reset_password`),fsys 不再直接读写用户表。
|
||||
- 统一品牌"树状菜单",支持深色模式(`css/theme.css`)与样式节点注入的自定义主题/样式(见 [docs/dnfs_custom_themes.md](docs/dnfs_custom_themes.md))。
|
||||
- 配置全部经 CLI 传递(CLI > 环境变量 > 默认值,原 `confs/configs/static.ini` 迁入 `core/services.yaml` 的 fsys 条目 args),不依赖 confs。
|
||||
|
||||
@@ -171,6 +175,12 @@
|
||||
- 安全机制:表名白名单(字母数字下划线)、`sqlite3_bind_text` 参数绑定、仅反射已有表无法建删表、RpcServer 异常脱敏。
|
||||
- 配置全部经 CLI 传递(CLI > 环境变量 > 默认值,原 `db.ini` 迁入 `core/services.yaml` 的 CLI args),不依赖 confs;详见 [db2/README.md](db2/README.md)。
|
||||
|
||||
### memdb(内存缓存数据库)
|
||||
|
||||
- 基于 [libmemdb](libmemdb/)(C++20 header-only 模板库:三索引存储——`std::list` 主存 + `unordered_map` 键索引 + `multiset` TTL 索引,TTL 惰性驱逐)+ [libmsg](libmsg/) RPC 服务,监听 127.0.0.1:8011。
|
||||
- 9 路由:`ping` / `set` / `get` / `del` / `update` / `exists` / `size` / `show` / `expire`;`std::mutex` 保护线程安全,SIGINT/SIGTERM 优雅退出。
|
||||
- 已纳入 `core/services.yaml` 自启动;暂无业务调用方,预留热点数据/会话缓存;详见 [memdb/README.md](memdb/README.md)。
|
||||
|
||||
### ai_review(AI 内容审核服务)
|
||||
|
||||
- 基于 DeepSeek LLM + ChromaDB 向量检索的 RAG 审核架构。
|
||||
@@ -225,7 +235,8 @@
|
||||
|
||||
- 一套相对独立的单智能体产品,基于 DeepSeek + WebSocket(:8765),由 `core` 拉起但与 waf / gateway / dnfs2 等主干微服务无调用关系。
|
||||
- 自带技能系统(`skills/`)、任务调度器、记忆管理、前端页面与语音客户端(`voice_client.py`)。
|
||||
- 默认模型 `deepseek-v4-flash`,配置在 `config.json`;详见 [aibox/README.md](aibox/README.md)。
|
||||
- 默认模型 `deepseek-v4-flash`,配置在 `config.json`;`models` 配置节可挂多个 OpenAI 兼容端点按档案名切换(档案级 reasoning/温度/展示名/`fallback` 失败转移),`model_config.py` 提供交互式接入点管理(列表/添加/删除/连通性验证,Key 仅记环境变量名不落盘),`--once` 支持无 WebSocket 的 CLI 一次性对话;详见 [aibox/README.md](aibox/README.md)。
|
||||
- **六轮安全审计(2026-08,Batch E/F/G/H/I/K/L)**:命令执行保护链(1MB/超时/中断)、凭据脱敏(代理密码不落盘、日志打码、Bearer 脱敏)、取消传播(子代理/情绪子代理/调度器超时)、DB 事务与 FTS5 转义、POSIX 进程组终止、任务系统入参校验(畸形时间/数字字段防调度器热循环)、前端 XSS 净化(DOMPurify + SRI)等;第六轮交叉验证 Batch K 零回归,R1~R3 低危已随 Batch L 落地(72 项发现全部修复);单测 490 项全过(Linux);审计报告见 [docs/reports/aibox_audit_20260817.md](docs/reports/aibox_audit_20260817.md)。
|
||||
|
||||
### BotOps(aibox 智能体 TUI 前端)
|
||||
|
||||
@@ -243,6 +254,21 @@
|
||||
- 附带数学库(矩阵运算、NN/CNN、Adam 优化器)、CUDA GPU 加速 kernel 和 Lua 硬件模块(ADS1115 / ADS1256 / SPI)。
|
||||
- 与主干微服务完全独立,需部署在树莓派上独立运行。
|
||||
|
||||
### pybot(机器人动作脚本解析器 · Python 迁移版)
|
||||
|
||||
- [bot4cj](#bot4cj机器人硬件脚本解析器) 的 Python 迁移版,零第三方依赖(纯标准库),与仓颉版逐语义对齐(含占位符 `-` 递归继承、微步分解、块序依赖)。
|
||||
- 三种运行模式:脚本执行 `run()`、`--record` 录制 actpred episode(images + actions.npy + imu.npy,≤4 个 MPU6050 按块行序拼接)、`--loop` 实时预测闭环(daemon 线程队列非阻塞取帧,overrun 告警;HEF 后端自动走异步推理——执行上一步动作与 NPU 推理并行,滞后一拍)。
|
||||
- 扬声器静音语义:`frequency <= 0` 不出声仅保持节拍(对齐仓颉版 sleep 占位语义)。
|
||||
- Windows 开发机可跑测试(fcntl mock + 符号链接 fallback);详见 [pybot/README.md](pybot/README.md)。
|
||||
|
||||
### mtc(.mt 脚本编译器)
|
||||
|
||||
- 命令式动作脚本(Python DSL)→ `.mt` 矩阵脚本编译器,零第三方依赖;产物可直接由 pybot / bot4cj 执行。
|
||||
- 管线:script.py `exec` 注入命名空间 → 事件收集 → 时间轴 IR(相对→绝对、重叠截断)→ 10ms 栅格化 → out.mt。
|
||||
- 非阻塞时间轴语义:`Action`/`SetTo`/`Beep` 在当前游标处注册运动段即返回,`Delay`/`Wait` 推进游标,并发运动天然表达;`dur` 须为 10 的倍数(与 pybot 微步节拍一致)。
|
||||
- 扬声器通道静止值回落 0(段前/段间空隙/段后 = 静音),与 pybot speaker 静音语义对齐;`spk.frequency`/`spk.delay` 列自动成对补齐。
|
||||
- 用法:`python -m mtc walk.py -o walk.mt`(须在 `mtc/` 目录或设 `PYTHONPATH=mtc` 运行);详见 [mtc/README.md](mtc/README.md)。
|
||||
|
||||
### lci(嵌入式帧缓冲图形界面)
|
||||
|
||||
- LCI(Linux Chinese Interface)是基于 Linux 帧缓冲设备(`/dev/fb0`)的图形界面系统,**无需 X11 或 Wayland**,直接 `mmap` 操作显存,专为树莓派等嵌入式环境设计。
|
||||
@@ -291,8 +317,12 @@
|
||||
| 8007 | ai_review | 127.0.0.1 | AI 审核(仅内部调用) |
|
||||
| 8009 | captcha2 | 127.0.0.1 | 验证码服务(pymsg/TCP,经 ca 代理) |
|
||||
| 8008 | smtp | 127.0.0.1 | 邮件服务(仅内部调用,注册验证码邮件后端) |
|
||||
| 8011 | memdb | 127.0.0.1 | 内存缓存数据库(libmsg RPC,仅内部调用,暂无业务调用方) |
|
||||
| 8765 | aibox | 127.0.0.1 | AI 智能体(WebSocket,独立产品) |
|
||||
| — | bot4cj | — (树莓派 I2C) | 机器人硬件脚本解析器(独立硬件项目,不联网) |
|
||||
| — | pybot | — (树莓派) | bot4cj 的 Python 迁移版(.mt 解析/录制/--loop 实时闭环,独立硬件项目) |
|
||||
| — | mtc | — | 命令式动作脚本 → .mt 编译器(开发期工具,无端口) |
|
||||
| — | actpred | — | 动作预测模型训练/部署(实验,无端口) |
|
||||
| — | lci | — (树莓派 /dev/fb0) | 嵌入式帧缓冲图形界面(独立硬件项目,不联网) |
|
||||
| — | BotOps | — | aibox 智能体 TUI 前端(WebSocket 客户端,连接 aibox:8765) |
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# Virtual environment
|
||||
venv/
|
||||
venv_dfc/
|
||||
.env/
|
||||
|
||||
# Python bytecode
|
||||
@@ -12,6 +13,14 @@ checkpoints/
|
||||
*.pt
|
||||
*.pth
|
||||
*.onnx
|
||||
*.onnx.data
|
||||
*.hef
|
||||
*.har
|
||||
|
||||
# DFC logs
|
||||
acceleras.log
|
||||
hailo_sdk.*.log
|
||||
hailort.log
|
||||
|
||||
# Dataset (too large for git)
|
||||
data/
|
||||
|
||||
+158
-40
@@ -10,15 +10,16 @@ Transformer(pre-LN)自注意力融合,取 [PRED] token 回归下一步动
|
||||
|
||||
```
|
||||
F 帧图像 (F×3×H×W) ─► 共享 CNN 编码器 ─► F 个 token ─► proj → d_model ┐
|
||||
├─► [img×F, act×H, PRED] + 位置编码
|
||||
动作历史 (H×D) ─► Linear proj → d_model ─► H 个 token ────────┤ │
|
||||
可学习 [PRED] token ────────────────────────────────────────┘ ▼
|
||||
TransformerEncoder (pre-LN, 2 层, 4 头, d_model=128)
|
||||
├─► [img×F, act×H, (IMU), PRED]
|
||||
动作历史 (H×D) ─► Linear proj → d_model ─► H 个 token ────────┤ │ + 位置编码
|
||||
IMU 姿态 (imu_dim) ─► Linear proj → 1 个 token(R2,可选)────────┤ ▼
|
||||
可学习 [PRED] token ────────────────────────────────────────┘ TransformerEncoder (pre-LN, 2 层, 4 头, d_model=128)
|
||||
│
|
||||
取 PRED 位置输出 ─► MLP 头 ─► 预测动作 (action_dim)
|
||||
```
|
||||
|
||||
token 序列长度固定为 F+H+1(默认 3+5+1=9),除 batch 外无动态 shape,对 NPU(如树莓派
|
||||
token 序列长度固定为 F+H+1(默认 3+5+1=9;`--imu-dim`>0 时 IMU token 插到
|
||||
PRED 之前,长度 F+H+2),除 batch 外无动态 shape,对 NPU(如树莓派
|
||||
AI 扩展板)部署友好。
|
||||
|
||||
## 快速开始
|
||||
@@ -33,6 +34,23 @@ bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18
|
||||
# 断点续训练(从 last.pt 恢复 model/optimizer/scheduler/epoch)
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --resume checkpoints/last.pt
|
||||
|
||||
# R1 改进:残差预测 + 多步展开训练(可叠加)
|
||||
# --predict-delta 预测 Δ=actions[t]-actions[t-1](惯性先验,平滑动作;推理端自动加回 a_{t-1})
|
||||
# --rollout-len K 多步展开训练(K>1 时 scheduled sampling,治 exposure bias)
|
||||
# --ss-prob P scheduled sampling 真实历史概率(默认 1.0=全真;0.5 表示半数步用自预测回传)
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 \
|
||||
--predict-delta --rollout-len 4 --ss-prob 0.7
|
||||
|
||||
# R2 姿态反馈:--imu-dim 8 启用独立 IMU token(数据需含 imu.npy,见"数据格式")
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 --imu-dim 8
|
||||
|
||||
# 归一化+逐列加权:固定物理范围归一化(免 stats 文件)+ 舵机列损失权重上浮
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 \
|
||||
--normalize-actions --normalize-mode fixed_range
|
||||
# 自定义列权重(逗号分隔 action_dim 个浮点;默认舵机1.0/spk0.2/保留位0.0)
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --normalize-actions \
|
||||
--normalize-mode fixed_range --column-weights "1,1,1,1,1,1,1,1,1,1,0.2,0.2,0,0,0,0,0,0"
|
||||
|
||||
# 推理(--images 按时间顺序旧→新传入 frame_len 个图像路径,冷启动时可重复传同一帧;
|
||||
# --action-history 格式:冒号分隔各步,每步 action_dim 个值,共 history_len 步)
|
||||
bash _run.sh --model checkpoints/best_inference.pt \
|
||||
@@ -42,6 +60,12 @@ bash _run.sh --model checkpoints/best_inference.pt \
|
||||
# 导出 ONNX(如安装了 onnxruntime,导出后自动验证一致性)
|
||||
bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
||||
|
||||
# R4 离线评估:开环滚动预测,量化闭环代理指标(误差发散/平滑度/越界率)
|
||||
python evaluate.py --model checkpoints/best_inference.pt --data-dir ./data
|
||||
# A/B 对比两个 checkpoint(负值 = 主模型更优)
|
||||
python evaluate.py --model checkpoints/r1.pt --compare checkpoints/baseline.pt \
|
||||
--data-dir ./data --episodes 3 --horizon 50 --output report.json
|
||||
|
||||
# 运行单元测试
|
||||
python test_actpred.py
|
||||
```
|
||||
@@ -62,9 +86,7 @@ bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0
|
||||
bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0 \
|
||||
--resume checkpoints/pretrain_last.pt
|
||||
|
||||
# 用预训练权重作为微调起点
|
||||
# 注意:预训练在非归一化动作空间进行,微调时请勿开启 --normalize-actions,
|
||||
# 否则输入分布与预训练权重不匹配,预训练收益会大打折扣
|
||||
# 用预训练权重作为微调起点(pretrain 也支持 --init-from,行为同 train)
|
||||
bash _run.sh train --data-dir ./data --epochs 50 \
|
||||
--init-from checkpoints/pretrain_best_inference.pt
|
||||
```
|
||||
@@ -72,9 +94,21 @@ bash _run.sh train --data-dir ./data --epochs 50 \
|
||||
| 参数 | 说明 |
|
||||
|------|------|
|
||||
| `--mt-dir` | .mt 脚本目录(如 `../bot4cj/scripts`) |
|
||||
| `--noise-std` | 高斯噪声标准差(数据增强,建议 0.5~2.0) |
|
||||
| `--noise-std` | 高斯噪声标准差(数据增强,建议 0.5~2.0;与 `--normalize-actions` 联用时按 z-score 单位,建议 0.05~0.2) |
|
||||
| `--random-images` | 用 ImageNet 归一化均匀噪声替代全零图像占位,**恢复 BN 统计**(详见下方警告) |
|
||||
| `--init-from` | 预训练 checkpoint 路径,作为微调起点(strict=False;与当前运行动作空间不一致时打印 WARNING) |
|
||||
|
||||
> **注意**:.mt 脚本列数(12)与 `action_dim`(18)不一致时自动补零。预训练保存到 `checkpoints/pretrain_best_inference.pt`。
|
||||
> 预训练与微调必须使用**同一动作空间**(`--normalize-actions` 同时开或同时关),
|
||||
> 否则输入分布与权重不匹配,预训练收益会大打折扣。
|
||||
|
||||
> **警告(2026-08-13 实测)**:全零图像占位会让 BN `running_var` 停在初始小值,
|
||||
> 推理时特征放大到 ±1e6,**HEF uint8 量化不可行**(输出对输入不敏感)。
|
||||
> 面向 NPU 部署的预训练/微调请务必使用 `--random-images`(或真实数据微调),
|
||||
> 导出时 `export_deploy.py` 的 `check_token_ranges` 会拦截退化的 checkpoint。
|
||||
> 详见 [根因分析报告](../docs/reports/actpred_hef_insensitivity_rootcause_20260813.md)。
|
||||
> 建议同时开启 `--normalize-actions`:z-score 后 act_tokens 与 img_tokens 同量级,
|
||||
> 避免 HEF merge 点共享量化 scale 粗量化图像路径。
|
||||
|
||||
## 数据格式
|
||||
|
||||
@@ -86,6 +120,7 @@ data/
|
||||
0001.jpg
|
||||
...
|
||||
actions.npy # shape: (T, action_dim) float32
|
||||
imu.npy # 可选,shape: (T, 8×N) float32(N 个 MPU6050 姿态拼接,R2)
|
||||
episode_002/
|
||||
...
|
||||
```
|
||||
@@ -93,6 +128,15 @@ data/
|
||||
每个 timestep 的样本:帧 `[t-frame_len+1, t]`(旧→新)+ 前 `history_len` 步动作 → 当前步动作。
|
||||
采样起点为 `max(history_len, frame_len-1)`。
|
||||
|
||||
`imu.npy` 每传感器占 8 槽 `[roll, pitch, ax, ay, az, gx, gy, gz]`(roll/pitch
|
||||
角度制,其余为 MPU6050 原始 int16),多传感器(最多 4 个)按脚本
|
||||
`@MPU6050_BIND` 行顺序拼接,由 pybot `--record` 逐帧同步落盘。**旧 episode
|
||||
没有姿态数据,训练 IMU 模型需用新版 pybot 重新录制**;`--imu-dim`>0 时缺失
|
||||
`imu.npy` 的 episode 按零向量兜底(向后兼容,但无姿态先验收益)。
|
||||
姿态归一化按物理量程固定缩放(roll/pitch ÷180,原始 int16 ÷32768,8 槽模式
|
||||
按传感器数平铺),训练与推理共用 `imu_norm.py`,无需统计文件。
|
||||
录制/训练/闭环脚本的传感器行顺序必须一致,否则槽位错位。
|
||||
|
||||
### 数据采集
|
||||
|
||||
```python
|
||||
@@ -110,19 +154,22 @@ recorder.flush()
|
||||
|
||||
| 参数 | 默认值 | 说明 |
|
||||
|------|--------|------|
|
||||
| `--image-width` | 1920 | 图像缩放宽度 |
|
||||
| `--image-height` | 1080 | 图像缩放高度 |
|
||||
| `--image-width` | 640 | 图像缩放宽度 |
|
||||
| `--image-height` | 480 | 图像缩放高度 |
|
||||
| `--action-dim` | 18 | 动作向量维度(对齐 bot4cj 18 列:10 舵机 + 2 云台 + 扬声器频率 + 延时 + 4 预留)|
|
||||
| `--history-len` | 5 | 历史动作窗口长度 |
|
||||
| `--frame-len` | 3 | 输入图像帧数(含当前帧,旧帧在前)|
|
||||
| `--imu-dim` | 0 | R2 姿态输入维度(>0 启用独立 IMU token;8×N,N=传感器数≤4)|
|
||||
| `--backbone` | simple | CNN 类型 (simple/mobilenet) |
|
||||
| `--attn-dim` | 128 | 注意力 token 维度 d_model(须能被 attn-heads 整除)|
|
||||
| `--attn-heads` | 4 | 注意力头数 |
|
||||
| `--attn-layers` | 2 | TransformerEncoder 层数 |
|
||||
| `--epochs` | 100 | 训练轮数 |
|
||||
| `--lr` | 0.001 | 学习率 |
|
||||
| `--batch-size` | 32 | 批次大小 |
|
||||
| `--normalize-actions` | False | 动作归一化(`--no-normalize-actions` 关闭) |
|
||||
| `--batch-size` | 4 | 批次大小(默认分辨率 640×480;更高分辨率需调小,小分辨率可调大) |
|
||||
| `--normalize-actions` | True | 动作归一化(`--no-normalize-actions` 关闭;2026-08-14 起默认开启——复现旧行为需显式关闭) |
|
||||
| `--normalize-mode` | zscore | 归一化方式:`zscore`=数据统计(依赖 `action_stats.npy`);`fixed_range`=固定物理范围 min-max→[0,1](常量表见 `action_norm.py`,免 stats 文件,越界 clamp 防御) |
|
||||
| `--column-weights` | 舵机1.0/spk0.2/保留位0.0 | 逐列损失权重(逗号分隔 action_dim 个浮点;仅影响训练损失,推理端无感知;**始终生效**——包括 `--no-normalize-actions` 的未归一化训练) |
|
||||
| `--dropout` | 0.1 | Dropout 率 |
|
||||
| `--weight-decay` | 0.0001 | 权重衰减 |
|
||||
| `--val-split` | 0.15 | 验证集比例 |
|
||||
@@ -132,14 +179,41 @@ recorder.flush()
|
||||
| `--num-workers` | 2 | DataLoader 工作进程数 |
|
||||
| `--checkpoint-dir` | ./checkpoints | 检查点保存目录 |
|
||||
| `--init-from` | 无 | 预训练 checkpoint 路径,作为微调起点(strict=False) |
|
||||
| `--resume` | 无 | 从 last.pt 断点续训练(恢复 model/optimizer/scheduler/epoch) |
|
||||
| `--resume` | 无 | 从 last.pt 断点续训练(恢复 model/optimizer/scheduler/epoch;scheduler 的 T_max 一并恢复,resume 时改 `--epochs` 不影响 LR 余弦曲线,只影响循环边界) |
|
||||
| `--save-config` | 无 | 保存配置到 JSON 文件 |
|
||||
| `--predict-delta` | False | R1-C 残差预测(target 改 Δ,推理端自动加回 a_{t-1}) |
|
||||
| `--rollout-len` | 1 | R1-B 多步展开步数 K(>1 时 scheduled sampling) |
|
||||
| `--ss-prob` | 1.0 | scheduled sampling 真实历史概率(仅 rollout-len>1 生效) |
|
||||
|
||||
## 离线评估(R4)
|
||||
|
||||
`evaluate.py` 在留出 episode 上**开环滚动预测**(复现 pybot 闭环口径:图像用真值滚动、
|
||||
动作用自预测回写),输出闭环代理指标。单步 MSE(teacher forcing)看不见闭环误差累积,
|
||||
本评估补上这一盲区:
|
||||
|
||||
| 指标 | 含义 | 判读 |
|
||||
|------|------|------|
|
||||
| `overall_mse` | 全部滚动步的平均 MSE | 越低越好 |
|
||||
| `error_growth` | 末段(后25%)误差 / 首段(前25%)误差(按 episode 分别计算后取中位数) | >1 表示误差随步数发散(闭环不稳) |
|
||||
| `error_curve` | 误差随步数的增长曲线 | 看发散拐点 |
|
||||
| `smoothness_ratio` | 预测平滑度 / 真值平滑度 | <1 表示比真值更平滑(残差模式的预期收益) |
|
||||
| `oob_rate` | 预测落在物理范围外的比例 | 越高越危险(舵机越界风险) |
|
||||
|
||||
越界判定用固定物理界(`action_norm.py` 常量表:舵机 [0,180]、spk.F [0,5000] 等),
|
||||
"越物理界"比"越经验界"更符合安全语义,且恒值列不会使经验界退化。`--compare` 用于 A/B 对比(如 R1 开/关)。
|
||||
滚动复现闭环口径:模型带 imu 输入时按 episode 的 `imu.npy` 真值姿态逐步喂入
|
||||
(缺失则零向量,与 pybot 闭环无 `@MPU6050_BIND` 块时的行为一致)。
|
||||
|
||||
## 部署注意
|
||||
|
||||
- 启用 `--normalize-actions` 训练的模型,其 `action_stats.npy` 保存在 checkpoint 同目录;
|
||||
部署时**必须将 stats 文件与 checkpoint 一起拷贝**(推理端优先按 checkpoint 内记录的路径查找,
|
||||
找不到时回退到 checkpoint 同目录的 `action_stats.npy`)。
|
||||
- 启用 `--normalize-actions` 且 `--normalize-mode zscore`(默认)训练的模型,其 `action_stats.npy`
|
||||
保存在 checkpoint 同目录;部署时**必须将 stats 文件与 checkpoint 一起拷贝**(推理端优先按
|
||||
checkpoint 内记录的路径查找,找不到时回退到 checkpoint 同目录的 `action_stats.npy`)。
|
||||
- `--normalize-mode fixed_range` 用固定物理范围 min-max→[0,1](常量表见 `action_norm.py`),
|
||||
**免 stats 文件**,消除 stats 路径类 bug 的整类风险;推理端(predict/onnx_infer/hailo_infer)
|
||||
从 checkpoint/ONNX 元数据自动识别归一化方式(含 predict_delta 残差模式),无需额外参数。
|
||||
注:旧版导出的 ONNX 无元数据,onnx_infer 回退 zscore + 非残差,需显式
|
||||
`--normalize-mode/--predict-delta` 或重新导出。
|
||||
- `actions.npy` 的列数必须与 `--action-dim` 一致;图像与动作帧数不一致时训练会打印 WARNING 并按较少者使用。
|
||||
|
||||
## 编程 API
|
||||
@@ -152,15 +226,19 @@ print(predictor.action_dim, predictor.history_len, predictor.frame_len) # 公
|
||||
# images: 长度为 frame_len 的 RGB 图像列表(旧→新,末帧为当前帧;冷启动可重复同一帧)
|
||||
action = predictor.predict(images, action_history) # np.ndarray (action_dim,)
|
||||
# action_history 形状须为 (history_len, action_dim)
|
||||
# 模型含 IMU 输入时(predictor.imu_dim 非 None)可传 imu= 原始姿态向量
|
||||
# (roll,pitch,ax,ay,az,gx,gy,gz;内部自动归一化,None 等价零向量)
|
||||
action = predictor.predict(images, action_history, imu=imu_vec)
|
||||
```
|
||||
|
||||
## ONNX Runtime 后端(CPU 加速)
|
||||
|
||||
`onnx_infer.py` 的 `OnnxPredictor` 走 ONNX Runtime CPU 推理,接口与
|
||||
`ActionPredictorInference` 对齐,维度从 ONNX 输入元数据自动推导(无需
|
||||
checkpoint)。Pi 5 实测 12ms/次含预处理(纯推理 5.9ms),比 PyTorch 快 ~5 倍,
|
||||
与 torch 前向输出差异 <1e-5。数据见
|
||||
[测速报告](../docs/reports/actpred_inference_speed_benchmark_20260813.md)。
|
||||
`ActionPredictorInference` 对齐(含可选 `imu=` 入参),维度从 ONNX 输入元数据
|
||||
自动推导(无需 checkpoint)。Pi 5 实测(224×224,2026-08-13 复测)11.4ms/次含预处理
|
||||
(纯推理 10.7ms),PyTorch 21ms,约 2 倍加速,与 torch 前向输出差异
|
||||
<1e-5。数据见
|
||||
[测速报告](../docs/reports/actpred_inference_speed_benchmark_20260813.md) §6。
|
||||
|
||||
```bash
|
||||
bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
||||
@@ -180,30 +258,70 @@ python3 onnx_infer.py --model model.onnx --images f1.jpg f2.jpg f3.jpg \
|
||||
|
||||
## Hailo 部署(AI HAT 2 / Hailo-10H)
|
||||
|
||||
树莓派 AI HAT 2 上通过 `hailo_infer.py` 的 `HailoPredictor` 走 NPU 推理,
|
||||
接口与 `ActionPredictorInference` 对齐。模型需先经 Hailo DFC 编译为 HEF:
|
||||
模型需先经 Hailo DFC(仅 x86 Linux,`venv_dfc`)编译为 HEF。
|
||||
训练态 ONNX(5D 输入、AdaptiveAvgPool、MultiheadAttention 的 head
|
||||
reshape)**无法被 DFC 解析**,必须用 `export_deploy.py` 的部署版导出
|
||||
(按帧展开 CNN、固定 kernel AvgPool、注意力按头展开为独立 Linear、
|
||||
全程 (1,1,S,d) 无 transpose 布局,权重数值不变)。
|
||||
|
||||
当前 NPU 目标分辨率为 **224×224**(pretrain checkpoint 配置),HEF 输入
|
||||
为 frame0~2 各 (1,224,224,3) NHWC + action_history (1,1,5,18) +
|
||||
pred_token (1,1,1,128)(`--imu-dim`>0 时追加 imu (1,1,1,imu_dim),共 6 输入),
|
||||
形状与量化参数由 `hailo_infer.py` 从 vstream
|
||||
info 自动推导,不硬编码。
|
||||
|
||||
```bash
|
||||
# 1. 导出 ONNX
|
||||
bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
||||
# 2. 在 x86 主机(或 Hailo 云端 DFC)上编译 + 量化校准 → model.hef,拷回树莓派
|
||||
# 3. 推理(零 torch 依赖,仅需 hailo_platform + numpy + PIL)
|
||||
# x86 主机(venv_dfc 内已装 hailo-sdk 5.3):
|
||||
# 1. 导出部署版 ONNX + 学习到的 pred_token(推理端喂入用)
|
||||
python3 export_deploy.py --model checkpoints/pretrain_best_inference.pt \
|
||||
--pred-token-out checkpoints/actpred_deploy_pred_token.npy
|
||||
# 2. 生成校准集(pred_token 按学习值 ±30% 抖动,避免常量量化退化)
|
||||
python3 gen_calibration.py --deploy --model checkpoints/pretrain_best_inference.pt
|
||||
# 3. 一键编译(parser → optimize → compiler,各步自动做输入连通性检查)
|
||||
python3 dfc_compile.py # 产物 checkpoints/actpred_pretrain.hef
|
||||
# 4. 将 HEF + actpred_deploy_pred_token.npy 拷回树莓派
|
||||
```
|
||||
|
||||
```python
|
||||
from hailo_infer import HailoPredictor
|
||||
已验证:完整 110 层网络(含 matmul/softmax/layernorm 注意力)解析、
|
||||
量化、编译全部通过;hailo10h 目标。
|
||||
编译要点与踩坑记录见 `dfc_compile.py` 文档字符串。
|
||||
|
||||
with HailoPredictor("model.hef") as predictor: # stats_path= 仅归一化模型需要
|
||||
action = predictor.predict(images, action_history)
|
||||
```
|
||||
> **验证记录(2026-08-13 树莓派实测)**:归一化轮次(`--normalize-actions`
|
||||
> + `--random-images` 预训练,224×224)重编译的 HEF 已通过 `_npu_test.py`
|
||||
> 实机复验:img 敏感性 22.3、hist 敏感性 38.9(阈值 0.1),HEF vs torch
|
||||
> max_diff 0.37(原始动作空间),约 12ms/次(含预处理)。此前两轮失败的
|
||||
> 根因与修复均已闭环:
|
||||
> - BN 统计退化 → `--random-images` 预训练 + `export_deploy.py`
|
||||
> `check_token_ranges` 导出前拦截
|
||||
> - pred_token 常量校准退化(scale=18571.7)→ 抖动校准,推理端喂学习值
|
||||
> - 输入层断连 → `dfc_compile.py` parser/optimize 后连通性检查
|
||||
> - vstream 报告顺序乱序 → `hailo_infer.py` 按 `input_layerN` 排序绑定
|
||||
> - 帧归一化判定错误(把归一化帧当原始像素喂,图像路径饱和破坏)→
|
||||
> 按量化网格 dequant(0) 判定,见 hailo_infer.py 注释
|
||||
>
|
||||
> 验证链:`export_deploy.py`(token 量级)→ `_onnx_validate.py`(ONNX
|
||||
> 数值)→ `dfc_compile.py`(输入连通性)→ `_npu_test.py`(4 角敏感性)。
|
||||
> 遗留:全零历史角(z≈-3.5)超出校准输入范围(-2.7)会被截断,
|
||||
> 该角仅作连通性验证,非真实部署场景。
|
||||
|
||||
```bash
|
||||
python3 hailo_infer.py --hef model.hef --images f1.jpg f2.jpg f3.jpg \
|
||||
--action-history "45,0,...:45,0,..." [--stats action_stats.npy]
|
||||
```
|
||||
|
||||
HEF 的 vstream 名称/形状/量化参数全部从 `HEF.get_input/output_vstream_infos()`
|
||||
动态推导,不硬编码;图像输入按 Hailo 惯例为 uint8 原始像素
|
||||
(ImageNet 归一化由 DFC 编译时折叠进网络)。
|
||||
`hailo_infer.py` 的 `HailoPredictor` 适配 5/6 输入 HEF(含可选 imu 输入,
|
||||
按形状自动归类绑定),输入布局与量化
|
||||
参数(帧归一化由量化网格 dequant(0) 自动判定)从 vstream info 自动推导,
|
||||
单元测试含 fake hailo_platform 的 5 输入与乱序报告顺序用例。CLI 支持
|
||||
`--pred-token` 喂入学习值、`--stats` 喂入归一化统计、`--imu` 喂入姿态。
|
||||
环境调研与 API 踩坑记录见
|
||||
[docs/reports/actpred_hailo_hat2_inference_research_20260813.md](../docs/reports/actpred_hailo_hat2_inference_research_20260813.md)。
|
||||
|
||||
### 异步推理(`predict_async`,run_async 流水线)
|
||||
|
||||
`HailoPredictor.predict_async(images, action_history, pred_token=None, imu=None)`
|
||||
提交异步推理并立即返回 `_AsyncJob`,需要结果时 `job.result(timeout)` 阻塞获取:
|
||||
|
||||
- 与同步 `predict` 语义/输出完全一致(公共 `_prepare_inputs/_fill_bindings/_postprocess`)
|
||||
- 内部维护 4 个绑定缓冲池 + 完成回调;在途任务超过缓冲数时提交会阻塞等待
|
||||
空闲缓冲,`job.result` 超时**不会**释放缓冲(缓冲由完成回调释放)
|
||||
- 连续提交(≥2 个在途)时 NPU 保持满负荷:实测连续吞吐 6.0ms vs 同步
|
||||
15.5ms(约 2.6×);扩展测试覆盖乱序取结果/混合调用/缓冲压力/300 次
|
||||
稳定性/4 线程并发,输出与同步 0 差异
|
||||
- pybot `--loop` 检测到 `predict_async` 即自动走异步闭环(推理期间并行
|
||||
执行上一步动作),见 [pybot/README.md](../pybot/README.md)
|
||||
@@ -0,0 +1,174 @@
|
||||
"""actpred HEF NPU 输入敏感性复验脚本(重编译后的验收门槛)
|
||||
|
||||
用法: venv/bin/python _npu_test.py
|
||||
|
||||
4 角测试(真实/全零帧 × 真实/全零历史)+ 与 PyTorch 参考对比。
|
||||
HEF 尺寸/布局从 vstream info 自动推导,checkpoint 用
|
||||
checkpoints/best_inference.pt(v4 onehot,fixed_range)。
|
||||
"""
|
||||
import sys
|
||||
import time
|
||||
sys.path.insert(0, "/opt/blackbird/actpred")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from PIL import Image
|
||||
from hailo_platform import HEF
|
||||
|
||||
from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
from hailo_infer import HailoPredictor
|
||||
|
||||
HEF_PATH = "checkpoints/actpred_v4.hef"
|
||||
CKPT = "checkpoints/best_inference.pt"
|
||||
IMG_DIR = "data/episode_0000/images"
|
||||
|
||||
hef = HEF(HEF_PATH)
|
||||
print("=== HEF vstream infos ===")
|
||||
for i in hef.get_input_vstream_infos():
|
||||
q = i.quant_info
|
||||
print(f"IN {i.name:16s} {tuple(i.shape)} {i.format.type} "
|
||||
f"zp={q.qp_zp} scale={q.qp_scale:.6g}")
|
||||
for o in hef.get_output_vstream_infos():
|
||||
q = o.quant_info
|
||||
print(f"OUT {o.name:16s} {tuple(o.shape)} {o.format.type} "
|
||||
f"zp={q.qp_zp} scale={q.qp_scale:.6g}")
|
||||
|
||||
imgs_real = [np.array(Image.open(f"{IMG_DIR}/000{i}.jpg").convert("RGB"))
|
||||
for i in range(3)]
|
||||
imgs_zero = [np.zeros_like(imgs_real[0]) for _ in range(3)]
|
||||
|
||||
hist_real = np.array([
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
], dtype=np.float32)
|
||||
hist_zero = np.zeros((5, 18), dtype=np.float32)
|
||||
|
||||
# PyTorch 参考(checkpoint 自带 config)
|
||||
ck = torch.load(CKPT, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(ck["config"])
|
||||
model = ActionPredictor(cfg.model, cfg.data, pretrained=False)
|
||||
model.load_state_dict(ck["model_state_dict"])
|
||||
model.eval()
|
||||
pred_learned = model.pred_token.data.unsqueeze(1).numpy() # (1,1,1,d)
|
||||
imu_dim = cfg.data.imu_dim
|
||||
|
||||
# 归一化模型:torch 参考在喂入前归一化;HailoPredictor 传 normalize_mode
|
||||
# 由封装内部归一化(hist_real/hist_zero 保持原始值,避免双重归一化)。
|
||||
# fixed_range 免 action_stats.npy(该模式训练不产出统计文件)。
|
||||
import os
|
||||
stats_path = None
|
||||
norm_mode = cfg.data.normalize_mode if cfg.data.normalize_actions else "none"
|
||||
if norm_mode == "zscore":
|
||||
stats_path = cfg.data.action_stats_path or os.path.join(
|
||||
os.path.dirname(os.path.abspath(CKPT)), "action_stats.npy")
|
||||
if not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(os.path.dirname(os.path.abspath(CKPT)),
|
||||
stats_path)
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(f"统计文件不存在: {stats_path}")
|
||||
print(f"HailoPredictor stats: {stats_path}")
|
||||
elif norm_mode == "fixed_range":
|
||||
print("fixed_range 归一化:免 action_stats.npy,按固定物理范围 min-max")
|
||||
|
||||
from torchvision import transforms
|
||||
tf = transforms.Compose([
|
||||
transforms.Resize((cfg.data.image_height, cfg.data.image_width)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
|
||||
def torch_pred(imgs, hist, use_learned_pred=True, imu=None):
|
||||
x = torch.stack([tf(Image.fromarray(im)) for im in imgs]).unsqueeze(0)
|
||||
h = torch.tensor(hist).unsqueeze(0)
|
||||
if norm_mode == "zscore" and stats_path is not None:
|
||||
stats = np.load(stats_path)
|
||||
h = (h - torch.tensor(stats[0])) / torch.tensor(stats[1])
|
||||
elif norm_mode == "fixed_range":
|
||||
from action_norm import normalize_action
|
||||
h = torch.from_numpy(normalize_action(h.numpy(), cfg.data.action_dim))
|
||||
ii = None
|
||||
if imu_dim > 0:
|
||||
# 与 HailoPredictor 同口径:原始姿态 → 内部归一化前需手动归一
|
||||
from imu_norm import normalize_imu
|
||||
ii = torch.from_numpy(normalize_imu(imu, imu_dim)).unsqueeze(0)
|
||||
if use_learned_pred:
|
||||
with torch.no_grad():
|
||||
out = model(x, h, ii).squeeze(0).numpy()
|
||||
else:
|
||||
saved = model.pred_token.data.clone()
|
||||
model.pred_token.data.zero_()
|
||||
with torch.no_grad():
|
||||
out = model(x, h, ii).squeeze(0).numpy()
|
||||
model.pred_token.data.copy_(saved)
|
||||
if cfg.data.onehot:
|
||||
# 多头 logits → 期望值解码 → 原始量纲(与 predict.py 同语义,
|
||||
# 含权重 0 列固定常量),对比前无需反归一化
|
||||
from bins import decode as bin_decode, fixed_zero_columns
|
||||
from action_norm import default_column_weights
|
||||
out = bin_decode(out, cfg.data.action_dim).astype(np.float32)
|
||||
w = (np.asarray(cfg.data.column_weights, dtype=np.float32)
|
||||
if cfg.data.column_weights is not None
|
||||
else default_column_weights(cfg.data.action_dim))
|
||||
fixed = fixed_zero_columns(w, cfg.data.action_dim)
|
||||
out[w <= 0] = fixed[w <= 0]
|
||||
return out
|
||||
|
||||
|
||||
# R2 方案一:imu_dim>0 时每步传原始姿态(HailoPredictor 内部归一化)
|
||||
imu_real = np.array([5.0, -3.0] + [0.0] * 6, dtype=np.float32)
|
||||
|
||||
# normalize_actions=False 时 normalize_mode 无意义(stats_path=None 即不归一化)
|
||||
with HailoPredictor(HEF_PATH, stats_path=stats_path,
|
||||
normalize_mode=(norm_mode if norm_mode != "none"
|
||||
else "zscore")) as p:
|
||||
print(f"\ndims: frame_len={p.frame_len} img={p.image_width}x{p.image_height} "
|
||||
f"history={p.history_len} action_dim={p.action_dim}")
|
||||
t0 = time.perf_counter()
|
||||
a = p.predict(imgs_real, hist_real, pred_token=pred_learned, imu=imu_real)
|
||||
b = p.predict(imgs_zero, hist_real, pred_token=pred_learned, imu=imu_real)
|
||||
c = p.predict(imgs_real, hist_zero, pred_token=pred_learned, imu=imu_real)
|
||||
d = p.predict(imgs_zero, hist_zero, pred_token=pred_learned)
|
||||
t1 = time.perf_counter()
|
||||
print(f"4x predict: {(t1 - t0) / 4 * 1000:.1f} ms/次")
|
||||
for name, r in [("real_img+real_hist", a), ("zero_img+real_hist", b),
|
||||
("real_img+zero_hist", c), ("zero_img+zero_hist", d)]:
|
||||
print(f"{name:20s} max_abs={np.abs(r).max():.4f} out={np.round(r[:6], 3)}...")
|
||||
|
||||
img_sens = np.abs(a - b).max()
|
||||
hist_sens = np.abs(a - c).max()
|
||||
print(f"\nimg sensitivity = {img_sens:.4f}(>0.1 视为输入连通)")
|
||||
print(f"hist sensitivity = {hist_sens:.4f}(>0.1 视为输入连通)")
|
||||
ok = img_sens > 0.1 and hist_sens > 0.1
|
||||
if imu_dim > 0:
|
||||
e = p.predict(imgs_real, hist_real, pred_token=pred_learned,
|
||||
imu=np.zeros_like(imu_real))
|
||||
print(f"imu sensitivity = {np.abs(a - e).max():.4f}(>0 视为姿态连通)")
|
||||
|
||||
ref = torch_pred(imgs_real, hist_real, True, imu_real)
|
||||
ref_zero_pred = torch_pred(imgs_real, hist_real, False, imu_real)
|
||||
# torch 参考是模型原始归一化空间输出;HailoPredictor 已反归一化
|
||||
# 到原始动作空间,比较前须同空间对齐,否则归一化模型下 max_diff
|
||||
# 无意义(onehot 模型 torch_pred 已 decode 为原始量纲,无需再处理)
|
||||
if not cfg.data.onehot:
|
||||
if norm_mode == "fixed_range":
|
||||
from action_norm import denormalize_action
|
||||
ref = denormalize_action(ref, cfg.data.action_dim)
|
||||
ref_zero_pred = denormalize_action(ref_zero_pred,
|
||||
cfg.data.action_dim)
|
||||
elif stats_path is not None:
|
||||
s = np.load(stats_path)
|
||||
ref = ref * s[1] + s[0]
|
||||
ref_zero_pred = ref_zero_pred * s[1] + s[0]
|
||||
print(f"torch learned_pred : {np.round(ref[:6], 3)}...")
|
||||
print(f"torch zero_pred : {np.round(ref_zero_pred[:6], 3)}...")
|
||||
print(f"HEF vs torch(learned) max_diff={np.abs(a - ref).max():.4f}"
|
||||
f"(原始动作空间)")
|
||||
|
||||
print("\nPASS: HEF 输入连通,可继续与 torch 对比量化误差"
|
||||
if ok else "\nFAIL: HEF 输出对输入不敏感,编译仍有问题")
|
||||
@@ -0,0 +1,138 @@
|
||||
"""部署版 ONNX 数值验证:输入敏感性 + 与 PyTorch 一致性(重编译前门槛)
|
||||
|
||||
用法: venv/bin/python _onnx_validate.py
|
||||
"""
|
||||
import sys
|
||||
sys.path.insert(0, "/opt/blackbird/actpred")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
import onnxruntime as ort
|
||||
|
||||
from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
|
||||
CKPT = "checkpoints/pretrain_best_inference.pt"
|
||||
ONNX = "checkpoints/actpred_deploy.onnx"
|
||||
IMG_DIR = "data/episode_0000/images"
|
||||
|
||||
ck = torch.load(CKPT, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(ck["config"])
|
||||
model = ActionPredictor(cfg.model, cfg.data, pretrained=False)
|
||||
model.load_state_dict(ck["model_state_dict"])
|
||||
model.eval()
|
||||
|
||||
action_stats = None
|
||||
norm_mode = cfg.data.normalize_mode if cfg.data.normalize_actions else "none"
|
||||
if norm_mode == "zscore":
|
||||
import os
|
||||
stats_path = cfg.data.action_stats_path or os.path.join(
|
||||
os.path.dirname(os.path.abspath(CKPT)), "action_stats.npy")
|
||||
if not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(os.path.dirname(os.path.abspath(CKPT)),
|
||||
stats_path)
|
||||
action_stats = np.load(stats_path)
|
||||
if action_stats.shape != (2, cfg.data.action_dim):
|
||||
raise ValueError(f"action stats 形状应为 (2, {cfg.data.action_dim}),"
|
||||
f"实际为 {action_stats.shape}: {stats_path}")
|
||||
print(f"history 将 z-score 归一化({stats_path})")
|
||||
elif norm_mode == "fixed_range":
|
||||
print("history 将 fixed_range min-max 归一化(action_norm.py,免 stats)")
|
||||
|
||||
sess = ort.InferenceSession(ONNX)
|
||||
in_names = [i.name for i in sess.get_inputs()]
|
||||
print("ONNX inputs:", in_names)
|
||||
|
||||
h, w = cfg.data.image_height, cfg.data.image_width
|
||||
tf = transforms.Compose([
|
||||
transforms.Resize((h, w)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
imgs_real = [np.array(Image.open(f"{IMG_DIR}/000{i}.jpg").convert("RGB"))
|
||||
for i in range(3)]
|
||||
imgs_zero = [np.zeros_like(imgs_real[0]) for _ in range(3)]
|
||||
|
||||
hist_real = np.array([
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 10, 200, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 20, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
], dtype=np.float32)
|
||||
hist_zero = np.zeros((5, 18), dtype=np.float32)
|
||||
|
||||
if norm_mode == "fixed_range":
|
||||
from action_norm import normalize_action
|
||||
hist_real = normalize_action(hist_real, cfg.data.action_dim)
|
||||
hist_zero = normalize_action(hist_zero, cfg.data.action_dim)
|
||||
elif action_stats is not None:
|
||||
hist_real = (hist_real - action_stats[0]) / action_stats[1]
|
||||
hist_zero = (hist_zero - action_stats[0]) / action_stats[1]
|
||||
|
||||
pred_learned = model.pred_token.data.unsqueeze(1).numpy() # (1,1,1,128)
|
||||
pred_none = np.zeros((1, 1, 1, 128), dtype=np.float32)
|
||||
|
||||
# R2 方案一:imu_dim>0 时模型/ONNX 多了归一化姿态输入
|
||||
imu_dim = cfg.data.imu_dim
|
||||
imu_real = None
|
||||
imu_zero = None
|
||||
if imu_dim > 0:
|
||||
from imu_norm import normalize_imu
|
||||
# 原始姿态(roll=5°, pitch=-3°, 其余 0)→ 部署版输入布局 (1,1,1,imu_dim)
|
||||
imu_real = normalize_imu([5.0, -3.0] + [0.0] * 6, imu_dim).reshape(
|
||||
1, 1, 1, imu_dim)
|
||||
imu_zero = np.zeros((1, 1, 1, imu_dim), dtype=np.float32)
|
||||
|
||||
|
||||
def frames_nhwc(imgs):
|
||||
"""列表帧 → (1,H,W,3) 归一化 float32 NHWC × frame_len"""
|
||||
chw = np.stack([tf(Image.fromarray(im)).numpy() for im in imgs])
|
||||
return [chw[i].transpose(1, 2, 0)[None] for i in range(len(imgs))]
|
||||
|
||||
|
||||
def run_onnx(imgs, hist, pred, imu=None):
|
||||
feeds = {f"frame{i}": f for i, f in enumerate(frames_nhwc(imgs))}
|
||||
feeds["action_history"] = hist[None, None]
|
||||
feeds["pred_token"] = pred
|
||||
if imu_dim > 0:
|
||||
feeds["imu"] = imu if imu is not None else imu_zero
|
||||
return sess.run(None, feeds)[0].reshape(-1)
|
||||
|
||||
|
||||
def run_torch(imgs, hist, imu=None):
|
||||
x = torch.stack([tf(Image.fromarray(im)) for im in imgs]).unsqueeze(0)
|
||||
hh = torch.tensor(hist).unsqueeze(0)
|
||||
ii = None
|
||||
if imu_dim > 0:
|
||||
ii = torch.tensor(
|
||||
imu if imu is not None else imu_zero).reshape(1, imu_dim).float()
|
||||
with torch.no_grad():
|
||||
return model(x, hh, ii).squeeze(0).numpy()
|
||||
|
||||
|
||||
a = run_onnx(imgs_real, hist_real, pred_learned, imu_real)
|
||||
b = run_onnx(imgs_zero, hist_real, pred_learned, imu_real)
|
||||
c = run_onnx(imgs_real, hist_zero, pred_learned, imu_real)
|
||||
d = run_onnx(imgs_zero, hist_zero, pred_learned)
|
||||
p0 = run_onnx(imgs_real, hist_real, pred_none, imu_real)
|
||||
|
||||
print(f"onnx real_img+real_hist: {np.round(a[:6], 3)}...")
|
||||
print(f"onnx pred_token=0 diff : max={np.abs(a - p0).max():.4f}")
|
||||
print(f"img sensitivity (a-b): max={np.abs(a - b).max():.4f}")
|
||||
print(f"hist sensitivity (a-c): max={np.abs(a - c).max():.4f}")
|
||||
|
||||
ok = np.abs(a - b).max() > 1e-4 and np.abs(a - c).max() > 1e-4
|
||||
if imu_dim > 0:
|
||||
e = run_onnx(imgs_real, hist_real, pred_learned, imu_zero)
|
||||
print(f"imu sensitivity (a-e): max={np.abs(a - e).max():.4f}")
|
||||
|
||||
ref = run_torch(imgs_real, hist_real, imu_real)
|
||||
print(f"torch reference : {np.round(ref[:6], 3)}...")
|
||||
print(f"onnx vs torch max_diff : {np.abs(a - ref).max():.6e}")
|
||||
|
||||
ok = ok and np.abs(a - ref).max() < 1e-4
|
||||
print("PASS" if ok else "FAIL")
|
||||
+5
-1
@@ -11,6 +11,10 @@ elif [ "${1:-}" = "pretrain" ]; then
|
||||
elif [ "${1:-}" = "export" ]; then
|
||||
shift
|
||||
exec python3 export_onnx.py "$@"
|
||||
else
|
||||
elif [ "${1:-}" = "predict" ]; then
|
||||
shift
|
||||
exec python3 predict.py "$@"
|
||||
else
|
||||
echo "未知子命令: ${1:-}(可用: train / pretrain / export / predict)" >&2
|
||||
exit 1
|
||||
fi
|
||||
@@ -16,6 +16,17 @@ pip install onnxscript onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple
|
||||
pip check
|
||||
mkdir -p checkpoints data
|
||||
|
||||
# R2-4 防呆:确认装到的是 CPU 版 torch——requirements.txt 含 torch,
|
||||
# 若被镜像源重装成 CUDA 版(无 GPU 部署会被 CUDA 依赖拖垮)在此拦截
|
||||
python3 -c "
|
||||
import torch
|
||||
assert torch.version.cuda is None, (
|
||||
f'检测到 CUDA 版 torch({torch.__version__}),'
|
||||
'请用 _set.sh 的 CPU 源重装: pip install torch torchvision '
|
||||
'--index-url https://download.pytorch.org/whl/cpu --force-reinstall')
|
||||
print('torch CPU build OK:', torch.__version__)
|
||||
"
|
||||
|
||||
# Pre-download MobileNetV2 weights for offline fine-tuning (optional, skip if backbone=simple)
|
||||
python3 -c "
|
||||
import torch, sys, os
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
"""动作向量固定物理范围归一化 —— 训练(dataset/pretrain_mt)与推理
|
||||
(predict / onnx_infer / hailo_infer / gen_calibration / export_deploy)共用
|
||||
|
||||
18 列动作向量的物理量纲(pybot 侧契约):
|
||||
0-9 舵机角度 [0, 180]°(pybot/hd/servo.py MAX_ANGLE=180.0 校准后全行程)
|
||||
10 spk.F 扬声器频率 [0, 5000] Hz(@Speaker max_f=5000)
|
||||
11 spk.delay 延时 [0, 2000] ms(脚本实用上限,std.mt 用 10-500)
|
||||
12-17 保留位(恒 0,hi-lo=1.0 使归一化后仍为 0;损失权重 0 不参与优化)
|
||||
|
||||
min-max 映射到 [0,1]:x_norm = clip((x - lo) / (hi - lo), 0, 1)。
|
||||
与 z-score 的区别:固定物理范围不依赖 action_stats.npy 数据统计文件,
|
||||
部署链零状态;clamp 防御真实数据越界(越界值截断到 [0,1],不扩散)。
|
||||
|
||||
仅依赖 numpy,保证 hailo_infer.py(HAT 部署环境无 torch)可直接复用。
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
# 各列物理下界(当前全 0;保留数组以便未来引入非零下界列)
|
||||
ACTION_LO = np.zeros(18, dtype=np.float32)
|
||||
# 各列物理上界(18 列标准布局;action_dim 不同用 action_bounds() 适配)
|
||||
ACTION_HI = np.array(
|
||||
[180.0] * 10 + # 0-9 舵机角
|
||||
[5000.0, 2000.0] + # 10 spk.F,11 spk.delay
|
||||
[1.0] * 6, # 12-17 保留位(恒 0,span=1.0 保持归一化后为 0)
|
||||
dtype=np.float32)
|
||||
|
||||
# 归一化方式:fixed_range 线性 min-max;log_range 在 spk 两列上用
|
||||
# log10(1+x) 对数尺度,修正低端聚集导致的精度错配
|
||||
# (10Hz 在线性尺度占 0.002、对数尺度占 0.27)
|
||||
RANGE_MODES = ("fixed_range", "log_range")
|
||||
LOG_COLS = (10, 11) # spk.F / spk.delay
|
||||
|
||||
|
||||
def is_range_mode(mode: str) -> bool:
|
||||
"""normalize_mode 是否为固定范围类(fixed_range / log_range)"""
|
||||
return mode in RANGE_MODES
|
||||
|
||||
|
||||
def action_bounds(action_dim: int = 18):
|
||||
"""(lo, hi) 物理界:action_dim ≤ 18 取前缀;> 18 的超出列按 [0, 1] 补足"""
|
||||
if action_dim <= len(ACTION_LO):
|
||||
return ACTION_LO[:action_dim], ACTION_HI[:action_dim]
|
||||
pad = action_dim - len(ACTION_LO)
|
||||
return (np.concatenate([ACTION_LO, np.zeros(pad, dtype=np.float32)]),
|
||||
np.concatenate([ACTION_HI, np.ones(pad, dtype=np.float32)]))
|
||||
|
||||
|
||||
def normalize_action(action, action_dim: int = None,
|
||||
mode: str = "fixed_range") -> np.ndarray:
|
||||
"""原始动作向量 → [0,1] 归一化向量(clamp 防御越界值)
|
||||
|
||||
action: (..., action_dim) 任意前导维;action_dim 缺省取末维长度。
|
||||
mode="fixed_range":全列线性 min-max;
|
||||
mode="log_range":LOG_COLS 列用 log10(1+x)(负值先截 0),其余列线性。
|
||||
"""
|
||||
a = np.asarray(action, dtype=np.float32)
|
||||
lo, hi = action_bounds(action_dim or a.shape[-1])
|
||||
out = np.clip((a - lo) / (hi - lo), 0.0, 1.0)
|
||||
if mode == "log_range":
|
||||
for c in LOG_COLS:
|
||||
if c < a.shape[-1]:
|
||||
hc = float(hi[c])
|
||||
col = np.clip(a[..., c], 0.0, None)
|
||||
out[..., c] = np.log10(col + 1.0) / np.log10(hc + 1.0)
|
||||
return out
|
||||
|
||||
|
||||
def denormalize_action(norm_action, action_dim: int = None,
|
||||
mode: str = "fixed_range") -> np.ndarray:
|
||||
"""[0,1] 归一化向量 → 原始动作向量(不 clamp,保留模型外推能力)
|
||||
|
||||
mode="log_range" 时 LOG_COLS 列指数还原:x = 10^(n·log10(hi+1)) - 1。
|
||||
"""
|
||||
n = np.asarray(norm_action, dtype=np.float32)
|
||||
lo, hi = action_bounds(action_dim or n.shape[-1])
|
||||
out = n * (hi - lo) + lo
|
||||
if mode == "log_range":
|
||||
for c in LOG_COLS:
|
||||
if c < n.shape[-1]:
|
||||
hc = float(hi[c])
|
||||
out[..., c] = (np.power(10.0, n[..., c] * np.log10(hc + 1.0))
|
||||
- 1.0)
|
||||
return out
|
||||
|
||||
|
||||
def default_column_weights(action_dim: int = 18) -> np.ndarray:
|
||||
"""逐列损失权重默认表:舵机 1.0(安全关键);spk.F / spk.delay 0.0
|
||||
(只影响声音/节奏不影响运动,推理端固定常量,见 bins.FIXED_OUTPUT);
|
||||
保留位 0.0(恒 0 纯噪声)。action_dim>18 的扩展列默认 1.0:超出
|
||||
当前 18 列语义布局的通道视为新执行器,静默给 0 会让其完全不参与
|
||||
训练(P2-a1)。"""
|
||||
w = np.zeros(action_dim, dtype=np.float32)
|
||||
w[:min(10, action_dim)] = 1.0
|
||||
if action_dim > 18:
|
||||
w[18:] = 1.0
|
||||
return w
|
||||
+134
@@ -0,0 +1,134 @@
|
||||
"""多头 one-hot 分箱 —— 18 维连续动作向量与多组 one-hot 分类目标互转
|
||||
|
||||
分箱设计:
|
||||
0-9 舵机角:均匀 1°/箱 → 180 箱(匹配 0.75° 微动作插值数据粒度)
|
||||
10 spk.F:对数分箱 log10(1+x) ∈ [0, log10(5001)] → 32 箱
|
||||
11 spk.delay:对数分箱 log10(1+x) ∈ [0, log10(2001)] → 32 箱
|
||||
12-17 保留位:无头(恒 0,不参与分类)
|
||||
|
||||
编码/解码与 action_norm 的关系:箱生活在原始物理空间(舵机箱=角度),
|
||||
decode 直接产出原始量纲,无需反归一化;输入历史仍按 normalize_mode
|
||||
归一化(与现状一致),仅监督目标换为 one-hot。
|
||||
|
||||
软标签:encode(x, sigma) 把真值 one-hot 高斯扩散到邻近箱,缓解
|
||||
804→1864 分类面的数据稀疏(多数箱极少被访问)。
|
||||
|
||||
仅依赖 numpy,推理/部署文件可直接复用。
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
# 每组 one-hot 的箱数(按列索引;保留位无头)
|
||||
BINS_PER_COL = [180] * 10 + [32, 32]
|
||||
# 有头的列
|
||||
HEAD_COLS = tuple(range(12))
|
||||
# 各组起始位置(组 i 占 [START[i], START[i+1]))
|
||||
START = np.concatenate([[0], np.cumsum(BINS_PER_COL)])
|
||||
TOTAL_BINS = int(START[-1]) # 1864
|
||||
|
||||
LOG_COLS = (10, 11) # 对数分箱列
|
||||
LOG_HI = {10: 5000.0, 11: 2000.0}
|
||||
|
||||
# 权重 0 列的输出固定值(原始量纲):spk.F 固定 0(只影响声音不影响
|
||||
# 运动,闭环中持续发声无意义);spk.delay 在步级数据中恒 10ms(微动作
|
||||
# 分解的隐式节拍)。无监督头在闭环中会输出中段垃圾污染历史,
|
||||
# 推理/反馈时直接固定为常量。其余权重 0 列(保留位)固定 0。
|
||||
FIXED_OUTPUT = {10: 0.0, 11: 10.0}
|
||||
|
||||
|
||||
def fixed_zero_columns(weights, action_dim: int = 18) -> np.ndarray:
|
||||
"""按列权重表返回权重<=0 列的固定输出向量(其余列 0)
|
||||
|
||||
与 default_column_weights 联动:列权重 0 → 推理输出固定为
|
||||
FIXED_OUTPUT 常量(缺省 0),避免无监督头在闭环反馈中输出垃圾。
|
||||
"""
|
||||
w = np.asarray(weights, dtype=np.float32)
|
||||
fixed = np.zeros(action_dim, dtype=np.float32)
|
||||
for i in range(min(action_dim, len(w))):
|
||||
if w[i] <= 0.0:
|
||||
fixed[i] = FIXED_OUTPUT.get(i, 0.0)
|
||||
return fixed
|
||||
|
||||
|
||||
def bin_edges(col: int) -> np.ndarray:
|
||||
"""列 col 的箱边界数组(长度 bins+1),原始量纲递增"""
|
||||
n = BINS_PER_COL[col]
|
||||
if col in LOG_COLS:
|
||||
hi = LOG_HI[col]
|
||||
lo_log, hi_log = 0.0, np.log10(hi + 1.0)
|
||||
# 对数空间均匀边界 → 原始空间箱边界
|
||||
return np.power(10.0, np.linspace(lo_log, hi_log, n + 1)) - 1.0
|
||||
lo, hi = 0.0, 180.0
|
||||
return np.linspace(lo, hi, n + 1)
|
||||
|
||||
|
||||
def bin_centers(col: int) -> np.ndarray:
|
||||
"""列 col 的箱中心(decode 期望值的坐标),原始量纲"""
|
||||
e = bin_edges(col)
|
||||
return 0.5 * (e[:-1] + e[1:])
|
||||
|
||||
|
||||
def encode(x: np.ndarray, action_dim: int = 18, sigma: float = 0.0,
|
||||
out: np.ndarray = None) -> np.ndarray:
|
||||
"""连续动作向量 (..., D) → 多组 one-hot (..., TOTAL_BINS)
|
||||
|
||||
sigma > 0 时启用软标签:真值箱高斯扩散 exp(-d²/2σ²),d=与真值箱的
|
||||
距离(按箱索引);归一化后总和为 1。保留位列(12+)无头,恒不编码。
|
||||
"""
|
||||
a = np.asarray(x, dtype=np.float32)
|
||||
flat = a.reshape(-1, a.shape[-1])
|
||||
n_rows = flat.shape[0]
|
||||
if out is None:
|
||||
out = np.zeros(a.shape[:-1] + (TOTAL_BINS,), dtype=np.float32)
|
||||
out_flat = out.reshape(n_rows, TOTAL_BINS)
|
||||
rows = np.arange(n_rows)
|
||||
for col in HEAD_COLS:
|
||||
if col >= flat.shape[1]:
|
||||
break
|
||||
edges = bin_edges(col)
|
||||
idx = np.clip(np.searchsorted(edges, flat[:, col], side="right") - 1,
|
||||
0, BINS_PER_COL[col] - 1)
|
||||
s, e = START[col], START[col + 1]
|
||||
if sigma <= 0.0:
|
||||
out_flat[rows, s + idx] = 1.0
|
||||
else:
|
||||
n = BINS_PER_COL[col]
|
||||
d = np.abs(np.arange(n)[None, :] - idx[:, None])
|
||||
g = np.exp(-(d * d) / (2.0 * sigma * sigma))
|
||||
g = g / g.sum(axis=1, keepdims=True)
|
||||
out_flat[:, s:s + n] = g
|
||||
return out
|
||||
|
||||
|
||||
def decode(logits: np.ndarray, action_dim: int = 18,
|
||||
temperature: float = 1.0) -> np.ndarray:
|
||||
"""多头 logits (..., TOTAL_BINS) → 连续动作向量 (..., D),原始量纲
|
||||
|
||||
逐头 softmax → 期望值 Σ pᵢ·箱中心(分辨率超越箱宽);保留位列填 0。
|
||||
temperature 用于探索采样(=1 为正常推理)。
|
||||
"""
|
||||
l = np.asarray(logits, dtype=np.float32)
|
||||
out = np.zeros(l.shape[:-1] + (action_dim,), dtype=np.float32)
|
||||
for col in HEAD_COLS:
|
||||
if col >= action_dim:
|
||||
break
|
||||
s, e = START[col], START[col + 1]
|
||||
z = l[..., s:e] / max(temperature, 1e-6)
|
||||
z = z - z.max(axis=-1, keepdims=True) # 数值稳定
|
||||
p = np.exp(z)
|
||||
p = p / p.sum(axis=-1, keepdims=True)
|
||||
out[..., col] = p @ bin_centers(col)
|
||||
return out
|
||||
|
||||
|
||||
def argmax_decode(logits: np.ndarray, action_dim: int = 18) -> np.ndarray:
|
||||
"""argmax 解码:每头取最大概率箱中心(确定性,量化到箱宽)"""
|
||||
l = np.asarray(logits, dtype=np.float32)
|
||||
out = np.zeros(l.shape[:-1] + (action_dim,), dtype=np.float32)
|
||||
for col in HEAD_COLS:
|
||||
if col >= action_dim:
|
||||
break
|
||||
s, e = START[col], START[col + 1]
|
||||
idx = np.argmax(l[..., s:e], axis=-1)
|
||||
out[..., col] = bin_centers(col)[idx]
|
||||
return out
|
||||
+181
-12
@@ -1,18 +1,48 @@
|
||||
import argparse
|
||||
from dataclasses import dataclass, asdict, field
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataConfig:
|
||||
image_width: int = 1920
|
||||
image_height: int = 1080
|
||||
image_width: int = 640
|
||||
image_height: int = 480
|
||||
action_dim: int = 18
|
||||
history_len: int = 5
|
||||
frame_len: int = 3
|
||||
data_dir: str = "./data"
|
||||
normalize_actions: bool = False
|
||||
normalize_actions: bool = True
|
||||
action_stats_path: Optional[str] = None
|
||||
# 归一化方式(仅 normalize_actions=True 时生效):
|
||||
# "zscore" 数据统计 z-score(现状,依赖 action_stats.npy)
|
||||
# "fixed_range" 固定物理范围 min-max → [0,1](免 stats 文件,零状态部署,
|
||||
# 常量表见 action_norm.py;越界值 clamp 防御)
|
||||
# "log_range" 同上,但 spk.F/spk.delay 两列用 log10(1+x) 对数尺度,
|
||||
# 修正低端聚集导致的精度错配(见实验报告)
|
||||
normalize_mode: str = "zscore"
|
||||
# 逐列损失权重(仅训练损失用,推理端无感知):None → action_norm
|
||||
# .default_column_weights(action_dim)(舵机 1.0 / spk 0.2 / 保留位 0.0)。
|
||||
# 显式指定时长度须等于 action_dim。
|
||||
column_weights: Optional[List[float]] = None
|
||||
# 多头 one-hot 分类输出(bins.py 分箱):True 时监督目标为多组 one-hot
|
||||
# (舵机 1°/箱 ×10、spk 对数 32 箱 ×2,共 1864 logits),损失用加权 CE,
|
||||
# 推理端期望值解码。要求 normalize_actions=True 且 normalize_mode 为
|
||||
# 范围类(fixed_range/log_range)——历史输入沿用范围归一化,免 stats。
|
||||
# 与 predict_delta 互斥(残差在连续空间定义)。
|
||||
onehot: bool = False
|
||||
# one-hot 软标签高斯宽度(箱):0=硬标签;>0 真值扩散到邻近箱,
|
||||
# 缓解小样本下分类面数据稀疏(1°/箱 × 804 样本多数箱极少被访问)
|
||||
onehot_sigma: float = 1.0
|
||||
# 动作历史角度抖动增强(原始量纲比例):>0 时对输入历史加
|
||||
# 高斯噪声,σ = action_jitter × 列物理范围(舵机 0.01 → ±1.8°)。
|
||||
# 制造虚拟邻近状态,对抗微步插值样本高度相关导致的有效多样性
|
||||
# 不足;且与闭环自预测历史噪声同源,增强反馈鲁棒性。目标不变。
|
||||
action_jitter: float = 0.0
|
||||
# R2 方案一:MPU6050 姿态作为独立 IMU token 注入(序列 F+H+2)。
|
||||
# 0 = 关闭(默认,序列保持 F+H+1,旧 checkpoint 完全兼容);
|
||||
# 8×N = N 个传感器的 roll,pitch,ax,ay,az,gx,gy,gz 按 pybot 脚本
|
||||
# @MPU6050_BIND 行顺序拼接(最多 4 个 → 32,pybot recorder 的 imu.npy 格式)。
|
||||
imu_dim: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -28,16 +58,32 @@ class ModelConfig:
|
||||
|
||||
@dataclass
|
||||
class TrainConfig:
|
||||
batch_size: int = 32
|
||||
# 默认输入分辨率 640×480(1080p 采集的预处理缩放在实机测速中主导
|
||||
# 端到端延迟,见测速报告 §6.1);更高分辨率训练时需调小 batch
|
||||
batch_size: int = 4
|
||||
epochs: int = 100
|
||||
lr: float = 1e-3
|
||||
weight_decay: float = 1e-4
|
||||
val_split: float = 0.15
|
||||
# 验证集划分方式:
|
||||
# "episode" 整集划分(防数据泄漏,但单 episode 验证噪声大)
|
||||
# "sample" 样本级随机划分(默认:多集混合、验证集更大更稳,
|
||||
# 代价是相邻高相关样本泄漏使 val 略乐观)
|
||||
val_split_mode: str = "sample"
|
||||
checkpoint_dir: str = "./checkpoints"
|
||||
patience: int = 15
|
||||
num_workers: int = 2
|
||||
seed: int = 42
|
||||
max_grad_norm: float = 1.0
|
||||
augment: bool = False
|
||||
# R1-C 残差预测:True 时 target 改为 Δ=actions[t]-actions[t-1],
|
||||
# 推理端 â_t = a_{t-1} + Δ̂(惯性先验,平滑动作)。默认 False 兼容旧 checkpoint。
|
||||
predict_delta: bool = False
|
||||
# R1-B 多步展开训练:rollout_len>1 时 dataset 返回 K 步连续序列,
|
||||
# 训练用 scheduled sampling(前 ss_prob 比例步用真历史,其余用自预测回传),
|
||||
# 治 exposure bias。rollout_len=1 时退化为现状单步。
|
||||
rollout_len: int = 1
|
||||
ss_prob: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -68,6 +114,25 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser:
|
||||
parser.add_argument("--data-dir", type=str, default=None)
|
||||
parser.add_argument("--normalize-actions", action=argparse.BooleanOptionalAction, default=None)
|
||||
parser.add_argument("--action-stats-path", type=str, default=None)
|
||||
parser.add_argument("--normalize-mode", type=str, default=None,
|
||||
choices=["zscore", "fixed_range", "log_range"],
|
||||
help="动作归一化方式:zscore=数据统计(依赖 action_stats.npy);"
|
||||
"fixed_range=固定物理范围 min-max→[0,1](免 stats 文件);"
|
||||
"log_range=spk 两列对数尺度(免 stats 文件)")
|
||||
parser.add_argument("--onehot", action=argparse.BooleanOptionalAction,
|
||||
default=None,
|
||||
help="多头 one-hot 分类输出(bins.py 分箱,舵机 1°/箱)")
|
||||
parser.add_argument("--onehot-sigma", type=float, default=None,
|
||||
help="one-hot 软标签高斯宽度(箱,0=硬标签,默认 1.0)")
|
||||
parser.add_argument("--action-jitter", type=float, default=None,
|
||||
help="动作历史角度抖动(原始量纲比例,0=关;"
|
||||
"如 0.01 = 舵机 ±1.8°)")
|
||||
parser.add_argument("--column-weights", type=str, default=None,
|
||||
help="逐列损失权重(逗号分隔 action_dim 个浮点),"
|
||||
"缺省用默认表(舵机1.0/spk0.2/保留位0.0);仅影响训练损失")
|
||||
parser.add_argument("--imu-dim", type=int, default=None,
|
||||
help="R2 方案一:MPU6050 姿态维度(0=关闭,默认;"
|
||||
"8×N=N 个传感器拼接(≤4 个),序列变 F+H+2)")
|
||||
parser.add_argument("--backbone", type=str, default=None, choices=["simple", "mobilenet"])
|
||||
parser.add_argument("--cnn-feature-dim", type=int, default=None)
|
||||
parser.add_argument("--fused-hidden-dim", type=int, default=None)
|
||||
@@ -80,11 +145,24 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser:
|
||||
parser.add_argument("--lr", type=float, default=None)
|
||||
parser.add_argument("--weight-decay", type=float, default=None)
|
||||
parser.add_argument("--val-split", type=float, default=None)
|
||||
parser.add_argument("--val-split-mode", type=str, default=None,
|
||||
choices=["episode", "sample"],
|
||||
help="验证集划分:episode=整集(防泄漏但噪声大);"
|
||||
"sample=样本级随机(默认,更大更稳)")
|
||||
parser.add_argument("--checkpoint-dir", type=str, default=None)
|
||||
parser.add_argument("--patience", type=int, default=None)
|
||||
parser.add_argument("--num-workers", type=int, default=None)
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument("--max-grad-norm", type=float, default=None)
|
||||
parser.add_argument("--augment", action=argparse.BooleanOptionalAction, default=None,
|
||||
help="训练时启用安全图像增强(ColorJitter:亮度/对比度/饱和度/色调)")
|
||||
parser.add_argument("--predict-delta", action=argparse.BooleanOptionalAction, default=None,
|
||||
help="R1-C 残差预测:target 改为 Δ=actions[t]-actions[t-1],"
|
||||
"推理端 â_t=a_{t-1}+Δ̂(惯性先验平滑动作,默认关兼容旧 checkpoint)")
|
||||
parser.add_argument("--rollout-len", type=int, default=None,
|
||||
help="R1-B 多步展开训练步数 K(默认 1=单步;K>1 时 scheduled sampling)")
|
||||
parser.add_argument("--ss-prob", type=float, default=None,
|
||||
help="scheduled sampling 真实历史概率(仅 rollout-len>1 生效,默认 1.0=全真)")
|
||||
parser.add_argument("--init-from", type=str, default=None,
|
||||
help="预训练 checkpoint 路径,作为微调起点(strict=False 加载兼容权重)")
|
||||
parser.add_argument("--resume", type=str, default=None,
|
||||
@@ -106,6 +184,7 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
("history_len", data_cfg.history_len, 1, 1024, int),
|
||||
("frame_len", data_cfg.frame_len, 1, 1024, int),
|
||||
("action_dim", data_cfg.action_dim, 1, 1024, int),
|
||||
("imu_dim", data_cfg.imu_dim, 0, 64, int),
|
||||
("cnn_feature_dim", model_cfg.cnn_feature_dim, 1, 4096, int),
|
||||
("fused_hidden_dim", model_cfg.fused_hidden_dim, 1, 4096, int),
|
||||
("dropout", model_cfg.dropout, 0.0, 1.0, float),
|
||||
@@ -127,6 +206,42 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
f"attn_dim={model_cfg.attn_dim} 必须能被 attn_heads="
|
||||
f"{model_cfg.attn_heads} 整除")
|
||||
|
||||
if not (data_cfg.normalize_mode in ("zscore", "fixed_range", "log_range")):
|
||||
raise ValueError(
|
||||
f"normalize_mode={data_cfg.normalize_mode!r} 非法,"
|
||||
f"应为 'zscore'、'fixed_range' 或 'log_range'")
|
||||
|
||||
if data_cfg.onehot:
|
||||
if not data_cfg.normalize_actions:
|
||||
raise ValueError("onehot=True 要求 normalize_actions=True"
|
||||
"(历史输入须范围归一化)")
|
||||
if not (data_cfg.normalize_mode in ("fixed_range", "log_range")):
|
||||
raise ValueError(
|
||||
f"onehot=True 要求 normalize_mode 为范围类,"
|
||||
f"实际为 {data_cfg.normalize_mode!r}")
|
||||
if train_cfg.predict_delta:
|
||||
raise ValueError("onehot 与 predict_delta 互斥"
|
||||
"(残差定义在连续空间)")
|
||||
if not (0.0 <= data_cfg.action_jitter <= 1.0):
|
||||
raise ValueError(
|
||||
f"action_jitter={data_cfg.action_jitter} 应在 [0, 1] 内")
|
||||
if not (0.0 <= data_cfg.onehot_sigma <= 10.0):
|
||||
raise ValueError(
|
||||
f"onehot_sigma={data_cfg.onehot_sigma} 应在 [0, 10] 内")
|
||||
|
||||
if data_cfg.column_weights is not None:
|
||||
w = data_cfg.column_weights
|
||||
if len(w) != data_cfg.action_dim:
|
||||
raise ValueError(
|
||||
f"column_weights 长度 {len(w)} 应等于 action_dim="
|
||||
f"{data_cfg.action_dim}")
|
||||
if any((not isinstance(v, (int, float))) or isinstance(v, bool)
|
||||
or not (0.0 <= v <= 1e6) for v in w):
|
||||
raise ValueError(
|
||||
f"column_weights 元素应为 [0, 1e6] 内的数值,实际: {w}")
|
||||
if not any(v > 0 for v in w):
|
||||
raise ValueError("column_weights 全 0:损失恒 0,训练无梯度")
|
||||
|
||||
if train_cfg is not None:
|
||||
t_checks = [
|
||||
("batch_size", train_cfg.batch_size, 1, 65536, int),
|
||||
@@ -137,6 +252,8 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
("patience", train_cfg.patience, 0, 100000, int),
|
||||
("num_workers", train_cfg.num_workers, 0, 256, int),
|
||||
("max_grad_norm", train_cfg.max_grad_norm, 0.0, 1e6, float),
|
||||
("rollout_len", train_cfg.rollout_len, 1, 1024, int),
|
||||
("ss_prob", train_cfg.ss_prob, 0.0, 1.0, float),
|
||||
]
|
||||
for name, val, lo, hi, typ in t_checks:
|
||||
if typ is int:
|
||||
@@ -148,24 +265,63 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
f"配置 {name}={val} 超出合理范围 [{lo}, {hi}]")
|
||||
|
||||
|
||||
def warn_init_from_mismatch(init_cfg: dict, current: Config):
|
||||
"""--init-from 加载预训练权重前的兼容性警告(train.py / pretrain_mt.py 共用)
|
||||
|
||||
normalize_actions 不一致时输入分布与权重不匹配,预训练收益大打折扣;
|
||||
action_dim 不一致时 act_proj 无法加载对应权重(strict=False 静默跳过,
|
||||
该层随机初始化)。
|
||||
"""
|
||||
if not isinstance(init_cfg, dict):
|
||||
return
|
||||
init_data = init_cfg.get("data") or {}
|
||||
init_norm = init_data.get("normalize_actions", False)
|
||||
if init_norm != current.data.normalize_actions:
|
||||
print(f"WARNING: init-from checkpoint 的 normalize_actions="
|
||||
f"{init_norm} 与当前运行 {current.data.normalize_actions} "
|
||||
f"不一致,输入分布不匹配会损失预训练收益(见 README 预训练说明)")
|
||||
init_mode = init_data.get("normalize_mode", "zscore") # 旧 checkpoint 无此键 → zscore
|
||||
if init_norm and current.data.normalize_actions and init_mode != current.data.normalize_mode:
|
||||
print(f"WARNING: init-from checkpoint 的 normalize_mode={init_mode} "
|
||||
f"与当前运行 {current.data.normalize_mode} 不一致,"
|
||||
f"归一化空间不同会损失预训练收益")
|
||||
init_dim = init_data.get("action_dim")
|
||||
if init_dim is not None and init_dim != current.data.action_dim:
|
||||
print(f"WARNING: init-from checkpoint 的 action_dim={init_dim} "
|
||||
f"与当前运行 {current.data.action_dim} 不一致,"
|
||||
f"act_proj 权重将随机初始化")
|
||||
init_imu = init_data.get("imu_dim", 0) # 旧 checkpoint 无此键 → 0
|
||||
if init_imu != current.data.imu_dim:
|
||||
print(f"WARNING: init-from checkpoint 的 imu_dim={init_imu} "
|
||||
f"与当前运行 {current.data.imu_dim} 不一致,"
|
||||
f"imu_proj/pos_embed 权重将随机初始化")
|
||||
|
||||
|
||||
def parse_checkpoint_config(cfg_dict: dict) -> Config:
|
||||
"""从 checkpoint 恢复 Config,校验键集合与值域
|
||||
|
||||
- 顶层缺 data/model 键 → ValueError(带明确信息)
|
||||
- data/model 段含未知键 → ValueError(列出未知键,避免裸 TypeError)
|
||||
- data/model/train 段含未知键 → ValueError(列出未知键,避免裸 TypeError)
|
||||
- 值域非法 → ValueError(复用 validate_config_values)
|
||||
"""
|
||||
if not isinstance(cfg_dict, dict):
|
||||
raise ValueError(f"checkpoint config 应为 dict,实际为 {type(cfg_dict).__name__}")
|
||||
for section in ("data", "model"):
|
||||
section_dataclasses = {
|
||||
"data": DataConfig,
|
||||
"model": ModelConfig,
|
||||
"train": TrainConfig,
|
||||
}
|
||||
for section, dataclass in section_dataclasses.items():
|
||||
if section not in cfg_dict:
|
||||
raise ValueError(f"checkpoint config 缺少 '{section}' 段")
|
||||
if section in ("data", "model"):
|
||||
raise ValueError(f"checkpoint config 缺少 '{section}' 段")
|
||||
continue
|
||||
section_dict = cfg_dict[section]
|
||||
if not isinstance(section_dict, dict):
|
||||
raise ValueError(f"checkpoint config['{section}'] 应为 dict,实际为 {type(section_dict).__name__}")
|
||||
known = set(DataConfig.__dataclass_fields__) if section == "data" \
|
||||
else set(ModelConfig.__dataclass_fields__)
|
||||
unknown = set(section_dict) - known
|
||||
raise ValueError(
|
||||
f"checkpoint config['{section}'] 应为 dict,"
|
||||
f"实际为 {type(section_dict).__name__}")
|
||||
unknown = set(section_dict) - set(dataclass.__dataclass_fields__)
|
||||
if unknown:
|
||||
raise ValueError(
|
||||
f"checkpoint config['{section}'] 含未知键: {sorted(unknown)}")
|
||||
@@ -184,6 +340,12 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
|
||||
"data_dir": ("data", "data_dir"),
|
||||
"normalize_actions": ("data", "normalize_actions"),
|
||||
"action_stats_path": ("data", "action_stats_path"),
|
||||
"normalize_mode": ("data", "normalize_mode"),
|
||||
"column_weights": ("data", "column_weights"),
|
||||
"onehot": ("data", "onehot"),
|
||||
"onehot_sigma": ("data", "onehot_sigma"),
|
||||
"action_jitter": ("data", "action_jitter"),
|
||||
"imu_dim": ("data", "imu_dim"),
|
||||
"backbone": ("model", "cnn_backbone"),
|
||||
"cnn_feature_dim": ("model", "cnn_feature_dim"),
|
||||
"fused_hidden_dim": ("model", "fused_hidden_dim"),
|
||||
@@ -196,15 +358,22 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
|
||||
"lr": ("train", "lr"),
|
||||
"weight_decay": ("train", "weight_decay"),
|
||||
"val_split": ("train", "val_split"),
|
||||
"val_split_mode": ("train", "val_split_mode"),
|
||||
"checkpoint_dir": ("train", "checkpoint_dir"),
|
||||
"patience": ("train", "patience"),
|
||||
"num_workers": ("train", "num_workers"),
|
||||
"seed": ("train", "seed"),
|
||||
"max_grad_norm": ("train", "max_grad_norm"),
|
||||
"augment": ("train", "augment"),
|
||||
"predict_delta": ("train", "predict_delta"),
|
||||
"rollout_len": ("train", "rollout_len"),
|
||||
"ss_prob": ("train", "ss_prob"),
|
||||
}
|
||||
args_dict = vars(args)
|
||||
for arg_name, (section, field) in mapping.items():
|
||||
val = args_dict.get(arg_name)
|
||||
if val is not None:
|
||||
if field == "column_weights" and isinstance(val, str):
|
||||
val = [float(v) for v in val.split(",")]
|
||||
setattr(getattr(cfg, section), field, val)
|
||||
return cfg
|
||||
+178
-18
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
@@ -7,6 +8,12 @@ from PIL import Image
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision import transforms
|
||||
from torchvision.transforms import functional as TF
|
||||
|
||||
from imu_norm import imu_scales
|
||||
from action_norm import (normalize_action, is_range_mode,
|
||||
action_bounds)
|
||||
from bins import encode as bin_encode
|
||||
|
||||
|
||||
def _natural_key(p):
|
||||
@@ -16,17 +23,42 @@ def _natural_key(p):
|
||||
|
||||
class ActionDataset(Dataset):
|
||||
def __init__(self, data_dir: str, history_len: int, action_dim: int,
|
||||
image_width: int = 1920, image_height: int = 1080,
|
||||
image_width: int = 640, image_height: int = 480,
|
||||
normalize_actions: bool = False,
|
||||
action_stats_path: str = None, episode_indices: list = None,
|
||||
frame_len: int = 3):
|
||||
frame_len: int = 3, sample_indices: list = None,
|
||||
augment: bool = False, action_jitter: float = 0.0,
|
||||
predict_delta: bool = False, rollout_len: int = 1,
|
||||
imu_dim: int = 0, normalize_mode: str = "zscore",
|
||||
onehot: bool = False, onehot_sigma: float = 1.0):
|
||||
self.history_len = history_len
|
||||
self.action_dim = action_dim
|
||||
# 固定范围类归一化(fixed_range / log_range,免 stats 文件);
|
||||
# 仅 normalize_actions=True 生效
|
||||
self._range_mode = (normalize_mode
|
||||
if (normalize_actions and is_range_mode(normalize_mode))
|
||||
else None)
|
||||
# 多头 one-hot:监督目标为多组 one-hot(bins.py),历史仍走 _norm
|
||||
self.onehot = onehot
|
||||
self.onehot_sigma = onehot_sigma
|
||||
# 动作历史角度抖动(原始空间加噪,σ ∝ 列物理范围)
|
||||
self.action_jitter = action_jitter
|
||||
# R2 方案一:>0 时样本带 "imu" 键(归一化姿态向量,缺失 episode 补零)
|
||||
self.imu_dim = imu_dim
|
||||
self.image_width = image_width
|
||||
self.image_height = image_height
|
||||
self.frame_len = frame_len
|
||||
self.augment = augment
|
||||
# R1-C 残差预测:target 改为 Δ=actions[t]-actions[t-1](归一化空间内求差)
|
||||
self.predict_delta = predict_delta
|
||||
# R1-B 多步展开:rollout_len>1 时每个样本返回 K 步连续 (images,history,target)
|
||||
self.rollout_len = max(1, rollout_len)
|
||||
|
||||
self.transform = transforms.Compose([
|
||||
# Resize 不保持宽高比(直接拉伸)是有意为之:推理端
|
||||
# (predict/onnx_infer/_onnx_validate)使用完全相同的预处理,
|
||||
# 训练/推理口径必须一致;改为 letterbox 会使既有 checkpoint
|
||||
# 输入分布失配,如需变更须训练/推理同步修改并重训(P2-d)
|
||||
transforms.Resize((image_height, image_width)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
@@ -35,7 +67,10 @@ class ActionDataset(Dataset):
|
||||
episodes = ActionDataset._valid_episode_dirs(data_dir)
|
||||
if episode_indices is not None:
|
||||
episodes = [episodes[i] for i in episode_indices]
|
||||
self._sample_indices = sample_indices
|
||||
|
||||
# 样本存 (image_files, actions, t_start):t_start 是 rollout 第一步的目标时刻,
|
||||
# __getitem__ 按 rollout_len 展开 K 步。存原始数组引用避免重复加载/拷贝。
|
||||
self.samples = []
|
||||
for ep_dir in episodes:
|
||||
actions_path = ep_dir / "actions.npy"
|
||||
@@ -55,20 +90,33 @@ class ActionDataset(Dataset):
|
||||
print(f"WARNING: {ep_dir} 动作帧数 {n_frames} 与图像数 {n_images} 不一致,"
|
||||
f"按较少者({min(n_frames, n_images)})使用,请检查录制完整性")
|
||||
n = min(n_frames, n_images)
|
||||
# 每个样本:帧 [t-frame_len+1, t](旧→新)+ 动作历史 [t-history_len, t) → 动作 t
|
||||
imus = self._load_imu(ep_dir, n)
|
||||
# 第 k 步(k=0..K-1)目标时刻 t=t_start+k:帧 [t-frame_len+1, t]、
|
||||
# 历史 [t-history_len, t)、目标 actions[t]。t_start 下界保证第 0 步窗口完整,
|
||||
# 上界保证第 K-1 步目标不越界。
|
||||
start = max(history_len, frame_len - 1)
|
||||
for t in range(start, n):
|
||||
self.samples.append((image_files[t - frame_len + 1: t + 1],
|
||||
actions[t - history_len:t], actions[t]))
|
||||
for t_start in range(start, n - (self.rollout_len - 1)):
|
||||
self.samples.append((image_files, actions, imus, t_start))
|
||||
|
||||
# 样本级随机划分:按样本索引过滤(相邻微步高度相关,随机划分
|
||||
# 会让 val 略乐观——高相关样本的近邻大量在训练集,见训练报告)
|
||||
if sample_indices is not None:
|
||||
idx_set = set(sample_indices)
|
||||
self.samples = [s for i, s in enumerate(self.samples)
|
||||
if i in idx_set]
|
||||
|
||||
self._action_mean = None
|
||||
self._action_std = None
|
||||
if normalize_actions:
|
||||
if normalize_actions and not self._range_mode:
|
||||
if action_stats_path:
|
||||
if not os.path.exists(action_stats_path):
|
||||
raise FileNotFoundError(
|
||||
f"指定的 action_stats_path 不存在: {action_stats_path}")
|
||||
stats = np.load(action_stats_path)
|
||||
if stats.shape != (2, self.action_dim):
|
||||
raise ValueError(
|
||||
f"action stats 形状应为 (2, {self.action_dim}),"
|
||||
f"实际为 {stats.shape}: {action_stats_path}")
|
||||
self._action_mean = stats[0]
|
||||
self._action_std = stats[1]
|
||||
else:
|
||||
@@ -76,26 +124,138 @@ class ActionDataset(Dataset):
|
||||
raise ValueError(
|
||||
f"数据集为空(0 样本),无法计算归一化统计量。"
|
||||
f"请检查 data_dir={data_dir} 中的 episode 数据。")
|
||||
all_actions = np.stack([s[2] for s in self.samples])
|
||||
# 归一化统计基于绝对动作(与 Δ 无关):覆盖所有样本 rollout 范围内
|
||||
# 出现的 history 与 target 动作,保持与旧版一致的统计口径。
|
||||
all_actions = np.stack([
|
||||
s[1][s[3] - self.history_len: s[3] + self.rollout_len]
|
||||
for s in self.samples
|
||||
]).reshape(-1, self.action_dim)
|
||||
self._action_mean = all_actions.mean(axis=0)
|
||||
self._action_std = all_actions.std(axis=0) + 1e-8
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_paths, action_history, target = self.samples[idx]
|
||||
def _jit_history(self, hist_raw):
|
||||
"""对原始动作历史加高斯抖动(σ = action_jitter × 列物理范围)
|
||||
|
||||
制造虚拟邻近状态对抗微步样本相关性;与闭环自预测历史噪声同源。
|
||||
目标(真值)不抖动——监督信号保持干净。
|
||||
"""
|
||||
if self.action_jitter <= 0.0:
|
||||
return hist_raw
|
||||
lo, hi = action_bounds(self.action_dim)
|
||||
sigma = self.action_jitter * (hi - lo)
|
||||
noise = np.random.normal(0.0, sigma, hist_raw.shape).astype(np.float32)
|
||||
return hist_raw + noise
|
||||
|
||||
def _norm(self, a):
|
||||
"""动作归一化(numpy → torch 前在 numpy 空间做,保持 float32)
|
||||
|
||||
fixed_range / log_range:固定物理范围 → [0,1](clamp 防御,免 stats);
|
||||
否则 z-score(_action_mean 为 None 时直通原始值)。
|
||||
"""
|
||||
if self._range_mode:
|
||||
return normalize_action(a, self.action_dim, mode=self._range_mode)
|
||||
if self._action_mean is None:
|
||||
return a
|
||||
return (a - self._action_mean) / self._action_std
|
||||
|
||||
def _load_frames(self, image_files, t):
|
||||
"""加载时刻 t 的 F 帧窗口 [t-frame_len+1, t](旧→新),augment 对整个窗口一致"""
|
||||
paths = image_files[t - self.frame_len + 1: t + 1]
|
||||
if self.augment:
|
||||
bf = random.uniform(0.8, 1.2)
|
||||
cf = random.uniform(0.8, 1.2)
|
||||
sf = random.uniform(0.8, 1.2)
|
||||
hf = random.uniform(-0.05, 0.05)
|
||||
frames = []
|
||||
for p in image_paths:
|
||||
for p in paths:
|
||||
image = Image.open(p).convert("RGB")
|
||||
if self.augment:
|
||||
image = TF.adjust_brightness(image, bf)
|
||||
image = TF.adjust_contrast(image, cf)
|
||||
image = TF.adjust_saturation(image, sf)
|
||||
image = TF.adjust_hue(image, hf)
|
||||
frames.append(self.transform(image))
|
||||
images = torch.stack(frames) # (F, 3, S, S)
|
||||
action_history = torch.tensor(action_history, dtype=torch.float32)
|
||||
target = torch.tensor(target, dtype=torch.float32)
|
||||
if self._action_mean is not None:
|
||||
action_history = (action_history - self._action_mean) / self._action_std
|
||||
target = (target - self._action_mean) / self._action_std
|
||||
return {"images": images, "action_history": action_history, "target": target}
|
||||
return torch.stack(frames) # (F, 3, H, W)
|
||||
|
||||
def _step_target(self, actions, t):
|
||||
"""时刻 t 的监督目标;onehot 时为多组 one-hot(原始空间分箱),
|
||||
predict_delta 时为 Δ=actions[t]-actions[t-1](归一化空间),
|
||||
否则为归一化连续值"""
|
||||
if self.onehot:
|
||||
return bin_encode(actions[t], self.action_dim,
|
||||
sigma=self.onehot_sigma)
|
||||
if self.predict_delta:
|
||||
return self._norm(actions[t]) - self._norm(actions[t - 1])
|
||||
return self._norm(actions[t])
|
||||
|
||||
def _load_imu(self, ep_dir, n):
|
||||
"""加载 episode 的 imu.npy 并按物理量程归一化(R2)
|
||||
|
||||
imu_dim=0(未启用)或文件缺失 → None(样本侧补零向量,等价无姿态
|
||||
先验,混合数据集可训)。形状 pad/truncate 到 (n, imu_dim)。
|
||||
"""
|
||||
if self.imu_dim <= 0:
|
||||
return None
|
||||
path = ep_dir / "imu.npy"
|
||||
if not path.exists():
|
||||
print(f"WARNING: {ep_dir} 无 imu.npy,该 episode 姿态按零向量处理")
|
||||
return None
|
||||
imus = np.load(path).astype(np.float32)
|
||||
if imus.ndim != 2:
|
||||
raise ValueError(
|
||||
f"{path} 形状应为 (T, imu_dim),实际为 {imus.shape}")
|
||||
if imus.shape[1] != self.imu_dim:
|
||||
fixed = np.zeros((imus.shape[0], self.imu_dim), dtype=np.float32)
|
||||
m = min(imus.shape[1], self.imu_dim)
|
||||
fixed[:, :m] = imus[:, :m]
|
||||
imus = fixed
|
||||
if len(imus) < n: # 与帧数对齐:不足补零
|
||||
imus = np.concatenate([
|
||||
imus, np.zeros((n - len(imus), self.imu_dim), dtype=np.float32)])
|
||||
return imus[:n] / imu_scales(self.imu_dim)
|
||||
|
||||
def _imu_at(self, imus, t):
|
||||
"""时刻 t 的归一化姿态向量(与最新帧对齐);无数据 → 零向量"""
|
||||
if imus is None:
|
||||
return torch.zeros(self.imu_dim, dtype=torch.float32)
|
||||
return torch.tensor(imus[t], dtype=torch.float32)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
image_files, actions, imus, t_start = self.samples[idx]
|
||||
if self.rollout_len == 1:
|
||||
# 单步:与旧版输出结构完全一致(imu_dim>0 时附加 "imu" 键)
|
||||
images = self._load_frames(image_files, t_start)
|
||||
history = torch.tensor(self._norm(self._jit_history(
|
||||
actions[t_start - self.history_len: t_start])),
|
||||
dtype=torch.float32)
|
||||
target = torch.tensor(self._step_target(actions, t_start),
|
||||
dtype=torch.float32)
|
||||
out = {"images": images, "action_history": history, "target": target}
|
||||
if self.imu_dim > 0:
|
||||
out["imu"] = self._imu_at(imus, t_start)
|
||||
return out
|
||||
|
||||
# 多步:返回 K 步序列,scheduled sampling 在 train 循环内按步取用
|
||||
steps = []
|
||||
for k in range(self.rollout_len):
|
||||
t = t_start + k
|
||||
step = {
|
||||
"images": self._load_frames(image_files, t),
|
||||
# 真实历史(scheduled sampling 的真值来源;首步必用)
|
||||
"action_history": torch.tensor(
|
||||
self._norm(self._jit_history(
|
||||
actions[t - self.history_len: t])), dtype=torch.float32),
|
||||
"target": torch.tensor(self._step_target(actions, t), dtype=torch.float32),
|
||||
# 上一步真实动作(Δ 残差的基准 a_{t-1},归一化空间)
|
||||
"prev_action": torch.tensor(self._norm(actions[t - 1]), dtype=torch.float32),
|
||||
}
|
||||
if self.imu_dim > 0:
|
||||
step["imu"] = self._imu_at(imus, t)
|
||||
steps.append(step)
|
||||
return {"steps": steps}
|
||||
|
||||
def get_action_stats(self):
|
||||
if self._action_mean is not None:
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
"""actpred 部署版 ONNX → HEF 编译脚本(Hailo DFC 5.3,仅 x86 Linux)
|
||||
|
||||
流程(对应 DFC 5.x CLI):
|
||||
hailo parser onnx → HAR(网络翻译)
|
||||
hailo optimize → 量化优化后的 HAR(需校准集 npz)
|
||||
hailo compiler → HEF
|
||||
|
||||
用法:
|
||||
# 1. 导出部署版 ONNX(DFC 友好图结构)
|
||||
python export_deploy.py --model checkpoints/pretrain_best_inference.pt
|
||||
# 2. 生成校准集(首次)
|
||||
python gen_calibration.py --deploy
|
||||
# 3. 编译
|
||||
python dfc_compile.py
|
||||
|
||||
关键点(踩坑记录):
|
||||
- 训练态 ONNX(5D 输入、AdaptiveAvgPool、MultiheadAttention 的 head
|
||||
reshape)DFC 无法解析,必须用 export_deploy.py 的改造版图
|
||||
- pred_token 必须作为图输入:常量 initializer 会被 DFC 误转置
|
||||
- 必须显式 --start/--end-node-names 与 --input-format,且禁用
|
||||
onnxsim(本图简化无收益,反而引入 reshape)
|
||||
- 校准集键名必须匹配 HN 层名 actpred_deploy/input_layerN
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tarfile
|
||||
import tempfile
|
||||
|
||||
|
||||
NET_NAME = "actpred_deploy"
|
||||
|
||||
|
||||
def start_nodes(frame_count: int, with_imu: bool = False) -> list:
|
||||
"""export_deploy.py 的输入命名约定:frame{i} × frame_count + action_history
|
||||
+ pred_token(R2 imu_dim>0 时末尾追加 imu)"""
|
||||
nodes = ([f"frame{i}" for i in range(frame_count)]
|
||||
+ ["action_history", "pred_token"])
|
||||
if with_imu:
|
||||
nodes.append("imu")
|
||||
return nodes
|
||||
|
||||
|
||||
def hailo(*args, venv_bin: str):
|
||||
cmd = [os.path.join(venv_bin, "hailo"), *args]
|
||||
env = dict(os.environ, PATH=venv_bin + os.pathsep + os.environ["PATH"])
|
||||
return subprocess.run(cmd, env=env, check=True)
|
||||
|
||||
|
||||
def _har_hn_layers(har_path: str) -> dict:
|
||||
"""解包 HAR 并读取主 hn 文件的 layers(优化产物含 fp/native 变体,取主图)"""
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
with tarfile.open(har_path) as tf:
|
||||
tf.extractall(tmp)
|
||||
hn_files = [f for f in os.listdir(tmp) if f.endswith(".hn")]
|
||||
if not hn_files:
|
||||
raise ValueError(f"{har_path} 中无 .hn 文件")
|
||||
main = [f for f in hn_files
|
||||
if not f.endswith(".fp.hn") and not f.endswith(".native.hn")]
|
||||
hn_file = main[0] if main else hn_files[0]
|
||||
with open(os.path.join(tmp, hn_file)) as f:
|
||||
return json.load(f)["layers"]
|
||||
|
||||
|
||||
def _input_layer_index(name: str) -> int:
|
||||
m = re.search(r"input_layer(\d+)$", name)
|
||||
return int(m.group(1)) if m else 0
|
||||
|
||||
|
||||
def har_input_layer_names(har_path: str) -> list:
|
||||
"""从 HAR 主 hn 图读输入层名,按 input_layerN 数字后缀排序"""
|
||||
layers = _har_hn_layers(har_path)
|
||||
names = [n for n, l in layers.items() if l["type"] == "input_layer"]
|
||||
names.sort(key=_input_layer_index)
|
||||
return names
|
||||
|
||||
|
||||
def check_input_connectivity(har_path: str, input_layers: list):
|
||||
"""检查 HAR 图中每个输入层是否可达任一输出层
|
||||
|
||||
输入层被优化断开时编译照常成功,但 HEF 输出恒定值且无任何报错
|
||||
(2026-08-13 实测踩坑),必须在编译前拦截。DFC 5.3 的 HAR 不含
|
||||
.edges 文件,连通性改从主 hn 图的 layer output 边 BFS 判定。
|
||||
"""
|
||||
layers = _har_hn_layers(har_path)
|
||||
adj = {n: l.get("output", []) for n, l in layers.items()}
|
||||
targets = {n for n, l in layers.items() if l["type"] == "output_layer"}
|
||||
if not targets:
|
||||
print(f"WARNING: {har_path} 图中无输出层,跳过连通性检查")
|
||||
return
|
||||
disconnected = []
|
||||
for name in input_layers:
|
||||
seen, stack = set(), [name]
|
||||
while stack:
|
||||
cur = stack.pop()
|
||||
if cur in seen:
|
||||
continue
|
||||
seen.add(cur)
|
||||
stack.extend(adj.get(cur, []))
|
||||
if not (seen & targets):
|
||||
disconnected.append(name)
|
||||
if disconnected:
|
||||
sys.exit(
|
||||
f"错误: {har_path} 中输入层不可达输出层(输入被优化断开,"
|
||||
f"编译出的 HEF 将输出恒定值): {disconnected}")
|
||||
print(f"输入连通性检查通过: {len(input_layers)} 个输入层均可达输出层")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Compile actpred deploy ONNX to HEF")
|
||||
parser.add_argument("--onnx", default="checkpoints/actpred_deploy.onnx")
|
||||
parser.add_argument("--calib-npz", default="data/calibration_deploy.npz")
|
||||
parser.add_argument("--output", default="checkpoints/actpred_pretrain.hef")
|
||||
parser.add_argument("--hw-arch", default="hailo10h",
|
||||
choices=["hailo15h", "hailo15m", "hailo15l",
|
||||
"hailo10h", "hailo10p"])
|
||||
parser.add_argument("--frame-count", type=int, default=3,
|
||||
help="帧输入个数(与导出 ONNX 的 frame_len 一致,"
|
||||
"默认 3)")
|
||||
parser.add_argument("--with-imu", action="store_true",
|
||||
help="R2:导出 ONNX 含 imu 输入时开启(start nodes "
|
||||
"末尾追加 imu,与 export_deploy.py 命名一致)")
|
||||
parser.add_argument("--venv-bin", default="venv_dfc/bin",
|
||||
help="安装了 hailo-sdk 的 venv bin 目录")
|
||||
args = parser.parse_args()
|
||||
|
||||
venv_bin = os.path.abspath(args.venv_bin)
|
||||
if not os.path.exists(os.path.join(venv_bin, "hailo")):
|
||||
sys.exit(f"错误: {venv_bin}/hailo 不存在,请先安装 hailo-sdk")
|
||||
|
||||
nodes = start_nodes(args.frame_count, args.with_imu)
|
||||
input_formats = ["NHWC"] * len(nodes)
|
||||
|
||||
workdir = tempfile.mkdtemp(prefix="dfc_")
|
||||
har_raw = os.path.join(workdir, "parsed.har")
|
||||
har_opt = os.path.join(workdir, "optimized.har")
|
||||
hef_dir = os.path.join(workdir, "hef")
|
||||
os.makedirs(hef_dir)
|
||||
|
||||
try:
|
||||
print("[1/3] parser: ONNX → HAR")
|
||||
hailo("parser", "onnx", args.onnx, "--hw-arch", args.hw_arch, "-y",
|
||||
"--har-path", har_raw,
|
||||
"--start-node-names", *nodes,
|
||||
"--end-node-names", "predicted_action",
|
||||
"--input-format",
|
||||
*[f"{n}={f}" for n, f in zip(nodes, input_formats)],
|
||||
"--disable-onnx-simplifier", venv_bin=venv_bin)
|
||||
|
||||
# 校验未被推荐机制截断:输出层应来自预测头(conv 而非 ew_add)
|
||||
names = har_input_layer_names(har_raw)
|
||||
if len(names) != len(nodes):
|
||||
sys.exit(f"错误: HAR 输入层数 {len(names)} != {len(nodes)},解析可能被截断")
|
||||
check_input_connectivity(har_raw, names)
|
||||
|
||||
print("[2/3] optimize: 量化校准")
|
||||
hailo("optimize", har_raw, "--hw-arch", args.hw_arch,
|
||||
"--calib-set-path", os.path.abspath(args.calib_npz),
|
||||
"--output-har-path", har_opt, venv_bin=venv_bin)
|
||||
check_input_connectivity(har_opt, har_input_layer_names(har_opt))
|
||||
|
||||
print("[3/3] compiler: HAR → HEF")
|
||||
hailo("compiler", har_opt, "--hw-arch", args.hw_arch,
|
||||
"--output-dir", hef_dir, venv_bin=venv_bin)
|
||||
|
||||
hef_src = os.path.join(hef_dir, f"{NET_NAME}.hef")
|
||||
shutil.copy(hef_src, args.output)
|
||||
shutil.copy(har_opt, args.output.replace(".hef", "_optimized.har"))
|
||||
print(f"HEF 已保存: {args.output} ({os.path.getsize(args.output)} bytes)")
|
||||
finally:
|
||||
shutil.rmtree(workdir, ignore_errors=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,263 @@
|
||||
"""R4 离线评估基准:开环滚动预测,量化闭环代理指标
|
||||
|
||||
现状只有开环单步 MSE(teacher forcing),无法反映真实闭环表现。本脚本在留出
|
||||
episode 上复现 pybot 闭环的开环滚动(图像用真值滚动、动作用自预测回写),
|
||||
输出三类闭环代理指标:
|
||||
|
||||
累积误差 逐步 MSE 及随步数的增长曲线(衡量误差累积/发散)
|
||||
平滑度 预测动作相邻差分方差(越低越平滑,对物理舵机越友好),与真值对比
|
||||
物理可行性 预测动作落在训练数据经验范围外的比例(越界率,越高越危险)
|
||||
|
||||
用法:
|
||||
python evaluate.py --model checkpoints/best_inference.pt --data-dir ./data
|
||||
python evaluate.py --model A.pt --compare B.pt --data-dir ./data # A/B 对比
|
||||
python evaluate.py --model m.pt --data-dir ./data --episodes 3 --horizon 50 \
|
||||
--output report.json
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from dataset import ActionDataset
|
||||
from predict import ActionPredictorInference
|
||||
|
||||
|
||||
def _episode_ids(data_dir):
|
||||
return [d.name for d in ActionDataset._valid_episode_dirs(data_dir)]
|
||||
|
||||
|
||||
def _load_episode(data_dir, ep_name):
|
||||
"""加载单个 episode 的图像路径(自然序)、动作真值 (T, D)、姿态真值 (T, imu_dim)
|
||||
|
||||
imu.npy 缺失时第三项返回 None(旧数据兼容,滚动时按零向量喂入)。
|
||||
姿态保持原始量纲——predict 端内部归一化,与 pybot 闭环口径一致。
|
||||
"""
|
||||
ep_dir = [d for d in ActionDataset._valid_episode_dirs(data_dir)
|
||||
if d.name == ep_name][0]
|
||||
from dataset import _natural_key
|
||||
images_dir = ep_dir / "images"
|
||||
image_files = sorted(list(images_dir.glob("*.jpg")) +
|
||||
list(images_dir.glob("*.png")), key=_natural_key)
|
||||
actions = np.load(ep_dir / "actions.npy").astype(np.float32)
|
||||
n = min(len(image_files), len(actions))
|
||||
imus = None
|
||||
imu_path = ep_dir / "imu.npy"
|
||||
if imu_path.exists():
|
||||
imus = np.load(imu_path).astype(np.float32)
|
||||
if imus.ndim != 2:
|
||||
raise ValueError(f"{imu_path} 形状应为 (T, imu_dim),实际为 {imus.shape}")
|
||||
imus = imus[:n]
|
||||
return image_files[:n], actions[:n], imus
|
||||
|
||||
|
||||
def _action_bounds(action_dim):
|
||||
"""固定物理界作为越界判据(与 action_norm.py 常量表共用)
|
||||
|
||||
旧版用数据经验界(min/max+margin),恒值列 span≈0 时任何非零预测都
|
||||
算越界(unnormalized 模型 oob_rate 曾恒 100%)。"越物理界"比"越经验界"
|
||||
更符合安全语义:舵机 [0,180]、spk.F [0,5000]、spk.delay [0,2000]。
|
||||
"""
|
||||
from action_norm import action_bounds
|
||||
return action_bounds(action_dim)
|
||||
|
||||
|
||||
@np.errstate(all="ignore")
|
||||
def rollout_episode(predictor, image_files, actions, horizon=None,
|
||||
lo=None, hi=None, imus=None):
|
||||
"""单 episode 开环滚动,复现闭环口径(图像真值滚动、动作自预测回写)。
|
||||
|
||||
从 t0 = max(history_len, frame_len-1) 起步(窗口完整处),每步:
|
||||
帧窗 = 真值图像 [t-frame_len+1, t](旧→新)
|
||||
历史 = 前 H 步动作(首步用真值填充,之后用自预测回写,复刻 run_loop)
|
||||
姿态 = imus[t](predictor.imu_dim 非 None 时传入;imus 为 None 或
|
||||
长度不足按零向量 = 无姿态先验)
|
||||
预测 â_t → 与真值 a_t 比较 → 回写历史
|
||||
返回逐步记录 dict。horizon=None 表示滚到 episode 末尾。
|
||||
"""
|
||||
F, H = predictor.frame_len, predictor.history_len
|
||||
use_imu = getattr(predictor, "imu_dim", None) is not None
|
||||
t0 = max(H, F - 1)
|
||||
n = len(actions)
|
||||
if horizon is None:
|
||||
horizon = n - t0
|
||||
horizon = min(horizon, n - t0)
|
||||
if horizon <= 0:
|
||||
return None
|
||||
|
||||
cache = {}
|
||||
|
||||
def load(t):
|
||||
if t not in cache:
|
||||
cache[t] = np.array(Image.open(image_files[t]).convert("RGB"))
|
||||
return cache[t]
|
||||
|
||||
history = [actions[t0 - H + k].copy() for k in range(H)] # 真值冷启动
|
||||
steps = []
|
||||
for k in range(horizon):
|
||||
t = t0 + k
|
||||
frames = [load(t - F + 1 + i) for i in range(F)]
|
||||
kwargs = {}
|
||||
if use_imu:
|
||||
kwargs["imu"] = (imus[t] if imus is not None and t < len(imus)
|
||||
else None)
|
||||
pred = np.asarray(
|
||||
predictor.predict(frames, np.array(history), **kwargs),
|
||||
dtype=np.float32).reshape(-1)
|
||||
true = actions[t]
|
||||
err = pred - true
|
||||
rec = {
|
||||
"t": int(t),
|
||||
"mse": float(np.mean(err ** 2)),
|
||||
"pred": pred,
|
||||
"true": true,
|
||||
}
|
||||
if lo is not None:
|
||||
rec["oob"] = bool(np.any((pred < lo) | (pred > hi)))
|
||||
steps.append(rec)
|
||||
history = history[1:] + [pred.copy()] # 自预测回写(复刻 run_loop)
|
||||
return steps
|
||||
|
||||
|
||||
def _smoothness(seq):
|
||||
"""相邻差分方差(多维取均值);序列 <2 返回 0。越低越平滑。"""
|
||||
seq = np.asarray(seq, dtype=np.float32)
|
||||
if len(seq) < 2:
|
||||
return 0.0
|
||||
return float(np.mean(np.var(np.diff(seq, axis=0), axis=0)))
|
||||
|
||||
|
||||
def summarize(all_steps, lo, hi):
|
||||
"""汇总多个 episode 的滚动记录为报告指标
|
||||
|
||||
error_growth 改为按 episode 归一化后取中位数:
|
||||
每 episode 内部算 (末25% MSE / 首25% MSE) → 中位数聚合。
|
||||
避免短 episode 尖峰被长 episode 均值淹没,也更鲁棒于长度差异。
|
||||
"""
|
||||
if not all_steps:
|
||||
raise ValueError("无可评估的滚动步(episode 太短或 horizon 过小)")
|
||||
flat = [s for steps in all_steps for s in steps]
|
||||
mses = np.array([s["mse"] for s in flat], dtype=np.float32)
|
||||
|
||||
# 误差随步数增长曲线:按相对步位置对齐各 episode,取每位置的平均 MSE
|
||||
max_len = max(len(s) for s in all_steps)
|
||||
curve = []
|
||||
for k in range(max_len):
|
||||
vals = [s[k]["mse"] for s in all_steps if k < len(s)]
|
||||
if vals:
|
||||
curve.append(float(np.mean(vals)))
|
||||
|
||||
# 逐 episode 归一化 error_growth:每 episode 自身首末段比值 → 中位数
|
||||
growths = []
|
||||
for steps in all_steps:
|
||||
n = len(steps)
|
||||
if n < 2:
|
||||
continue
|
||||
m = np.array([s["mse"] for s in steps], dtype=np.float32)
|
||||
front = np.mean(m[:max(1, n // 4)])
|
||||
back = np.mean(m[max(0, n * 3 // 4):])
|
||||
growths.append(back / max(front, 1e-12))
|
||||
error_growth = float(np.median(growths)) if growths else 1.0
|
||||
|
||||
preds = np.array([s["pred"] for s in flat], dtype=np.float32)
|
||||
trues = np.array([s["true"] for s in flat], dtype=np.float32)
|
||||
oob = [s.get("oob", False) for s in flat]
|
||||
|
||||
return {
|
||||
"n_episodes": len(all_steps),
|
||||
"n_steps": len(flat),
|
||||
"overall_mse": float(np.mean(mses)),
|
||||
# 末段(后 25% 步)相对首段(前 25%)的误差比,>1 表示误差随步数发散
|
||||
"error_growth": error_growth,
|
||||
"error_curve": curve,
|
||||
"pred_smoothness": _smoothness(preds),
|
||||
"true_smoothness": _smoothness(trues),
|
||||
# 平滑度比:<1 表示预测比真值还平滑(好),>1 表示预测更抖(坏)
|
||||
"smoothness_ratio": _smoothness(preds) / max(_smoothness(trues), 1e-12),
|
||||
"oob_rate": float(np.mean(oob)) if lo is not None else None,
|
||||
}
|
||||
|
||||
|
||||
def evaluate(checkpoint, data_dir, episodes=None, horizon=None, device=None):
|
||||
predictor = ActionPredictorInference(checkpoint, device)
|
||||
ep_names = _episode_ids(data_dir)
|
||||
if episodes is not None:
|
||||
ep_names = ep_names[-episodes:] # 取尾部 episode 作为留出集
|
||||
lo, hi = _action_bounds(predictor.action_dim)
|
||||
all_steps = []
|
||||
for name in ep_names:
|
||||
image_files, actions, imus = _load_episode(data_dir, name)
|
||||
steps = rollout_episode(predictor, image_files, actions, horizon,
|
||||
lo, hi, imus=imus)
|
||||
if steps:
|
||||
all_steps.append(steps)
|
||||
rep = summarize(all_steps, lo, hi)
|
||||
rep["checkpoint"] = checkpoint
|
||||
rep["frame_len"] = predictor.frame_len
|
||||
rep["history_len"] = predictor.history_len
|
||||
rep["imu_dim"] = getattr(predictor, "imu_dim", None)
|
||||
rep["predict_delta"] = getattr(predictor, "_predict_delta", False)
|
||||
return rep
|
||||
|
||||
|
||||
def _fmt(rep):
|
||||
lines = [
|
||||
f"checkpoint : {rep['checkpoint']}",
|
||||
f"predict_delta : {rep['predict_delta']}",
|
||||
f"imu_dim : {rep.get('imu_dim')}",
|
||||
f"episodes/steps : {rep['n_episodes']} / {rep['n_steps']}",
|
||||
f"overall_mse : {rep['overall_mse']:.6f}",
|
||||
f"error_growth : {rep['error_growth']:.3f} (>1 = 误差随步数发散)",
|
||||
f"smoothness_ratio: {rep['smoothness_ratio']:.3f} (<1 = 比真值更平滑)",
|
||||
f" pred/true : {rep['pred_smoothness']:.4f} / {rep['true_smoothness']:.4f}",
|
||||
]
|
||||
if rep["oob_rate"] is not None:
|
||||
lines.append(f"oob_rate : {rep['oob_rate']:.4f} (越界率, 越低越安全)")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description="actpred 离线闭环代理评估(R4)")
|
||||
p.add_argument("--model", required=True, help="主 checkpoint 路径")
|
||||
p.add_argument("--data-dir", required=True, help="episode 数据目录")
|
||||
p.add_argument("--compare", default=None, help="可选对比 checkpoint(A/B)")
|
||||
p.add_argument("--episodes", type=int, default=None,
|
||||
help="评估用尾部 episode 数(默认全部)")
|
||||
p.add_argument("--horizon", type=int, default=None,
|
||||
help="每 episode 滚动步数上限(默认滚到末尾)")
|
||||
p.add_argument("--device", default=None)
|
||||
p.add_argument("--output", default=None, help="保存 JSON 报告路径")
|
||||
args = p.parse_args()
|
||||
|
||||
rep = evaluate(args.model, args.data_dir, args.episodes, args.horizon,
|
||||
args.device)
|
||||
print("== 主模型 ==")
|
||||
print(_fmt(rep))
|
||||
|
||||
result = {"main": rep}
|
||||
if args.compare:
|
||||
rep_b = evaluate(args.compare, args.data_dir, args.episodes,
|
||||
args.horizon, args.device)
|
||||
print("\n== 对比模型 ==")
|
||||
print(_fmt(rep_b))
|
||||
# 关键指标对比(负值=主模型更优)
|
||||
result["compare"] = rep_b
|
||||
result["delta"] = {
|
||||
"overall_mse": rep["overall_mse"] - rep_b["overall_mse"],
|
||||
"error_growth": rep["error_growth"] - rep_b["error_growth"],
|
||||
"smoothness_ratio": rep["smoothness_ratio"] - rep_b["smoothness_ratio"],
|
||||
}
|
||||
d = result["delta"]
|
||||
print("\n== 主 − 对比 (负值=主模型更优) ==")
|
||||
print(f"overall_mse : {d['overall_mse']:+.6f}")
|
||||
print(f"error_growth : {d['error_growth']:+.3f}")
|
||||
print(f"smoothness_ratio: {d['smoothness_ratio']:+.3f}")
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||||
print(f"\n报告已保存: {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,317 @@
|
||||
"""导出 Hailo 部署版 ONNX:DFC 友好图结构
|
||||
|
||||
训练态 ONNX 对 DFC 不友好:5D 图像输入、AdaptiveAvgPool2d(被
|
||||
torch.export 展开成 Loop 节点)、MultiheadAttention 的 head 拆分
|
||||
Reshape 序列。此脚本逐项改造(权重数值完全不变):
|
||||
1. 图像输入拆为 F 个独立 (1,H,W,3) NHWC 输入,按帧展开 CNN 分支,
|
||||
各分支持有独立权重副本(共享 initializer 会被 DFC 误融合)
|
||||
2. AdaptiveAvgPool2d 换成等价的固定 kernel AvgPool2d
|
||||
3. 自注意力手工展开为每头独立 Linear(128→32),无 head reshape;
|
||||
pre-LN 残差结构与训练态 nn.TransformerEncoderLayer(norm_first=True)
|
||||
逐位等价
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
|
||||
# 一致性校验阈值:wrapper 改造/ORT 转换与原模型输出差异超过即拦截退出,
|
||||
# 防止 ManualAttnBlock 依赖的 PyTorch 布局约定变更静默穿透(v3-P2-2)
|
||||
MAX_TOLERATED_DIFF = 1e-4
|
||||
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
MAX_TOKEN_ABS = 100.0
|
||||
MERGE_RATIO_WARN = 32.0
|
||||
|
||||
|
||||
class ManualAttnBlock(nn.Module):
|
||||
"""等价替换 nn.TransformerEncoderLayer(d=128, nhead=4, ff=256, norm_first=True)"""
|
||||
|
||||
def __init__(self, layer: nn.TransformerEncoderLayer):
|
||||
super().__init__()
|
||||
d = layer.self_attn.embed_dim
|
||||
nh = layer.self_attn.num_heads
|
||||
self.head_dim = d // nh
|
||||
self.nh = nh
|
||||
self.scale = self.head_dim ** -0.5
|
||||
|
||||
ipw = layer.self_attn.in_proj_weight # (3d, d) = [Wq; Wk; Wv]
|
||||
ipb = layer.self_attn.in_proj_bias
|
||||
self.q_proj = nn.ModuleList()
|
||||
self.k_proj = nn.ModuleList()
|
||||
self.v_proj = nn.ModuleList()
|
||||
for h in range(nh):
|
||||
sl = slice(h * self.head_dim, (h + 1) * self.head_dim)
|
||||
slices = (sl, slice(d + sl.start, d + sl.stop),
|
||||
slice(2 * d + sl.start, 2 * d + sl.stop))
|
||||
for proj, ws in ((self.q_proj, slices[0]), (self.k_proj, slices[1]),
|
||||
(self.v_proj, slices[2])):
|
||||
lin = nn.Linear(d, self.head_dim)
|
||||
lin.weight.data.copy_(ipw[ws])
|
||||
proj.append(lin)
|
||||
for h in range(nh):
|
||||
self.q_proj[h].bias.data.copy_(ipb[h * self.head_dim:(h + 1) * self.head_dim])
|
||||
self.k_proj[h].bias.data.copy_(ipb[d + h * self.head_dim:d + (h + 1) * self.head_dim])
|
||||
self.v_proj[h].bias.data.copy_(ipb[2 * d + h * self.head_dim:2 * d + (h + 1) * self.head_dim])
|
||||
|
||||
self.out_proj = layer.self_attn.out_proj
|
||||
self.norm1 = layer.norm1
|
||||
self.norm2 = layer.norm2
|
||||
self.fc1 = layer.linear1
|
||||
self.fc2 = layer.linear2
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# x: (1, 1, S, d),seq 占 W 维,pre-LN 自注意力 + 残差
|
||||
n = self.norm1(x)
|
||||
heads = []
|
||||
for h in range(self.nh):
|
||||
q = self.q_proj[h](n)
|
||||
k = self.k_proj[h](n)
|
||||
v = self.v_proj[h](n)
|
||||
s = torch.matmul(q, k.transpose(-1, -2)) * self.scale
|
||||
a = torch.softmax(s, dim=-1)
|
||||
heads.append(torch.matmul(a, v))
|
||||
attn = self.out_proj(torch.cat(heads, dim=3)) # feature 维 concat
|
||||
x = x + attn
|
||||
x = x + self.fc2(torch.relu(self.fc1(self.norm2(x))))
|
||||
return x
|
||||
|
||||
|
||||
class DeployWrapper(nn.Module):
|
||||
def __init__(self, model: ActionPredictor):
|
||||
super().__init__()
|
||||
import copy
|
||||
m = model
|
||||
if not hasattr(m.image_encoder, "net"):
|
||||
raise ValueError(
|
||||
"部署版导出仅支持 backbone=simple(MobileNet 的 DFC 改造未实现)")
|
||||
ref_cnn = nn.Sequential(*list(m.image_encoder.net.children())[:-1])
|
||||
self.cnns = nn.ModuleList(
|
||||
[copy.deepcopy(ref_cnn) for _ in range(m._frame_len)])
|
||||
# img_proj ∘ AvgPool == AvgPool ∘ Conv1x1(均为线性算子,可交换)。
|
||||
# 把 img_proj 的权重折进每分支末尾的 1×1 卷积,偏置单独用加法,
|
||||
# 避免 DFC 对 3D Linear 的错误卷积融合
|
||||
self.proj_convs = nn.ModuleList()
|
||||
for _ in range(m._frame_len):
|
||||
c = nn.Conv2d(m.img_proj.in_features, m.img_proj.out_features, 1, bias=False)
|
||||
c.weight.data.copy_(m.img_proj.weight.data.view(
|
||||
m.img_proj.out_features, m.img_proj.in_features, 1, 1))
|
||||
self.proj_convs.append(c)
|
||||
self.proj_bias = nn.Parameter(m.img_proj.bias.data.clone())
|
||||
# 固定 kernel AvgPool 须与训练态 AdaptiveAvgPool2d(1) 逐值等价:
|
||||
# kernel 取 ref_cnn 的真实输出尺寸。H//16 只在分辨率能被 16 整除
|
||||
# 时成立(1080 高度下卷积输出 68 行,按 67 取 kernel 会静默丢弃
|
||||
# 末行,破坏与训练图的数值一致)
|
||||
with torch.no_grad():
|
||||
probe = ref_cnn(torch.zeros(
|
||||
1, 3, m._image_height, m._image_width))
|
||||
self.pool = nn.AvgPool2d((probe.shape[2], probe.shape[3]))
|
||||
# act_proj 权重注册为常量,用显式 MatMul+Add 代替 Linear,
|
||||
# 避免 DFC 对 3D Linear 的错误卷积融合
|
||||
self.act_w = nn.Parameter(m.act_proj.weight.data.t().clone()) # (D, d)
|
||||
self.act_b = nn.Parameter(m.act_proj.bias.data.clone())
|
||||
# R2 方案一:IMU 独立 token(imu_dim>0 时追加图输入,同 MatMul+Add 手法)
|
||||
self._imu_dim = m._imu_dim
|
||||
if self._imu_dim > 0:
|
||||
self.imu_w = nn.Parameter(m.imu_proj.weight.data.t().clone())
|
||||
self.imu_b = nn.Parameter(m.imu_proj.bias.data.clone())
|
||||
# pred_token 改为图输入(常量 initializer 会被 DFC 误转置),
|
||||
# 位置编码经 concat 后的 ew_add 广播可被正确推断
|
||||
self.pos_embed4 = nn.Parameter(m.pos_embed.data.unsqueeze(1).clone()) # (1, 1, S, d)
|
||||
self.blocks = nn.ModuleList(
|
||||
[ManualAttnBlock(l) for l in model.encoder.layers])
|
||||
self.head = m.head
|
||||
self._frame_len = m._frame_len
|
||||
|
||||
def _frame_feat(self, frame_nhwc: torch.Tensor, i: int) -> torch.Tensor:
|
||||
x = self.cnns[i](frame_nhwc.permute(0, 3, 1, 2))
|
||||
x = self.proj_convs[i](x) # 1x1 conv ≡ 池化后的 img_proj
|
||||
return self.pool(x) # (1, d, 1, 1) NCHW
|
||||
|
||||
def forward(self, *inputs: torch.Tensor) -> torch.Tensor:
|
||||
# inputs = frame0..frame{F-1}, action_history, pred_token, [imu]
|
||||
# 全程 (1, 1, S, d) 布局:seq 占 W 维,避免任何 transpose
|
||||
frames = list(inputs[:self._frame_len])
|
||||
action_history = inputs[self._frame_len]
|
||||
pred_token = inputs[self._frame_len + 1]
|
||||
imu = inputs[self._frame_len + 2] if self._imu_dim > 0 else None
|
||||
feats = [self._frame_feat(frames[i], i).permute(0, 2, 3, 1)
|
||||
for i in range(self._frame_len)] # (1, 1, 1, d) NHWC
|
||||
img_tokens = torch.cat(feats, dim=2) + self.proj_bias # (1, 1, F, d)
|
||||
act_tokens = torch.matmul(action_history, self.act_w) + self.act_b # (1, 1, H, d)
|
||||
parts = [img_tokens, act_tokens]
|
||||
if self._imu_dim > 0:
|
||||
# imu: (1, 1, 1, imu_dim) 归一化姿态 → (1, 1, 1, d),位于 PRED 之前
|
||||
parts.append(torch.matmul(imu, self.imu_w) + self.imu_b)
|
||||
parts.append(pred_token)
|
||||
tokens = torch.cat(parts, dim=2) # (1, 1, S, d)
|
||||
tokens = tokens + self.pos_embed4
|
||||
for blk in self.blocks:
|
||||
tokens = blk(tokens)
|
||||
return self.head(tokens[:, :, -1:, :]).squeeze(1) # (1, 1, action_dim)
|
||||
|
||||
|
||||
def check_token_ranges(wrapper: "DeployWrapper", frame_len: int, height: int,
|
||||
width: int, history_len: int, action_dim: int,
|
||||
pred_token: torch.Tensor, max_abs: float = MAX_TOKEN_ABS,
|
||||
action_stats_path: str = None,
|
||||
normalize_mode: str = "zscore"):
|
||||
"""导出前检查 token 量级:幅度过大时 uint8 量化会把小信号压成 0
|
||||
|
||||
预训练用全零图像占位 → BN running_var 停在初始小值 → 推理时每级
|
||||
BN 放大 ~50 倍,图像 token 可达 ±2e5。DFC 在 token merge 点按最大
|
||||
幅度分配共享量化 scale(实测 scale=1965.92,pred_token/act_tokens
|
||||
全部量化成 0),HEF 输出对任何输入恒定。必须在编译前拦截。
|
||||
|
||||
归一化模型(normalize_actions)的 HEF 输入是归一化后的历史,
|
||||
检查输入须与推理端一致(zscore → z 空间采样;fixed_range /
|
||||
log_range → [0,1])。
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
raw = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
|
||||
norm = (raw.astype(np.float32) / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
|
||||
frames = [torch.tensor(norm[None]) for _ in range(frame_len)]
|
||||
history = rng.uniform(0.0, 500.0, size=(1, 1, history_len, action_dim))
|
||||
from action_norm import is_range_mode
|
||||
if is_range_mode(normalize_mode):
|
||||
# fixed_range / log_range 归一化空间即 [0,1],直接在归一化空间采样
|
||||
history = rng.uniform(0.0, 1.0, size=(1, 1, history_len, action_dim))
|
||||
elif action_stats_path is not None:
|
||||
stats = np.load(action_stats_path)
|
||||
if stats.shape != (2, action_dim):
|
||||
raise ValueError(
|
||||
f"action stats 形状应为 (2, {action_dim}),"
|
||||
f"实际为 {stats.shape}: {action_stats_path}")
|
||||
# z-score 归一化模型直接在 z 空间采样:raw 均匀采样会在零方差列上
|
||||
# 产生天文数字 z 值(实际推理中这些列恒为 0,z 恒为 0)
|
||||
history = rng.normal(0.0, 1.5, size=(1, 1, history_len, action_dim))
|
||||
history = torch.tensor(history, dtype=torch.float32)
|
||||
with torch.no_grad():
|
||||
feats = [wrapper._frame_feat(f, i).permute(0, 2, 3, 1)
|
||||
for i, f in enumerate(frames)]
|
||||
img_tokens = torch.cat(feats, dim=2) + wrapper.proj_bias
|
||||
act_tokens = torch.matmul(history, wrapper.act_w) + wrapper.act_b
|
||||
img_max = img_tokens.abs().max().item()
|
||||
act_max = act_tokens.abs().max().item()
|
||||
pred_max = pred_token.abs().max().item()
|
||||
print(f"token 量级检查: |img|={img_max:.4g} |act|={act_max:.4g} "
|
||||
f"|pred|={pred_max:.4g}(图像上限 {max_abs})")
|
||||
if img_max > max_abs:
|
||||
raise SystemExit(
|
||||
f"错误: 图像 token 幅度 {img_max:.3g} 超过可量化上限 {max_abs}。"
|
||||
f"checkpoint 的 BN 统计退化(预训练用全零图像占位,running_var "
|
||||
f"停在初始小值),uint8 量化会把小信号全部压成 0、HEF 输出恒定。"
|
||||
f"请先用真实/随机图像数据微调后再导出。")
|
||||
if act_max > MERGE_RATIO_WARN * max(img_max, 1.0):
|
||||
print(f"WARNING: act_tokens 幅度 {act_max:.3g} 远大于 img_tokens "
|
||||
f"{img_max:.3g},merge 点共享 scale 会粗量化图像路径;"
|
||||
f"建议用 --normalize-actions 训练(动作归一化)压小 act_tokens")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Export Hailo-friendly ONNX")
|
||||
parser.add_argument("--model", required=True, help="checkpoint path")
|
||||
parser.add_argument("--output", default="checkpoints/actpred_deploy.onnx")
|
||||
parser.add_argument("--pred-token-out", default=None,
|
||||
help="保存学习到的 pred_token 到 .npy(HEF 推理端喂入用,"
|
||||
"配合抖动校准避免常量量化退化)")
|
||||
args = parser.parse_args()
|
||||
|
||||
checkpoint = torch.load(args.model, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(checkpoint["config"])
|
||||
|
||||
model = ActionPredictor(cfg.model, cfg.data, pretrained=False)
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
model.eval()
|
||||
wrapper = DeployWrapper(model).eval()
|
||||
|
||||
h, w = cfg.data.image_height, cfg.data.image_width
|
||||
f = cfg.data.frame_len
|
||||
stats_path = None
|
||||
from action_norm import is_range_mode
|
||||
if cfg.data.normalize_actions and not is_range_mode(cfg.data.normalize_mode):
|
||||
# 与 predict.py 语义一致:相对路径按 checkpoint 目录解析
|
||||
# (fixed_range / log_range 模式免 stats 文件,跳过)
|
||||
stats_path = cfg.data.action_stats_path
|
||||
if not stats_path:
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)),
|
||||
"action_stats.npy")
|
||||
elif not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)), stats_path)
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(
|
||||
f"模型训练时启用了动作归一化,但未找到统计文件: {stats_path}")
|
||||
check_token_ranges(wrapper, f, h, w, cfg.data.history_len,
|
||||
cfg.data.action_dim, model.pred_token.data.unsqueeze(1),
|
||||
action_stats_path=stats_path,
|
||||
normalize_mode=cfg.data.normalize_mode)
|
||||
dummy_frames = [torch.randn(1, h, w, 3) for _ in range(f)]
|
||||
dummy_history = torch.randn(1, 1, cfg.data.history_len, cfg.data.action_dim)
|
||||
dummy_pred = model.pred_token.data.unsqueeze(1).clone() # (1, 1, 1, d)
|
||||
imu_dim = cfg.data.imu_dim
|
||||
# imu_dim>0 时追加归一化姿态输入(均匀采样的归一化值域)
|
||||
dummy_imu = torch.randn(1, 1, 1, imu_dim) if imu_dim > 0 else None
|
||||
|
||||
extra_inputs = (dummy_imu,) if imu_dim > 0 else ()
|
||||
extra_names = ["imu"] if imu_dim > 0 else []
|
||||
|
||||
with torch.no_grad():
|
||||
ref = wrapper(*dummy_frames, dummy_history, dummy_pred, *extra_inputs)
|
||||
|
||||
torch.onnx.export(
|
||||
wrapper, (*dummy_frames, dummy_history, dummy_pred, *extra_inputs),
|
||||
args.output,
|
||||
input_names=[f"frame{i}" for i in range(f)] + ["action_history", "pred_token"] + extra_names,
|
||||
output_names=["predicted_action"],
|
||||
opset_version=18,
|
||||
)
|
||||
print(f"部署版 ONNX 已导出: {args.output}({f} 个帧输入 (1,{h},{w},3) NHWC + action_history + pred_token"
|
||||
f"{' + imu' if imu_dim > 0 else ''})")
|
||||
|
||||
if args.pred_token_out:
|
||||
import numpy as np
|
||||
np.save(args.pred_token_out, model.pred_token.data.unsqueeze(1).numpy())
|
||||
print(f"pred_token 已保存: {args.pred_token_out} (1,1,1,{model.pred_token.shape[-1]})")
|
||||
|
||||
stacked = torch.stack([fr.permute(0, 3, 1, 2) for fr in dummy_frames], dim=1)
|
||||
with torch.no_grad():
|
||||
orig_out = model(stacked, dummy_history.squeeze(1),
|
||||
dummy_imu.squeeze(1).squeeze(1) if imu_dim > 0 else None)
|
||||
wrap_diff = (ref.squeeze(1) - orig_out).abs().max().item()
|
||||
print(f"wrapper 与原模型最大差异: {wrap_diff:.6e}")
|
||||
if wrap_diff > MAX_TOLERATED_DIFF:
|
||||
sys.exit(f"错误: wrapper 与原模型差异 {wrap_diff:.3e} > "
|
||||
f"{MAX_TOLERATED_DIFF:.0e}(ManualAttnBlock 布局约定可能已失效)")
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
import numpy as np
|
||||
sess = ort.InferenceSession(args.output)
|
||||
feeds = {f"frame{i}": dummy_frames[i].numpy() for i in range(f)}
|
||||
feeds["action_history"] = dummy_history.numpy()
|
||||
feeds["pred_token"] = dummy_pred.numpy()
|
||||
if imu_dim > 0:
|
||||
feeds["imu"] = dummy_imu.numpy()
|
||||
ort_out = sess.run(None, feeds)[0]
|
||||
ort_diff = np.abs(ort_out - ref.numpy()).max()
|
||||
if ort_diff > MAX_TOLERATED_DIFF:
|
||||
sys.exit(f"错误: ONNX 与 wrapper 差异 {ort_diff:.3e} > "
|
||||
f"{MAX_TOLERATED_DIFF:.0e}")
|
||||
print(f"ONNX 一致性验证通过(最大差异: {ort_diff:.6e})")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+45
-11
@@ -1,9 +1,14 @@
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
import torch
|
||||
from torch.export import Dim
|
||||
from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
|
||||
# 一致性校验阈值:ORT 转换与原模型输出差异超过即拦截退出(v3-P2-2)
|
||||
MAX_TOLERATED_DIFF = 1e-4
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Export actpred model to ONNX")
|
||||
@@ -23,32 +28,61 @@ def main():
|
||||
dummy_images = torch.randn(
|
||||
1, data_cfg.frame_len, 3, data_cfg.image_height, data_cfg.image_width)
|
||||
dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
|
||||
# R2 方案一:imu_dim>0 时追加归一化姿态输入 (1, imu_dim)
|
||||
imu_dim = data_cfg.imu_dim
|
||||
extra_inputs = (torch.randn(1, imu_dim),) if imu_dim > 0 else ()
|
||||
extra_names = ["imu"] if imu_dim > 0 else []
|
||||
|
||||
batch = Dim("batch")
|
||||
dynamic = {"images": {0: batch}, "action_history": {0: batch}}
|
||||
if imu_dim > 0:
|
||||
dynamic["imu"] = {0: batch}
|
||||
torch.onnx.export(
|
||||
model, (dummy_images, dummy_history),
|
||||
model, (dummy_images, dummy_history, *extra_inputs),
|
||||
args.output,
|
||||
input_names=["images", "action_history"],
|
||||
input_names=["images", "action_history"] + extra_names,
|
||||
output_names=["predicted_action"],
|
||||
dynamic_shapes={
|
||||
"images": {0: batch},
|
||||
"action_history": {0: batch},
|
||||
},
|
||||
dynamic_shapes=dynamic,
|
||||
opset_version=18,
|
||||
)
|
||||
print(f"ONNX model exported to: {args.output}")
|
||||
print(f"ONNX model exported to: {args.output}"
|
||||
f"{'(含 imu 输入)' if imu_dim > 0 else ''}")
|
||||
|
||||
# 写入训练配置元数据:onnx_infer 据此自动识别归一化方式与残差模式,
|
||||
# 避免推理端配置与训练不一致导致静默错误(normalize_actions=False 时
|
||||
# 记为 "none",推理端不做任何归一化)
|
||||
import onnx
|
||||
m = onnx.load(args.output)
|
||||
meta = {
|
||||
"normalize_mode": (data_cfg.normalize_mode
|
||||
if data_cfg.normalize_actions else "none"),
|
||||
"predict_delta": ("1" if getattr(cfg.train, "predict_delta", False)
|
||||
else "0"),
|
||||
"imu_dim": str(imu_dim),
|
||||
"action_dim": str(data_cfg.action_dim),
|
||||
}
|
||||
for k, v in meta.items():
|
||||
p = m.metadata_props.add()
|
||||
p.key, p.value = k, v
|
||||
onnx.save(m, args.output)
|
||||
print(f"ONNX 元数据已写入: {meta}")
|
||||
|
||||
# 可选:验证导出的 ONNX 模型一致性
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
import numpy as np
|
||||
sess = ort.InferenceSession(args.output)
|
||||
img_np = dummy_images.numpy()
|
||||
hist_np = dummy_history.numpy()
|
||||
ort_out = sess.run(None, {"images": img_np, "action_history": hist_np})[0]
|
||||
feeds = {"images": dummy_images.numpy(),
|
||||
"action_history": dummy_history.numpy()}
|
||||
for name, t in zip(extra_names, extra_inputs):
|
||||
feeds[name] = t.numpy()
|
||||
ort_out = sess.run(None, feeds)[0]
|
||||
with torch.no_grad():
|
||||
torch_out = model(dummy_images, dummy_history).numpy()
|
||||
torch_out = model(dummy_images, dummy_history, *extra_inputs).numpy()
|
||||
max_diff = np.abs(ort_out - torch_out).max()
|
||||
if max_diff > MAX_TOLERATED_DIFF:
|
||||
sys.exit(f"错误: ONNX 与 PyTorch 输出差异 {max_diff:.3e} > "
|
||||
f"{MAX_TOLERATED_DIFF:.0e}")
|
||||
print(f"ONNX 一致性验证通过(最大差异: {max_diff:.6e})")
|
||||
except ImportError:
|
||||
print("提示: 安装 onnxruntime 可验证导出模型一致性 (pip install onnxruntime)")
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
"""生成 DFC 量化校准数据集
|
||||
|
||||
两种模式:
|
||||
默认模式(data/calibration.npz):为旧版 dfc_compile.py 的 ClientModel
|
||||
流程准备,目前保留兼容。
|
||||
--deploy 模式(data/calibration_deploy.npz):为 export_deploy.py 导出
|
||||
的部署版 ONNX 生成校准集,键名匹配 HN 层名 actpred_deploy/input_layerN,
|
||||
形状匹配 HEF 输入:frame0~2 (N,H,W,3) 已归一化 float32、
|
||||
action_history (N,1,H_len,D)、pred_token (N,1,1,d) 学习值 ± 抖动
|
||||
(常量校准会使量化范围退化,见 --pred-jitter)。
|
||||
图像分辨率与 --model checkpoint 的配置一致。
|
||||
|
||||
数据源:
|
||||
- 图像:无真实 episode 时用 [0,255] 均匀采样像素 + ImageNet 归一化
|
||||
- 动作历史:从 bot4cj/scripts/*.mt 的真实 @Action 矩阵滑窗采样
|
||||
(列不足 18 补零,与 pretrain_mt.py 语义一致)
|
||||
|
||||
真实数据到位后应替换为真实帧重新校准。
|
||||
|
||||
用法:
|
||||
python gen_calibration.py --deploy --model checkpoints/pretrain_best_inference.pt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
|
||||
from config import DataConfig, parse_checkpoint_config
|
||||
from mt_parser import parse_mt_directory
|
||||
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
|
||||
def build_action_history_samples(mt_dir, n_samples, history_len, action_dim, rng):
|
||||
sequences = []
|
||||
for _path, matrix in parse_mt_directory(mt_dir):
|
||||
n_cols = matrix.shape[1]
|
||||
if n_cols < action_dim:
|
||||
matrix = np.pad(matrix, ((0, 0), (0, action_dim - n_cols)))
|
||||
elif n_cols > action_dim:
|
||||
matrix = matrix[:, :action_dim]
|
||||
if len(matrix) >= history_len:
|
||||
sequences.append(matrix)
|
||||
if not sequences:
|
||||
raise ValueError(f"{mt_dir} 中没有可用的 @Action 序列")
|
||||
|
||||
windows = []
|
||||
for matrix in sequences:
|
||||
for i in range(len(matrix) - history_len + 1):
|
||||
windows.append(matrix[i:i + history_len])
|
||||
idx = rng.integers(0, len(windows), size=n_samples)
|
||||
return np.stack([windows[i] for i in idx]).astype(np.float32)
|
||||
|
||||
|
||||
def build_image_samples(n_samples, frame_len, height, width, rng,
|
||||
chunk_samples: int = 4):
|
||||
"""生成 (n_samples, frame_len, 3, H, W) uint8 校准图像
|
||||
|
||||
存 uint8 原始像素值(0~255),加载时再做 ToTensor+Normalize,节省磁盘。
|
||||
经 memmap 分块填充后返回:高分辨率(如 1080p)下整块数组可达 ~600MB,
|
||||
一次性分配容易挤爆内存,分块写盘把峰值压到 chunk_samples 帧量级。
|
||||
"""
|
||||
if chunk_samples < 1:
|
||||
raise ValueError(f"chunk_samples 应 >= 1,实际为 {chunk_samples}")
|
||||
chunk_samples = min(chunk_samples, n_samples)
|
||||
fd, tmp_path = tempfile.mkstemp(suffix=".npy")
|
||||
os.close(fd)
|
||||
images = np.lib.format.open_memmap(
|
||||
tmp_path, mode="w+", dtype=np.uint8,
|
||||
shape=(n_samples, frame_len, 3, height, width))
|
||||
try:
|
||||
for start in range(0, n_samples, chunk_samples):
|
||||
stop = min(start + chunk_samples, n_samples)
|
||||
images[start:stop] = rng.integers(
|
||||
0, 256, size=(stop - start, frame_len, 3, height, width),
|
||||
dtype=np.uint8)
|
||||
except Exception:
|
||||
images._mmap.close()
|
||||
os.unlink(images.filename)
|
||||
raise
|
||||
return images
|
||||
|
||||
|
||||
def build_deploy_set(args):
|
||||
import torch
|
||||
from action_norm import is_range_mode, normalize_action
|
||||
from model import ActionPredictor
|
||||
|
||||
ckpt = torch.load(args.model, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(ckpt["config"]).data
|
||||
rng = np.random.default_rng(args.seed)
|
||||
|
||||
h, w = cfg.image_height, cfg.image_width
|
||||
raw = rng.integers(0, 256, size=(args.samples, h, w, 3),
|
||||
dtype=np.uint8).astype(np.float32) / 255.0
|
||||
norm = ((raw - IMAGENET_MEAN) / IMAGENET_STD).astype(np.float32)
|
||||
frames = [norm.copy() for _ in range(cfg.frame_len)]
|
||||
|
||||
history = build_action_history_samples(
|
||||
args.mt_dir, args.samples, cfg.history_len, cfg.action_dim, rng)
|
||||
history = history[:, None, :, :] # (N, 1, H_len, D)
|
||||
if cfg.normalize_actions and is_range_mode(cfg.normalize_mode):
|
||||
# 固定范围类归一化:校准数据与推理端同一变换(免 stats 文件)
|
||||
from action_norm import normalize_action, is_range_mode
|
||||
history = normalize_action(history, cfg.action_dim,
|
||||
mode=cfg.normalize_mode)
|
||||
print(f" action_history: 已 {cfg.normalize_mode} 归一化"
|
||||
f"(固定物理范围)")
|
||||
elif cfg.normalize_actions:
|
||||
# z-score 归一化模型的 HEF 输入是 z-score 历史(hailo_infer --stats 同语义),
|
||||
# 校准数据必须走同一变换,否则量化范围与推理分布不匹配
|
||||
stats_path = cfg.action_stats_path
|
||||
if not stats_path:
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)),
|
||||
"action_stats.npy")
|
||||
elif not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)), stats_path)
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(
|
||||
f"模型训练时启用了动作归一化,但未找到统计文件: {stats_path}")
|
||||
stats = np.load(stats_path)
|
||||
if stats.shape != (2, cfg.action_dim):
|
||||
raise ValueError(
|
||||
f"action stats 形状应为 (2, {cfg.action_dim}),"
|
||||
f"实际为 {stats.shape}: {stats_path}")
|
||||
history = (history - stats[0]) / stats[1]
|
||||
print(f" action_history: 已 z-score 归一化({stats_path})")
|
||||
|
||||
model = ActionPredictor(parse_checkpoint_config(ckpt["config"]).model,
|
||||
cfg, pretrained=False)
|
||||
model.load_state_dict(ckpt["model_state_dict"])
|
||||
d = model.pred_token.shape[-1]
|
||||
v = model.pred_token.data.unsqueeze(1).numpy() # (1, 1, d) 学习到的常量
|
||||
# 常量校准会使量化范围退化(上次实测 scale=18571.7,任何输入都被
|
||||
# 量化成 0);按学习值 ±相对抖动生成校准样本,运行时喂学习值即可
|
||||
# 无损还原(量化误差 ~1e-4)
|
||||
jitter = v * rng.uniform(-args.pred_jitter, args.pred_jitter,
|
||||
size=(args.samples, 1, 1, d)).astype(np.float32)
|
||||
jitter = jitter + rng.normal(0.0, args.pred_jitter_abs,
|
||||
size=jitter.shape).astype(np.float32)
|
||||
pred = np.broadcast_to(v, (args.samples, 1, 1, d)).astype(np.float32) + jitter
|
||||
|
||||
inputs = [*frames, history, pred]
|
||||
if cfg.imu_dim > 0:
|
||||
# R2 方案一:imu 输入为归一化姿态(roll,pitch/180,ax..gz/32768)。
|
||||
# 校准覆盖典型工作范围:归一化值域 ±0.5(≈±90°/±16384)均匀采样
|
||||
imu = rng.uniform(-0.5, 0.5,
|
||||
size=(args.samples, 1, 1, cfg.imu_dim)).astype(np.float32)
|
||||
inputs.append(imu)
|
||||
|
||||
np.savez_compressed(
|
||||
args.output,
|
||||
**{f"actpred_deploy/input_layer{i + 1}": arr
|
||||
for i, arr in enumerate(inputs)})
|
||||
print(f"部署校准集已保存: {args.output}")
|
||||
print(f" frame0~2: {frames[0].shape}(分辨率 {w}x{h},与 checkpoint 一致)")
|
||||
print(f" action_history: {history.shape} pred_token: {pred.shape}"
|
||||
f"(学习值 ±{args.pred_jitter:.0%} 抖动)")
|
||||
if cfg.imu_dim > 0:
|
||||
print(f" imu: {inputs[-1].shape}(归一化值域 ±0.5 均匀采样)")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate DFC calibration dataset")
|
||||
parser.add_argument("--deploy", action="store_true",
|
||||
help="生成部署版(export_deploy.py)校准集")
|
||||
parser.add_argument("--model", default="checkpoints/pretrain_best_inference.pt",
|
||||
help="--deploy 模式的 checkpoint(决定分辨率与 pred_token)")
|
||||
parser.add_argument("--mt-dir", default="../bot4cj/scripts")
|
||||
parser.add_argument("--output", default=None)
|
||||
parser.add_argument("--samples", type=int, default=32)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--pred-jitter", type=float, default=0.3,
|
||||
help="pred_token 校准相对抖动幅度(默认 ±30%)")
|
||||
parser.add_argument("--pred-jitter-abs", type=float, default=0.005,
|
||||
help="pred_token 校准绝对抖动标准差(近零元素保底)")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.deploy:
|
||||
args.output = args.output or "data/calibration_deploy.npz"
|
||||
build_deploy_set(args)
|
||||
return
|
||||
|
||||
args.output = args.output or "data/calibration.npz"
|
||||
cfg = DataConfig()
|
||||
rng = np.random.default_rng(args.seed)
|
||||
|
||||
history = build_action_history_samples(
|
||||
args.mt_dir, args.samples, cfg.history_len, cfg.action_dim, rng)
|
||||
images = build_image_samples(
|
||||
args.samples, cfg.frame_len, cfg.image_height, cfg.image_width, rng)
|
||||
|
||||
np.savez_compressed(args.output, images=images, action_history=history)
|
||||
# 释放 memmap 并删除临时文件(数据已写入 npz)
|
||||
images._mmap.close()
|
||||
os.unlink(images.filename)
|
||||
print(f"校准集已保存: {args.output}")
|
||||
print(f" images: {images.shape} dtype={images.dtype}")
|
||||
print(f" action_history: {history.shape} dtype={history.dtype}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+461
-85
@@ -3,57 +3,130 @@
|
||||
接口对齐 predict.ActionPredictorInference:predict(images, action_history)。
|
||||
|
||||
HEF 由 actpred ONNX 经 Hailo DFC 编译生成(量化校准后)。本封装不硬编码
|
||||
vstream 名称/形状/量化参数,全部从 HEF 的 input/output vstream info 动态推导:
|
||||
vstream 名称/形状/量化参数,全部从 HEF 的 input/output vstream info 动态推导,
|
||||
支持两种布局:
|
||||
|
||||
图像输入: (1, F, 3, H, W) → NHWC uint8 (1, F, H, W, 3)
|
||||
或 (1, H, W, F*3) → 帧沿通道拼接
|
||||
动作输入: (1, H, D) float32,uint8 边缘时按 quant_info 量化
|
||||
动作输出: (1, D),uint8 边缘时按 quant_info 反量化
|
||||
新版 5/6 输入(export_deploy.py + dfc_compile.py 产物):
|
||||
frame0~2 各 (H,W,3) NHWC uint8(ImageNet 归一化校准,负零偏自动识别)
|
||||
+ action_history (1,H_len,D) + pred_token (1,1,d)
|
||||
+ imu (1,1,1,imu_dim)(R2 方案一可选输入,原始姿态内部归一化)
|
||||
旧版单图像输入:
|
||||
(1, F, 3, H, W) → NHWC uint8 (1, F, H, W, 3)
|
||||
或 (1, H, W, F*3) → 帧沿通道拼接;action_history (1, H, D)
|
||||
|
||||
uint8 边缘按 quant_info 量化/反量化;pred_token 输入默认喂零,
|
||||
新编译 HEF(抖动校准,见 gen_calibration.py --pred-jitter)应喂
|
||||
checkpoint 的学习值(--pred-token npy 或 predict(pred_token=...)),
|
||||
否则该 token 语义与训练时不一致。
|
||||
|
||||
背景与踩坑记录见
|
||||
docs/reports/actpred_hailo_hat2_inference_research_20260813.md:
|
||||
- hailort 5.x 对 Hailo-10 仅支持新 API(create_infer_model),
|
||||
旧版 VDevice.configure(hef) 报 HAILO_NOT_IMPLEMENTED;
|
||||
- 绑定缓冲区 dtype/大小必须匹配 HEF 边缘格式(量化 uint8);
|
||||
- 图像输入按 Hailo 惯例为 uint8 原始像素(ImageNet 归一化由
|
||||
DFC 编译时折叠进网络)。
|
||||
- 绑定缓冲区 dtype/大小必须匹配 HEF 边缘格式(量化 uint8)。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# 与 predict.parse_action_history 相同逻辑;此处内联以保持
|
||||
# 本模块零 torch 依赖(HAT 部署环境只装 hailo_platform + numpy + PIL)
|
||||
try:
|
||||
import cv2 # 可选:帧归一化+量化合并查表(加速预处理)
|
||||
HAS_CV2 = True
|
||||
except ImportError:
|
||||
cv2 = None
|
||||
HAS_CV2 = False
|
||||
|
||||
from imu_norm import normalize_imu
|
||||
from action_norm import (normalize_action, denormalize_action,
|
||||
is_range_mode)
|
||||
from utils import parse_action_history
|
||||
from bins import TOTAL_BINS, decode as bin_decode, fixed_zero_columns
|
||||
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
# 异步推理缓冲数:≥4 时流水线吞吐趋于饱和(实测 4 缓冲 ~2.7×、6 缓冲 ~2.9×)
|
||||
_ASYNC_BUFFERS = 4
|
||||
|
||||
|
||||
def parse_action_history(s: str, action_dim: int) -> np.ndarray:
|
||||
steps = s.split(":")
|
||||
history = []
|
||||
for step in steps:
|
||||
values = [float(v) for v in step.split(",")]
|
||||
if len(values) != action_dim:
|
||||
raise ValueError(
|
||||
f"Expected {action_dim} values per step, got {len(values)}")
|
||||
history.append(values)
|
||||
return np.array(history, dtype=np.float32)
|
||||
class _AsyncJob:
|
||||
"""predict_async 的推理句柄:result() 阻塞返回 (action_dim,) 动作向量"""
|
||||
|
||||
def __init__(self):
|
||||
self._done = threading.Event()
|
||||
self._result = None
|
||||
self._error = None
|
||||
|
||||
def result(self, timeout: float = None) -> np.ndarray:
|
||||
if not self._done.wait(timeout):
|
||||
raise TimeoutError("async HEF inference timeout")
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return self._result
|
||||
|
||||
def _complete(self, result=None, error=None):
|
||||
self._result = result
|
||||
self._error = error
|
||||
self._done.set()
|
||||
|
||||
|
||||
def _input_layer_index(name: str) -> int:
|
||||
m = re.search(r"input_layer(\d+)$", name)
|
||||
return int(m.group(1)) if m else 0
|
||||
|
||||
|
||||
class HailoPredictor:
|
||||
"""HEF 动作预测推理器(Hailo NPU 后端)
|
||||
|
||||
支持新版 5 输入 HEF(逐帧输入 + pred_token)与旧版单图像输入 HEF,
|
||||
布局与量化参数从 vstream info 自动推导。
|
||||
|
||||
用法:
|
||||
with HailoPredictor("actpred.hef") as predictor:
|
||||
action = predictor.predict(images, action_history)
|
||||
"""
|
||||
|
||||
def __init__(self, hef_path: str, stats_path: str = None, vdevice=None):
|
||||
def __init__(self, hef_path: str, stats_path: str = None, vdevice=None,
|
||||
predict_delta: bool = False, pred_token=None,
|
||||
normalize_mode: str = "zscore", column_weights=None):
|
||||
"""HEF 推理器。
|
||||
|
||||
predict_delta: 模型用 --predict-delta 训练时须置 True,推理输出
|
||||
(归一化空间)会加回 history 末步再反归一化,与 predict.py 对齐。
|
||||
pred_token: np.ndarray 或 .npy 路径;export_deploy.py --pred-token-out
|
||||
产物。缺省时喂零(仅适合 pred_token 恰好全零的模型)。
|
||||
normalize_mode: 与训练时 DataConfig.normalize_mode 一致。
|
||||
"zscore"(默认)需传 stats_path;"fixed_range" / "log_range"
|
||||
免 stats 文件,按固定物理范围归一化(action_norm.py)。
|
||||
column_weights: 逐列损失权重(onehot HEF 的权重 0 列按
|
||||
bins.FIXED_OUTPUT 固定,防闭环垃圾回写);None 时用
|
||||
action_norm.default_column_weights。
|
||||
"""
|
||||
from hailo_platform import HEF, VDevice
|
||||
|
||||
self._predict_delta = predict_delta
|
||||
self._range_mode = (normalize_mode
|
||||
if is_range_mode(normalize_mode) else None)
|
||||
self._column_weights = None
|
||||
if column_weights is not None:
|
||||
self._column_weights = np.asarray(column_weights,
|
||||
dtype=np.float32)
|
||||
self._default_pred_token = None
|
||||
if pred_token is not None:
|
||||
if isinstance(pred_token, str):
|
||||
self._default_pred_token = np.load(pred_token)
|
||||
else:
|
||||
self._default_pred_token = np.asarray(pred_token,
|
||||
dtype=np.float32)
|
||||
|
||||
self._stack = contextlib.ExitStack()
|
||||
if vdevice is None:
|
||||
vdevice = self._stack.enter_context(VDevice())
|
||||
@@ -68,14 +141,14 @@ class HailoPredictor:
|
||||
if len(self._out_infos) != 1:
|
||||
raise ValueError(
|
||||
f"HEF 应只有一个输出 vstream,实际 {len(self._out_infos)}")
|
||||
self._image_info = self._find_image_input()
|
||||
self._history_info = self._find_history_input()
|
||||
self._out_info = self._out_infos[0]
|
||||
self._frame_infos, self._history_info, self._pred_info, \
|
||||
self._imu_info = self._classify_inputs()
|
||||
self._derive_dims()
|
||||
|
||||
self._action_mean = None
|
||||
self._action_std = None
|
||||
if stats_path is not None:
|
||||
if stats_path is not None and not self._range_mode:
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(f"统计文件不存在: {stats_path}")
|
||||
stats = np.load(stats_path)
|
||||
@@ -89,49 +162,118 @@ class HailoPredictor:
|
||||
self._stack.close() # 初始化失败也释放设备,避免 fd 泄漏
|
||||
raise
|
||||
|
||||
# ---------------- 形状推导 ----------------
|
||||
# ---------------- 输入分类与形状推导 ----------------
|
||||
|
||||
def _find_image_input(self):
|
||||
cands = [i for i in self._in_infos if len(i.shape) in (4, 5)]
|
||||
if len(cands) != 1:
|
||||
raise ValueError(
|
||||
"HEF 输入中无法唯一识别图像 vstream(期望一个 4/5 维输入),"
|
||||
f"实际: {[(i.name, tuple(i.shape)) for i in self._in_infos]}")
|
||||
return cands[0]
|
||||
def _classify_inputs(self):
|
||||
"""将 HEF 输入分为 帧列表 / 动作历史 / pred_token / imu(后两者可为 None)
|
||||
|
||||
def _find_history_input(self):
|
||||
cands = [i for i in self._in_infos
|
||||
if len(i.shape) == 3 and i is not self._image_info]
|
||||
if len(cands) != 1:
|
||||
多帧输入的绑定顺序按 input_layerN 数字后缀排序,而非 vstream
|
||||
报告顺序:hailort 5.3.0 实测报告顺序被打乱(1,5,4,3,2),
|
||||
按报告顺序绑定会导致 frame1/frame2 错位。
|
||||
imu 为 R2 方案一新增输入(1,1,1,imu_dim),旧 HEF 无此输入。
|
||||
"""
|
||||
ins = list(self._in_infos)
|
||||
# 新版:每帧独立 (H,W,3) 输入
|
||||
frames = [i for i in ins if len(i.shape) == 3
|
||||
and i.shape[2] == 3 and i.shape[0] > 3]
|
||||
if frames:
|
||||
frames.sort(key=lambda i: _input_layer_index(i.name))
|
||||
else:
|
||||
# 旧版:单一 4/5 维图像输入
|
||||
frames = [i for i in ins if len(i.shape) in (4, 5)]
|
||||
rest = [i for i in ins if i not in frames]
|
||||
history, pred, imu = None, None, None
|
||||
for i in rest:
|
||||
s = tuple(i.shape)
|
||||
if len(s) == 3 and s[1] == 1:
|
||||
pred = i # pred_token (1,1,d)
|
||||
elif len(s) == 3:
|
||||
history = i # 动作历史 (1,H_len,D)
|
||||
elif len(s) == 4 and s[0] == 1 and s[1] == 1 and s[2] == 1:
|
||||
imu = i # R2:imu (1,1,1,imu_dim)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"无法识别 HEF 输入 {i.name} 形状 {s}(期望动作历史 "
|
||||
f"(1,H_len,D)、pred_token (1,1,d) 或 imu (1,1,1,imu_dim))")
|
||||
if not frames:
|
||||
raise ValueError("HEF 输入中未找到图像输入")
|
||||
if history is None:
|
||||
raise ValueError(
|
||||
"HEF 输入中无法唯一识别动作历史 vstream(期望一个 3 维输入),"
|
||||
f"实际: {[(i.name, tuple(i.shape)) for i in self._in_infos]}")
|
||||
return cands[0]
|
||||
f"HEF 输入中未找到动作历史输入(输出维度 "
|
||||
f"{tuple(self._out_info.shape)}),"
|
||||
f"实际: {[(i.name, tuple(i.shape)) for i in ins]}")
|
||||
return frames, history, pred, imu
|
||||
|
||||
def _derive_dims(self):
|
||||
shape = tuple(self._image_info.shape)
|
||||
if len(shape) == 5:
|
||||
_, self.frame_len, c, self.image_height, self.image_width = shape
|
||||
if c != 3:
|
||||
raise ValueError(f"图像输入通道维应为 3,实际形状 {shape}")
|
||||
self._image_layout = "ncfhw"
|
||||
else: # (1, H, W, F*3) NHWC 帧拼通道
|
||||
_, self.image_height, self.image_width, c = shape
|
||||
if c % 3 != 0 or c == 3:
|
||||
raise ValueError(
|
||||
f"无法从 4 维图像输入推导帧数(期望 (1,H,W,F*3)),"
|
||||
f"实际形状 {shape}")
|
||||
self.frame_len = c // 3
|
||||
self._image_layout = "nhwc_fused"
|
||||
# 帧输入布局
|
||||
if len(self._frame_infos) == 1 and len(self._frame_infos[0].shape) in (4, 5):
|
||||
shape = tuple(self._frame_infos[0].shape)
|
||||
if len(shape) == 5:
|
||||
_, self.frame_len, c, self.image_height, self.image_width = shape
|
||||
if c != 3:
|
||||
raise ValueError(f"图像输入通道维应为 3,实际形状 {shape}")
|
||||
self._image_layout = "ncfhw"
|
||||
else: # (1, H, W, F*3) NHWC 帧拼通道
|
||||
_, self.image_height, self.image_width, c = shape
|
||||
if c % 3 != 0 or c == 3:
|
||||
raise ValueError(
|
||||
f"无法从 4 维图像输入推导帧数(期望 (1,H,W,F*3)),"
|
||||
f"实际形状 {shape}")
|
||||
self.frame_len = c // 3
|
||||
self._image_layout = "nhwc_fused"
|
||||
else:
|
||||
# 新版:每帧独立 (H,W,3) 输入,形状须一致
|
||||
shapes = [tuple(i.shape) for i in self._frame_infos]
|
||||
if any(len(s) != 3 or s[2] != 3
|
||||
or s[0] != shapes[0][0] or s[1] != shapes[0][1]
|
||||
for s in shapes):
|
||||
raise ValueError(f"帧输入形状不一致: {shapes}")
|
||||
self.image_height, self.image_width, _ = shapes[0]
|
||||
self.frame_len = len(self._frame_infos)
|
||||
self._image_layout = "nhwc_frames"
|
||||
self._frame_quant = self._frame_infos[0].quant_info
|
||||
# 帧输入是否需 ImageNet 归一化,由量化网格判定:DFC 的 zp 恒非负
|
||||
# (zp = round(-min/scale)),码 0 反量化值 dequant(0) = -zp*scale。
|
||||
# 原始像素输入(min=0, zp=0)下为 0;归一化输入(min≈-2.1)下
|
||||
# ≈ -2.1(实测 HEF zp=114, scale=0.0187)。旧启发式 zp*scale<0
|
||||
# 把归一化输入误判为原始像素直通,图像路径被饱和破坏。
|
||||
self._normalize_frames = \
|
||||
(0 - self._frame_quant.qp_zp) * self._frame_quant.qp_scale < -0.5
|
||||
|
||||
# cv2.LUT 查表(归一化+量化合并一步):仅逐帧独立输入布局且
|
||||
# 帧边缘为 uint8 时可用;无 cv2 或旧布局回退 float 路径
|
||||
from hailo_platform import FormatType
|
||||
self._frame_luts = None
|
||||
if HAS_CV2 and self._image_layout == "nhwc_frames" and \
|
||||
all(f.format.type == FormatType.UINT8
|
||||
for f in self._frame_infos):
|
||||
self._frame_luts = self._build_frame_luts()
|
||||
|
||||
# 异步推理缓冲池(predict_async 复用)
|
||||
self._async_buffers = [self._configured.create_bindings()
|
||||
for _ in range(_ASYNC_BUFFERS)]
|
||||
self._async_flags = [[True] for _ in range(_ASYNC_BUFFERS)]
|
||||
|
||||
h_shape = tuple(self._history_info.shape)
|
||||
self.history_len, self.action_dim = h_shape[1], h_shape[2]
|
||||
|
||||
# R2:imu 输入存在时暴露 imu_dim(pybot 闭环鸭子类型据此传 imu=),
|
||||
# 无则 None(旧 HEF 保持旧两参调用路径)
|
||||
self.imu_dim = tuple(self._imu_info.shape)[-1] \
|
||||
if self._imu_info is not None else None
|
||||
|
||||
out_shape = tuple(self._out_info.shape)
|
||||
out_dim = out_shape[-1]
|
||||
if out_dim != self.action_dim:
|
||||
raise ValueError(
|
||||
f"输出动作维度 {out_dim} 与动作历史维度 {self.action_dim} 不一致")
|
||||
if out_dim == TOTAL_BINS:
|
||||
# onehot 模型:输出为多头 logits,predict 时经 bins.decode
|
||||
# 解码为 action_dim 原始量纲(无需反归一化)
|
||||
self._onehot = True
|
||||
else:
|
||||
self._onehot = False
|
||||
if out_dim != self.action_dim:
|
||||
raise ValueError(
|
||||
f"输出动作维度 {out_dim} 与动作历史维度 "
|
||||
f"{self.action_dim} 不一致")
|
||||
|
||||
# ---------------- 生命周期 ----------------
|
||||
|
||||
@@ -147,11 +289,8 @@ class HailoPredictor:
|
||||
|
||||
# ---------------- 推理 ----------------
|
||||
|
||||
def predict(self, images, action_history: np.ndarray) -> np.ndarray:
|
||||
"""images: 长度为 frame_len 的 RGB 图像列表(旧→新,末帧为当前帧)
|
||||
|
||||
返回 (action_dim,) float32 动作向量。
|
||||
"""
|
||||
def _prepare_inputs(self, images, action_history):
|
||||
"""输入校验 + 动作历史归一化(predict/predict_async 共用)"""
|
||||
if not isinstance(images, (list, tuple)):
|
||||
images = [images]
|
||||
if len(images) != self.frame_len:
|
||||
@@ -164,15 +303,26 @@ class HailoPredictor:
|
||||
raise ValueError(
|
||||
f"action_history 形状应为 {expected},实际为 {action_history.shape}")
|
||||
|
||||
if self._action_mean is not None:
|
||||
if self._range_mode:
|
||||
action_history = normalize_action(action_history, self.action_dim,
|
||||
mode=self._range_mode)
|
||||
elif self._action_mean is not None:
|
||||
action_history = (action_history - self._action_mean) / self._action_std
|
||||
return images, action_history
|
||||
|
||||
def _fill_bindings(self, bindings, images, action_history,
|
||||
pred_token=None, imu=None):
|
||||
"""按 HEF 布局把(已归一化的)输入填充到 bindings 缓冲"""
|
||||
from hailo_platform import FormatType
|
||||
bindings = self._configured.create_bindings()
|
||||
|
||||
frames = np.stack([self._preprocess(img) for img in images])
|
||||
bindings.input(self._image_info.name).set_buffer(
|
||||
self._image_buffer(frames))
|
||||
if self._frame_luts is not None:
|
||||
# cv2.LUT 快路径:逐帧独立输入,直接绑定量化缓冲
|
||||
for i, (info, img) in enumerate(zip(self._frame_infos, images)):
|
||||
buf = self._frame_lut_quant(img, i)
|
||||
bindings.input(info.name).set_buffer(buf)
|
||||
else:
|
||||
frame_float = np.stack([self._frame_float(img) for img in images])
|
||||
self._bind_frames(bindings, frame_float)
|
||||
|
||||
h_shape = tuple(self._history_info.shape)
|
||||
if self._history_info.format.type == FormatType.UINT8:
|
||||
@@ -183,6 +333,30 @@ class HailoPredictor:
|
||||
buf = action_history.astype(np.float32).reshape(h_shape)
|
||||
bindings.input(self._history_info.name).set_buffer(buf)
|
||||
|
||||
if self._pred_info is not None:
|
||||
p_shape = tuple(self._pred_info.shape)
|
||||
if pred_token is None:
|
||||
pred_token = self._default_pred_token
|
||||
if pred_token is None:
|
||||
values = np.zeros(p_shape, dtype=np.float32)
|
||||
else:
|
||||
values = np.asarray(pred_token, dtype=np.float32).reshape(p_shape)
|
||||
if self._pred_info.format.type == FormatType.UINT8:
|
||||
q = self._pred_info.quant_info
|
||||
values = np.clip(np.round(values / q.qp_scale + q.qp_zp),
|
||||
0, 255).astype(np.uint8)
|
||||
bindings.input(self._pred_info.name).set_buffer(values)
|
||||
|
||||
if self._imu_info is not None:
|
||||
i_shape = tuple(self._imu_info.shape)
|
||||
# 原始姿态 → 物理量程归一化(与训练 dataset 同口径)
|
||||
values = normalize_imu(imu, i_shape[-1]).reshape(i_shape)
|
||||
if self._imu_info.format.type == FormatType.UINT8:
|
||||
q = self._imu_info.quant_info
|
||||
values = np.clip(np.round(values / q.qp_scale + q.qp_zp),
|
||||
0, 255).astype(np.uint8)
|
||||
bindings.input(self._imu_info.name).set_buffer(values)
|
||||
|
||||
out_shape = tuple(self._out_info.shape)
|
||||
if self._out_info.format.type == FormatType.UINT8:
|
||||
out_buf = np.zeros(out_shape, dtype=np.uint8)
|
||||
@@ -190,32 +364,194 @@ class HailoPredictor:
|
||||
out_buf = np.zeros(out_shape, dtype=np.float32)
|
||||
bindings.output(self._out_info.name).set_buffer(out_buf)
|
||||
|
||||
self._configured.run([bindings], timeout=10000)
|
||||
def _postprocess(self, raw_result, history_tail=None) -> np.ndarray:
|
||||
"""反量化 + 解码/反归一化 → (action_dim,)(predict/predict_async 共用)
|
||||
|
||||
result = np.asarray(
|
||||
bindings.output(self._out_info.name).get_buffer(),
|
||||
dtype=np.float32)
|
||||
history_tail: 归一化空间的历史末步(predict_delta 残差模式用)。
|
||||
"""
|
||||
from hailo_platform import FormatType
|
||||
if self._out_info.format.type == FormatType.UINT8:
|
||||
q = self._out_info.quant_info
|
||||
result = (result - q.qp_zp) * q.qp_scale
|
||||
result = result.reshape(-1)
|
||||
if self._action_mean is not None:
|
||||
result = result * self._action_std + self._action_mean
|
||||
raw_result = (raw_result - q.qp_zp) * q.qp_scale
|
||||
result = raw_result.reshape(-1)
|
||||
if self._onehot:
|
||||
# 多头 logits → 期望值解码 → 原始量纲(箱生活在原始空间,
|
||||
# 无需反归一化);权重 0 列固定为常量(与 predict.py 对齐),
|
||||
# 防闭环垃圾回写
|
||||
result = bin_decode(result, self.action_dim).reshape(-1)
|
||||
if self._column_weights is None:
|
||||
from action_norm import default_column_weights
|
||||
self._column_weights = default_column_weights(
|
||||
self.action_dim)
|
||||
w = np.asarray(self._column_weights, dtype=np.float32)
|
||||
fixed = fixed_zero_columns(w, self.action_dim)
|
||||
result[w[:self.action_dim] <= 0.0] = (
|
||||
fixed[w[:self.action_dim] <= 0.0])
|
||||
else:
|
||||
if self._predict_delta:
|
||||
# 残差模式:Δ̂ + a_{t-1}(归一化空间),与 predict.py 对齐
|
||||
result = result + history_tail
|
||||
if self._range_mode:
|
||||
result = denormalize_action(result, self.action_dim,
|
||||
mode=self._range_mode)
|
||||
elif self._action_mean is not None:
|
||||
result = result * self._action_std + self._action_mean
|
||||
if not np.all(np.isfinite(result)):
|
||||
print("WARNING: HEF output contains NaN/Inf, replacing with zeros")
|
||||
result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
|
||||
return result
|
||||
|
||||
def _preprocess(self, image) -> np.ndarray:
|
||||
def predict(self, images, action_history: np.ndarray,
|
||||
pred_token: np.ndarray = None,
|
||||
imu: np.ndarray = None) -> np.ndarray:
|
||||
"""images: 长度为 frame_len 的 RGB 图像列表(旧→新,末帧为当前帧)
|
||||
imu: 可选原始姿态向量(roll,pitch,ax,ay,az,gx,gy,gz,HEF 含 imu
|
||||
输入时生效,内部按物理量程归一化;None 等价零向量)
|
||||
|
||||
返回 (action_dim,) float32 动作向量。pred_token 仅新版 5 输入 HEF
|
||||
需要;默认喂零,抖动校准的新 HEF 建议喂 checkpoint 学习值。
|
||||
"""
|
||||
images, action_history = self._prepare_inputs(images, action_history)
|
||||
bindings = self._configured.create_bindings()
|
||||
self._fill_bindings(bindings, images, action_history,
|
||||
pred_token=pred_token, imu=imu)
|
||||
self._configured.run([bindings], timeout=10000)
|
||||
raw = np.asarray(
|
||||
bindings.output(self._out_info.name).get_buffer(),
|
||||
dtype=np.float32)
|
||||
return self._postprocess(raw, action_history[-1])
|
||||
|
||||
# ---------------- 异步推理 ----------------
|
||||
|
||||
def predict_async(self, images, action_history: np.ndarray,
|
||||
pred_token: np.ndarray = None,
|
||||
imu: np.ndarray = None) -> "_AsyncJob":
|
||||
"""提交异步推理(run_async 流水线,多缓冲复用)
|
||||
|
||||
与 predict() 语义/输出完全一致,但立即返回 job;调用方可在
|
||||
NPU 推理期间并行做其他工作(如抓帧/执行上一步动作),需要结果
|
||||
时 job.result() 阻塞获取。连续提交时(≥2 个在途)NPU 保持满负荷,
|
||||
实测吞吐可达同步 predict 的 ~2.9×。
|
||||
|
||||
注意:绑定缓冲有限(_ASYNC_BUFFERS 个),在途任务超过缓冲数时
|
||||
本调用会阻塞等待空闲缓冲;job.result(timeout) 超时不会释放缓冲,
|
||||
缓冲由完成回调释放。
|
||||
"""
|
||||
images, action_history = self._prepare_inputs(images, action_history)
|
||||
history_tail = action_history[-1]
|
||||
idx = self._acquire_async_buffer()
|
||||
bindings = self._async_buffers[idx]
|
||||
job = _AsyncJob()
|
||||
try:
|
||||
self._fill_bindings(bindings, images, action_history,
|
||||
pred_token=pred_token, imu=imu)
|
||||
except Exception:
|
||||
self._release_async_buffer(idx)
|
||||
raise
|
||||
|
||||
def cb(completion_info):
|
||||
try:
|
||||
exc = getattr(completion_info, "exception", None)
|
||||
if exc:
|
||||
raise exc
|
||||
raw = np.asarray(
|
||||
bindings.output(self._out_info.name).get_buffer(),
|
||||
dtype=np.float32)
|
||||
job._complete(result=self._postprocess(raw, history_tail))
|
||||
except Exception as e: # noqa: BLE001 —— 回调内兜底
|
||||
job._complete(error=e)
|
||||
finally:
|
||||
self._release_async_buffer(idx)
|
||||
|
||||
self._configured.run_async([bindings], callback=cb)
|
||||
return job
|
||||
|
||||
def _acquire_async_buffer(self) -> int:
|
||||
"""取一个空闲绑定缓冲(在途任务超限时阻塞等待)"""
|
||||
while True:
|
||||
for i in range(len(self._async_buffers)):
|
||||
if self._async_flags[i][0]:
|
||||
self._async_flags[i][0] = False
|
||||
return i
|
||||
time.sleep(0.0001)
|
||||
|
||||
def _release_async_buffer(self, idx: int):
|
||||
self._async_flags[idx][0] = True
|
||||
|
||||
def _frame_float(self, image) -> np.ndarray:
|
||||
"""Resize → float32 HWC;归一化 HEF(负零偏)再走 ImageNet 归一化"""
|
||||
img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
|
||||
img = img.convert("RGB").resize(
|
||||
(self.image_width, self.image_height), Image.BILINEAR)
|
||||
return np.asarray(img, dtype=np.uint8) # HWC uint8
|
||||
x = np.asarray(img, dtype=np.float32) # HWC 0..255
|
||||
if self._normalize_frames:
|
||||
x = (x / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
|
||||
return x
|
||||
|
||||
def _build_frame_luts(self) -> list:
|
||||
"""按帧 quant_info 构建 256 项单通道查表(归一化+量化合并一步)
|
||||
|
||||
OpenCV 5.0 的 cv2.LUT 只接受单通道表,故每帧返回 [R,G,B] 三张表。
|
||||
数学上与 float 路径(_frame_float + quantize)等价(实测 0 LSB)。
|
||||
"""
|
||||
luts = []
|
||||
for info in self._frame_infos:
|
||||
q = info.quant_info
|
||||
c = np.arange(256, dtype=np.float32)[:, None]
|
||||
if self._normalize_frames:
|
||||
v = (c / 255.0 - IMAGENET_MEAN[None, :]) / IMAGENET_STD[None, :]
|
||||
else:
|
||||
v = np.broadcast_to(c, (256, 3))
|
||||
v = v / q.qp_scale + q.qp_zp
|
||||
lut = np.clip(np.round(v), 0, 255).astype(np.uint8) # (256,3)
|
||||
luts.append([lut[:, ch] for ch in range(3)])
|
||||
return luts
|
||||
|
||||
def _frame_lut_quant(self, image, frame_idx: int = 0) -> np.ndarray:
|
||||
"""cv2.LUT 直通:原始 uint8 像素 → 量化 uint8 缓冲(免 float 中间态)
|
||||
|
||||
非 uint8 / 非 3 通道 / 尺寸不匹配输入走与 _frame_float 相同的 PIL
|
||||
转换,保证两条路径语义一致(RGBA 经 convert 合成、灰度转 RGB)。
|
||||
"""
|
||||
x = np.asarray(image)
|
||||
if x.dtype != np.uint8 or x.ndim != 3 or x.shape[2] != 3 or \
|
||||
x.shape[:2] != (self.image_height, self.image_width):
|
||||
img = Image.fromarray(x).convert("RGB")
|
||||
if img.size != (self.image_width, self.image_height):
|
||||
img = img.resize((self.image_width, self.image_height),
|
||||
Image.BILINEAR)
|
||||
x = np.asarray(img, dtype=np.uint8)
|
||||
luts = self._frame_luts[frame_idx]
|
||||
outs = [cv2.LUT(x[..., ch], lut) for ch, lut in enumerate(luts)]
|
||||
return np.ascontiguousarray(cv2.merge(outs))
|
||||
|
||||
def _bind_frames(self, bindings, frame_float: np.ndarray):
|
||||
"""(F,H,W,3) float 帧 → HEF 帧输入缓冲区(按布局与量化)"""
|
||||
from hailo_platform import FormatType
|
||||
|
||||
def quantize(buf, q):
|
||||
return np.clip(np.round(buf / q.qp_scale + q.qp_zp),
|
||||
0, 255).astype(np.uint8)
|
||||
|
||||
def _image_buffer(self, frames: np.ndarray) -> np.ndarray:
|
||||
"""(F, H, W, 3) → HEF 图像输入缓冲区(按编译后布局)"""
|
||||
if self._image_layout == "ncfhw":
|
||||
return frames.transpose(0, 3, 1, 2)[np.newaxis] # (1,F,3,H,W)
|
||||
# nhwc_fused: 帧沿通道拼接 (1, H, W, F*3)
|
||||
return frames.reshape(1, self.image_height, self.image_width,
|
||||
self.frame_len * 3)
|
||||
buf = frame_float.transpose(0, 3, 1, 2)[np.newaxis] # (1,F,3,H,W)
|
||||
info = self._frame_infos[0]
|
||||
if info.format.type == FormatType.UINT8:
|
||||
buf = quantize(buf, info.quant_info)
|
||||
bindings.input(info.name).set_buffer(buf)
|
||||
elif self._image_layout == "nhwc_fused":
|
||||
buf = frame_float.reshape(
|
||||
1, self.image_height, self.image_width, self.frame_len * 3)
|
||||
info = self._frame_infos[0]
|
||||
if info.format.type == FormatType.UINT8:
|
||||
buf = quantize(buf, info.quant_info)
|
||||
bindings.input(info.name).set_buffer(buf)
|
||||
else: # nhwc_frames:每帧独立输入
|
||||
for i, info in enumerate(self._frame_infos):
|
||||
buf = frame_float[i]
|
||||
if info.format.type == FormatType.UINT8:
|
||||
buf = quantize(buf, info.quant_info)
|
||||
bindings.input(info.name).set_buffer(buf)
|
||||
|
||||
|
||||
def main():
|
||||
@@ -226,19 +562,59 @@ def main():
|
||||
parser.add_argument("--action-history", required=True,
|
||||
help="冒号分隔各步,每步 action_dim 个值")
|
||||
parser.add_argument("--stats", default=None,
|
||||
help="action_stats.npy 路径(训练时启用归一化才需要)")
|
||||
help="action_stats.npy 路径(zscore 归一化训练才需要)")
|
||||
parser.add_argument("--normalize-mode", choices=["zscore", "fixed_range"],
|
||||
default=None,
|
||||
help="与训练时 normalize_mode 一致;fixed_range 免 --stats。"
|
||||
"缺省按 --ckpt 配置自动判定,无 --ckpt 时为 zscore")
|
||||
parser.add_argument("--pred-token", default=None,
|
||||
help="pred_token .npy 路径(export_deploy.py "
|
||||
"--pred-token-out 产物;缺省喂零)")
|
||||
parser.add_argument("--ckpt", default=None,
|
||||
help="checkpoint .pt 路径(自动检测 predict_delta;"
|
||||
"HEF 部署时需要以恢复训练配置)")
|
||||
parser.add_argument("--imu", type=str, default=None,
|
||||
help="可选姿态向量(逗号分隔,原始量纲:roll,pitch 度 / "
|
||||
"ax..gz int16 原始值);HEF 含 imu 输入且缺省时按零向量")
|
||||
parser.add_argument("--output-format", choices=["json", "raw"],
|
||||
default="json")
|
||||
args = parser.parse_args()
|
||||
|
||||
with HailoPredictor(args.hef, stats_path=args.stats) as predictor:
|
||||
pred_token = np.load(args.pred_token) if args.pred_token else None
|
||||
imu = None
|
||||
if args.imu is not None:
|
||||
imu = np.array([float(v) for v in args.imu.split(",")],
|
||||
dtype=np.float32)
|
||||
|
||||
predict_delta = False
|
||||
normalize_mode = args.normalize_mode or "zscore"
|
||||
if args.ckpt:
|
||||
import torch
|
||||
ck = torch.load(args.ckpt, map_location="cpu", weights_only=True)
|
||||
cfg = ck.get("config", {})
|
||||
predict_delta = cfg.get("train", {}).get("predict_delta", False)
|
||||
if predict_delta:
|
||||
print("predict_delta=True: HEF 输出将加回历史末步(残差模式)")
|
||||
if args.normalize_mode is None:
|
||||
# 未显式指定时按 checkpoint 配置判定归一化方式(旧 checkpoint 无
|
||||
# normalize_mode 键,默认 zscore,行为不变)
|
||||
data_cfg = cfg.get("data", {})
|
||||
if (data_cfg.get("normalize_actions")
|
||||
and data_cfg.get("normalize_mode") == "fixed_range"):
|
||||
normalize_mode = "fixed_range"
|
||||
print("normalize_mode=fixed_range: 固定物理范围归一化(免 stats)")
|
||||
|
||||
with HailoPredictor(args.hef, stats_path=args.stats,
|
||||
predict_delta=predict_delta,
|
||||
normalize_mode=normalize_mode) as predictor:
|
||||
if len(args.images) != predictor.frame_len:
|
||||
raise SystemExit(
|
||||
f"--images 需传入 {predictor.frame_len} 个路径(旧→新),"
|
||||
f"实际为 {len(args.images)} 个")
|
||||
images = [np.array(Image.open(p).convert("RGB")) for p in args.images]
|
||||
history = parse_action_history(args.action_history, predictor.action_dim)
|
||||
predicted = predictor.predict(images, history)
|
||||
predicted = predictor.predict(images, history, pred_token=pred_token,
|
||||
imu=imu)
|
||||
|
||||
if args.output_format == "json":
|
||||
print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
"""IMU 姿态归一化 —— 训练(dataset)与推理(predict / hailo_infer)共用
|
||||
|
||||
量纲差异大(pybot MPU6050 原始输出,见 pybot/hd/mpu6050.py):
|
||||
roll/pitch 角度制,±180°
|
||||
ax/ay/az 加速度计原始 int16(±8g 量程,1g≈4096 LSB)
|
||||
gx/gy/gz 陀螺仪原始 int16(±500°/s 量程)
|
||||
|
||||
按物理量程做固定缩放归一化(不依赖数据统计文件,部署链零状态):
|
||||
索引 0,1 /180 → ≈[-1, 1]
|
||||
索引 2..7 /32768 → ≈[-1, 1]
|
||||
|
||||
单传感器键序与 pybot robot.IMU_KEYS / recorder 一致:
|
||||
[roll, pitch, ax, ay, az, gx, gy, gz]
|
||||
多传感器(最多 4 个)按 pybot 脚本 @MPU6050_BIND 行顺序拼接,
|
||||
每传感器占 8 槽,imu_dim = 8×N,归一化除数按 8 槽模式平铺。
|
||||
|
||||
仅依赖 numpy,保证 hailo_infer.py(HAT 部署环境无 torch)可直接复用。
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
# 各索引归一化除数(单传感器 8 槽;多传感器用 imu_scales() 平铺)
|
||||
IMU_SCALES = np.array([180.0, 180.0,
|
||||
32768.0, 32768.0, 32768.0,
|
||||
32768.0, 32768.0, 32768.0], dtype=np.float32)
|
||||
|
||||
|
||||
def imu_scales(imu_dim: int) -> np.ndarray:
|
||||
"""归一化除数向量:8 槽模式按传感器数平铺(imu_dim = 8×N,N≤4)
|
||||
|
||||
imu_dim ≤ 8 时等价 IMU_SCALES 前缀(roll/pitch 子集仍可用)。
|
||||
"""
|
||||
reps = (imu_dim + len(IMU_SCALES) - 1) // len(IMU_SCALES)
|
||||
return np.tile(IMU_SCALES, reps)[:imu_dim]
|
||||
|
||||
|
||||
def normalize_imu(imu, imu_dim: int = 8) -> np.ndarray:
|
||||
"""原始姿态向量 → 归一化向量(pad/truncate 到 imu_dim)
|
||||
|
||||
imu: list/np.ndarray,长度任意(None 等价全零)。
|
||||
返回 (imu_dim,) float32。
|
||||
"""
|
||||
vec = np.zeros(imu_dim, dtype=np.float32)
|
||||
if imu is not None:
|
||||
src = np.asarray(imu, dtype=np.float32).ravel()
|
||||
n = min(len(src), imu_dim)
|
||||
vec[:n] = src[:n]
|
||||
return vec / imu_scales(imu_dim)
|
||||
+21
-5
@@ -89,7 +89,8 @@ class PredictionHead(nn.Module):
|
||||
class ActionPredictor(nn.Module):
|
||||
"""多帧图像 + 历史动作的 Transformer 注意力融合预测器
|
||||
|
||||
token 序列:[img×F, act×H, PRED],长度固定为 F+H+1,
|
||||
token 序列:[img×F, act×H, (IMU), PRED],长度固定为 F+H+1(imu_dim=0)
|
||||
或 F+H+2(imu_dim>0,R2 方案一:IMU 独立 token 位于 PRED 之前),
|
||||
取 PRED 位置输出进预测头回归下一时间步动作。
|
||||
除 batch 维外全部维度静态,保证 ONNX/NPU 部署友好。
|
||||
"""
|
||||
@@ -107,7 +108,12 @@ class ActionPredictor(nn.Module):
|
||||
self.img_proj = nn.Linear(model_cfg.cnn_feature_dim, d)
|
||||
self.act_proj = nn.Linear(data_cfg.action_dim, d)
|
||||
|
||||
seq_len = data_cfg.frame_len + data_cfg.history_len + 1
|
||||
self._imu_dim = data_cfg.imu_dim
|
||||
if self._imu_dim > 0:
|
||||
self.imu_proj = nn.Linear(self._imu_dim, d)
|
||||
|
||||
seq_len = data_cfg.frame_len + data_cfg.history_len + 1 \
|
||||
+ (1 if self._imu_dim > 0 else 0)
|
||||
self.pred_token = nn.Parameter(torch.zeros(1, 1, d))
|
||||
self.pos_embed = nn.Parameter(torch.zeros(1, seq_len, d))
|
||||
nn.init.trunc_normal_(self.pred_token, std=0.02)
|
||||
@@ -121,8 +127,10 @@ class ActionPredictor(nn.Module):
|
||||
encoder_layer, num_layers=model_cfg.attn_layers,
|
||||
enable_nested_tensor=False)
|
||||
|
||||
from bins import TOTAL_BINS
|
||||
head_out = TOTAL_BINS if data_cfg.onehot else data_cfg.action_dim
|
||||
self.head = PredictionHead(
|
||||
d, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout)
|
||||
d, model_cfg.fused_hidden_dim, head_out, model_cfg.dropout)
|
||||
|
||||
self._frame_len = data_cfg.frame_len
|
||||
self._history_len = data_cfg.history_len
|
||||
@@ -131,8 +139,10 @@ class ActionPredictor(nn.Module):
|
||||
self._image_height = data_cfg.image_height
|
||||
self._cnn_feature_dim = model_cfg.cnn_feature_dim
|
||||
|
||||
def forward(self, images: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor:
|
||||
def forward(self, images: torch.Tensor, action_history: torch.Tensor,
|
||||
imu: torch.Tensor = None) -> torch.Tensor:
|
||||
# images: (B, F, 3, H, W) action_history: (B, H_len, D)
|
||||
# imu: (B, imu_dim) 已归一化;None 时填零(无姿态先验)
|
||||
b = images.size(0)
|
||||
frames = images.reshape(
|
||||
b * self._frame_len, 3, self._image_height, self._image_width)
|
||||
@@ -141,7 +151,13 @@ class ActionPredictor(nn.Module):
|
||||
img_tokens = self.img_proj(img_feat)
|
||||
act_tokens = self.act_proj(action_history)
|
||||
pred_token = self.pred_token.expand(b, -1, -1)
|
||||
tokens = torch.cat([img_tokens, act_tokens, pred_token], dim=1)
|
||||
parts = [img_tokens, act_tokens]
|
||||
if self._imu_dim > 0:
|
||||
if imu is None:
|
||||
imu = images.new_zeros(b, self._imu_dim)
|
||||
parts.append(self.imu_proj(imu).unsqueeze(1))
|
||||
parts.append(pred_token)
|
||||
tokens = torch.cat(parts, dim=1)
|
||||
tokens = tokens + self.pos_embed
|
||||
out = self.encoder(tokens)
|
||||
return self.head(out[:, -1])
|
||||
+10
-7
@@ -2,8 +2,7 @@
|
||||
|
||||
从 .mt 文件中提取 @Action 块,解析为 (n_rows, n_cols) 的 float32 矩阵。
|
||||
支持 `-` 占位符(沿用上一行同列的值);首行 `-` 填 0,
|
||||
与 pybot/bot4cj 运行时语义一致(2026-08 统一,见
|
||||
docs/reports/bot4cj_python_migration_dev_doc.md §2.3)。
|
||||
与 pybot/bot4cj 运行时语义一致(2026-08 统一)。
|
||||
"""
|
||||
|
||||
import re
|
||||
@@ -30,15 +29,19 @@ def parse_mt_action(path) -> np.ndarray:
|
||||
Raises:
|
||||
ValueError: 文件中无 @Action 块、数据行不足、行列不一致或数值非法
|
||||
"""
|
||||
text = Path(path).read_text(encoding="utf-8")
|
||||
text = Path(path).read_text(encoding="utf-8-sig")
|
||||
# 先去除整行注释(# 开头),避免注释中的 @Action 干扰正则匹配
|
||||
text = re.sub(r"(?m)^[ \t]*#.*$", "", text)
|
||||
|
||||
match = re.search(r"@Action\s+(.*?)\s*@", text, re.DOTALL)
|
||||
if not match:
|
||||
start_m = re.search(r"@Action\s+", text)
|
||||
if not start_m:
|
||||
raise ValueError(f"未找到 @Action 块: {path}")
|
||||
|
||||
block = match.group(1)
|
||||
# 块结束于下一个 @ 段标记;@Action 为文件最后一个块时取到文末
|
||||
# (旧正则 @Action\s+(.*?)\s*@ 要求后续必有 @,末位块会解析失败)
|
||||
block = text[start_m.end():]
|
||||
end_m = re.search(r"\s*@", block)
|
||||
if end_m:
|
||||
block = block[:end_m.start()]
|
||||
|
||||
rows = []
|
||||
for line in block.split("\n"):
|
||||
|
||||
+96
-26
@@ -1,10 +1,11 @@
|
||||
"""actpred ONNX Runtime 推理封装 —— CPU 后端(Pi 5 实测 5.9ms/次,PyTorch 的 ~5 倍)
|
||||
|
||||
接口对齐 predict.ActionPredictorInference:predict(images, action_history)。
|
||||
接口对齐 predict.ActionPredictorInference:predict(images, action_history, imu=None)。
|
||||
|
||||
模型维度从 ONNX session 输入元数据动态推导(images: [1,F,3,H,W],
|
||||
action_history: [1,H,D]),无需 checkpoint。
|
||||
图像预处理复现训练时的 dataset.py 流水线:Resize → /255 → ImageNet 归一化。
|
||||
action_history: [1,H,D];R2 imu_dim>0 的模型另有 imu: [1,imu_dim]),
|
||||
无需 checkpoint。图像预处理复现训练时的 dataset.py 流水线:
|
||||
Resize → /255 → ImageNet 归一化。
|
||||
|
||||
背景见 docs/reports/actpred_inference_speed_benchmark_20260813.md。
|
||||
"""
|
||||
@@ -16,20 +17,15 @@ import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from imu_norm import normalize_imu
|
||||
from action_norm import (normalize_action, denormalize_action,
|
||||
is_range_mode)
|
||||
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
|
||||
def parse_action_history(s: str, action_dim: int) -> np.ndarray:
|
||||
steps = s.split(":")
|
||||
history = []
|
||||
for step in steps:
|
||||
values = [float(v) for v in step.split(",")]
|
||||
if len(values) != action_dim:
|
||||
raise ValueError(
|
||||
f"Expected {action_dim} values per step, got {len(values)}")
|
||||
history.append(values)
|
||||
return np.array(history, dtype=np.float32)
|
||||
from utils import parse_action_history
|
||||
|
||||
|
||||
class OnnxPredictor:
|
||||
@@ -41,7 +37,19 @@ class OnnxPredictor:
|
||||
"""
|
||||
|
||||
def __init__(self, onnx_path: str, stats_path: str = None,
|
||||
providers=None):
|
||||
providers=None, normalize_mode: str = None,
|
||||
predict_delta: bool = None):
|
||||
"""ONNX 推理器。
|
||||
|
||||
normalize_mode: 与训练时 DataConfig.normalize_mode 一致。
|
||||
"zscore" 需传 stats_path;"fixed_range" 免 stats 文件,
|
||||
按固定物理范围 min-max 归一化(action_norm.py);
|
||||
"none" 表示训练未归一化,输入输出均不处理。
|
||||
缺省 None 时从 ONNX 元数据自动识别(export_onnx 写入),
|
||||
旧模型无元数据回退 "zscore"。
|
||||
predict_delta: 残差模式(输出 Δ̂ 加回 history 末步得绝对动作)。
|
||||
缺省 None 时同样从 ONNX 元数据识别,无元数据回退 False。
|
||||
"""
|
||||
import onnxruntime as ort
|
||||
|
||||
if providers is None:
|
||||
@@ -54,15 +62,36 @@ class OnnxPredictor:
|
||||
|
||||
inputs = {i.name: [d if isinstance(d, int) else -1
|
||||
for d in i.shape] for i in self._session.get_inputs()}
|
||||
img_shape, hist_shape = self._locate_inputs(inputs)
|
||||
img_shape, hist_shape, imu_shape = self._locate_inputs(inputs)
|
||||
_, self.frame_len, _, self.image_height, self.image_width = img_shape
|
||||
_, self.history_len, self.action_dim = hist_shape
|
||||
self._image_name = next(n for n, s in inputs.items() if len(s) == 5)
|
||||
self._history_name = next(n for n, s in inputs.items() if len(s) == 3)
|
||||
# R2 方案一:2 维输入 [1, imu_dim] 为姿态(无则 None)。
|
||||
# pybot 闭环鸭子类型约定:None 时按旧两参调用。
|
||||
self.imu_dim = imu_shape[-1] if imu_shape is not None else None
|
||||
self._imu_name = (next(n for n, s in inputs.items() if len(s) == 2)
|
||||
if imu_shape is not None else None)
|
||||
|
||||
meta = self._read_metadata()
|
||||
if normalize_mode is None:
|
||||
normalize_mode = meta.get("normalize_mode") or "zscore"
|
||||
if normalize_mode not in ("zscore", "fixed_range", "none",
|
||||
"log_range"):
|
||||
raise ValueError(f"normalize_mode={normalize_mode!r} 非法,"
|
||||
"应为 zscore/fixed_range/none/log_range")
|
||||
self._range_mode = (normalize_mode
|
||||
if is_range_mode(normalize_mode) else None)
|
||||
self._no_norm = normalize_mode == "none"
|
||||
self.normalize_mode = normalize_mode # 解析后的最终模式(供调用方判断)
|
||||
# R1-C 残差预测:元数据 predict_delta=1 时模型输出 Δ̂,
|
||||
# 需在归一化空间加回 a_{t-1}(与 predict.py/hailo_infer 同语义)
|
||||
self._predict_delta = (predict_delta if predict_delta is not None
|
||||
else meta.get("predict_delta") == "1")
|
||||
|
||||
self._action_mean = None
|
||||
self._action_std = None
|
||||
if stats_path is not None:
|
||||
if stats_path is not None and not (self._range_mode or self._no_norm):
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(f"统计文件不存在: {stats_path}")
|
||||
stats = np.load(stats_path)
|
||||
@@ -73,19 +102,28 @@ class OnnxPredictor:
|
||||
self._action_mean = stats[0].astype(np.float32)
|
||||
self._action_std = stats[1].astype(np.float32)
|
||||
|
||||
def _read_metadata(self) -> dict:
|
||||
"""读 ONNX custom metadata(export_onnx 写入);旧模型无则空 dict"""
|
||||
try:
|
||||
return dict(self._session.get_modelmeta().custom_metadata_map)
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def _locate_inputs(inputs):
|
||||
"""按维度数识别图像(5维)/动作历史(3维)输入,返回各自形状"""
|
||||
"""按维度数识别图像(5维)/动作历史(3维)/姿态(2维,可选)输入,返回各自形状"""
|
||||
img = [s for s in inputs.values() if len(s) == 5]
|
||||
hist = [s for s in inputs.values() if len(s) == 3]
|
||||
if len(img) != 1 or len(hist) != 1:
|
||||
imu = [s for s in inputs.values() if len(s) == 2]
|
||||
if len(img) != 1 or len(hist) != 1 or len(imu) > 1:
|
||||
raise ValueError(
|
||||
"ONNX 输入应为 1 个 5 维图像 + 1 个 3 维动作历史,"
|
||||
"ONNX 输入应为 1 个 5 维图像 + 1 个 3 维动作历史"
|
||||
"(R2 可再加 1 个 2 维 imu),"
|
||||
f"实际: {inputs}")
|
||||
if img[0][2] != 3:
|
||||
raise ValueError(
|
||||
f"图像输入形状应为 [1,F,3,H,W],实际为 {img[0]}")
|
||||
return img[0], hist[0]
|
||||
return img[0], hist[0], (imu[0] if imu else None)
|
||||
|
||||
# ---------------- 生命周期 ----------------
|
||||
|
||||
@@ -101,8 +139,11 @@ class OnnxPredictor:
|
||||
|
||||
# ---------------- 推理 ----------------
|
||||
|
||||
def predict(self, images, action_history: np.ndarray) -> np.ndarray:
|
||||
def predict(self, images, action_history: np.ndarray,
|
||||
imu: np.ndarray = None) -> np.ndarray:
|
||||
"""images: 长度为 frame_len 的 RGB 图像列表(旧→新,末帧为当前帧)
|
||||
imu: 可选原始姿态向量(roll,pitch,ax,ay,az,gx,gy,gz,模型含 imu
|
||||
输入时生效,内部按物理量程归一化;None 等价零向量)
|
||||
|
||||
返回 (action_dim,) float32 动作向量。
|
||||
"""
|
||||
@@ -120,7 +161,10 @@ class OnnxPredictor:
|
||||
raise ValueError(
|
||||
f"action_history 形状应为 {expected},实际为 {action_history.shape}")
|
||||
|
||||
if self._action_mean is not None:
|
||||
if self._range_mode:
|
||||
action_history = normalize_action(action_history, self.action_dim,
|
||||
mode=self._range_mode)
|
||||
elif self._action_mean is not None:
|
||||
action_history = (action_history - self._action_mean) / self._action_std
|
||||
|
||||
frames = np.stack([self._preprocess(img) for img in images])
|
||||
@@ -128,8 +172,17 @@ class OnnxPredictor:
|
||||
self._image_name: frames[np.newaxis], # (1,F,3,S,S)
|
||||
self._history_name: action_history[np.newaxis], # (1,H,D)
|
||||
}
|
||||
if self._imu_name is not None:
|
||||
feeds[self._imu_name] = normalize_imu(
|
||||
imu, self.imu_dim)[np.newaxis] # (1,imu_dim)
|
||||
result = self._session.run(None, feeds)[0].reshape(-1).astype(np.float32)
|
||||
if self._action_mean is not None:
|
||||
if self._predict_delta:
|
||||
# 残差:Δ̂ + a_{t-1}(归一化空间,与 predict.py 对齐)
|
||||
result = result + action_history[-1]
|
||||
if self._range_mode:
|
||||
result = denormalize_action(result, self.action_dim,
|
||||
mode=self._range_mode)
|
||||
elif self._action_mean is not None:
|
||||
result = result * self._action_std + self._action_mean
|
||||
return result
|
||||
|
||||
@@ -151,19 +204,36 @@ def main():
|
||||
parser.add_argument("--action-history", required=True,
|
||||
help="冒号分隔各步,每步 action_dim 个值")
|
||||
parser.add_argument("--stats", default=None,
|
||||
help="action_stats.npy 路径(训练时启用归一化才需要)")
|
||||
help="action_stats.npy 路径(zscore 归一化训练才需要)")
|
||||
parser.add_argument("--normalize-mode",
|
||||
choices=["zscore", "fixed_range", "none", "log_range"],
|
||||
default=None,
|
||||
help="与训练时 normalize_mode 一致;缺省从 ONNX 元数据"
|
||||
"自动识别(旧模型回退 zscore)")
|
||||
parser.add_argument("--predict-delta", action="store_true", default=None,
|
||||
help="残差模式(输出加回 history 末步);缺省从 ONNX "
|
||||
"元数据自动识别(旧模型回退关闭)")
|
||||
parser.add_argument("--imu", type=str, default=None,
|
||||
help="可选姿态向量(逗号分隔,原始量纲:roll,pitch 度 / "
|
||||
"ax..gz int16 原始值);模型含 imu 输入且缺省时按零向量")
|
||||
parser.add_argument("--output-format", choices=["json", "raw"],
|
||||
default="json")
|
||||
args = parser.parse_args()
|
||||
|
||||
with OnnxPredictor(args.model, stats_path=args.stats) as predictor:
|
||||
with OnnxPredictor(args.model, stats_path=args.stats,
|
||||
normalize_mode=args.normalize_mode,
|
||||
predict_delta=args.predict_delta) as predictor:
|
||||
if len(args.images) != predictor.frame_len:
|
||||
raise SystemExit(
|
||||
f"--images 需传入 {predictor.frame_len} 个路径(旧→新),"
|
||||
f"实际为 {len(args.images)} 个")
|
||||
images = [np.array(Image.open(p).convert("RGB")) for p in args.images]
|
||||
history = parse_action_history(args.action_history, predictor.action_dim)
|
||||
predicted = predictor.predict(images, history)
|
||||
imu = None
|
||||
if args.imu is not None:
|
||||
imu = np.array([float(v) for v in args.imu.split(",")],
|
||||
dtype=np.float32)
|
||||
predicted = predictor.predict(images, history, imu=imu)
|
||||
|
||||
if args.output_format == "json":
|
||||
print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))
|
||||
|
||||
+77
-24
@@ -7,7 +7,12 @@ from PIL import Image
|
||||
from torchvision import transforms
|
||||
|
||||
from config import parse_checkpoint_config
|
||||
from imu_norm import normalize_imu
|
||||
from action_norm import (normalize_action, denormalize_action,
|
||||
is_range_mode)
|
||||
from bins import decode as bin_decode, fixed_zero_columns
|
||||
from model import ActionPredictor
|
||||
from utils import parse_action_history
|
||||
|
||||
|
||||
class ActionPredictorInference:
|
||||
@@ -29,6 +34,15 @@ class ActionPredictorInference:
|
||||
self.action_dim = self._data_cfg.action_dim
|
||||
self.history_len = self._data_cfg.history_len
|
||||
self.frame_len = self._data_cfg.frame_len
|
||||
# R2 方案一:imu_dim>0 时 predict 接受 imu= 原始姿态向量(内部归一化)。
|
||||
# pybot 闭环鸭子类型约定:此属性为 None 时按旧两参调用,
|
||||
# 故 imu_dim=0(旧 checkpoint)必须映射为 None。
|
||||
self.imu_dim = self._data_cfg.imu_dim or None
|
||||
self._onehot = self._data_cfg.onehot
|
||||
# R1-C 残差预测:checkpoint 记录 predict_delta=True 时,模型输出 Δ̂,
|
||||
# 需在归一化空间加回 a_{t-1}(history 末步)得到绝对动作。
|
||||
# 旧 checkpoint 无该字段,getattr 默认 False,行为与旧版一致。
|
||||
self._predict_delta = getattr(cfg.train, "predict_delta", False)
|
||||
|
||||
self._transform = transforms.Compose([
|
||||
transforms.Resize((self._data_cfg.image_height,
|
||||
@@ -39,7 +53,14 @@ class ActionPredictorInference:
|
||||
|
||||
self._action_mean = None
|
||||
self._action_std = None
|
||||
if self._data_cfg.normalize_actions:
|
||||
# fixed_range / log_range:固定物理范围 → [0,1](免 stats 文件);
|
||||
# 旧 checkpoint 无 normalize_mode 键,parse_checkpoint_config
|
||||
# 默认 "zscore",行为不变
|
||||
self._range_mode = (self._data_cfg.normalize_mode
|
||||
if (self._data_cfg.normalize_actions
|
||||
and is_range_mode(self._data_cfg.normalize_mode))
|
||||
else None)
|
||||
if self._data_cfg.normalize_actions and not self._range_mode:
|
||||
# 优先使用 checkpoint config 记录的 stats 路径(相对路径按 checkpoint 目录解析),
|
||||
# 旧 checkpoint 无该字段时回退到 checkpoint 同目录的 action_stats.npy
|
||||
stats_path = self._data_cfg.action_stats_path
|
||||
@@ -63,8 +84,12 @@ class ActionPredictorInference:
|
||||
self._action_std = torch.from_numpy(stats[1].astype(np.float32)).to(self._device)
|
||||
|
||||
@torch.no_grad()
|
||||
def predict(self, images, action_history: np.ndarray) -> np.ndarray:
|
||||
"""images: 长度为 frame_len 的 RGB 图像序列(旧→新,末帧为当前帧)"""
|
||||
def predict(self, images, action_history: np.ndarray,
|
||||
imu: np.ndarray = None) -> np.ndarray:
|
||||
"""images: 长度为 frame_len 的 RGB 图像序列(旧→新,末帧为当前帧)
|
||||
imu: 可选原始姿态向量(roll,pitch,ax,ay,az,gx,gy,gz,imu_dim>0 时
|
||||
生效,内部按物理量程归一化;None 等价零向量 = 无姿态先验)
|
||||
"""
|
||||
if not isinstance(images, (list, tuple)):
|
||||
images = [images]
|
||||
if len(images) != self._data_cfg.frame_len:
|
||||
@@ -81,25 +106,49 @@ class ActionPredictorInference:
|
||||
frames.append(self._transform(img))
|
||||
img_tensor = torch.stack(frames).unsqueeze(0).to(self._device) # (1, F, 3, S, S)
|
||||
history = torch.tensor(action_history, dtype=torch.float32)
|
||||
if self._action_mean is not None:
|
||||
if self._range_mode:
|
||||
history = torch.from_numpy(
|
||||
normalize_action(action_history, self.action_dim,
|
||||
mode=self._range_mode))
|
||||
elif self._action_mean is not None:
|
||||
history = (history - self._action_mean) / self._action_std
|
||||
history = history.unsqueeze(0).to(self._device)
|
||||
pred = self._model(img_tensor, history).squeeze(0)
|
||||
if self._action_mean is not None:
|
||||
pred = pred * self._action_std + self._action_mean
|
||||
return pred.cpu().numpy()
|
||||
|
||||
|
||||
def parse_action_history(s: str, action_dim: int) -> np.ndarray:
|
||||
steps = s.split(":")
|
||||
history = []
|
||||
for step in steps:
|
||||
values = [float(v) for v in step.split(",")]
|
||||
if len(values) != action_dim:
|
||||
raise ValueError(
|
||||
f"Expected {action_dim} values per step, got {len(values)}")
|
||||
history.append(values)
|
||||
return np.array(history, dtype=np.float32)
|
||||
imu_tensor = None
|
||||
if self._data_cfg.imu_dim > 0:
|
||||
imu_tensor = torch.from_numpy(
|
||||
normalize_imu(imu, self._data_cfg.imu_dim)
|
||||
).unsqueeze(0).to(self._device)
|
||||
pred = self._model(img_tensor, history, imu_tensor).squeeze(0) # 归一化空间
|
||||
if self._onehot:
|
||||
# 多头 logits → 期望值解码 → 原始量纲(箱生活在原始空间);
|
||||
# 权重 0 列固定为常量(delay=10ms 等),防闭环垃圾回写
|
||||
result = bin_decode(pred.cpu().numpy(), self.action_dim)
|
||||
col_w = self._data_cfg.column_weights
|
||||
if col_w is None:
|
||||
from action_norm import default_column_weights
|
||||
col_w = default_column_weights(self.action_dim)
|
||||
fixed = fixed_zero_columns(col_w, self.action_dim)
|
||||
# 按权重掩码覆盖(固定 0 的列也须覆盖——掩码不能只看常量值)
|
||||
w = np.asarray(col_w, dtype=np.float32)
|
||||
result[w[:self.action_dim] <= 0.0] = (
|
||||
fixed[w[:self.action_dim] <= 0.0])
|
||||
if not np.all(np.isfinite(result)):
|
||||
result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
|
||||
return result
|
||||
if self._predict_delta:
|
||||
# 残差:Δ̂ 加回 a_{t-1}(history 末步,同为归一化空间)→ 绝对动作
|
||||
pred = history[0, -1] + pred
|
||||
if self._range_mode:
|
||||
result = denormalize_action(pred.cpu().numpy(), self.action_dim,
|
||||
mode=self._range_mode)
|
||||
else:
|
||||
if self._action_mean is not None:
|
||||
pred = pred * self._action_std + self._action_mean
|
||||
result = pred.cpu().numpy()
|
||||
if not np.all(np.isfinite(result)):
|
||||
print("WARNING: predict output contains NaN/Inf, replacing with zeros")
|
||||
result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
|
||||
return result
|
||||
|
||||
|
||||
def main():
|
||||
@@ -111,14 +160,15 @@ def main():
|
||||
parser.add_argument("--action-history", required=True,
|
||||
help="冒号分隔各步,每步 action_dim 个值,共 history_len 步。"
|
||||
"e.g. '0,0,0,0:1,0,0,0' (action_dim=4, history_len=2)")
|
||||
parser.add_argument("--action-dim", type=int, default=None,
|
||||
help="override action dim (auto-detected from checkpoint if omitted)")
|
||||
parser.add_argument("--imu", type=str, default=None,
|
||||
help="可选姿态向量(逗号分隔,原始量纲:roll,pitch 度 / "
|
||||
"ax..gz int16 原始值);模型 imu_dim>0 且缺省时按零向量")
|
||||
parser.add_argument("--output-format", choices=["json", "raw"], default="json")
|
||||
parser.add_argument("--device", type=str, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
predictor = ActionPredictorInference(args.model, args.device)
|
||||
action_dim = args.action_dim or predictor.action_dim
|
||||
action_dim = predictor.action_dim
|
||||
|
||||
if len(args.images) != predictor.frame_len:
|
||||
raise SystemExit(
|
||||
@@ -126,7 +176,10 @@ def main():
|
||||
f"实际为 {len(args.images)} 个。冷启动时可重复传同一帧路径。")
|
||||
images = [np.array(Image.open(p).convert("RGB")) for p in args.images]
|
||||
history = parse_action_history(args.action_history, action_dim)
|
||||
predicted = predictor.predict(images, history)
|
||||
imu = None
|
||||
if args.imu is not None:
|
||||
imu = np.array([float(v) for v in args.imu.split(",")], dtype=np.float32)
|
||||
predicted = predictor.predict(images, history, imu=imu)
|
||||
|
||||
if args.output_format == "json":
|
||||
print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))
|
||||
|
||||
+166
-46
@@ -12,6 +12,12 @@
|
||||
数据增强:
|
||||
--noise-std 控制高斯噪声标准差。因 .mt 脚本数据量有限,
|
||||
适度噪声(如 0.5~2.0)可有效防止过拟合。
|
||||
|
||||
图像占位:
|
||||
--random-images 用 ImageNet 归一化的均匀噪声替代全零占位。
|
||||
全零图像会让 BN running_var 停在初始小值(推理时每级放大 ~50 倍,
|
||||
特征达 ±1e6),导致 HEF uint8 量化不可行;微调阶段必须用噪声或
|
||||
真实图像恢复 BN 统计。
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -23,9 +29,16 @@ import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from config import Config, build_argparser, apply_overrides, validate_config_values
|
||||
from config import (Config, build_argparser, apply_overrides,
|
||||
validate_config_values, warn_init_from_mismatch)
|
||||
from model import ActionPredictor
|
||||
from mt_parser import parse_mt_directory
|
||||
from action_norm import (normalize_action, default_column_weights,
|
||||
is_range_mode)
|
||||
from train import _weighted_mse
|
||||
|
||||
_IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
||||
_IMAGENET_STD = [0.229, 0.224, 0.225]
|
||||
|
||||
|
||||
class MTActionDataset(Dataset):
|
||||
@@ -35,9 +48,19 @@ class MTActionDataset(Dataset):
|
||||
"""
|
||||
|
||||
def __init__(self, mt_dir: str, history_len: int, action_dim: int,
|
||||
file_indices: list = None, files: list = None):
|
||||
file_indices: list = None, files: list = None,
|
||||
action_mean=None, action_std=None,
|
||||
normalize_mode: str = "zscore"):
|
||||
self.history_len = history_len
|
||||
self.action_dim = action_dim
|
||||
self._action_mean = action_mean
|
||||
self._action_std = action_std
|
||||
# 固定范围类归一化(fixed_range / log_range,免统计量;
|
||||
# action_mean 传入时优先 z-score)
|
||||
self._range_mode = (normalize_mode
|
||||
if (is_range_mode(normalize_mode)
|
||||
and action_mean is None)
|
||||
else None)
|
||||
self.samples = []
|
||||
|
||||
# 复用调用方已解析的结果(避免 train/val 各自重复解析同一目录)
|
||||
@@ -70,33 +93,56 @@ class MTActionDataset(Dataset):
|
||||
|
||||
def __getitem__(self, idx):
|
||||
history, target = self.samples[idx]
|
||||
if self._range_mode:
|
||||
history = normalize_action(history, self.action_dim,
|
||||
mode=self._range_mode)
|
||||
target = normalize_action(target, self.action_dim,
|
||||
mode=self._range_mode)
|
||||
history = torch.from_numpy(history.copy())
|
||||
target = torch.from_numpy(target.copy())
|
||||
if self._action_mean is not None:
|
||||
history = (history - self._action_mean) / self._action_std
|
||||
target = (target - self._action_mean) / self._action_std
|
||||
return {
|
||||
"action_history": torch.from_numpy(history.copy()),
|
||||
"target": torch.from_numpy(target.copy()),
|
||||
"action_history": history,
|
||||
"target": target,
|
||||
}
|
||||
|
||||
|
||||
def _placeholder_images(batch, frame_len, image_height, image_width,
|
||||
device, random_images: bool):
|
||||
"""构建 (B,F,3,H,W) 图像占位:全零或 ImageNet 归一化均匀噪声"""
|
||||
if not random_images:
|
||||
return torch.zeros(batch, frame_len, 3, image_height, image_width,
|
||||
device=device)
|
||||
x = torch.rand(batch, frame_len, 3, image_height, image_width,
|
||||
device=device)
|
||||
mean = torch.tensor(_IMAGENET_MEAN, device=device).view(1, 1, 3, 1, 1)
|
||||
std = torch.tensor(_IMAGENET_STD, device=device).view(1, 1, 3, 1, 1)
|
||||
return (x - mean) / std
|
||||
|
||||
|
||||
def pretrain_one_epoch(model, loader, optimizer, device, image_width,
|
||||
image_height, frame_len,
|
||||
max_grad_norm=None, noise_std=0.0):
|
||||
max_grad_norm=None, noise_std=0.0,
|
||||
random_images: bool = False, weights=None):
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
zero_images = None
|
||||
images = None
|
||||
for batch in tqdm(loader, desc=" pretrain", leave=False):
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
if noise_std > 0:
|
||||
history = history + torch.randn_like(history) * noise_std
|
||||
target = target + torch.randn_like(target) * noise_std
|
||||
if zero_images is None or zero_images.size(0) != history.size(0):
|
||||
zero_images = torch.zeros(
|
||||
history.size(0), frame_len, 3, image_height, image_width,
|
||||
device=device,
|
||||
)
|
||||
if images is None or images.size(0) != history.size(0) or random_images:
|
||||
images = _placeholder_images(
|
||||
history.size(0), frame_len, image_height, image_width,
|
||||
device, random_images)
|
||||
optimizer.zero_grad()
|
||||
pred = model(zero_images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
pred = model(images, history)
|
||||
loss = _weighted_mse(pred, target, weights)
|
||||
loss.backward()
|
||||
if max_grad_norm:
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
|
||||
@@ -107,21 +153,21 @@ def pretrain_one_epoch(model, loader, optimizer, device, image_width,
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate(model, loader, device, image_width, image_height, frame_len):
|
||||
def validate(model, loader, device, image_width, image_height, frame_len,
|
||||
random_images: bool = False, weights=None):
|
||||
model.eval()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
zero_images = None
|
||||
images = None
|
||||
for batch in tqdm(loader, desc=" val ", leave=False):
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
if zero_images is None or zero_images.size(0) != history.size(0):
|
||||
zero_images = torch.zeros(
|
||||
history.size(0), frame_len, 3, image_height, image_width,
|
||||
device=device,
|
||||
)
|
||||
pred = model(zero_images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
if images is None or images.size(0) != history.size(0) or random_images:
|
||||
images = _placeholder_images(
|
||||
history.size(0), frame_len, image_height, image_width,
|
||||
device, random_images)
|
||||
pred = model(images, history)
|
||||
loss = _weighted_mse(pred, target, weights)
|
||||
total_loss += loss.item() * history.size(0)
|
||||
count += history.size(0)
|
||||
if count == 0:
|
||||
@@ -135,11 +181,13 @@ def main():
|
||||
help=".mt 脚本目录(如 ../bot4cj/scripts)")
|
||||
parser.add_argument("--noise-std", type=float, default=0.0,
|
||||
help="高斯噪声标准差(数据增强,建议 0.5~2.0)")
|
||||
parser.add_argument("--random-images", action="store_true",
|
||||
help="用 ImageNet 归一化均匀噪声替代全零图像占位"
|
||||
"(微调时恢复 BN 统计,避免 HEF 量化退化)")
|
||||
args = parser.parse_args()
|
||||
|
||||
cfg = apply_overrides(Config(), args)
|
||||
validate_config_values(cfg.data, cfg.model, cfg.train)
|
||||
cfg.data.data_dir = args.mt_dir
|
||||
|
||||
os.makedirs(cfg.train.checkpoint_dir, exist_ok=True)
|
||||
|
||||
@@ -171,10 +219,56 @@ def main():
|
||||
train_ds = MTActionDataset(
|
||||
args.mt_dir, cfg.data.history_len, cfg.data.action_dim,
|
||||
file_indices=train_file_indices, files=all_files,
|
||||
normalize_mode=cfg.data.normalize_mode,
|
||||
)
|
||||
|
||||
# 动作归一化(与 train.py 语义一致):zscore 模式无 stats 文件时按训练集
|
||||
# targets 计算并保存到 checkpoint_dir,验证集复用同一 stats;
|
||||
# fixed_range 模式免统计量,跳过此段
|
||||
action_mean = action_std = None
|
||||
if cfg.data.normalize_actions and is_range_mode(cfg.data.normalize_mode):
|
||||
# 构造时 action_mean=None 已置 _range_mode,这里只打印口径
|
||||
print(f"动作归一化: {cfg.data.normalize_mode}"
|
||||
f"(固定物理范围→[0,1],免 stats)")
|
||||
elif cfg.data.normalize_actions:
|
||||
if cfg.data.action_stats_path:
|
||||
if not os.path.exists(cfg.data.action_stats_path):
|
||||
raise FileNotFoundError(
|
||||
f"指定的 action_stats_path 不存在: "
|
||||
f"{cfg.data.action_stats_path}")
|
||||
stats = np.load(cfg.data.action_stats_path)
|
||||
else:
|
||||
# 统计量必须覆盖 history 与 target 并集:只按 targets 算时,
|
||||
# targets 中恒定的列 std 落到 1e-8 地板值,而 history 窗口
|
||||
# 可能含该列的其他取值(如 .mt 脚本前几行),z-score 会爆炸
|
||||
# (实测 (90-45)/1e-8=4.5e9),后续 act_tokens ±1e9 直接
|
||||
# 打爆 DFC 量化(NegativeSlopeExponentNonFixable)
|
||||
all_actions = np.stack(
|
||||
[np.concatenate([s[0], s[1][None]]) for s in train_ds.samples]
|
||||
).reshape(-1, cfg.data.action_dim)
|
||||
stats = np.stack([all_actions.mean(axis=0),
|
||||
all_actions.std(axis=0) + 1e-8])
|
||||
stats_path = os.path.join(cfg.train.checkpoint_dir,
|
||||
"action_stats.npy")
|
||||
np.save(stats_path, stats)
|
||||
# 只存裸文件名:推理端按 checkpoint 目录解析相对路径
|
||||
# (predict.py/export_deploy.py/gen_calibration.py 同语义),
|
||||
# 存 "./checkpoints/..." 会被二次拼接成 checkpoints/./checkpoints
|
||||
cfg.data.action_stats_path = "action_stats.npy"
|
||||
print(f"action_stats 已保存: {stats_path}(训练集 targets 统计)")
|
||||
if stats.shape != (2, cfg.data.action_dim):
|
||||
raise ValueError(
|
||||
f"action stats 形状应为 (2, {cfg.data.action_dim}),"
|
||||
f"实际为 {stats.shape}")
|
||||
action_mean, action_std = stats[0], stats[1]
|
||||
train_ds._action_mean = action_mean
|
||||
train_ds._action_std = action_std
|
||||
|
||||
val_ds = MTActionDataset(
|
||||
args.mt_dir, cfg.data.history_len, cfg.data.action_dim,
|
||||
file_indices=val_file_indices, files=all_files,
|
||||
action_mean=action_mean, action_std=action_std,
|
||||
normalize_mode=cfg.data.normalize_mode,
|
||||
)
|
||||
|
||||
print(f"Files: {len(train_file_indices)} train, {len(val_file_indices)} val")
|
||||
@@ -199,13 +293,33 @@ def main():
|
||||
model = ActionPredictor(cfg.model, cfg.data, pretrained=False).to(device)
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
print(f"Model params: {n_params:,}")
|
||||
|
||||
# 逐列损失权重(与 train.py 同口径;仅训练损失用)
|
||||
col_w = cfg.data.column_weights
|
||||
if col_w is None:
|
||||
col_w = default_column_weights(cfg.data.action_dim)
|
||||
weights_t = torch.tensor(np.asarray(col_w, dtype=np.float32), device=device)
|
||||
print(f"Column weights: {[round(float(v), 4) for v in weights_t.tolist()]}")
|
||||
print(f"Config: action_dim={cfg.data.action_dim}, "
|
||||
f"history_len={cfg.data.history_len}, "
|
||||
f"frame_len={cfg.data.frame_len}, "
|
||||
f"backbone={cfg.model.cnn_backbone}, "
|
||||
f"attn={cfg.model.attn_layers}x{cfg.model.attn_heads}h"
|
||||
f"@{cfg.model.attn_dim}")
|
||||
print(f"Pre-training: noise_std={args.noise_std}")
|
||||
print(f"Pre-training: noise_std={args.noise_std}, "
|
||||
f"random_images={args.random_images}")
|
||||
|
||||
if args.init_from:
|
||||
init_ckpt = torch.load(args.init_from, map_location="cpu",
|
||||
weights_only=True)
|
||||
warn_init_from_mismatch(init_ckpt.get("config", {}), cfg)
|
||||
missing, unexpected = model.load_state_dict(
|
||||
init_ckpt["model_state_dict"], strict=False)
|
||||
if missing:
|
||||
print(f" init-from missing keys: {missing}")
|
||||
if unexpected:
|
||||
print(f" init-from unexpected keys: {unexpected}")
|
||||
print(f"Loaded pretrained weights from: {args.init_from}")
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
model.parameters(), lr=cfg.train.lr,
|
||||
@@ -241,35 +355,41 @@ def main():
|
||||
train_loss = pretrain_one_epoch(
|
||||
model, train_loader, optimizer, device,
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.frame_len,
|
||||
cfg.train.max_grad_norm, args.noise_std,
|
||||
cfg.train.max_grad_norm, args.noise_std, args.random_images,
|
||||
weights=weights_t,
|
||||
)
|
||||
val_loss = validate(model, val_loader, device,
|
||||
cfg.data.image_width, cfg.data.image_height,
|
||||
cfg.data.frame_len)
|
||||
cfg.data.frame_len, args.random_images,
|
||||
weights=weights_t)
|
||||
scheduler.step()
|
||||
print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}")
|
||||
|
||||
if val_loss < best_val_loss:
|
||||
best_val_loss = val_loss
|
||||
patience_counter = 0
|
||||
torch.save(
|
||||
{
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"val_loss": val_loss,
|
||||
"config": cfg.to_dict(),
|
||||
},
|
||||
best_path,
|
||||
)
|
||||
torch.save(
|
||||
{
|
||||
"model_state_dict": model.state_dict(),
|
||||
"config": cfg.to_dict(),
|
||||
},
|
||||
inference_path,
|
||||
)
|
||||
print(f" -> saved best (val_loss={val_loss:.6f})")
|
||||
if not np.isfinite(val_loss):
|
||||
print(f"WARNING: val_loss 非有限值({val_loss}),跳过 best 更新")
|
||||
patience_counter += 1
|
||||
else:
|
||||
best_val_loss = val_loss
|
||||
patience_counter = 0
|
||||
torch.save(
|
||||
{
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"val_loss": val_loss,
|
||||
"config": cfg.to_dict(),
|
||||
},
|
||||
best_path,
|
||||
)
|
||||
torch.save(
|
||||
{
|
||||
"model_state_dict": model.state_dict(),
|
||||
"config": cfg.to_dict(),
|
||||
},
|
||||
inference_path,
|
||||
)
|
||||
print(f" -> saved best (val_loss={val_loss:.6f})")
|
||||
else:
|
||||
patience_counter += 1
|
||||
if patience_counter >= cfg.train.patience:
|
||||
|
||||
@@ -4,5 +4,9 @@ torchvision>=0.15,<1.0
|
||||
numpy>=1.24,<3.0
|
||||
Pillow>=9.0,<14.0 # torchvision 依赖,上限放宽避免与 torch 源安装冲突
|
||||
|
||||
# 推理(Hailo 部署):帧归一化+量化合并查表(hailo_infer.py cv2.LUT 快路径,
|
||||
# 无 cv2 自动回退 float 路径——本依赖缺失仅损失 3.3× 端到端加速)
|
||||
opencv-python-headless>=4.9,<6.0
|
||||
|
||||
# Training utilities
|
||||
tqdm>=4.60,<5.0
|
||||
+2009
-4
File diff suppressed because it is too large.
Load diff
+256
-43
@@ -6,9 +6,14 @@ import torch
|
||||
from torch.utils.data import DataLoader, Subset
|
||||
from tqdm import tqdm
|
||||
|
||||
from config import Config, build_argparser, apply_overrides, validate_config_values
|
||||
from config import (Config, build_argparser, apply_overrides,
|
||||
validate_config_values, warn_init_from_mismatch)
|
||||
from action_norm import normalize_action, is_range_mode
|
||||
from bins import (encode as bin_encode, decode as bin_decode,
|
||||
fixed_zero_columns, START, HEAD_COLS)
|
||||
from model import ActionPredictor
|
||||
from dataset import ActionDataset
|
||||
from action_norm import default_column_weights
|
||||
|
||||
|
||||
def split_episodes(data_dir: str, val_ratio: float, seed: int = 42):
|
||||
@@ -20,39 +25,164 @@ def split_episodes(data_dir: str, val_ratio: float, seed: int = 42):
|
||||
return indices[val_count:], indices[:val_count]
|
||||
|
||||
|
||||
def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = None):
|
||||
def _weighted_mse(pred: torch.Tensor, target: torch.Tensor,
|
||||
weights: torch.Tensor = None) -> torch.Tensor:
|
||||
"""逐列加权 MSE:mean(w·(pred-target)²)/mean(w),weights None 时退化为普通 MSE。
|
||||
|
||||
除以 mean(w) 保持损失量级与未加权口径可比(w 全 1 时与普通 MSE 完全一致)。
|
||||
"""
|
||||
if weights is None:
|
||||
return torch.nn.functional.mse_loss(pred, target)
|
||||
return (weights * (pred - target) ** 2).mean() / weights.mean()
|
||||
|
||||
|
||||
def _collate(batch):
|
||||
"""多步样本({"steps":[...]})逐 k 聚合成 batch 列表;单步样本走默认 collate。
|
||||
|
||||
rollout_len==1 时 __getitem__ 返回标准 dict,直接默认 collate;
|
||||
rollout_len>1 时返回 {"steps":[step_0..step_{K-1}]},把每个 step_k 跨样本堆叠,
|
||||
返回 {"steps":[batch_0..batch_{K-1}]},供 train_one_epoch 按步 scheduled sampling。
|
||||
"""
|
||||
from torch.utils.data._utils.collate import default_collate
|
||||
if batch and isinstance(batch[0], dict) and "steps" in batch[0]:
|
||||
k = len(batch[0]["steps"])
|
||||
return {"steps": [default_collate([b["steps"][i] for b in batch])
|
||||
for i in range(k)]}
|
||||
return default_collate(batch)
|
||||
|
||||
|
||||
def _onehot_ce(logits: torch.Tensor, targets: torch.Tensor,
|
||||
weights: torch.Tensor, action_dim: int) -> torch.Tensor:
|
||||
"""多头加权交叉熵:每组 softmax 与 one-hot(或软标签)求 CE,按列权重加权
|
||||
|
||||
logits/targets: (B, TOTAL_BINS);weights: (action_dim,) 列权重向量。
|
||||
软标签(sigma>0)各组概率和恒 1,CE 定义一致。
|
||||
"""
|
||||
ls = None
|
||||
head_ce = []
|
||||
for c in HEAD_COLS:
|
||||
if c >= action_dim:
|
||||
break
|
||||
s, e = START[c], START[c + 1]
|
||||
# 逐头 log_softmax(不能整行:各头 bin 数不同,且整行 max 会
|
||||
# 把其他头压到 -1e9 量级,CE 爆炸)
|
||||
head_ce.append(-(targets[:, s:e]
|
||||
* torch.log_softmax(logits[:, s:e], dim=-1)).sum(-1))
|
||||
head_ce = torch.stack(head_ce, dim=1) # (B, n_heads)
|
||||
w = weights[:head_ce.size(1)]
|
||||
return (head_ce * w).sum() / (w.sum() * head_ce.size(0))
|
||||
|
||||
|
||||
def _shift_history(history: torch.Tensor, new_action: torch.Tensor) -> torch.Tensor:
|
||||
"""动作历史左移一步、追加 new_action 到末尾(scheduled sampling 回传自预测用)。
|
||||
|
||||
history: (B, H, D),new_action: (B, D) → (B, H, D)
|
||||
"""
|
||||
return torch.cat([history[:, 1:], new_action.unsqueeze(1)], dim=1)
|
||||
|
||||
|
||||
def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = None,
|
||||
rollout_len: int = 1, ss_prob: float = 1.0,
|
||||
predict_delta: bool = False, weights: torch.Tensor = None,
|
||||
onehot: bool = False, range_mode: str = None,
|
||||
action_dim: int = 18):
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
rng = np.random.RandomState()
|
||||
for batch in tqdm(loader, desc=" train", leave=False):
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
optimizer.zero_grad()
|
||||
pred = model(images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
if rollout_len == 1:
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
imu = batch["imu"].to(device) if "imu" in batch else None
|
||||
pred = model(images, history, imu)
|
||||
loss = (_onehot_ce(pred, target, weights, action_dim) if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
else:
|
||||
# 多步 scheduled sampling:首步必用真历史;后续步以 ss_prob 用真历史、
|
||||
# 否则用上一步模型预测(detach 截断,避免梯度跨步爆炸)回传构造历史。
|
||||
loss = 0.0
|
||||
bs = 0
|
||||
history = None
|
||||
for k, step in enumerate(batch["steps"]):
|
||||
images = step["images"].to(device)
|
||||
true_history = step["action_history"].to(device)
|
||||
target = step["target"].to(device)
|
||||
prev_action = step["prev_action"].to(device)
|
||||
imu = step["imu"].to(device) if "imu" in step else None
|
||||
if k == 0 or rng.random_sample() < ss_prob:
|
||||
history = true_history
|
||||
# 否则沿用上一轮回传的 history(自预测构造)
|
||||
pred = model(images, history, imu)
|
||||
loss = loss + (_onehot_ce(pred, target, weights, action_dim)
|
||||
if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
# 回传:onehot 先期望值解码(原始空间)再按范围归一化;
|
||||
# 残差模式把 Δ̂ 加回 a_{t-1} 得到绝对动作;再移入历史
|
||||
with torch.no_grad():
|
||||
if onehot:
|
||||
raw = bin_decode(pred.detach().cpu().numpy(),
|
||||
action_dim)
|
||||
# 权重 0 列固定为常量,防自预测垃圾进入历史
|
||||
fixed = fixed_zero_columns(
|
||||
weights.cpu().numpy(), action_dim)
|
||||
w = weights.cpu().numpy()[:action_dim]
|
||||
raw[w <= 0.0] = fixed[w <= 0.0]
|
||||
abs_action = torch.from_numpy(normalize_action(
|
||||
raw, action_dim, mode=range_mode)).to(device)
|
||||
else:
|
||||
abs_action = (prev_action + pred
|
||||
if predict_delta else pred)
|
||||
history = _shift_history(true_history, abs_action.detach())
|
||||
loss = loss / len(batch["steps"])
|
||||
loss.backward()
|
||||
if max_grad_norm:
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
|
||||
optimizer.step()
|
||||
total_loss += loss.item() * images.size(0)
|
||||
count += images.size(0)
|
||||
total_loss += loss.item() * bs
|
||||
count += bs
|
||||
return total_loss / max(count, 1)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate(model, loader, device):
|
||||
def validate(model, loader, device, rollout_len: int = 1, predict_delta: bool = False,
|
||||
weights: torch.Tensor = None, onehot: bool = False,
|
||||
action_dim: int = 18):
|
||||
model.eval()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
for batch in tqdm(loader, desc=" val ", leave=False):
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
pred = model(images, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
total_loss += loss.item() * images.size(0)
|
||||
count += images.size(0)
|
||||
if rollout_len == 1:
|
||||
images = batch["images"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
imu = batch["imu"].to(device) if "imu" in batch else None
|
||||
pred = model(images, history, imu)
|
||||
loss = (_onehot_ce(pred, target, weights, action_dim) if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
else:
|
||||
# 验证固定全真历史(确定性,不带 scheduled sampling 随机性),
|
||||
# 对 K 步累加,保证 best 模型选择口径稳定可比。
|
||||
loss = 0.0
|
||||
bs = 0
|
||||
for step in batch["steps"]:
|
||||
images = step["images"].to(device)
|
||||
history = step["action_history"].to(device)
|
||||
target = step["target"].to(device)
|
||||
imu = step["imu"].to(device) if "imu" in step else None
|
||||
pred = model(images, history, imu)
|
||||
loss = loss + (_onehot_ce(pred, target, weights, action_dim)
|
||||
if onehot
|
||||
else _weighted_mse(pred, target, weights))
|
||||
bs = images.size(0)
|
||||
loss = loss / len(batch["steps"])
|
||||
total_loss += loss.item() * bs
|
||||
count += bs
|
||||
if count == 0:
|
||||
return float("inf") # 空验证集不参与比较,避免假 best
|
||||
return total_loss / count
|
||||
@@ -76,24 +206,64 @@ def main():
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Device: {device}")
|
||||
|
||||
train_indices, val_indices = split_episodes(
|
||||
cfg.data.data_dir, cfg.train.val_split, seed=cfg.train.seed)
|
||||
print(f"Episodes: {len(train_indices)} train, {len(val_indices)} val")
|
||||
if cfg.train.val_split_mode == "sample":
|
||||
# 样本级随机划分:先建全量数据集取样本索引,再按 seed 随机分成
|
||||
# train/val(zscore 统计量仍只由训练样本计算,经 stats 文件复用)
|
||||
probe_ds = ActionDataset(
|
||||
cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
|
||||
cfg.data.image_width, cfg.data.image_height, False,
|
||||
cfg.data.action_stats_path, None, cfg.data.frame_len,
|
||||
predict_delta=cfg.train.predict_delta,
|
||||
rollout_len=cfg.train.rollout_len, imu_dim=cfg.data.imu_dim,
|
||||
normalize_mode=cfg.data.normalize_mode)
|
||||
n = len(probe_ds)
|
||||
rng = np.random.RandomState(cfg.train.seed)
|
||||
perm = rng.permutation(n)
|
||||
val_count = max(1, int(n * cfg.train.val_split))
|
||||
val_idx = perm[:val_count]
|
||||
train_idx = perm[val_count:]
|
||||
train_indices = None
|
||||
val_indices = None
|
||||
print(f"Sample split: {n - val_count} train, {val_count} val "
|
||||
f"(mode=sample)")
|
||||
else:
|
||||
train_indices, val_indices = split_episodes(
|
||||
cfg.data.data_dir, cfg.train.val_split, seed=cfg.train.seed)
|
||||
print(f"Episodes: {len(train_indices)} train, {len(val_indices)} val")
|
||||
|
||||
train_ds = ActionDataset(
|
||||
cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions,
|
||||
cfg.data.action_stats_path, train_indices, cfg.data.frame_len)
|
||||
if cfg.data.normalize_actions and not cfg.data.action_stats_path:
|
||||
cfg.data.action_stats_path, train_indices, cfg.data.frame_len,
|
||||
sample_indices=train_idx if cfg.train.val_split_mode == "sample" else None,
|
||||
augment=cfg.train.augment,
|
||||
predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len,
|
||||
imu_dim=cfg.data.imu_dim, normalize_mode=cfg.data.normalize_mode,
|
||||
onehot=cfg.data.onehot, onehot_sigma=cfg.data.onehot_sigma,
|
||||
action_jitter=cfg.data.action_jitter)
|
||||
auto_stats = False
|
||||
if (cfg.data.normalize_actions and cfg.data.normalize_mode == "zscore"
|
||||
and not cfg.data.action_stats_path):
|
||||
# 归一化统计量必须只用训练集计算,再供验证集复用,否则 val loss 不可信
|
||||
# (fixed_range 模式免 stats 文件,跳过此段)
|
||||
stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy")
|
||||
np.save(stats_path, train_ds.get_action_stats())
|
||||
print(f"Action stats saved to: {stats_path}")
|
||||
# 验证集构建需要完整路径(相对路径按 cwd 解析会找不到文件)
|
||||
cfg.data.action_stats_path = stats_path
|
||||
auto_stats = True
|
||||
val_ds = ActionDataset(
|
||||
cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions,
|
||||
cfg.data.action_stats_path, val_indices, cfg.data.frame_len)
|
||||
cfg.data.action_stats_path, val_indices, cfg.data.frame_len,
|
||||
sample_indices=val_idx if cfg.train.val_split_mode == "sample" else None,
|
||||
predict_delta=cfg.train.predict_delta, rollout_len=cfg.train.rollout_len,
|
||||
imu_dim=cfg.data.imu_dim, normalize_mode=cfg.data.normalize_mode,
|
||||
onehot=cfg.data.onehot, onehot_sigma=cfg.data.onehot_sigma,
|
||||
action_jitter=cfg.data.action_jitter)
|
||||
if auto_stats:
|
||||
# checkpoint 配置只存裸文件名:推理端按 checkpoint 目录解析相对路径
|
||||
cfg.data.action_stats_path = "action_stats.npy"
|
||||
|
||||
print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")
|
||||
if len(train_ds) == 0:
|
||||
@@ -108,16 +278,33 @@ def main():
|
||||
loader_gen.manual_seed(cfg.train.seed)
|
||||
train_loader = DataLoader(train_ds, batch_size=cfg.train.batch_size,
|
||||
shuffle=True, num_workers=cfg.train.num_workers,
|
||||
generator=loader_gen)
|
||||
generator=loader_gen, collate_fn=_collate)
|
||||
val_loader = DataLoader(val_ds, batch_size=cfg.train.batch_size,
|
||||
shuffle=False, num_workers=cfg.train.num_workers)
|
||||
shuffle=False, num_workers=cfg.train.num_workers,
|
||||
collate_fn=_collate)
|
||||
|
||||
model = ActionPredictor(cfg.model, cfg.data).to(device)
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
print(f"Model params: {n_params:,}")
|
||||
|
||||
# 逐列损失权重(仅训练/验证损失用,推理端无感知)
|
||||
col_w = cfg.data.column_weights
|
||||
if col_w is None:
|
||||
col_w = default_column_weights(cfg.data.action_dim)
|
||||
weights_t = torch.tensor(np.asarray(col_w, dtype=np.float32), device=device)
|
||||
print(f"Column weights: {[round(float(v), 4) for v in weights_t.tolist()]}")
|
||||
range_mode = (cfg.data.normalize_mode
|
||||
if (cfg.data.normalize_actions
|
||||
and is_range_mode(cfg.data.normalize_mode))
|
||||
else None)
|
||||
if cfg.data.onehot:
|
||||
print(f"One-hot 输出: {sum(1 for c in HEAD_COLS if c < cfg.data.action_dim)} 头"
|
||||
f" / {sum(START[c + 1] - START[c] for c in HEAD_COLS if c < cfg.data.action_dim)} bins"
|
||||
f"(sigma={cfg.data.onehot_sigma})")
|
||||
|
||||
if args.init_from:
|
||||
init_ckpt = torch.load(args.init_from, map_location="cpu", weights_only=True)
|
||||
warn_init_from_mismatch(init_ckpt.get("config", {}), cfg)
|
||||
missing, unexpected = model.load_state_dict(
|
||||
init_ckpt["model_state_dict"], strict=False)
|
||||
if missing:
|
||||
@@ -141,6 +328,9 @@ def main():
|
||||
model.load_state_dict(resume_ckpt["model_state_dict"])
|
||||
optimizer.load_state_dict(resume_ckpt["optimizer_state_dict"])
|
||||
scheduler.load_state_dict(resume_ckpt["scheduler_state_dict"])
|
||||
# 注意:scheduler 状态(含 T_max)随 checkpoint 恢复,LR 余弦曲线
|
||||
# 按原调度继续;resume 时改 --epochs 只影响循环边界,不重塑 LR
|
||||
# 曲线(如需新曲线请重新训练,勿 resume)
|
||||
start_epoch = resume_ckpt["epoch"] + 1
|
||||
best_val_loss = resume_ckpt.get("best_val_loss", float("inf"))
|
||||
patience_counter = resume_ckpt.get("patience_counter", 0)
|
||||
@@ -149,28 +339,42 @@ def main():
|
||||
|
||||
for epoch in range(start_epoch, cfg.train.epochs + 1):
|
||||
print(f"\nEpoch {epoch}/{cfg.train.epochs} lr={optimizer.param_groups[0]['lr']:.6f}")
|
||||
train_loss = train_one_epoch(model, train_loader, optimizer, device,
|
||||
cfg.train.max_grad_norm)
|
||||
val_loss = validate(model, val_loader, device)
|
||||
train_loss = train_one_epoch(
|
||||
model, train_loader, optimizer, device,
|
||||
cfg.train.max_grad_norm,
|
||||
rollout_len=cfg.train.rollout_len,
|
||||
ss_prob=cfg.train.ss_prob,
|
||||
predict_delta=cfg.train.predict_delta,
|
||||
weights=weights_t, onehot=cfg.data.onehot,
|
||||
range_mode=range_mode, action_dim=cfg.data.action_dim)
|
||||
val_loss = validate(model, val_loader, device,
|
||||
rollout_len=cfg.train.rollout_len,
|
||||
predict_delta=cfg.train.predict_delta,
|
||||
weights=weights_t, onehot=cfg.data.onehot,
|
||||
action_dim=cfg.data.action_dim)
|
||||
scheduler.step()
|
||||
print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}")
|
||||
|
||||
if val_loss < best_val_loss:
|
||||
best_val_loss = val_loss
|
||||
patience_counter = 0
|
||||
torch.save({
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"val_loss": val_loss,
|
||||
"config": cfg.to_dict(),
|
||||
}, best_path)
|
||||
inference_path = os.path.join(cfg.train.checkpoint_dir, "best_inference.pt")
|
||||
torch.save({
|
||||
"model_state_dict": model.state_dict(),
|
||||
"config": cfg.to_dict(),
|
||||
}, inference_path)
|
||||
print(f" -> saved best (val_loss={val_loss:.6f})")
|
||||
if not np.isfinite(val_loss):
|
||||
print(f"WARNING: val_loss 非有限值({val_loss}),跳过 best 更新")
|
||||
patience_counter += 1
|
||||
else:
|
||||
best_val_loss = val_loss
|
||||
patience_counter = 0
|
||||
torch.save({
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"val_loss": val_loss,
|
||||
"config": cfg.to_dict(),
|
||||
}, best_path)
|
||||
inference_path = os.path.join(cfg.train.checkpoint_dir, "best_inference.pt")
|
||||
torch.save({
|
||||
"model_state_dict": model.state_dict(),
|
||||
"config": cfg.to_dict(),
|
||||
}, inference_path)
|
||||
print(f" -> saved best (val_loss={val_loss:.6f})")
|
||||
else:
|
||||
patience_counter += 1
|
||||
if patience_counter >= cfg.train.patience:
|
||||
@@ -196,6 +400,15 @@ def main():
|
||||
"model_state_dict": model.state_dict(),
|
||||
"config": cfg.to_dict(),
|
||||
}, os.path.join(cfg.train.checkpoint_dir, "best_inference.pt"))
|
||||
torch.save({
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"scheduler_state_dict": scheduler.state_dict(),
|
||||
"best_val_loss": best_val_loss,
|
||||
"patience_counter": patience_counter,
|
||||
"config": cfg.to_dict(),
|
||||
}, os.path.join(cfg.train.checkpoint_dir, "best.pt"))
|
||||
print("WARNING: saved final weights as best_inference.pt (no valid validation loss)")
|
||||
else:
|
||||
print(f"Best model saved to: {best_path}")
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
"""actpred 共享工具函数(无 torch 依赖,推理/部署文件可复用)"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def parse_action_history(s: str, action_dim: int) -> np.ndarray:
|
||||
"""解析冒号分隔的动作历史字符串为 (H, D) float32 数组"""
|
||||
steps = s.split(":")
|
||||
history = []
|
||||
for step in steps:
|
||||
values = [float(v) for v in step.split(",")]
|
||||
if len(values) != action_dim:
|
||||
raise ValueError(
|
||||
f"Expected {action_dim} values per step, got {len(values)}")
|
||||
history.append(values)
|
||||
return np.array(history, dtype=np.float32)
|
||||
+320
-17
@@ -17,12 +17,14 @@
|
||||
- 🧬 **智能体自我进化**:通过 `update_soul` / `update_identity` 工具动态更新角色定位与用户画像
|
||||
- 🤖 **子代理**:通过 `create_subagent` 工具派生出拥有独立上下文/角色设定的轻量子代理,主代理可分解长任务、隔离上下文,子代理可选用不同模型
|
||||
- 🔄 **技能热重载**:无需重启即可添加、修改或删除扩展技能
|
||||
- 🔐 **SOCKS5 代理**:通过 `set_proxy` / `get_proxy` 工具为网页访问与 API 调用配置代理,可持久化
|
||||
- 🔐 **SOCKS5 代理**:通过 `set_proxy` / `get_proxy` 工具为网页访问与 API 调用配置代理,可持久化;密码明文仅存内存不落盘,跨重启经 `password_env` 环境变量注入(对齐 `api_key_env` 惯例)
|
||||
- 💾 **对话管理**:历史对话归档、开始新对话(重置 AI 记忆并重新加载提示词)
|
||||
- 📦 **对话压缩**:消息数或 Token 数达到阈值时自动压缩,保持上下文质量
|
||||
- 🛡️ **内存安全防护**:文件读取按需加载不全载、命令输出 1MB 限制、单行超长截断、工具结果多层防护防止 OOM
|
||||
- 🔒 **错误自动恢复**:API 请求失败自动回滚消息,工具结果安全清洗(控制字符过滤、非字符串转换)
|
||||
- 🛑 **可靠中断机制**:支持 Ctrl+C 实时中断,通过 `threading.Event` + HTTP 连接关闭 + 子进程 kill 三层取消,流式响应和长时间运行命令均可立即终止
|
||||
- 💻 **CLI 一次性对话**:`--once "问题"` 不启动 WebSocket 直接执行并退出,stdout/stderr 分流可管道组合,支持 `--quiet` / `--show-reasoning` / `--worker-only` / `--model` / `--temperature` / `--once-timeout`
|
||||
- 🔀 **多模型档案**:`config.json` 的 `models` 节可配置多个 OpenAI Chat Completions 兼容端点(DeepSeek/Kimi/Qwen/本地 Ollama 等),按名称切换,档案级 api_url/model/api_key_env/reasoning_enabled/timeout/temperature/display_name;`--model` 支持 `model:`/`profile:` 前缀逃逸
|
||||
|
||||
## 🚀 快速开始
|
||||
|
||||
@@ -55,21 +57,79 @@ python aibox.py
|
||||
|
||||
启动后,WebSocket 服务器将运行在 `ws://localhost:8765`(地址与端口可在 `config.json` 的 `system.websocket_host` / `system.websocket_port` 中修改)。
|
||||
|
||||
> 注意:本模块入口文件为 `aibox.py`(单文件约 8450 行,集成 WebSocket 服务器、AI 引擎、工具/技能系统、任务调度器于一体)。`client.py` 为示例 WebSocket 客户端,`voice_client.py` 为语音客户端,`_run.sh` 为被 `core` 拉起的启动脚本。
|
||||
### 管理大模型接入点(不用手写配置)
|
||||
|
||||
不想手写 `config.json` 的 `models` 节?用交互式菜单工具:
|
||||
|
||||
```bash
|
||||
python model_config.py # 交互式菜单:列出 / 添加 / 删除 / 连通性验证
|
||||
python model_config.py --list # 非交互:列出当前接入点
|
||||
```
|
||||
|
||||
- **添加接入点**:按提示填写档案名、API 地址(附常见厂商端点参考)、模型名;Key 两种方式——环境变量名(推荐,Key 不落盘)或免密本地端点(Ollama 等);可选显示名/温度/超时/思考模式/失败转移(从已有接入点多选)/设为默认
|
||||
- **删除接入点**:编号选择 + 确认,自动清理其他档案的失败转移引用
|
||||
- **连通性验证**:对选中接入点发一次最小请求,直接给出成败原因
|
||||
- 添加后可立即使用:`python aibox.py --model <档案名>` 指定,或设为默认后无需 `--model`
|
||||
|
||||
> 注意:本模块入口文件为 `aibox.py`(当前约 3408 行,集成 WebSocket 服务器、AI 引擎、工具/技能系统、任务调度器于一体;自 v3.0 起工具类与子系统已抽离至 `tools/`、`subsystems/`、`common/`、`tests/` 包,aibox.py 仅保留核心类与编排逻辑)。`client.py` 为示例 WebSocket 客户端,`voice_client.py` 为语音客户端,`llm_client.py` 为多厂商大模型客户端抽象层,`_run.sh` 为被 `core` 拉起的启动脚本。
|
||||
|
||||
## 📁 目录结构
|
||||
|
||||
```
|
||||
aibox/
|
||||
├── aibox.py # 服务主程序(WebSocket 服务器 + AI 引擎 + 工具/技能系统 + 任务调度器)
|
||||
├── aibox.py # 服务主程序(WebSocket 服务器 + AI 引擎 + 任务调度器 + 核心类 DeepSeekChat/WebSocketHandler/EnhancedAIManager)
|
||||
├── model_config.py # 大模型接入点交互式管理工具(列表/添加/删除/连通性验证,不用手写配置)
|
||||
├── client.py # 示例 WebSocket 客户端
|
||||
├── voice_client.py # 语音客户端
|
||||
├── llm_client.py # 大模型客户端抽象层(多厂商切换,DeepSeekClient 兼容 OpenAI 协议)
|
||||
├── entry.html # 前端入口页
|
||||
├── entry1.html # 前端入口页(变体)
|
||||
├── frontend.html # 前端主页面
|
||||
├── llm_client.py # 大模型客户端抽象层(多厂商切换)
|
||||
├── common/ # 基础设施包(i18n / 配置 / 常量 / 日志)
|
||||
├── tools/ # 工具系统基础包(Tool 基类 + ToolRegistry 注册器)
|
||||
│ ├── config.py # ConfigManager 配置管理
|
||||
│ ├── constants.py # 常量定义
|
||||
│ ├── i18n.py # 多语言
|
||||
│ └── logger.py # 日志器
|
||||
├── subsystems/ # 子系统包(Wave 1/2:从 aibox.py 抽离的基础设施)
|
||||
│ ├── executors.py # CommandExecutor + SubprocessCommandExecutor
|
||||
│ ├── prompt_manager.py # PromptManager(单例,提示词按数字前缀排序加载)
|
||||
│ ├── db.py # MemoryDatabase + MemoDatabase(SQLite,FTS5 trigram 中文检索 + LIKE 兜底)
|
||||
│ ├── skill_manager.py # SkillManager(技能发现/加载/热重载)
|
||||
│ ├── conversation_history.py # ConversationHistory(对话历史管理)
|
||||
│ ├── conversation_compressor.py # ConversationCompressor(对话压缩,注入 llm_client 解耦)
|
||||
│ └── process_manager.py # ProcessManager(后台进程管理,全局单例 + 环形输出缓冲 + atexit 收尾)
|
||||
├── tools/ # 工具系统包(A 批 + B 批 + 后台进程共 22 个工具类)
|
||||
│ ├── base.py # Tool 抽象基类
|
||||
│ ├── registry.py # ToolRegistry 工具注册器
|
||||
│ ├── basic_tools.py # A 批:ReadFile/WriteFile/ListFiles/SimpleWebCrawler/GetCurrentTime/Calculator
|
||||
│ ├── command_tool.py # B 批:CommandLineTool(run_command)
|
||||
│ ├── skill_tools.py # B 批:SearchSkills/ReloadSkills/InstallSkill/CreateSkill
|
||||
│ ├── db_tools.py # B 批:MemoryTool/MemoTool
|
||||
│ ├── prompt_tools.py # B 批:UpdateIdentity/UpdateSoul/SaveMemoryAndEndConversation
|
||||
│ ├── proxy_tools.py # B 批:SetProxy/GetProxy
|
||||
│ ├── misc_tools.py # B 批:AskUserTool
|
||||
│ └── process_tools.py # 后台进程:BgRun/BgWrite/BgList/BgOutput/BgKill(bg_run/bg_write/bg_list/bg_output/bg_kill)
|
||||
├── tests/ # 单元测试套件(598 项全过)
|
||||
│ ├── conftest.py # sys.path 配置 + websockets 桩
|
||||
│ ├── test_basic_tools.py # A 批工具 43 项
|
||||
│ ├── test_b_tools.py # B 批工具 58 项
|
||||
│ ├── test_executors.py # 命令执行器
|
||||
│ ├── test_prompt_manager.py # 提示词管理器
|
||||
│ ├── test_db.py # 记忆/备忘录数据库(含中文检索/迁移)
|
||||
│ ├── test_conversation_history.py # 对话历史
|
||||
│ ├── test_conversation_compressor.py # 对话压缩器
|
||||
│ ├── test_skill_manager.py # 技能管理器
|
||||
│ ├── test_once_mode.py # CLI 一次性对话模式(Batch A/B/C 15 项)
|
||||
│ ├── test_models_profile.py # 多模型档案(Batch D 17 项)
|
||||
│ ├── test_config.py # 配置脱敏(F1 3 项)
|
||||
│ ├── test_f2_fixes.py # 超时清理/线程收尾(F2 5 项)
|
||||
│ ├── test_g_fixes.py # 第三轮审计 Batch G(20 项)
|
||||
│ ├── test_h_fixes.py # 第四轮审计 Batch H(26 项)
|
||||
│ ├── test_i_fixes.py # 遗留收尾 Batch I(16 项)
|
||||
│ ├── test_subagent_max_turns.py # 子代理最大轮次行为回归(14 项)
|
||||
│ ├── test_k_fixes.py # 第五轮审计 Batch K(35 项)
|
||||
│ ├── test_model_config.py # 模型接入点管理工具(17 项)
|
||||
│ └── test_process_manager.py # 后台进程管理器与 bg_* 工具(39 项:生命周期/交互 I/O/收尾)
|
||||
├── skills/ # 内置扩展技能(*.py 可执行脚本 / *.md Markdown 技能)
|
||||
│ ├── hello-world.py
|
||||
│ ├── aes-crypto.py
|
||||
@@ -79,6 +139,7 @@ aibox/
|
||||
│ ├── ssh-exec.py
|
||||
│ ├── sftp-transfer.py
|
||||
│ └── ...
|
||||
├── prompts/ # 系统提示词源文件(运行期复制到 ~/.aibox/prompts/)
|
||||
├── SKILLS_FORMAT.md # 技能文件格式说明
|
||||
├── _run.sh # 启动脚本(被 core 守护时调用)
|
||||
├── _runcli.sh # 命令行启动脚本
|
||||
@@ -87,6 +148,26 @@ aibox/
|
||||
└── README.md
|
||||
```
|
||||
|
||||
**包依赖方向**(单向,禁止反向引用):
|
||||
|
||||
```
|
||||
common ← subsystems ← tools ← aibox.py
|
||||
```
|
||||
|
||||
- `common`:仅依赖标准库,不引用任何其他包
|
||||
- `subsystems`:依赖 `common`,不引用 `tools` 或 `aibox`
|
||||
- `tools`:依赖 `common` + `subsystems`,不引用 `aibox`
|
||||
- `aibox.py`:聚合层,导入上述三个包
|
||||
|
||||
**运行测试**:
|
||||
|
||||
```bash
|
||||
# 从仓库根目录运行
|
||||
py -m pytest aibox/tests -q
|
||||
# 当前:598 项全过(Linux;test_execute_timeout 为 Windows 向用例,本机无 ping 预存失败)
|
||||
# Windows 开发机:532 passed / 4 failed(平台环境项:symlink 权限、python3 缺失、POSIX 可执行位语义 ×2)/ 1 skipped
|
||||
```
|
||||
|
||||
运行期数据存放于 `~/.aibox/`:
|
||||
|
||||
```
|
||||
@@ -136,7 +217,7 @@ aibox/
|
||||
| `memory` | 保存、搜索(支持分页)、删除记忆 |
|
||||
| `memo` | 创建、查看、完成、搜索任务 |
|
||||
| `express_emotion` | 表达情绪并记录当前情绪状态(anger/joy/sadness/surprise),可指定 `persona` 从 `~/.aibox/skills/` 加载人格技能(如 angry/coquettish/tsundere)派生携带该人格的情绪子代理,保证人格独立、不污染主对话上下文 |
|
||||
| `create_subagent` | 派生子代理:独立上下文/角色设定,仅拥有执行面工具(无记忆/人格/派生/提问能力),可选独立模型,用于长任务分解、上下文隔离、独立角色专注处理 |
|
||||
| `create_subagent` | 派生子代理:独立上下文/角色设定,仅拥有执行面工具(无记忆/人格/派生/提问能力),可选独立模型(工具描述动态列出 `models` 节可用档案名,防 LLM 编造名称),用于长任务分解、上下文隔离、独立角色专注处理 |
|
||||
| `search_skills` | 搜索扩展技能(支持 Markdown 和可执行技能),查到后派生子代理执行 |
|
||||
| `reload_skills` | 热重载扩展技能 |
|
||||
| `create_skill` | 创建新技能(AI 自主编写 Python 脚本,实际写入经子代理完成) |
|
||||
@@ -146,7 +227,7 @@ aibox/
|
||||
| `update_soul` | 更新 AI 自我定义(`99_SOUL.md`,智能体进化) |
|
||||
| `save_memory_and_end_conversation` | 保存记忆并结束对话 |
|
||||
| `get_current_time` | 获取当前时间 |
|
||||
| `set_proxy` | 设置 SOCKS5 代理(host/port/可选账号密码,可持久化,即时生效) |
|
||||
| `set_proxy` | 设置 SOCKS5 代理(host/port/可选账号密码,可持久化,即时生效);`password` 明文仅本次会话有效不落盘,跨重启用 `password_env` 指定存有密码的环境变量名 |
|
||||
| `get_proxy` | 查看当前 SOCKS5 代理配置 |
|
||||
|
||||
### 子代理(worker)工具集
|
||||
@@ -160,6 +241,13 @@ aibox/
|
||||
| `web` | 访问网页、获取标题/文本/链接(返回 markdown 指引,改用 `curl` 命令行访问) |
|
||||
| `calculator` | 数学计算 |
|
||||
| `get_current_time` | 获取当前时间 |
|
||||
| `bg_run` | 后台启动长驻进程(dev server / 训练任务 / watch),立即返回句柄不阻塞;stdin 为管道可交互;默认 `detach=false`(aibox 退出时终止),`detach=true` 豁免但脱离句柄管理;`log_file=true` 同步镜像全程输出到 `workspace/bg_<句柄>.log`(无界全档,可被 run_command 的 tail/cat 原生读取);名称参数为 `proc_name`(≤64 字符,防与 `execute_tool` 首参撞名) |
|
||||
| `bg_write` | 向后台进程 stdin 写入一行(行协议交互:REPL / 问答式 CLI);返回 **watermark 水位**——`bg_output(since=watermark)` 得到的必然且仅是本条输入引发的响应(写读对齐,不含 `[stdin]` 回显行);`wait_lines=N` 同步等待 N 行响应(expect 式一步问答);单次上限 64KB,超时写入不阻塞工具线程;无 TTY,全屏程序(vim/top)不可用,建议启动时强制行缓冲(`python -i -u` / `PYTHONUNBUFFERED=1` / `stdbuf -oL`) |
|
||||
| `bg_list` | 列出全部后台进程(运行中在前,已退出保留尾部输出供验尸,LRU 淘汰) |
|
||||
| `bg_output` | 拉取输出:`tail=N` 尾部 N 行 / `since=游标` 增量续读(游标为上次返回的 `next_cursor`);单次默认 64KB、上限 1MB;环形缓冲(默认 256KB/进程)淘汰的早期行如实提示 |
|
||||
| `bg_kill` | 终止后台进程(整棵进程组):三级优雅终止——关 stdin(行协议进程见 EOF 常自行退出)→ SIGTERM 宽限 → SIGKILL 兜底;`force=true` 直接 SIGKILL;对已退出进程幂等 |
|
||||
|
||||
> **后台进程语义(v3.8,T22 属主隔离)**:注册表为全局单例(跨会话/跨子代理可见——list/output 观察面全局共享;写/杀按属主隔离,仅启动者所在会话可 `bg_write`/`bg_kill` 自己的进程,防提示注入跨会话破坏他人 dev server;未标识属主的调用不受限);用户中断(Ctrl+C/取消事件)**不**影响后台进程;aibox 进程退出经 atexit 统一收尾(非 detach 全杀,SIGKILL/SIGTERM 硬杀不覆盖);单行输出截断 16KB;运行中上限 16 / 总条目上限 32 / 单次 stdin 写入上限 64KB(`background_processes` 配置节可调);POSIX 按进程组终止(`start_new_session`,孙进程不悬挂),Windows `taskkill /T /F` 收割进程树。
|
||||
|
||||
## 💬 WebSocket API
|
||||
|
||||
@@ -452,15 +540,69 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
||||
"broadcast_enabled": true
|
||||
},
|
||||
"api": {
|
||||
"provider": "deepseek",
|
||||
"provider": "openai_compatible",
|
||||
"deepseek_api_url": "https://api.deepseek.com/chat/completions",
|
||||
"deepseek_model": "deepseek-v4-flash",
|
||||
"reasoning_enabled": false,
|
||||
"api_key_env": "DEEPSEEK_API_KEY",
|
||||
"timeout": 120,
|
||||
"temperature": 0.7
|
||||
"temperature": 0.7,
|
||||
"default_model": "deepseek"
|
||||
},
|
||||
"context": {
|
||||
"models": {},
|
||||
```
|
||||
|
||||
### 🔀 多模型档案(models 节)
|
||||
|
||||
`models` 可配置多个 OpenAI Chat Completions 协议兼容端点,按**档案名**切换(v3.2):
|
||||
|
||||
```json
|
||||
{
|
||||
"models": {
|
||||
"kimi": {
|
||||
"api_url": "https://api.moonshot.cn/v1/chat/completions",
|
||||
"model": "kimi-k2",
|
||||
"api_key_env": "MOONSHOT_API_KEY",
|
||||
"reasoning_enabled": false,
|
||||
"timeout": 120
|
||||
},
|
||||
"local": {
|
||||
"api_url": "http://localhost:11434/v1/chat/completions",
|
||||
"model": "qwen3:32b",
|
||||
"api_key_env": ""
|
||||
}
|
||||
},
|
||||
"api": { "default_model": "kimi" }
|
||||
}
|
||||
```
|
||||
|
||||
**档案字段**:
|
||||
|
||||
| 字段 | 说明 |
|
||||
|---|---|
|
||||
| `api_url` | **必填**。端点 URL(OpenAI Chat Completions 协议) |
|
||||
| `model` | **必填**(除非调用时传裸模型名)。该档案使用的模型名 |
|
||||
| `api_key_env` | API Key 环境变量名;**未配置回退 `api.api_key_env`;显式空串 `""` 为本地免密端点**(Ollama/vLLM)。档案端点环境变量取不到时**直接报错,绝不回退 `~/.cert/deepseek.key`**——该文件为 DeepSeek 专属凭据,跨厂商复用属于泄漏(key 文件回退仅对无档案的旧 `api` 扁平节生效) |
|
||||
| `reasoning_enabled` | 档案级思考模式开关(缺省回退 `api.reasoning_enabled`) |
|
||||
| `timeout` | 请求超时秒数 |
|
||||
| `temperature` | 档案级默认采样温度(可选)。优先级:CLI/调用方显式传入 > 本字段 > 全局默认;reasoning 模型建议按厂商文档设置 |
|
||||
| `display_name` | 展示名(可选,缺省用档案名)。`client.name` 用于对话上下文中的 assistant 名,不再恒显示 "DeepSeek" |
|
||||
| `fallback` | 失败转移档案名列表(可选,如 `["kimi", "local"]`,一层不递归)。主档案瞬时故障(HTTP 429/5xx、连接失败/超时)时按序转移;认证/参数类 4xx **不转移**(配置问题响亮暴露);失败档案进入 60s 冷却期(429 响应带 Retry-After 数字秒时按其延长)。`chat_completion`(对话压缩/轻量子代理)无条件转移;`stream_chat`(主对话流式)仅**首事件前**转移(已输出内容/用户中断绝不转移) |
|
||||
| `failover_timeout_budget` | 整条转移链的总时间预算秒数(可选,缺省不限)。耗尽即停止尝试剩余档案,防止长链故障时超时叠加(3 档案最坏 3×timeout);`FailoverClient.get_stats()` 暴露 requests/failovers/cooldown_skips/exhausted 计数,`on_failover(from, to, reason)` 回调供管理面/前端通知接入 |
|
||||
|
||||
**选择优先级**:`--model <档案名>`(命中 `models` 节)→ `api.default_model` 指向档案 → 回退旧 `api` 扁平节。`--model` 为**进程级覆盖**:WebSocket 服务模式(主对话、每个会话实例、定时任务实例)与 `--once` 一次性模式均生效。
|
||||
|
||||
`--model` 传**裸模型名**(未命中档案)时,沿用默认端点与 Key,仅覆盖模型名——即 "同端点换模型" 语义。
|
||||
|
||||
**前缀逃逸**(档案名与裸模型名歧义时):`--model model:NAME` 强制按裸模型名解析(档案恰好叫 NAME 时仍走裸名语义);`--model profile:NAME` 强制按档案名解析(不存在则明确报错,不再静默落回裸名语义)。
|
||||
|
||||
**兼容性**:`models` 留空 `{}` 时行为与旧版完全一致(走 `api` 节),已有配置无需任何改动。
|
||||
|
||||
其余配置节:
|
||||
|
||||
```json
|
||||
{
|
||||
"context": {
|
||||
"max_history_messages": 500,
|
||||
"max_context_messages": 120,
|
||||
"max_tokens": 8192,
|
||||
@@ -483,6 +625,14 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
||||
"commands": {
|
||||
"enabled": true
|
||||
},
|
||||
"background_processes": {
|
||||
"enabled": true,
|
||||
"max_processes": 16,
|
||||
"max_entries": 32,
|
||||
"buffer_per_process": 262144,
|
||||
"max_lifetime": 0,
|
||||
"max_write_bytes": 65536
|
||||
},
|
||||
"web": {
|
||||
"timeout": 10,
|
||||
"delay": 0,
|
||||
@@ -512,11 +662,14 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
||||
"llm_max_tokens": 800
|
||||
},
|
||||
"proxy": {
|
||||
"type": "",
|
||||
"host": "",
|
||||
"port": 0,
|
||||
"username": "",
|
||||
"password": ""
|
||||
"socks5": {
|
||||
"enabled": false,
|
||||
"host": "127.0.0.1",
|
||||
"port": 1080,
|
||||
"username": "",
|
||||
"password": null,
|
||||
"password_env": "SOCKS5_PASSWORD"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -524,6 +677,8 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
||||
> 默认模型为 `deepseek-v4-flash`(通过 `api.deepseek_model` 修改)。设置 `api.reasoning_enabled: true` 可启用思考模式,此时会在请求中添加 `thinking: {type: "enabled"}` 参数,模型会先进行思考再回答,实时展示思考过程(灰色斜体区域),回答质量更高且**支持工具调用**(文件操作、命令执行、记忆搜索等工具均可在思考模式下正常使用)。
|
||||
>
|
||||
> API Key 从环境变量 `DEEPSEEK_API_KEY` 读取(亦支持 `~/.cert/deepseek.key` 文件)。
|
||||
>
|
||||
> **代理密码不落盘**:`proxy.socks5.password` 仅在运行期内存中有效,保存配置时强制置空;跨重启持久化请将密码存入环境变量(如 `SOCKS5_PASSWORD`),配置 `password_env` 指向该变量名。加载时检测到旧版明文密码会自动迁移(本会话保留可用,文件立即重写脱敏)。
|
||||
>
|
||||
> **多厂商切换**:大模型调用已通过 `aibox/llm_client.py` 抽象层解耦,业务层不再直接依赖具体厂商。配置 `api.provider`(默认 `deepseek`)即可切换厂商,当前内置 `DeepSeekClient`(兼容 OpenAI Chat Completions 协议,Kimi / Qwen / GLM / Doubao 等绝大多数厂商均可复用)。新增厂商只需在此文件新增一个 `LLMClient` 子类并在 `create_llm_client` 工厂中注册,无需改动业务代码,为后续子代理/多模型共发预留了扩展点。
|
||||
|
||||
@@ -545,6 +700,14 @@ python aibox.py [选项]
|
||||
--memo-add "标题|内容|时间" 添加任务
|
||||
--memo-list 列出任务
|
||||
--memo-complete ID 完成任务
|
||||
--once QUERY 一次性对话:执行 QUERY 后退出(不启动 WebSocket)
|
||||
--quiet 一次性模式静默输出(stdout 仅回答正文)
|
||||
--show-reasoning 一次性模式输出思考流(默认静默,走 stderr)
|
||||
--no-save-memory 一次性模式禁用记忆工具(干净跳过,不进请求 schema)
|
||||
--worker-only 一次性模式纯执行面(worker profile,无人格/记忆工具)
|
||||
--model NAME 模型档案名(models 节)或裸模型名(同端点换模型);歧义时 model:NAME 强制裸名 / profile:NAME 强制档案;服务模式与 --once 均生效(进程级覆盖)
|
||||
--temperature T 覆盖配置的对话温度
|
||||
--once-timeout SEC 一次性模式全局超时(默认 300 秒)
|
||||
```
|
||||
|
||||
### 使用示例
|
||||
@@ -566,6 +729,30 @@ python aibox.py --memo-add "提醒开会|下午3点会议室A|+1h"
|
||||
python aibox.py --memo-list
|
||||
```
|
||||
|
||||
### 💻 CLI 一次性对话(--once)
|
||||
|
||||
不启动 WebSocket,单次执行后退出。回答正文走 **stdout**,工具过程/思考流走 **stderr**,可直接管道组合:
|
||||
|
||||
```bash
|
||||
# 基础问答
|
||||
python aibox.py --once "帮我看下当前日期"
|
||||
|
||||
# 管道组合(--quiet 静默模式)
|
||||
python aibox.py --once --quiet "列出当前目录 py 文件" | grep ".py" | wc -l
|
||||
|
||||
# 按档案名选模型 / 同端点换模型
|
||||
python aibox.py --once --model kimi "总结这段代码"
|
||||
python aibox.py --once --model deepseek-v4-lite "1+1?"
|
||||
|
||||
# 纯执行面(无人格/记忆,单次收敛,适合脚本调用)
|
||||
python aibox.py --once --worker-only "执行 echo hello 并返回输出"
|
||||
|
||||
# 显式开启思考流 / 禁用记忆写入
|
||||
python aibox.py --once --show-reasoning --no-save-memory "分析这个问题"
|
||||
```
|
||||
|
||||
**退出码**:`0` 成功 / `1` 工具执行失败 / `2` API 调用失败 / `3` 用户中断(Ctrl+C)/ `4` 全局超时。
|
||||
|
||||
## 🔄 架构流程
|
||||
|
||||
```
|
||||
@@ -632,10 +819,19 @@ Server -> Client: {type: "dialogue_reset", content: "✅ 对话状态已完全
|
||||
|
||||
### 添加新工具
|
||||
|
||||
1. 继承 `Tool` 基类(位于 `tools/base.py`)
|
||||
1. 继承 `Tool` 基类(位于 [tools/base.py](tools/base.py))
|
||||
2. 实现 `get_name()`、`get_description()`、`get_parameters()`、`execute()`
|
||||
3. 可选实现 `execute_streaming()` 支持流式输出(通过 `output_queue` 发送 `("line", content)` 和 `("complete", final_result)`)
|
||||
4. 在 `DeepSeekChat._register_tools()` 中注册
|
||||
4. 按工具类别放入对应文件:
|
||||
- 基础工具(文件/网页/计算/时间)→ `tools/basic_tools.py`
|
||||
- 命令执行 → `tools/command_tool.py`
|
||||
- 技能管理 → `tools/skill_tools.py`
|
||||
- 数据库工具 → `tools/db_tools.py`
|
||||
- 提示词/身份更新 → `tools/prompt_tools.py`
|
||||
- 代理配置 → `tools/proxy_tools.py`
|
||||
- 其他 → `tools/misc_tools.py`
|
||||
5. 在 [tools/__init__.py](tools/__init__.py) 的 `__all__` 中导出新类
|
||||
6. 在 `DeepSeekChat._register_tools()` 中按 `tool_profile`(orchestrator/worker)注册
|
||||
|
||||
```python
|
||||
class MyTool(Tool):
|
||||
@@ -656,6 +852,8 @@ class MyTool(Tool):
|
||||
return f"执行结果: {param}"
|
||||
```
|
||||
|
||||
> **依赖约束**:`tools/` 包只能依赖 `common` + `subsystems` + 标准库,不得反向引用 `aibox.py` 任何符号;如需访问 aibox 内部状态,应通过构造函数注入(参考 `MemoryTool(memory_db=...)` 模式)。
|
||||
|
||||
> **注意**:`execute()` 必须返回字符串。工具结果会经过多层安全防护:
|
||||
> 1. **工具层源头限制**:文件/网页工具返回 markdown 指引,实际 I/O 由 `run_command` 完成(`run_command` 输出≤1MB,超出立即终止进程;`read_file` 指引命令限制 500 行、单行≤10KB)
|
||||
> 2. **`_sanitize_tool_result()`**:None/非字符串转为安全字符串,过滤控制字符
|
||||
@@ -734,6 +932,111 @@ elif command == "/mycommand":
|
||||
|
||||
## 📄 更新日志
|
||||
|
||||
### v3.8
|
||||
- ✅ **后台进程池(智能体运行/管理长驻进程)**:新增 `subsystems/process_manager.py`(`ProcessManager` 全局单例)+ `tools/process_tools.py` 5 个 worker 执行面工具(`bg_run` / `bg_write` / `bg_list` / `bg_output` / `bg_kill`),弥补 `run_command` 同步阻塞、超时即杀的不足——dev server / 训练任务 / watch 等长驻进程可后台启动、跨会话管理:
|
||||
- **生命周期语义(三决策)**:用户中断(Ctrl+C/取消事件)不影响后台进程;aibox 退出经 atexit 统一收尾(非 detach 全杀,`detach=true` 豁免但脱离句柄管理);注册表跨会话/跨子代理共享(观察面全局可见;写/杀按属主隔离——T22,仅启动者所在会话可操作,定时任务 AI 为精简 `task` profile 且无 update_identity/update_soul/set_proxy)
|
||||
- **交互式 I/O(stdin 管道 + watermark 写读对齐)**:stdin 挂管道,`bg_write` 向行协议进程(REPL / 问答式 CLI)写入一行,返回 **watermark**(`[stdin]` 回显行序号)——`bg_output(since=watermark)` 得到的必然且仅是该输入引发的响应(写读配对);`wait_lines=N` expect 式同步等待;回显行入环形缓冲与 log 镜像(tail 可见完整交互时序);写入 watchdog 线程兜底(进程不读 stdin 管道写满时不阻塞工具线程);单次写入上限 64KB
|
||||
- **输出缓冲**:每进程有界环形缓冲(默认 256KB,`buffer_per_process` 可调),单行截断 16KB;`bg_output` 支持 `tail` 尾读与 `since` 游标增量续读(淘汰行如实提示),单次默认 64KB、上限 1MB;`bg_run log_file=true` 输出同步 tee 到 `workspace/bg_<句柄>.log`(无界全档,`[stdin]` 回显同入档,可被 run_command 的 tail/cat 原生读取)
|
||||
- **三级优雅终止**:`bg_kill` 关 stdin(行协议进程见 EOF 常自行退出)→ SIGTERM 宽限(grace 比例 4:6)→ SIGKILL 兜底;`force=true` 直接 SIGKILL;shutdown 收尾同样先关 stdin
|
||||
- **安全边界**:运行中上限 16 / 总条目上限 32(LRU 淘汰已退出条目);`background_processes` 配置节(含 `enabled` 双闸:`commands.enabled` 关闭时一并拒绝);`proc_name` 标签 ≤64 字符且拒绝控制字符(M5 惯例);POSIX `start_new_session` + 进程组终止(对齐执行器 L9 语义),Windows `taskkill /T /F` 收割进程树;TTL `max_lifetime` 惰性清理(默认 0 不启用);凭据类交互(sudo -S 喂密码)会留对话记录,敏感凭据走环境变量注入(工具描述已指引)
|
||||
- **实现要点**:H2 reader 线程模式复用(EOF 即 poll 收割防僵尸);`list_procs` 惰性刷新全表状态;`proc_name` 命名规避 `ToolRegistry.execute_tool(name=...)` 首参撞名;编排者误调拦截表 `WORKER_ONLY_TOOLS` 同步纳入;无 TTY 的适用边界(全屏程序不可用、行缓冲指引)写入工具描述供模型自检
|
||||
- **测试**:新增 `tests/test_process_manager.py` 39 项(生命周期/游标续读/环形淘汰/单行截断/上限/TTL/收尾/工具分发/禁用闸/watermark 对齐/回显/sync 等待/超时/EOF 优雅终止/log tee/注册表分发),Windows 开发机 532 passed / 4 failed(平台环境项,与代码无关)
|
||||
|
||||
### v3.7
|
||||
- ✅ **第六轮安全审计 + Batch L 修复(2026-08-20,交叉验证轮:Batch K 回归验证 / 并发与生命周期 / 降级路径 / 修复间相互作用,4 路并行审读 + 主审逐条复核)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md) §16):
|
||||
- **Batch K 回归验证**:J1~J16 逐项在位、零回归,无修复引入的缺陷;历届 Batch E~I 修复抽查(executors reader 收尾、1MB 输出上限、registry 流式预算、broadcast 快照等)均在位
|
||||
- **Batch L(R1~R3 低危全落地)**:R1 `/memos` WS 命令 strptime 容错(畸形 reminder_time 按"待执行"展示,不再整体"处理失败");R2 移除 `handle_interrupt` 对全局 `ai_manager.command_executor` 的取消(WS 中断不再误杀 CLI 主 AI 子进程);R3 `unregister` 补齐回收(关闭 websocket + 取消在途处理任务,broadcast 超时路径与连接关闭路径行为对齐,半开连接不再泄漏)
|
||||
- **测试**:新增 `tests/test_l_fixes.py` 6 例(R1 容错与语义不回归 ×2 / R2 范围收窄与中断能力不回归 ×2 / R3 回收与幂等 ×2);Linux 维持 490 passed,Windows 开发机 493 passed / 4 failed(均为平台环境项,与代码无关)
|
||||
|
||||
### v3.6
|
||||
- ✅ **交互式大模型接入点管理([model_config.py](model_config.py))**:菜单式 列表 / 添加 / 删除 / 连通性验证,不用手写 `config.json`;Key 仅记录环境变量名不落盘(空 = 免密本地端点);删除自动清理其他档案的 fallback 引用;默认档案经 `api.default_model` 生效(详见 [模型接入点管理](#管理大模型接入点不用手写配置))
|
||||
- ✅ **第五轮安全审计(2026-08-19,4 路并行审读 + Linux PoC 2/2 实证,发现 J1~J16:高危 2 / 中危 5 / 低危 6 / 信息 3)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md) §13~§14):
|
||||
- **K1(J1 高危·任务系统投毒)**:畸形时间/数字字段经 CLI/WS/memo 工具三入口零校验入库 → 调度器失败热循环 + 每个新连接握手即崩——入库端统一校验(`reminder_time`/`repeat_end_time` strptime 预检、`repeat_interval` 强转 int、`repeat_type` 白名单)、调度器预检失败按到期执行(Q5 语义)、`get_pending_tasks_for_user` 解析失败按到期展示、`register` 移入 try/finally(会话登记不再泄漏);附带修复 `add_user_task` 缺失 `repeat_end_time` 参数透传
|
||||
- **K2(J2 高危·前端 XSS)**:三页面 7 处 `marked.parse` 直渲(AI 输出/命令结果全链路 innerHTML)统一改经 `renderMarkdown()` + DOMPurify(缺失时退化为转义文本防 XSS 兜底);`entry1.addSystemMessage` 改 DOM 构建 + textContent;任务标题改 createTextNode;CDN 脚本(marked/hljs/socket.io/DOMPurify)补 SRI integrity(J13)
|
||||
- **K3(中危)**:FTS 外部内容表删除换序——先删 `memories_fts` 再删 `memories`,索引残留消除(J3,PoC 实证);config/prompt/对话写路径统一原子写(`common/atomic.py`,mkstemp + fsync + os.replace)+ 损坏备份 `config.json.corrupt.<时间戳>` 拒绝覆盖(J5);PromptManager/ConfigManager 单例初始化纳入锁(J6,并发首例不再双重加载系统提示词)
|
||||
- **K4(中危)**:`model_config.py` 连通性验证异常消息复用 `_sanitize_log_msg` 脱敏(J4,Bearer 令牌不再进终端);`broadcast` 改每客户端并发发送 + `asyncio.wait_for` 超时(`broadcast_send_timeout` 默认 5s),超时即 unregister——僵死客户端(半开 TCP)不再卡死广播与调度器循环(J7)
|
||||
- **K5(低危随批)**:J8 交互输入容错(非法温度/超时友好报错);J9 MemoDatabase WAL + busy_timeout=10000 覆盖全部写路径;J10 memo 批量删除补实现(`memo_ids`);J11 `/new` 取消-重置竞态(排队 reset 等待 is_processing 清零);J12 skills 注入面(nmap hosts `shlex.quote`、screenshot PowerShell 单引号转义、SKILLS_FORMAT.md 新增"安全边界"节);J14 voice_client 死代码/僵尸连接/录音残留清理;J15 死代码移除(`full_response_queue` 等)+ 未知 task action 明确报错 + 广播门控对齐 + 热路径日志降级;J16 `--help` 输出 2000 上限 + 同名技能键确定性解决
|
||||
- ✅ **测试**:490 项全过(Linux 环境;第五轮新增 Batch K 35 项 + 子代理 max_turns 回归 14 项;`test_execute_timeout`/`test_flat_api_still_reads_key_file` 为 Windows 向用例,Linux 预存失败与本轮无关)
|
||||
|
||||
### v3.5
|
||||
- ✅ **安全审计与修复(第三轮 Batch G + 第四轮 Batch H + 遗留收尾 Batch I)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md)):
|
||||
- **Batch G(llm 加固 + 线程生命周期 + 取消传播 + DB 事务)**:
|
||||
- llm_client:非瞬时 HTTP 错误体截断 `[:200]`(密钥/错误页不进异常栈)、Retry-After 上限 cap 3600s、fallback 档案缺 model 跳过并告警、fallback 去重、删除 `LLMClient._cancelled` 死状态、FailoverClient 并发契约文档化
|
||||
- 子代理取消传播:`run_full_subagent` 共享父 `_cancel_event` + 0.2s 观察者线程转发深度中断(强关在途 HTTP、kill 子进程);`SubAgentTool` 构造注入父事件
|
||||
- drain 线程重构:批次级单线程 + stop 事件 + try/finally 统一回收(消除每次工具调用泄漏一个常驻线程);`/new` 的 `reset_conversation` 移入线程池(不再同步获取 `_msg_lock` 冻结事件循环);broadcast 收窄 except + unregister 改 discard;`_processing_tasks` 改按连接的任务集合登记(关闭时取消全部在途任务)
|
||||
- DB:FTS5 关键词双引号翻倍转义(`"` → `""`)、两处 `LIMIT` f-string 参数化、`mark_as_triggered` 读写收进单连接 `BEGIN IMMEDIATE` 事务(消灭并发触发竞态);历史索引裸 except 收窄 + 归档文件名同秒加序号
|
||||
- **Batch H(第四轮审计,3 路并行审读 + Linux PoC 6/6 实证)**:
|
||||
- **H1(高危)**:修复 N5 回归——子代理共享父取消事件后首轮 think 无条件 `clear()` 会吞掉父级已置位的取消信号(用户中断 / --once 超时被绕过),共享事件时跳过 clear + 父已取消直接早退
|
||||
- `/reload_prompts` 移入线程池(N3 同类遗漏);技能扫描(串行 `--help` 子进程)启动前预热,不再在首条消息/首个任务时冻结事件循环
|
||||
- 调度器任务级超时(`scheduler.task_timeout`,默认 600s)——AI think 卡死不再停摆整个调度循环
|
||||
- `get_recent_memories(limit=None)` 默认上限(不再全表加载);`memo delete` 改物理删除(此前只标记完成);`PromptManager.reload` 本地构建后原子替换(并发读者不再读到空/半截系统提示词)
|
||||
- read_file 的 sed 参数强转 int(防引号注入);`set_proxy` 落盘失败如实报告;SSE 非 UTF-8 行容错;persona 路径 realpath 校验(软链越界拒绝);1MB 截断无半个字符;compressor 除零防护;InstallSkillTool 指引命令路径引用;ask_user 缺参明确报错
|
||||
- **POSIX 进程组终止**:executor 的 `Popen` 加 `start_new_session` + 组信号 kill(组 ID 在 Popen 后立即捕获)——`(sleep 300 &)` 形态的孙进程持管道不再悬挂 reader 线程/泄漏 fd(实证修复前孙进程残留)
|
||||
- web 工具 scheme 白名单(http/https,`file://` 等拒绝);registry 日志打码改包含匹配(apiKey/passwd 变体);FailoverClient 日志 Bearer 令牌正则脱敏(网关回显 Authorization 头不再进日志)
|
||||
- **Batch I(遗留收尾)**:
|
||||
- **N9**:prompt 写入内容长度上限(`MAX_CONTENT_LENGTH=10000`,identity/soul/status 超长拒绝)——约束持久化注入面
|
||||
- **F5**:情绪子代理取消传播——`run_light_subagent` 改走 stream_chat + 观察者线程,用户中断可即时终止(此前阻塞至请求超时)
|
||||
- 提醒时间格式异常任务按"到期"返回(此前被 SQL 层字符串比较静默丢弃,任务永不执行);`mark_as_triggered` 加 `busy_timeout=10000`(多进程同库写竞争等待而非立即报 locked)
|
||||
- 对话文件原子写(tmp + os.replace,崩溃不留半截 JSON);`load_conversation` realpath 越界防护;WS 命令专用线程池(`/new`、`/reload_prompts` 不再占用默认池);--once 硬超时兜底(软超时后宽限一个周期强杀进程)
|
||||
- PromptManager 单例加锁;`limit=0` 语义修正(None 才不加 LIMIT)
|
||||
- ✅ **测试**:424 项全过(Linux 环境;Batch G 20 + Batch H 26 + Batch I 16 项新增;`test_execute_timeout`/`test_flat_api_still_reads_key_file` 为 Windows 向用例,Linux 预存失败与本轮无关)
|
||||
|
||||
### v3.4
|
||||
- ✅ **多模型档案加固(H1/H2/H3/M1/M2/M3 六项修复)**:
|
||||
- **H1(凭据隔离)**:档案端点环境变量取不到时直接报错,**绝不回退读 `~/.cert/deepseek.key`**(该文件为 DeepSeek 专属凭据,跨厂商复用属于泄漏);key 文件回退仅对无档案的旧 `api` 扁平节生效
|
||||
- **H2(schema 校验前移)**:档案缺 `api_url`/`model` 或类型错误时,解析期即抛出指明档案名与缺失字段的 `ValueError`——不再把档案名当模型名发出、不再延迟到 `requests InvalidURL` 或 API 400
|
||||
- **H3(--model 服务模式生效)**:`--model` 升级为进程级覆盖,WebSocket 服务模式(主对话/每个会话实例/定时任务实例)与 `--once` 均生效,此前服务模式静默忽略
|
||||
- **M2(防模型名幻觉)**:`create_subagent` 工具描述与 JSON Schema 动态携带 `models` 节可用档案名列表(默认档案标注),LLM 无从编造名称;`models` 为空时描述保持原样(零破坏)
|
||||
- **M1(换模型同步派生引用)**:新增 `DeepSeekChat.set_client()` 统一替换 LLM 客户端——同步刷新 `compressor.llm_client` 与 `api_key`/`name` 快照;`create_subagent` 子代理换模型经此入口,压缩链不再残留默认模型
|
||||
- **M3(reasoning 契约显式化)**:档案级 `reasoning_enabled` 从"构造后挂实例属性"改为 `LLMClient` 构造参数(三态:`True`/`False` 档案级显式、`None` 未配置回退全局),外部误设属性无从静默覆盖全局开关
|
||||
- ✅ **L 批体验与兼容修复(L1/L2/L3/L5 四项)**:
|
||||
- **L1(展示名按档案)**:`client.name` = 档案 `display_name` > 档案名,用 Kimi/Qwen 时对话上下文不再误标 "DeepSeek";旧扁平节保持 "DeepSeek"(零破坏)
|
||||
- **L2(前缀逃逸)**:`--model model:NAME` 强制裸模型名语义、`--model profile:NAME` 强制档案语义(不存在时明确报错并列出可用档案),档案名与模型名歧义可逃逸
|
||||
- **L3(档案级温度)**:档案 `temperature` 字段作为默认采样温度,优先级:CLI/调用方显式传入 > 档案字段 > 全局默认(0.7/0.3),主对话/轻量子代理/compressor 全链路生效
|
||||
- **L5(思考流字段别名)**:`_parse_delta` 识别 `reasoning_content`/`reasoning`/`thinking`/`thinking_content` 四种常见思考字段(首个命中生效),其他厂商思考流不再静默丢失
|
||||
- ✅ **L4a(chat_completion 跨档案失败转移)**:档案新字段 `fallback`(档案名列表,一层不递归);`create_llm_client` 构造 `FailoverClient` 包装链(主档案 + fallback)。瞬时故障(HTTP 429/5xx 抛 `TransientLLMError`、连接失败/超时)按链转移,认证/参数类 4xx 不转移(配置问题响亮暴露);失败档案 60s 冷却;fallback 档案缺 Key 等环境问题跳过 + warning 不阻断主档案;配置结构错误(类型/不存在/自引用)启动期硬性拦截。当前覆盖对话压缩与轻量子代理链路,流式(`stream_chat`)转移属 L4b
|
||||
- ✅ **L4b(stream_chat 首事件前转移 + 中断联动)**:流式仅在**未产出任何事件**(reasoning/content/tool_call)时转移——重放对用户不可见;已产出事件后失败直接上抛(防前端重复输出/工具增量错乱)。`LLMClient.cancel()` 置 `_cancelled` 标志,`FailoverClient.cancel()` 广播整条链并置包装器标志——用户中断引发的连接异常与真实故障不可区分,凭标志跳过转移(否则用户越中断、系统越换厂商重试);`chat_completion` 同受标志保护。链耗尽抛最后异常(既有回滚/报错处理接住),全链冷却抛 RuntimeError
|
||||
- ✅ **L4c(转移可观测性 + 预算控制 + Retry-After)**:`FailoverClient.get_stats()` 暴露 requests/failovers/cooldown_skips/exhausted 聚合计数;`on_failover(from, to, reason)` 回调(to=None 表示链耗尽,回调异常不影响主流程)供管理面/前端通知接入;档案字段 `failover_timeout_budget` 限制整链尝试总秒数(缺省不限,<=0 视为未配置),耗尽停止尝试剩余档案防超时叠加;429 响应 Retry-After(数字秒)经 `TransientLLMError.retry_after` 延长该档案冷却期(非数字格式忽略)
|
||||
- ✅ **测试**:355 项全过(H1 隔离 4 + H2 校验 6 + H3 传导 6 + M2 提示 5 + M1 替换 2 + M3 契约 4 + L 批 16 + L4a 转移 15 + L4b 流式 12 + L4c 可观测 13 项新增)
|
||||
|
||||
### v3.3
|
||||
- ✅ **安全审计与修复**(两轮审计报告见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md)):
|
||||
- **Batch E(命令执行保护链修复)**:`CommandLineTool.execute_streaming` 纯委托注入的 executor(删除绕过 1MB/超时/中断保护的 150 行重复 Popen 实现);`SubprocessCommandExecutor` 读取循环线程化(后台 reader 线程 + 0.1s 超时轮询,静默命令超时不再失效);两处 `Popen` 改 `stdin=DEVNULL`(等待输入的命令不再永久挂起工具线程);write_file 指引 heredoc 定界符唯一化;registry 流式转发循环加 3600s 总超时兜底。PoC 回归 5/5 + 全量单测通过
|
||||
- **Batch F(凭据保护 + 线程收尾)**:
|
||||
- **F1**:SOCKS5 代理密码不再明文落盘——`_save_config` 落盘前强制脱敏(纵深防御),新增 `password_env` 环境变量注入(对齐 `api_key_env` 惯例),旧版明文配置加载时自动迁移清除;`registry` debug 日志对 `password`/`api_key`/`token`/`secret`/`authorization` 参数打码
|
||||
- **F2**:registry 流式转发超时路径统一清理 `tool.streaming_output`/`tool.output_queue`(消除僵尸 daemon 线程经实例属性把输出混入后续会话队列的串流路径),`stream_tool_timeout` 提为实例属性(默认 3600s,测试可调);`SubprocessCommandExecutor` reader 线程收尾——正常 EOF 后显式 close 释放管道 fd,仍阻塞时有限 join 后放弃(daemon 兜底);实证修正:CPython 流锁下阻塞中的 readline 无法被 `close()` 中断(首版实现致超时被拖满命令时长,已按 PoC 回归修正)
|
||||
- ✅ **测试**:272 项全过(Batch E PoC 回归 + F1 配置脱敏 3 项 + F2 超时清理/线程收尾 5 项新增)
|
||||
|
||||
### v3.2
|
||||
- ✅ **CLI 一次性对话模式(`--once`)**(分三批落地):
|
||||
- **Batch A**:`--once QUERY` / `--quiet` / `--once-timeout`;不启动 WebSocket,执行后退出;回答走 stdout、过程走 stderr,可管道组合;SIGINT 与超时复用三层取消机制
|
||||
- **Batch B**:`--no-save-memory`(`enable_memory` 开关干净禁用 MemoryTool)、`--show-reasoning`(思考流默认静默);压缩提示改走 output_callback 不再污染 stdout;退出码 1(工具执行失败)
|
||||
- **Batch C**:SQLite WAL(多 CLI 进程并发写锁缓解)、`--worker-only`(纯执行面 + 精简 system prompt 防人格污染)、`--model` / `--temperature` 覆盖、子代理工具失败标记透传(`_SUBAGENT_TOOL_FAILED_MARKER` → orchestrator stats)
|
||||
- ✅ **多模型档案(`models` 配置节)**:配置多个 OpenAI Chat Completions 兼容端点(DeepSeek/Kimi/Qwen/本地 Ollama 等)按档案名切换;`--model` 支持档案名与裸模型名(同端点换模型)双语义;档案级 `api_key_env` 三级语义(回退 / 显式 / 空串免密);档案级 `reasoning_enabled` 优先于全局;`models: {}` 时零破坏回退旧 `api` 节;`llm_client.py` 增加 `openai_compatible` / `openai` provider 别名
|
||||
- ✅ **测试**:264 项全过(once 模式 15 项 + 多模型档案 17 项新增)
|
||||
|
||||
### v3.1
|
||||
- ✅ **修复 FTS5 中文搜索失效**([subsystems/db.py](subsystems/db.py)):
|
||||
- FTS5 虚拟表改用 `tokenize='trigram'` 分词器(需 SQLite >= 3.34,环境不支持时自动退回 unicode61)
|
||||
- 旧 unicode61 表在启动时自动迁移:DROP 重建后从 memories 表回灌全量数据
|
||||
- `search_memories` 增加 LIKE 模糊检索兜底(trigram 对 <3 字符中文短词无法命中;关键词中 `%`/`_` 已转义为字面字符;MATCH 语法错误不再抛异常)
|
||||
- ✅ **修复对话压缩 Token 阈值不生效**:
|
||||
- `should_compress` 补齐 Token 数判断:消息条数或 Token 数任一达到阈值即触发压缩(此前 `compress_token_threshold` 为死配置)
|
||||
- `compress_conversation` 配套:Token 触发(消息数未超阈值)时按比例压缩早期消息,不再因目标消息数计算而跳过
|
||||
- ✅ **ConversationCompressor 抽离至 [subsystems/conversation_compressor.py](subsystems/conversation_compressor.py)**(Wave 3):
|
||||
- 构造注入 `llm_client`(客户端抽象层)替代 AI 引擎实例,打破与 `DeepSeekChat` 的循环引用
|
||||
- aibox.py 3517 → 3018 行;配套 17 项单元测试(估算/阈值判断/压缩流程/LLM 失败回退/孤立 tool 消息清理)
|
||||
|
||||
### v3.0
|
||||
- ✅ **模块化重构(aibox.py 8450 行 → 3706 行,缩减 56%)**:
|
||||
- **A 批基础工具抽离**([tools/basic_tools.py](tools/basic_tools.py)):`FileToolBase` / `ReadFileTool` / `WriteFileTool` / `ListFilesTool` / `SimpleWebCrawlerTool` / `GetCurrentTimeTool` / `CalculatorTool` 共 7 类,无外部子系统依赖;配套 43 项单元测试
|
||||
- **B 批管理面工具抽离**([tools/](tools/) 目录):`CommandLineTool` / `SearchSkillsTool` / `ReloadSkillsTool` / `InstallSkillTool` / `CreateSkillTool` / `MemoryTool` / `MemoTool` / `UpdateIdentityTool` / `UpdateSoulTool` / `SaveMemoryAndEndConversationTool` / `SetProxyTool` / `GetProxyTool` / `AskUserTool` 共 13 类,按职责分 7 个文件;配套 58 项单元测试
|
||||
- **Wave 1 子系统抽离**([subsystems/](subsystems/) 目录):`CommandExecutor` + `SubprocessCommandExecutor`(executors.py)、`PromptManager`(prompt_manager.py)、`MemoryDatabase` + `MemoDatabase`(db.py)、`SkillManager`(skill_manager.py)共 4 子系统;配套 70 项单元测试
|
||||
- **测试套件**:[tests/](tests/) 目录 171 项全过,覆盖 A/B 批工具与 Wave 1 子系统
|
||||
- **修复**:`subsystems/db.py` 的 `get_recent_memories(limit=None)` SQL datatype mismatch bug(`LIMIT ?` 传 None 时移除 LIMIT 子句)
|
||||
- **架构约束**:建立 `common ← subsystems ← tools ← aibox` 单向依赖,禁止反向引用;`tools/__init__.py` 与 `subsystems/__init__.py` 采用"重导出"策略,调用点无需关心子模块位置
|
||||
- **未抽离**:`TaskScheduler`、C 批工具(`ExpressEmotionTool` / `SubAgentTool`)、核心类(`DeepSeekChat` / `WebSocketHandler` / `EnhancedAIManager`)仍保留在 `aibox.py` 内(`ConversationHistory` 于 Wave 2、`ConversationCompressor` 于 v3.1 抽离)
|
||||
|
||||
### v2.9
|
||||
- ✅ **编排者-执行者工具分层(工具集 profile)**:
|
||||
- `DeepSeekChat` 新增 `tool_profile` 参数(`orchestrator` / `worker`),`_register_tools()` 按 profile 分组装配工具集
|
||||
|
||||
@@ -227,3 +227,14 @@ except ImportError:
|
||||
```
|
||||
|
||||
**`--help` 要完备**。包含参数说明和使用示例两部分。AI 通过阅读 `--help` 输出学会如何调用技能,信息越完整,AI 调用越准确。
|
||||
|
||||
---
|
||||
|
||||
## 安全边界(J12,审计 2026-08-19)
|
||||
|
||||
**技能即可信代码**。可执行技能由 AI(可被网页内容/用户消息诱导)以子进程方式运行,参数为不可信输入:
|
||||
|
||||
- **参数必须经 shell 引号转义**:凡技能内部把参数拼进 shell 命令/脚本串(subprocess shell、PowerShell、python-nmap 等内部拼串库),一律先 `shlex.quote`(POSIX)或对应语言转义(PowerShell 单引号翻倍)
|
||||
- **禁止把凭据嵌入命令行参数**:命令行参数对同机进程可见(`ps`);密码/令牌应经环境变量或配置文件传递
|
||||
- **ssh/sftp 类技能**:`AutoAddPolicy` 信任任意主机密钥存在 MITM 风险,仅限可信网络使用;生产环境改用 `known_hosts` 校验
|
||||
- 内置技能示例(`skills/`)已按此规范处理;新增技能请遵守同一条例
|
||||
+946
-5164
File diff suppressed because it is too large.
Load diff
@@ -815,6 +815,14 @@ async def main():
|
||||
await cli.run()
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n{Colors.SYSTEM}程序已退出{Colors.RESET}")
|
||||
finally:
|
||||
# M8(审计 2026-08-20):正常退出(EOF 等)此前不关 websocket,
|
||||
# 服务端见异常断连——统一 finally 通知
|
||||
if cli.client:
|
||||
try:
|
||||
await cli.client.disconnect()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
common.atomic - 原子写文件工具(J5,审计 2026-08-19)
|
||||
|
||||
同目录 mkstemp 临时文件 + fsync + os.replace:进程中断/断电不留半截文件;
|
||||
mkstemp 随机文件名避免多进程并发写同路径互踩(conversation_history 旧实现
|
||||
的固定 .tmp 名在并发场景会互相覆盖)。仅依赖标准库,供 common/subsystems
|
||||
各写路径复用。
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
|
||||
def atomic_write_text(path: str, content: str) -> None:
|
||||
"""原子写文本文件。失败时清理临时文件并抛出原异常。"""
|
||||
path = os.path.expanduser(path)
|
||||
directory = os.path.dirname(path) or "."
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".aibox_tmp_", dir=directory)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
# T27(审计 2026-08-21):Windows 下目标被其他进程打开(杀毒扫描/
|
||||
# 编辑器/并发读者)时 os.replace 抛 PermissionError。POSIX 无条件
|
||||
# 成功不受影响;Windows 做短退避重试,高频写路径(每条消息都原子
|
||||
# 写 current_conversation.json)一次冲突不再直接丢该条持久化。
|
||||
last_exc = None
|
||||
for attempt in range(3):
|
||||
try:
|
||||
os.replace(tmp_path, path)
|
||||
last_exc = None
|
||||
break
|
||||
except PermissionError as e:
|
||||
last_exc = e
|
||||
if attempt < 2:
|
||||
time.sleep(0.05 * (attempt + 1))
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
except BaseException:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
+113
-25
@@ -8,8 +8,13 @@ common.config - aibox 配置管理器
|
||||
|
||||
import os
|
||||
import json
|
||||
import copy
|
||||
import shutil
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from typing import Dict, Optional
|
||||
|
||||
from .atomic import atomic_write_text
|
||||
from .i18n import i18n
|
||||
class ConfigManager:
|
||||
"""配置管理器 - 单例模式,负责加载和保存配置"""
|
||||
@@ -17,15 +22,20 @@ class ConfigManager:
|
||||
_instance = None
|
||||
_config = None
|
||||
_config_path = None
|
||||
_instance_lock = threading.Lock()
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
# J6(审计 2026-08-19):此前单例创建与初始化均无锁——多线程首例
|
||||
# 并发可双实例、双重加载配置
|
||||
with cls._instance_lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if self._config is None:
|
||||
self._load_config()
|
||||
with self._instance_lock:
|
||||
if self._config is None:
|
||||
self._load_config()
|
||||
|
||||
def _get_default_config_path(self) -> str:
|
||||
env_config = os.environ.get("AI_CONFIG_PATH")
|
||||
@@ -64,14 +74,36 @@ class ConfigManager:
|
||||
|
||||
# 大模型调用接口
|
||||
"api": {
|
||||
"provider": "openai_compatible", # 提供商:deepseek / openai_compatible(OpenAI Chat Completions 协议兼容)
|
||||
"deepseek_api_url": "https://api.deepseek.com/chat/completions", # deepseek API URL
|
||||
"deepseek_model": "deepseek-v4-flash", # 模型名字
|
||||
"reasoning_enabled": False, # 开启思维链
|
||||
"api_key_env": "DEEPSEEK_API_KEY", # API-KEY所在的环境变量
|
||||
"timeout": 120, # 调用大模型请求超时时间
|
||||
"temperature": 0.7 # 大模型对话温度(类似思维发散程度)
|
||||
"temperature": 0.7, # 大模型对话温度(类似思维发散程度)
|
||||
"default_model": "deepseek" # 默认模型档案名(对应 models 节的 key)
|
||||
},
|
||||
|
||||
# 多模型档案:可配置多个 OpenAI Chat Completions 兼容端点,按名称选择
|
||||
# 示例(在 ~/.aibox/config.json 中添加):
|
||||
# "models": {
|
||||
# "kimi": {
|
||||
# "api_url": "https://api.moonshot.cn/v1/chat/completions",
|
||||
# "model": "kimi-k2",
|
||||
# "api_key_env": "MOONSHOT_API_KEY", # 未配置回退 api.api_key_env;空串 "" 为免密端点
|
||||
# "reasoning_enabled": false,
|
||||
# "timeout": 120
|
||||
# },
|
||||
# "local": {
|
||||
# "api_url": "http://localhost:11434/v1/chat/completions",
|
||||
# "model": "qwen3:32b",
|
||||
# "api_key_env": ""
|
||||
# }
|
||||
# }
|
||||
# 选择优先级:--model <档案名> > api.default_model > 回退旧 api 扁平节。
|
||||
# 默认留空:用户未配置档案时走旧 api 节(零破坏)。
|
||||
"models": {},
|
||||
|
||||
# 大模型上下文
|
||||
"context": {
|
||||
"max_history_messages": 500, # 最大历史消息数量
|
||||
@@ -160,12 +192,26 @@ class ConfigManager:
|
||||
print(i18n.get('config_loaded', config_path))
|
||||
except Exception as e:
|
||||
print(f"⚠️ {i18n.get('error_prefix')}{e},使用默认配置")
|
||||
# J5(审计 2026-08-19):损坏文件先备份再继续——此前直接改用
|
||||
# 默认配置,任一后续保存会用"默认+本次改动"整体覆盖原文件,
|
||||
# proxy/api/context 等全部手工配置永久丢失
|
||||
self._backup_corrupted_config(config_path)
|
||||
self._config = self._get_default_config()
|
||||
else:
|
||||
print(i18n.get('config_not_found', config_path))
|
||||
self._config = self._get_default_config()
|
||||
self._save_default_config(config_path)
|
||||
self._migrate_legacy_proxy_password()
|
||||
|
||||
def _backup_corrupted_config(self, config_path: str):
|
||||
"""J5:损坏配置先复制备份(带时间戳),防止后续保存整体覆盖丢失"""
|
||||
try:
|
||||
backup = f"{config_path}.corrupt.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||||
shutil.copy2(config_path, backup)
|
||||
print(f"⚠️ 已备份损坏的配置文件到 {backup}")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 备份损坏配置文件失败: {e}")
|
||||
|
||||
def _merge_configs(self, default: Dict, user: Dict) -> Dict:
|
||||
result = default.copy()
|
||||
for key, value in user.items():
|
||||
@@ -177,9 +223,9 @@ class ConfigManager:
|
||||
|
||||
def _save_default_config(self, config_path: str):
|
||||
try:
|
||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
||||
with open(config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self._get_default_config(), f, ensure_ascii=False, indent=2)
|
||||
atomic_write_text(
|
||||
config_path,
|
||||
json.dumps(self._get_default_config(), ensure_ascii=False, indent=2))
|
||||
print(i18n.get('config_created', config_path))
|
||||
except:
|
||||
pass
|
||||
@@ -198,6 +244,28 @@ class ConfigManager:
|
||||
self._load_config()
|
||||
|
||||
|
||||
def _resolve_proxy_password(self, proxy_conf: Dict) -> Optional[str]:
|
||||
"""解析代理密码:内存明文优先,否则按 password_env 查环境变量。
|
||||
|
||||
对齐 api_key_env 惯例:配置文件只持久化环境变量名,明文密码不落盘。
|
||||
"""
|
||||
password = proxy_conf.get("password")
|
||||
if password:
|
||||
return password
|
||||
env_name = proxy_conf.get("password_env")
|
||||
if env_name:
|
||||
return os.environ.get(env_name) or None
|
||||
return None
|
||||
|
||||
def _migrate_legacy_proxy_password(self):
|
||||
"""一次性迁移:旧版把代理密码明文落盘。加载后明文保留在内存供本会话
|
||||
使用,同时立即重写脱敏配置文件,清除盘上存量明文。"""
|
||||
socks = (self._config or {}).get("proxy", {}).get("socks5")
|
||||
if socks and socks.get("password"):
|
||||
print("⚠️ 检测到配置文件中的代理明文密码:已保留供本次会话使用,不再落盘。"
|
||||
"跨重启请将密码存入环境变量,并把 proxy.socks5.password_env 指向该变量名")
|
||||
self._save_config()
|
||||
|
||||
def get_proxy(self) -> Optional[Dict[str, str]]:
|
||||
"""获取当前代理配置(返回 requests 可用的 proxies 字典)"""
|
||||
proxy_conf = self.get("proxy.socks5")
|
||||
@@ -206,15 +274,20 @@ class ConfigManager:
|
||||
host = proxy_conf.get("host", "127.0.0.1")
|
||||
port = proxy_conf.get("port", 1080)
|
||||
username = proxy_conf.get("username")
|
||||
password = proxy_conf.get("password")
|
||||
password = self._resolve_proxy_password(proxy_conf)
|
||||
if username and password:
|
||||
proxy_url = f"socks5://{username}:{password}@{host}:{port}"
|
||||
else:
|
||||
proxy_url = f"socks5://{host}:{port}"
|
||||
return {"http": proxy_url, "https": proxy_url}
|
||||
|
||||
def set_proxy(self, host: str, port: int, username: str = None, password: str = None, enabled: bool = True, save: bool = True):
|
||||
"""动态设置代理配置(内存中更新,并可选择写入配置文件)"""
|
||||
def set_proxy(self, host: str, port: int, username: str = None, password: str = None,
|
||||
password_env: str = None, enabled: bool = True, save: bool = True):
|
||||
"""动态设置代理配置(内存中更新,并可选择写入配置文件)。
|
||||
|
||||
password 明文仅存内存(_save_config 落盘时强制脱敏);
|
||||
跨重启持久化请用 password_env 指向存有密码的环境变量名。
|
||||
"""
|
||||
# 确保 proxy 节点存在
|
||||
if "proxy" not in self._config:
|
||||
self._config["proxy"] = {}
|
||||
@@ -223,31 +296,46 @@ class ConfigManager:
|
||||
"host": host,
|
||||
"port": port,
|
||||
"username": username,
|
||||
"password": password
|
||||
"password": password,
|
||||
"password_env": password_env
|
||||
}
|
||||
if save:
|
||||
self._save_config()
|
||||
return self._save_config()
|
||||
return True
|
||||
|
||||
def _save_config(self):
|
||||
"""将当前配置保存到配置文件(覆盖)"""
|
||||
def _save_config(self) -> bool:
|
||||
"""将当前配置保存到配置文件(覆盖)。返回是否成功落盘。
|
||||
|
||||
落盘前对代理密码做纵深防御脱敏:无论调用路径如何,明文密码一律不写入文件。
|
||||
L2(审计 2026-08-18):此前异常仅 print,set_proxy 无条件报"已保存",
|
||||
磁盘只读等场景下用户以为已持久化,重启后配置丢失。
|
||||
"""
|
||||
config_path = self._config_path
|
||||
if config_path:
|
||||
try:
|
||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
||||
with open(config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self._config, f, ensure_ascii=False, indent=2)
|
||||
print(f"✅ 配置已保存到 {config_path}")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 保存配置失败: {e}")
|
||||
if not config_path:
|
||||
return False
|
||||
try:
|
||||
safe_config = copy.deepcopy(self._config)
|
||||
socks = safe_config.get("proxy", {}).get("socks5")
|
||||
if socks and socks.get("password"):
|
||||
socks["password"] = None
|
||||
# J5:原子写(mkstemp + fsync + os.replace),崩溃不留半截 JSON
|
||||
atomic_write_text(
|
||||
config_path,
|
||||
json.dumps(safe_config, ensure_ascii=False, indent=2))
|
||||
print(f"✅ 配置已保存到 {config_path}")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 保存配置失败: {e}")
|
||||
return False
|
||||
|
||||
def get_proxy_info(self) -> Dict:
|
||||
"""获取代理配置的详细信息(用于显示)"""
|
||||
"""获取代理配置的详细信息(用于显示,不回传任何密码值)"""
|
||||
proxy_conf = self.get("proxy.socks5", {})
|
||||
return {
|
||||
"enabled": proxy_conf.get("enabled", False),
|
||||
"host": proxy_conf.get("host", ""),
|
||||
"port": proxy_conf.get("port", 0),
|
||||
"username": proxy_conf.get("username"),
|
||||
"password": proxy_conf.get("password") is not None # 只显示是否存在
|
||||
"password": self._resolve_proxy_password(proxy_conf) is not None, # 只显示是否存在
|
||||
"password_env": proxy_conf.get("password_env") or ""
|
||||
}
|
||||
@@ -236,6 +236,31 @@ class Constants:
|
||||
def COMPRESSOR_LLM_MAX_TOKENS(cls):
|
||||
return cls._get_config().get("compressor.llm_max_tokens", 1000)
|
||||
|
||||
# 后台进程管理(bg_run/bg_list/bg_output/bg_kill)
|
||||
@property
|
||||
def BG_PROCESSES_ENABLED(cls):
|
||||
return cls._get_config().get("background_processes.enabled", True)
|
||||
|
||||
@property
|
||||
def BG_MAX_PROCESSES(cls):
|
||||
return cls._get_config().get("background_processes.max_processes", 16)
|
||||
|
||||
@property
|
||||
def BG_MAX_ENTRIES(cls):
|
||||
return cls._get_config().get("background_processes.max_entries", 32)
|
||||
|
||||
@property
|
||||
def BG_BUFFER_PER_PROCESS(cls):
|
||||
return cls._get_config().get("background_processes.buffer_per_process", 262144)
|
||||
|
||||
@property
|
||||
def BG_MAX_LIFETIME(cls):
|
||||
return cls._get_config().get("background_processes.max_lifetime", 0)
|
||||
|
||||
@property
|
||||
def BG_MAX_WRITE_BYTES(cls):
|
||||
return cls._get_config().get("background_processes.max_write_bytes", 65536)
|
||||
|
||||
|
||||
# 创建常量实例
|
||||
constants = Constants()
|
||||
@@ -98,6 +98,7 @@ class I18n:
|
||||
'memo_overdue': "⚠️",
|
||||
'memo_pending': "⏰",
|
||||
'search_results': "🔍 搜索 '{}' 找到 {} 条记忆:",
|
||||
'search_usage': "用法: /search 关键词",
|
||||
'current_config': "⚙️ 当前配置:",
|
||||
'config_debug': " 调试级别: {}",
|
||||
'config_data_dir': " 数据目录: {}",
|
||||
@@ -110,6 +111,7 @@ class I18n:
|
||||
'enabled': "启用",
|
||||
'disabled': "禁用",
|
||||
'config_reloaded': "✅ 配置已重新加载",
|
||||
'config_reloaded_cli_kept': "(命令行进程级覆盖已保留)",
|
||||
'prompts_reloaded': "✅ 提示词已重新加载 ({} 个)",
|
||||
'checking_memos': "🔍 检查任务...",
|
||||
'overdue_found': "发现 {} 个到期任务:",
|
||||
@@ -157,6 +159,7 @@ class I18n:
|
||||
'memo_add_failed': "❌ 添加失败: {}",
|
||||
'memo_completed': "✅ 任务 #{} 已完成",
|
||||
'memo_complete_failed': "❌ 无法完成任务 #{}",
|
||||
'memo_complete_missing_id': "❌ 缺少 task_id,无法完成任务",
|
||||
|
||||
# 即时任务相关
|
||||
'immediate_memo_added': "⚡ 已添加为即时任务并立即检查",
|
||||
@@ -258,6 +261,7 @@ class I18n:
|
||||
'memo_overdue': "⚠️",
|
||||
'memo_pending': "⏰",
|
||||
'search_results': "🔍 Searching '{}' found {} memories:",
|
||||
'search_usage': "Usage: /search <keyword>",
|
||||
'current_config': "⚙️ Current configuration:",
|
||||
'config_debug': " Debug level: {}",
|
||||
'config_data_dir': " Data directory: {}",
|
||||
@@ -270,6 +274,7 @@ class I18n:
|
||||
'enabled': "enabled",
|
||||
'disabled': "disabled",
|
||||
'config_reloaded': "✅ Configuration reloaded",
|
||||
'config_reloaded_cli_kept': " (CLI process-level overrides preserved)",
|
||||
'prompts_reloaded': "✅ Prompts reloaded ({})",
|
||||
'checking_memos': "🔍 Checking tasks...",
|
||||
'overdue_found': "Found {} overdue tasks:",
|
||||
@@ -317,6 +322,7 @@ class I18n:
|
||||
'memo_add_failed': "❌ Failed to add: {}",
|
||||
'memo_completed': "✅ Task #{} completed",
|
||||
'memo_complete_failed': "❌ Cannot complete task #{}",
|
||||
'memo_complete_missing_id': "❌ Missing task_id, cannot complete task",
|
||||
|
||||
# Immediate task related
|
||||
'immediate_memo_added': "⚡ Added as immediate task and checked now",
|
||||
|
||||
+31
-13
@@ -4,12 +4,14 @@
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>AI智能助手 - Socket.IO客户端</title>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||
<!-- Socket.IO 客户端库 -->
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js"></script>
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js" integrity="sha384-7EyYLQZgWBi67fBtVxw60/OWl1kjsfrPFcaU0pp0nAh+i8FD068QogUvg85Ewy1k" crossorigin="anonymous"></script>
|
||||
<style>
|
||||
/* 样式保持不变,完全使用原有的样式 */
|
||||
* {
|
||||
@@ -982,6 +984,15 @@
|
||||
</template>
|
||||
|
||||
<script>
|
||||
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||
function renderMarkdown(content) {
|
||||
const html = marked.parse(content);
|
||||
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||
return html.replace(/</g, '<');
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', function() {
|
||||
// 配置marked
|
||||
marked.setOptions({
|
||||
@@ -1124,7 +1135,7 @@
|
||||
textDiv.setAttribute('data-raw-content', streamData.content);
|
||||
|
||||
if (markdownEnabled) {
|
||||
textDiv.innerHTML = marked.parse(streamData.content);
|
||||
textDiv.innerHTML = renderMarkdown(streamData.content);
|
||||
textDiv.querySelectorAll('pre code').forEach((block) => {
|
||||
hljs.highlightElement(block);
|
||||
});
|
||||
@@ -1419,7 +1430,7 @@
|
||||
const rawContent = msgDiv.getAttribute('data-raw-content');
|
||||
if (rawContent) {
|
||||
if (markdownEnabled) {
|
||||
msgDiv.innerHTML = marked.parse(rawContent);
|
||||
msgDiv.innerHTML = renderMarkdown(rawContent);
|
||||
msgDiv.classList.add('markdown-body');
|
||||
addCopyButtonsToCodeBlocks(msgDiv);
|
||||
} else {
|
||||
@@ -1438,6 +1449,7 @@
|
||||
let reconnectAttempts = 0;
|
||||
const MAX_RECONNECT_ATTEMPTS = 5;
|
||||
let targetWsUrl = null;
|
||||
let reconnectTimer = null; // F2(审计 2026-08-20):重连定时器句柄,手动断开时清除
|
||||
|
||||
function createSocketConnection() {
|
||||
targetWsUrl = document.getElementById('ws-url').value.trim();
|
||||
@@ -1503,7 +1515,8 @@
|
||||
if (reason !== 'io client disconnect' && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
||||
reconnectAttempts++;
|
||||
showSystemMessage(`⚠️ 连接断开,${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} 次重试...`);
|
||||
setTimeout(() => {
|
||||
reconnectTimer = setTimeout(() => {
|
||||
reconnectTimer = null;
|
||||
if (!socket || !socket.connected) {
|
||||
console.log('尝试重连...');
|
||||
socket = createSocketConnection();
|
||||
@@ -1676,7 +1689,8 @@
|
||||
|
||||
// 事件绑定
|
||||
document.getElementById('connect-btn').addEventListener('click', function() {
|
||||
if (socket && socket.connected) {
|
||||
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要断开,否则重复点击产生孤儿连接/双后端会话
|
||||
if (socket) {
|
||||
socket.disconnect();
|
||||
}
|
||||
reconnectAttempts = 0;
|
||||
@@ -1684,6 +1698,8 @@
|
||||
});
|
||||
|
||||
document.getElementById('disconnect-btn').addEventListener('click', function() {
|
||||
// F2:清除已排程的重连定时器,防止倒计时到期把连接"复活"
|
||||
if (reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
|
||||
if (socket) {
|
||||
socket.disconnect();
|
||||
socket = null;
|
||||
@@ -1717,11 +1733,13 @@
|
||||
|
||||
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
||||
const container = document.getElementById('messages-container');
|
||||
// F1(审计 2026-08-20):welcome-message 在 container 内,
|
||||
// 必须先 cloneNode 缓存再清空,否则 querySelector 返回 null 抛 TypeError
|
||||
const welcomeEl = document.querySelector('.welcome-message');
|
||||
const welcomeHtml = welcomeEl ? welcomeEl.cloneNode(true) : null;
|
||||
container.innerHTML = '';
|
||||
|
||||
const welcomeHtml = document.querySelector('.welcome-message').cloneNode(true);
|
||||
container.appendChild(welcomeHtml);
|
||||
|
||||
if (welcomeHtml) container.appendChild(welcomeHtml);
|
||||
|
||||
clearAllStreamingMessages();
|
||||
});
|
||||
|
||||
|
||||
+35
-11
@@ -5,13 +5,15 @@
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=no">
|
||||
<title>小猫AI · 智能喵语伙伴</title>
|
||||
<!-- Font Awesome 6 (免费图标库) -->
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||
<!-- highlight.js 代码高亮 + marked markdown解析 -->
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||
<!-- Socket.IO 客户端库 -->
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js"></script>
|
||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js" integrity="sha384-7EyYLQZgWBi67fBtVxw60/OWl1kjsfrPFcaU0pp0nAh+i8FD068QogUvg85Ewy1k" crossorigin="anonymous"></script>
|
||||
<style>
|
||||
/* ========== 融合小猫风格 + WebSocket增强样式 ========== */
|
||||
* {
|
||||
@@ -498,6 +500,15 @@
|
||||
|
||||
<script>
|
||||
(function(){
|
||||
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||
function renderMarkdown(content) {
|
||||
const html = marked.parse(content);
|
||||
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||
return html.replace(/</g, '<');
|
||||
}
|
||||
|
||||
// ---------- DOM 元素 ----------
|
||||
const bgFile = document.getElementById('bgFile');
|
||||
const roleImg = document.getElementById('roleImg');
|
||||
@@ -533,7 +544,12 @@
|
||||
const div = document.createElement('div');
|
||||
div.className = isError ? 'error-msg' : 'system-msg';
|
||||
div.style.cssText = 'background:rgba(0,0,0,0.6); border-radius:16px; padding:6px 12px; margin:6px 0; font-size:12px;';
|
||||
div.innerHTML = `<i class="fas ${isError ? 'fa-exclamation-triangle' : 'fa-info-circle'}"></i> ${text}`;
|
||||
// J2:text 为服务端内容(命令执行结果/AI 输出),禁止拼 innerHTML——
|
||||
// 原实现 `${text}` 模板拼接使任意 HTML 注入执行;图标为静态字面量
|
||||
const icon = document.createElement('i');
|
||||
icon.className = `fas ${isError ? 'fa-exclamation-triangle' : 'fa-info-circle'}`;
|
||||
div.appendChild(icon);
|
||||
div.appendChild(document.createTextNode(' ' + text));
|
||||
chatBox.appendChild(div);
|
||||
chatBox.scrollTop = chatBox.scrollHeight;
|
||||
}
|
||||
@@ -545,7 +561,7 @@
|
||||
if(textDiv) {
|
||||
textDiv.setAttribute('data-raw', stream.content);
|
||||
if(markdownEnabled) {
|
||||
textDiv.innerHTML = marked.parse(stream.content);
|
||||
textDiv.innerHTML = renderMarkdown(stream.content);
|
||||
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
||||
} else textDiv.innerText = stream.content;
|
||||
}
|
||||
@@ -572,7 +588,12 @@
|
||||
if(!streamingTasks[eid]) {
|
||||
const sep = document.createElement('div');
|
||||
sep.className = 'task-separator';
|
||||
sep.innerHTML = `<i class="fas fa-bolt"></i> 定时任务 #${taskMeta.taskId} (${taskMeta.taskTitle || '执行'}) 第${taskMeta.triggerCount||1}次执行`;
|
||||
// J2:taskMeta.taskTitle 为服务端可控内容,禁止拼 innerHTML
|
||||
const bolt = document.createElement('i');
|
||||
bolt.className = 'fas fa-bolt';
|
||||
sep.appendChild(bolt);
|
||||
sep.appendChild(document.createTextNode(
|
||||
` 定时任务 #${taskMeta.taskId} (${taskMeta.taskTitle || '执行'}) 第${taskMeta.triggerCount||1}次执行`));
|
||||
chatBox.appendChild(sep);
|
||||
const taskDiv = document.createElement('div');
|
||||
taskDiv.className = 'task-msg streaming';
|
||||
@@ -590,7 +611,7 @@
|
||||
if(st && st.element) {
|
||||
const textDiv = st.element.querySelector('.message-text');
|
||||
if(textDiv) {
|
||||
if(markdownEnabled) textDiv.innerHTML = marked.parse(st.content);
|
||||
if(markdownEnabled) textDiv.innerHTML = renderMarkdown(st.content);
|
||||
else textDiv.innerText = st.content;
|
||||
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
||||
}
|
||||
@@ -613,7 +634,7 @@
|
||||
if(task.timer) clearTimeout(task.timer);
|
||||
if(task.element) {
|
||||
const txtDiv = task.element.querySelector('.message-text');
|
||||
if(txtDiv && markdownEnabled) txtDiv.innerHTML = marked.parse(task.content);
|
||||
if(txtDiv && markdownEnabled) txtDiv.innerHTML = renderMarkdown(task.content);
|
||||
task.element.classList.remove('streaming');
|
||||
}
|
||||
globalMsgCount++;
|
||||
@@ -652,7 +673,8 @@
|
||||
}
|
||||
|
||||
function connectWebSocket() {
|
||||
if(socket && socket.connected) {
|
||||
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要断开,否则重复点击产生孤儿连接/双后端会话
|
||||
if(socket) {
|
||||
socket.disconnect();
|
||||
}
|
||||
|
||||
@@ -731,6 +753,8 @@
|
||||
|
||||
function disconnectWebSocket() {
|
||||
manualDisconnect = true;
|
||||
// F2(审计 2026-08-20):清除已排程的重连定时器,防止倒计时到期把连接"复活"
|
||||
if(reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
|
||||
if(socket) {
|
||||
socket.disconnect();
|
||||
socket = null;
|
||||
|
||||
+31
-13
@@ -4,10 +4,12 @@
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>AI智能助手 - WebSocket直连</title>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||
<style>
|
||||
/* 样式保持不变,完全使用原有的样式 */
|
||||
* {
|
||||
@@ -951,6 +953,15 @@
|
||||
</template>
|
||||
|
||||
<script>
|
||||
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||
function renderMarkdown(content) {
|
||||
const html = marked.parse(content);
|
||||
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||
return html.replace(/</g, '<');
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', function() {
|
||||
// 配置marked
|
||||
marked.setOptions({
|
||||
@@ -1048,7 +1059,7 @@
|
||||
textDiv.setAttribute('data-raw-content', streamData.content);
|
||||
|
||||
if (markdownEnabled) {
|
||||
textDiv.innerHTML = marked.parse(streamData.content);
|
||||
textDiv.innerHTML = renderMarkdown(streamData.content);
|
||||
textDiv.querySelectorAll('pre code').forEach((block) => {
|
||||
hljs.highlightElement(block);
|
||||
});
|
||||
@@ -1336,7 +1347,7 @@
|
||||
const rawContent = msgDiv.getAttribute('data-raw-content');
|
||||
if (rawContent) {
|
||||
if (markdownEnabled) {
|
||||
msgDiv.innerHTML = marked.parse(rawContent);
|
||||
msgDiv.innerHTML = renderMarkdown(rawContent);
|
||||
msgDiv.classList.add('markdown-body');
|
||||
addCopyButtonsToCodeBlocks(msgDiv);
|
||||
} else {
|
||||
@@ -1413,9 +1424,12 @@
|
||||
|
||||
if (!manualDisconnect && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
||||
reconnectAttempts++;
|
||||
showSystemMessage(`⚠️ 连接断开,${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} 秒后重试...`);
|
||||
// F6(审计 2026-08-20):实际延迟为 3000*N ms,文案与此对齐
|
||||
showSystemMessage(`⚠️ 连接断开,${3 * reconnectAttempts} 秒后重试(${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS})...`);
|
||||
setTimeout(() => {
|
||||
if (!ws || ws.readyState !== WebSocket.OPEN) {
|
||||
// F2(审计 2026-08-20):闭包到期时再判 manualDisconnect,
|
||||
// 防止手动断开后已排程定时器把连接"复活"
|
||||
if (!manualDisconnect && (!ws || ws.readyState !== WebSocket.OPEN)) {
|
||||
console.log('尝试重连...');
|
||||
ws = connectWebSocket();
|
||||
}
|
||||
@@ -1585,7 +1599,8 @@
|
||||
|
||||
// 事件绑定
|
||||
document.getElementById('connect-btn').addEventListener('click', function() {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要关闭,否则重复点击产生孤儿连接/双后端会话
|
||||
if (ws && ws.readyState !== WebSocket.CLOSED) {
|
||||
ws.close();
|
||||
}
|
||||
reconnectAttempts = 0;
|
||||
@@ -1628,11 +1643,14 @@
|
||||
|
||||
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
||||
const container = document.getElementById('messages-container');
|
||||
// F1(审计 2026-08-20):welcome-message 在 container 内,
|
||||
// 必须先 cloneNode 缓存再清空,否则 querySelector 返回 null 抛 TypeError,
|
||||
// 且后续 clearAllStreamingMessages/计数清零全部中断不执行
|
||||
const welcomeEl = document.querySelector('.welcome-message');
|
||||
const welcomeHtml = welcomeEl ? welcomeEl.cloneNode(true) : null;
|
||||
container.innerHTML = '';
|
||||
|
||||
const welcomeHtml = document.querySelector('.welcome-message').cloneNode(true);
|
||||
container.appendChild(welcomeHtml);
|
||||
|
||||
if (welcomeHtml) container.appendChild(welcomeHtml);
|
||||
|
||||
clearAllStreamingMessages();
|
||||
messageCount = 0;
|
||||
toolCallCount = 0;
|
||||
|
||||
+561
-39
@@ -22,27 +22,63 @@ from typing import Dict, Iterator, List, Optional
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import requests
|
||||
|
||||
# L12(审计 2026-08-18):瞬时错误体可能回显网关/WAF 错误页中的
|
||||
# Authorization 头,抹掉 Bearer 令牌后再进日志,防密钥落盘
|
||||
_AUTH_HEADER_RE = re.compile(r"Bearer\s+[A-Za-z0-9._~+/=-]+")
|
||||
# T13:错误响应体可能回显裸 API Key(如 sk-xxx),一并脱敏
|
||||
_API_KEY_RE = re.compile(r"sk-[A-Za-z0-9_-]{16,}")
|
||||
|
||||
# T19(审计 2026-08-21):瞬时网络异常白名单。此前 FailoverClient 捕获
|
||||
# requests.exceptions.RequestException 基类——档案 api_url 手写错(漏
|
||||
# scheme 抛 MissingSchema)等永久性配置错误也被当瞬时:标记冷却 60s、
|
||||
# 触发转移、计入 failovers,每次冷却到期再撞一次坏端点,统计与冷却表
|
||||
# 被永久错误污染。显式枚举瞬时异常,其余 requests 异常直接上抛响亮暴露
|
||||
_TRANSIENT_NET_ERRORS = (
|
||||
requests.exceptions.ConnectionError,
|
||||
requests.exceptions.Timeout,
|
||||
requests.exceptions.ChunkedEncodingError,
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_log_msg(msg: str) -> str:
|
||||
"""日志脱敏:替换错误消息中的 Bearer 令牌与裸 API Key。"""
|
||||
if not msg:
|
||||
return msg
|
||||
msg = _AUTH_HEADER_RE.sub("Bearer ***", msg)
|
||||
msg = _API_KEY_RE.sub("sk-***", msg)
|
||||
return msg
|
||||
|
||||
|
||||
class LLMClient(ABC):
|
||||
"""大模型客户端抽象基类。子类需实现流式/非流式两种调用,并归一化为统一事件格式。"""
|
||||
|
||||
def __init__(self, api_key: str, model: str, base_url: str,
|
||||
timeout: int, config=None, logger=None):
|
||||
timeout: int, config=None, logger=None,
|
||||
reasoning_enabled: bool = None,
|
||||
display_name: str = None, default_temperature: float = None):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.timeout = timeout
|
||||
self.config = config # 可选,用于读取代理等动态配置
|
||||
self.logger = logger # 可选,用于日志输出
|
||||
# M3:档案级思考模式开关为构造契约;None = 未配置(消费方回退全局 api.reasoning_enabled)
|
||||
self.reasoning_enabled = reasoning_enabled
|
||||
# L1:展示名(默认取档案名),业务层上下文(如 history 的 assistant 名)不再误标 "DeepSeek"
|
||||
self._display_name = display_name
|
||||
# L3:档案级默认采样温度;None = 未配置(方法参数未显式传入时回退硬编码默认值)
|
||||
self.default_temperature = default_temperature
|
||||
# 当前活动的 HTTP 响应引用,供 cancel() 强制关闭底层连接以中断阻塞读取
|
||||
self._current_response = None
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""厂商/模型展示名,业务层可用于上下文(如 history 里的 assistant 名)。"""
|
||||
return "LLM"
|
||||
return self._display_name or "LLM"
|
||||
|
||||
# ---------- 工具方法 ----------
|
||||
|
||||
@@ -55,12 +91,16 @@ class LLMClient(ABC):
|
||||
def _log(self, level: str, msg: str):
|
||||
if self.logger is not None and hasattr(self.logger, level):
|
||||
try:
|
||||
getattr(self.logger, level)(msg)
|
||||
getattr(self.logger, level)(_sanitize_log_msg(msg))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def cancel(self):
|
||||
"""中断当前流式请求。线程安全,可从任意线程调用。"""
|
||||
"""中断当前流式请求。线程安全,可从任意线程调用。
|
||||
|
||||
注:用户中断标志由 FailoverClient 维护(其 cancel() 广播整条链并置
|
||||
自身 _cancelled,据此跳过转移)——底层 client 只负责强关连接。
|
||||
"""
|
||||
resp = self._current_response
|
||||
self._current_response = None
|
||||
if resp is not None:
|
||||
@@ -78,16 +118,18 @@ class LLMClient(ABC):
|
||||
# ---------- 抽象接口 ----------
|
||||
|
||||
@abstractmethod
|
||||
def chat_completion(self, messages: List[Dict], temperature: float = 0.3,
|
||||
def chat_completion(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 1024, **kwargs) -> str:
|
||||
"""非流式对话,返回完整文本。用于对话压缩等离线摘要场景。"""
|
||||
"""非流式对话,返回完整文本。用于对话压缩等离线摘要场景。
|
||||
temperature=None 时按 default_temperature(档案级)→ 0.3 回退。"""
|
||||
|
||||
@abstractmethod
|
||||
def stream_chat(self, messages: List[Dict], temperature: float = 0.7,
|
||||
def stream_chat(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||
"""流式对话,逐条 yield 统一事件格式的 dict。业务层据此驱动输出/工具调用。"""
|
||||
"""流式对话,逐条 yield 统一事件格式的 dict。业务层据此驱动输出/工具调用。
|
||||
temperature=None 时按 default_temperature(档案级)→ 0.7 回退。"""
|
||||
|
||||
|
||||
class DeepSeekClient(LLMClient):
|
||||
@@ -99,12 +141,18 @@ class DeepSeekClient(LLMClient):
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str, model: str, base_url: str,
|
||||
timeout: int = 30, config=None, logger=None):
|
||||
super().__init__(api_key, model, base_url, timeout, config, logger)
|
||||
timeout: int = 30, config=None, logger=None,
|
||||
reasoning_enabled: bool = None,
|
||||
display_name: str = None, default_temperature: float = None):
|
||||
super().__init__(api_key, model, base_url, timeout, config, logger,
|
||||
reasoning_enabled=reasoning_enabled,
|
||||
display_name=display_name,
|
||||
default_temperature=default_temperature)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "DeepSeek"
|
||||
# L1:档案命中时展示档案名(或档案 display_name),旧扁平节保持 "DeepSeek"
|
||||
return self._display_name or "DeepSeek"
|
||||
|
||||
def _headers(self) -> Dict[str, str]:
|
||||
return {
|
||||
@@ -112,8 +160,11 @@ class DeepSeekClient(LLMClient):
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def chat_completion(self, messages: List[Dict], temperature: float = 0.3,
|
||||
def chat_completion(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 1024, **kwargs) -> str:
|
||||
# L3:None → 档案级 default_temperature → 硬编码默认(与旧版行为一致)
|
||||
if temperature is None:
|
||||
temperature = self.default_temperature if self.default_temperature is not None else 0.3
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
@@ -126,22 +177,44 @@ class DeepSeekClient(LLMClient):
|
||||
self.base_url,
|
||||
headers=self._headers(),
|
||||
json=payload,
|
||||
timeout=30,
|
||||
timeout=self.timeout,
|
||||
proxies=self._resolve_proxies(),
|
||||
)
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
return result["choices"][0]["message"]["content"].strip()
|
||||
self._log("error", f"LLM 压缩请求失败: HTTP {response.status_code}")
|
||||
try:
|
||||
result = response.json()
|
||||
return result["choices"][0]["message"]["content"].strip()
|
||||
finally:
|
||||
response.close()
|
||||
# T9+T13:非 200 立即关连接(stream=True 才需要,但统一防御),
|
||||
# 错误体脱敏后再入异常,防 API Key 经错误页回显进历史/日志
|
||||
status = response.status_code
|
||||
body = _sanitize_log_msg(response.text[:200])
|
||||
response.close()
|
||||
if status == 429 or status >= 500:
|
||||
# L4a:瞬时故障显式抛出——FailoverClient 据此转移下一档案;
|
||||
# 无包装器时由调用方 catch-all 接住,下游仍走失败分支(语义等价)
|
||||
raise TransientLLMError(
|
||||
f"LLM 请求瞬时失败: HTTP {status}: {body}",
|
||||
retry_after=_retry_after_of(response))
|
||||
self._log("error", f"LLM 压缩请求失败: HTTP {status}")
|
||||
return None
|
||||
except TransientLLMError:
|
||||
raise # L4a:瞬时故障上抛供 FailoverClient 转移,不吞
|
||||
except _TRANSIENT_NET_ERRORS:
|
||||
raise # T19:连接类瞬时故障同样上抛(原被吞成 None,无法转移;
|
||||
# ChunkedEncodingError 一并对齐——响应体中途损坏属瞬时网络故障)
|
||||
except Exception as e:
|
||||
self._log("error", f"LLM 压缩请求异常: {e}")
|
||||
return None
|
||||
|
||||
def stream_chat(self, messages: List[Dict], temperature: float = 0.7,
|
||||
def stream_chat(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||
# L3:None → 档案级 default_temperature → 硬编码默认(与旧版行为一致)
|
||||
if temperature is None:
|
||||
temperature = self.default_temperature if self.default_temperature is not None else 0.7
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
@@ -170,18 +243,32 @@ class DeepSeekClient(LLMClient):
|
||||
self._current_response = response
|
||||
|
||||
if response.status_code != 200:
|
||||
# T9:非 200 立即释放连接(stream=True 下不关闭会阻塞同域名
|
||||
# in-flight 请求直到 GC);T13:错误体脱敏后再入异常
|
||||
status = response.status_code
|
||||
retry_after = _retry_after_of(response)
|
||||
body = _sanitize_log_msg(response.text[:200])
|
||||
response.close()
|
||||
self._current_response = None
|
||||
if status == 429 or status >= 500:
|
||||
# L4b:瞬时故障显式分类(此时尚未产出任何事件,可安全转移重放)
|
||||
raise TransientLLMError(
|
||||
f"LLM 流式请求瞬时失败: HTTP {status}: {body}",
|
||||
retry_after=retry_after)
|
||||
raise RuntimeError(
|
||||
f"HTTP {response.status_code}: {response.text}")
|
||||
f"HTTP {status}: {body}")
|
||||
|
||||
saw_done = False
|
||||
try:
|
||||
for line in response.iter_lines():
|
||||
if not line:
|
||||
continue
|
||||
line_text = line.decode("utf-8")
|
||||
line_text = line.decode("utf-8", errors="replace") # L3:非 UTF-8 行不中断流
|
||||
if not line_text.startswith("data: "):
|
||||
continue
|
||||
data = line_text[6:]
|
||||
if data == "[DONE]":
|
||||
saw_done = True
|
||||
yield {"type": "done"}
|
||||
break
|
||||
try:
|
||||
@@ -197,12 +284,28 @@ class DeepSeekClient(LLMClient):
|
||||
yield event
|
||||
finally:
|
||||
self._current_response = None
|
||||
# T20(审计 2026-08-21):iter_lines 耗尽但从未收到 [DONE](代理缓冲
|
||||
# 后提前 FIN / 服务端中途断开,末行半截 JSON 解析失败被 continue)
|
||||
# 此前被静默视为成功——回答无声截断无任何信号。抛瞬时异常:首事件
|
||||
# 前 FailoverClient 可转移重放;已产出事件则上抛(既有语义:已输出
|
||||
# 内容绝不重放),由上层错误处理回滚/报错。用户中断(cancel 关连接
|
||||
# 抛异常)不会走到这里。
|
||||
if not saw_done:
|
||||
raise TransientLLMError("LLM 流式响应不完整:未收到 [DONE] 结束标记(连接中途断开)")
|
||||
|
||||
# L5:思考内容的常见厂商字段别名(首个命中的非空字段生效)
|
||||
REASONING_FIELDS = ("reasoning_content", "reasoning", "thinking", "thinking_content")
|
||||
|
||||
def _parse_delta(self, delta: Dict, reasoning_enabled: bool) -> Iterator[Dict]:
|
||||
"""将单条 delta 归一化为统一事件。子类覆写以适配不同厂商字段。"""
|
||||
# 思考内容
|
||||
if reasoning_enabled and delta.get("reasoning_content"):
|
||||
yield {"type": "reasoning", "text": delta["reasoning_content"]}
|
||||
# 思考内容:DeepSeek reasoning_content / OpenRouter reasoning /
|
||||
# 部分 GLM thinking(_content) —— 此前仅识别 DeepSeek 字段,其他厂商思考流静默丢失
|
||||
if reasoning_enabled:
|
||||
for field in self.REASONING_FIELDS:
|
||||
value = delta.get(field)
|
||||
if value:
|
||||
yield {"type": "reasoning", "text": value}
|
||||
break
|
||||
# 正文
|
||||
if delta.get("content"):
|
||||
yield {"type": "content", "text": delta["content"]}
|
||||
@@ -219,19 +322,279 @@ class DeepSeekClient(LLMClient):
|
||||
}
|
||||
|
||||
|
||||
def resolve_api_key(api_key: Optional[str], config=None) -> str:
|
||||
"""解析 API Key:优先显式传入,其次环境变量,最后 ~/.cert/deepseek.key。"""
|
||||
class TransientLLMError(RuntimeError):
|
||||
"""瞬时 LLM 故障(HTTP 429/5xx)——FailoverClient 据此转移下一档案(L4a)。
|
||||
|
||||
L4c:retry_after 携带服务端 Retry-After 响应头(数字秒)建议冷却时长,
|
||||
FailoverClient 以其替代默认冷却期(429 限流时尊重服务端节奏)。
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, retry_after: Optional[float] = None):
|
||||
super().__init__(message)
|
||||
self.retry_after = retry_after
|
||||
|
||||
|
||||
def _retry_after_of(response) -> Optional[float]:
|
||||
"""从 Retry-After 响应头解析冷却秒数(仅支持数字秒格式,其余忽略返回 None)。
|
||||
|
||||
P2:上限 cap 3600s——异常/恶意服务端返回超大值(如 999999)时,
|
||||
避免档案被单条响应头冷却数天。
|
||||
"""
|
||||
try:
|
||||
v = response.headers.get("Retry-After")
|
||||
if v is None:
|
||||
return None
|
||||
return min(max(1.0, float(v)), 3600.0)
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
return None
|
||||
|
||||
|
||||
class FailoverClient:
|
||||
"""跨档案失败转移包装器(L4a)。
|
||||
|
||||
持有序 client 链(主档案 + fallback 档案),对外实现与 LLMClient 相同的
|
||||
消费面(未显式定义的属性经 __getattr__ 委托主 client)。
|
||||
|
||||
L4a 范围——chat_completion(非流式幂等)按链转移:
|
||||
- 转移:HTTP 429/5xx(TransientLLMError)、连接失败/超时(requests 连接类异常)
|
||||
- 不转移:认证/参数类 4xx(底层返回 None,直接透传——配置问题应响亮暴露)
|
||||
失败档案进入冷却期(cooldown_seconds,默认 60s),期间直接跳过,
|
||||
避免每次请求都撞刚失败的端点。链耗尽返回 None(与 chat_completion 失败契约一致)。
|
||||
|
||||
L4b 范围——stream_chat 仅**首事件前**故障转移(连接失败/HTTP 429/5xx,
|
||||
尚未产出任何 reasoning/content/tool_call 事件,重放对用户不可见);
|
||||
已产出事件后失败直接上抛(重放会导致前端重复输出/工具增量错乱);
|
||||
用户中断(cancel() 置位)引发的连接异常绝不转移。
|
||||
|
||||
L4c 范围——可观测性与预算控制:
|
||||
- get_stats() 暴露 requests/failovers/cooldown_skips/exhausted 聚合计数
|
||||
- on_failover(from_name, to_name, reason) 回调(to_name=None 表示链耗尽),
|
||||
供管理面/前端通知接入,回调异常不影响主流程
|
||||
- timeout_budget 限制整条链尝试的总秒数(缺省不限),耗尽即停止尝试剩余档案,
|
||||
防止长链故障时超时叠加(3 档案最坏 3×timeout)
|
||||
- 429 响应带 Retry-After(数字秒)时按其延长该档案冷却期
|
||||
|
||||
并发契约(P4):实例非线程安全(_cancelled/_call_started/_stats 无锁)。
|
||||
当前调用面满足"单线程串行调用 + cancel() 可从其他线程"的模式(compressor
|
||||
在 _msg_lock 内同步调用、子代理各持独立实例);如需多线程共享调用,
|
||||
需由调用方加锁或为实例补锁。
|
||||
"""
|
||||
|
||||
def __init__(self, clients: List, logger=None, cooldown_seconds: float = 60.0,
|
||||
timeout_budget: Optional[float] = None, on_failover=None):
|
||||
self._clients = list(clients)
|
||||
self.logger = logger
|
||||
self.cooldown_seconds = cooldown_seconds
|
||||
self._fail_at = {} # client.name → (失败时刻, 冷却秒数)
|
||||
self._cancelled = False
|
||||
self._stats = {"requests": 0, "failovers": 0, "cooldown_skips": 0, "exhausted": 0}
|
||||
if timeout_budget is not None and timeout_budget <= 0:
|
||||
self._log("warning", f"timeout_budget 非法({timeout_budget} <= 0),按未配置处理")
|
||||
timeout_budget = None
|
||||
self.timeout_budget = timeout_budget
|
||||
self.on_failover = on_failover
|
||||
self._call_started = None
|
||||
|
||||
def __getattr__(self, item):
|
||||
# 未显式定义的属性/方法委托主 client(name/model/api_key/cancel 等);
|
||||
# 下划线属性直接报错,防 __init__ 前访问 self._clients 的无限递归
|
||||
if item.startswith("_"):
|
||||
raise AttributeError(item)
|
||||
return getattr(self._clients[0], item)
|
||||
|
||||
def _log(self, level: str, msg: str):
|
||||
if self.logger is not None and hasattr(self.logger, level):
|
||||
try:
|
||||
getattr(self.logger, level)(_sanitize_log_msg(msg))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _cool_key(client) -> int:
|
||||
"""T6:冷却键按客户端实例(id)——client.name 取 display_name 优先且
|
||||
无唯一性校验,两个 profile_name 不同但 display_name 相同的档案
|
||||
(如主备都叫 "Kimi")冷却状态互相污染:A 失败冷却 → B 被误判
|
||||
跳过(该转移时不转移)。P6 去重保证每档案恰好一个实例,实例被
|
||||
_clients 持有生命周期内 id 稳定,与去重键语义一致。"""
|
||||
return id(client)
|
||||
|
||||
def _cooling(self, client) -> bool:
|
||||
entry = self._fail_at.get(self._cool_key(client))
|
||||
if entry is None:
|
||||
return False
|
||||
t, cd = entry
|
||||
return (time.monotonic() - t) < cd
|
||||
|
||||
def _mark_fail(self, client, cooldown: Optional[float] = None):
|
||||
# L4c:429 Retry-After(经 TransientLLMError.retry_after)覆盖默认冷却时长
|
||||
self._fail_at[self._cool_key(client)] = (time.monotonic(),
|
||||
cooldown if cooldown and cooldown > 0
|
||||
else self.cooldown_seconds)
|
||||
|
||||
def _budget_exhausted(self) -> bool:
|
||||
if self.timeout_budget is None or self._call_started is None:
|
||||
return False
|
||||
return (time.monotonic() - self._call_started) >= self.timeout_budget
|
||||
|
||||
def get_stats(self) -> Dict:
|
||||
"""L4c:转移统计快照(requests/failovers/cooldown_skips/exhausted)。"""
|
||||
return dict(self._stats)
|
||||
|
||||
def _notify_failover(self, from_name: str, to_name, reason: str):
|
||||
self._stats["failovers"] += 1
|
||||
try:
|
||||
if self.on_failover is not None:
|
||||
self.on_failover(from_name, to_name, _sanitize_log_msg(reason))
|
||||
except Exception as e:
|
||||
self._log("warning", f"on_failover 回调异常(忽略): {e}")
|
||||
|
||||
def cancel(self):
|
||||
"""L4b:中断链上全部 client(任一可能正在流式中)并置自身中断标志。
|
||||
|
||||
显式定义以覆盖 __getattr__ 委托——包装器必须同时通知整条链。
|
||||
"""
|
||||
self._cancelled = True
|
||||
for client in self._clients:
|
||||
try:
|
||||
client.cancel()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def chat_completion(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 1024, **kwargs):
|
||||
self._cancelled = False # 新一轮请求,清除上一轮中断标志
|
||||
self._call_started = time.monotonic()
|
||||
self._stats["requests"] += 1
|
||||
last_exc = None
|
||||
tried = 0
|
||||
for idx, client in enumerate(self._clients):
|
||||
if self._cooling(client):
|
||||
self._stats["cooldown_skips"] += 1
|
||||
self._log("info", f"模型档案 {client.name} 处于失败冷却期,跳过")
|
||||
continue
|
||||
if self._budget_exhausted():
|
||||
self._log("warning",
|
||||
f"失败转移时间预算({self.timeout_budget}s)已耗尽,停止尝试剩余档案")
|
||||
break
|
||||
tried += 1
|
||||
try:
|
||||
result = client.chat_completion(messages, temperature=temperature,
|
||||
max_tokens=max_tokens, **kwargs)
|
||||
if result is not None:
|
||||
return result
|
||||
# None:认证/参数类 4xx 或响应异常——不可转移,直接透传
|
||||
return None
|
||||
except (TransientLLMError,) + _TRANSIENT_NET_ERRORS as e:
|
||||
if self._cancelled:
|
||||
return None # 用户中断引发的异常:不转移
|
||||
last_exc = e
|
||||
self._mark_fail(client, cooldown=getattr(e, "retry_after", None))
|
||||
to_name = (self._clients[idx + 1].name
|
||||
if idx + 1 < len(self._clients) else None)
|
||||
self._notify_failover(client.name, to_name, str(e))
|
||||
self._log("warning", f"模型档案 {client.name} 瞬时故障({e}),转移下一档案")
|
||||
self._stats["exhausted"] += 1 # 末尾统一计数(链耗尽/预算耗尽/全冷却)
|
||||
self._log("error", f"全部 {len(self._clients)} 个模型档案均不可用(尝试 {tried} 个):{last_exc}")
|
||||
return None
|
||||
|
||||
def stream_chat(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||
"""L4b:流式转移——仅首事件前故障可安全重放(对用户不可见)。
|
||||
|
||||
T7:本方法是普通函数(非生成器)——生成器函数体延迟到首次 next()
|
||||
才执行,若 cancel() 在创建生成器与首迭代之间被其他线程调用,
|
||||
延迟执行的 `_cancelled = False` 会把用户的取消清掉(请求照常发出)。
|
||||
标志重置与计数在返回生成器前同步完成,取消契约(类文档)恢复成立。
|
||||
"""
|
||||
self._cancelled = False # 新一轮请求,清除上一轮中断标志
|
||||
self._call_started = time.monotonic()
|
||||
self._stats["requests"] += 1
|
||||
return self._stream_chat_impl(messages, temperature=temperature,
|
||||
max_tokens=max_tokens, tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
reasoning_enabled=reasoning_enabled, **extra)
|
||||
|
||||
def _stream_chat_impl(self, messages: List[Dict], temperature: float = None,
|
||||
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||
last_exc = None
|
||||
for idx, client in enumerate(self._clients):
|
||||
if self._cooling(client):
|
||||
self._stats["cooldown_skips"] += 1
|
||||
self._log("info", f"模型档案 {client.name} 处于失败冷却期,跳过")
|
||||
continue
|
||||
if self._budget_exhausted():
|
||||
self._log("warning",
|
||||
f"失败转移时间预算({self.timeout_budget}s)已耗尽,停止尝试剩余档案")
|
||||
break
|
||||
produced_any = False
|
||||
try:
|
||||
for event in client.stream_chat(
|
||||
messages, temperature=temperature, max_tokens=max_tokens,
|
||||
tools=tools, tool_choice=tool_choice,
|
||||
reasoning_enabled=reasoning_enabled, **extra):
|
||||
produced_any = True
|
||||
yield event
|
||||
return # 自然结束(含 done 事件)视为成功
|
||||
except (TransientLLMError,) + _TRANSIENT_NET_ERRORS as e:
|
||||
if self._cancelled:
|
||||
raise # 用户中断(cancel 强关连接):绝不转移
|
||||
if produced_any:
|
||||
raise # 已向消费方产出事件,重放会重复输出/工具增量错乱
|
||||
last_exc = e
|
||||
self._mark_fail(client, cooldown=getattr(e, "retry_after", None))
|
||||
to_name = (self._clients[idx + 1].name
|
||||
if idx + 1 < len(self._clients) else None)
|
||||
self._notify_failover(client.name, to_name, str(e))
|
||||
self._log("warning",
|
||||
f"模型档案 {client.name} 流式瞬时故障(首事件前,{e}),转移下一档案")
|
||||
self._stats["exhausted"] += 1
|
||||
if last_exc is not None:
|
||||
raise last_exc # 链耗尽:抛最后异常,由既有错误处理(回滚/报错)接住
|
||||
raise RuntimeError(
|
||||
f"全部 {len(self._clients)} 个模型档案不可用(冷却期或转移时间预算耗尽)")
|
||||
|
||||
|
||||
def resolve_api_key(api_key: Optional[str], config=None,
|
||||
env_name: Optional[str] = None,
|
||||
allow_missing: bool = False,
|
||||
allow_key_file: bool = True) -> str:
|
||||
"""解析 API Key:优先显式传入,其次环境变量,最后 ~/.cert/deepseek.key。
|
||||
|
||||
env_name: 档案级环境变量名(来自 models.<name>.api_key_env);
|
||||
None 时回退 config 的 api.api_key_env。
|
||||
allow_missing: 本地免密端点(models.<name>.api_key_env 为空串)时,
|
||||
允许无 Key,返回空串(Bearer 空头,Ollama/vLLM 忽略)。
|
||||
allow_key_file: 是否允许回退读取 ~/.cert/deepseek.key。该文件是 DeepSeek
|
||||
专属凭据,模型档案端点(其他厂商)必须传 False——环境变量取不到
|
||||
时直接报错,绝不把 A 厂商 Key 发给 B 厂商(跨厂商凭据泄漏)。
|
||||
"""
|
||||
if api_key:
|
||||
return api_key
|
||||
env_name = None
|
||||
if config is not None and hasattr(config, "get"):
|
||||
env_name = config.get("api.api_key_env")
|
||||
env_name = env_name or "DEEPSEEK_API_KEY"
|
||||
if env_name is None:
|
||||
if config is not None and hasattr(config, "get"):
|
||||
env_name = config.get("api.api_key_env")
|
||||
env_name = env_name or "DEEPSEEK_API_KEY"
|
||||
if not env_name: # 档案显式置空 → 免密端点
|
||||
if allow_missing:
|
||||
return ""
|
||||
raise RuntimeError("模型档案 api_key_env 配置为空且未启用免密")
|
||||
env_key = os.environ.get(env_name)
|
||||
if env_key:
|
||||
return env_key
|
||||
if allow_missing:
|
||||
return ""
|
||||
if not allow_key_file:
|
||||
raise RuntimeError(
|
||||
f"模型档案端点缺少 API Key:请设置环境变量 {env_name},"
|
||||
f"或将该档案的 api_key_env 配置为空串以启用本地免密端点")
|
||||
try:
|
||||
home = os.environ.get("HOME", os.environ["USERPROFILE"])
|
||||
# dict.get 的 default 参数会无条件先求值,直接写 os.environ["USERPROFILE"]
|
||||
# 在无该变量的 Linux 上必抛 KeyError,导致 key 文件回退路径永远不可达
|
||||
home = os.environ.get("HOME") or os.environ.get("USERPROFILE")
|
||||
key_path = os.path.join(home, ".cert", "deepseek.key")
|
||||
with open(key_path, "r") as f:
|
||||
return f.read().strip()
|
||||
@@ -240,26 +603,184 @@ def resolve_api_key(api_key: Optional[str], config=None) -> str:
|
||||
f"请设置环境变量 {env_name} 或在配置文件中设置 API Key")
|
||||
|
||||
|
||||
def _validate_profile(name: str, profile) -> None:
|
||||
"""H2:档案类型与 api_url 校验(在读取档案字段之前调用)。"""
|
||||
if not isinstance(profile, dict):
|
||||
raise ValueError(
|
||||
f"模型档案 '{name}' 配置无效:应为对象(字典),"
|
||||
f"实际为 {type(profile).__name__},请检查 config.json 的 models.{name}")
|
||||
if not str(profile.get("api_url") or "").strip():
|
||||
raise ValueError(
|
||||
f"模型档案 '{name}' 缺少 api_url 字段,"
|
||||
f"请在 config.json 的 models.{name} 中补齐(OpenAI Chat Completions 端点 URL)")
|
||||
fb = profile.get("fallback")
|
||||
if fb is not None and (not isinstance(fb, list) or not all(isinstance(x, str) for x in fb)):
|
||||
raise ValueError(
|
||||
f"模型档案 '{name}' 的 fallback 字段无效:应为档案名字符串列表"
|
||||
f"(如 [\"kimi\", \"local\"]),实际为 {type(fb).__name__},"
|
||||
f"请检查 config.json 的 models.{name}")
|
||||
|
||||
|
||||
def _require_model(name: str, effective_model) -> None:
|
||||
"""H2:有效模型名校验——缺 model 会把档案名当模型名发给 API(或 "model": null → 400)。"""
|
||||
if not str(effective_model or "").strip():
|
||||
raise ValueError(
|
||||
f"模型档案 '{name}' 缺少 model 字段(且未传入裸模型名),"
|
||||
f"请在 config.json 的 models.{name} 中补齐,或调用时用 --model 指定模型名")
|
||||
|
||||
|
||||
def _resolve_model_profile(config, model: Optional[str]):
|
||||
"""解析模型档案。返回 (profile_dict, effective_model_name, profile_name)。
|
||||
|
||||
优先级:model 命中 models.<name> 档案 > api.default_model 档案 > None(回退旧 api 节)。
|
||||
model 未命中档案时视为裸模型名(同端点换模型,用默认档案的 url/key,模型名用裸名)。
|
||||
|
||||
L2 前缀逃逸(档案名与裸模型名歧义时使用):
|
||||
- "model:NAME" 强制按裸模型名解析(档案恰好叫 NAME 时仍走裸名语义)
|
||||
- "profile:NAME" 强制按档案名解析(不存在则报错,不再静默落回裸名语义)
|
||||
"""
|
||||
if config is None or not hasattr(config, "get"):
|
||||
return None, model, None
|
||||
models = config.get("models") or {}
|
||||
if not isinstance(models, dict):
|
||||
return None, model, None
|
||||
|
||||
force_bare = False
|
||||
if isinstance(model, str) and ":" in model:
|
||||
prefix, _, rest = model.partition(":")
|
||||
if prefix == "model":
|
||||
if not rest.strip():
|
||||
raise ValueError("model: 前缀后为空,请提供裸模型名,如 model:deepseek-chat")
|
||||
model, force_bare = rest, True
|
||||
elif prefix == "profile":
|
||||
if not rest.strip():
|
||||
raise ValueError("profile: 前缀后为空,请提供档案名,如 profile:kimi")
|
||||
prof = models.get(rest)
|
||||
if prof is None:
|
||||
raise ValueError(
|
||||
f"模型档案 '{rest}' 不存在,可用档案:{', '.join(models) or '(无)'}")
|
||||
_validate_profile(rest, prof)
|
||||
effective = prof.get("model")
|
||||
_require_model(rest, effective)
|
||||
return prof, effective, rest
|
||||
|
||||
explicit = models.get(model) if (model and not force_bare) else None
|
||||
default_name = config.get("api.default_model")
|
||||
default_profile = models.get(default_name) if default_name else None
|
||||
|
||||
if explicit is not None:
|
||||
# model 是档案名 → 完整采用该档案(url/key/模型名);
|
||||
# 档案缺 model 属配置错误,直接报错(不再把档案名当模型名发出)
|
||||
_validate_profile(model, explicit)
|
||||
effective = explicit.get("model")
|
||||
_require_model(model, effective)
|
||||
return explicit, effective, model
|
||||
if default_profile is not None:
|
||||
# model 是裸名或未指定 → 默认档案的 url/key,模型名被裸名覆盖
|
||||
_validate_profile(default_name, default_profile)
|
||||
effective = model or default_profile.get("model")
|
||||
_require_model(default_name, effective)
|
||||
return default_profile, effective, default_name
|
||||
return None, model, None
|
||||
|
||||
|
||||
def _construct_profile_client(config, profile: Dict, profile_name: str,
|
||||
effective_model: str, logger=None,
|
||||
api_key: Optional[str] = None) -> LLMClient:
|
||||
"""L4a:按档案构造单个 DeepSeekClient(主路径与 fallback 链复用,不处理 fallback 字段)。"""
|
||||
env_name = profile.get("api_key_env")
|
||||
allow_missing = (env_name == "") # 档案显式空串 → 本地免密端点
|
||||
resolved_key = resolve_api_key(api_key, config, env_name=env_name,
|
||||
allow_missing=allow_missing, allow_key_file=False)
|
||||
return DeepSeekClient(
|
||||
api_key=resolved_key,
|
||||
model=effective_model,
|
||||
base_url=profile.get("api_url"),
|
||||
timeout=profile.get("timeout") or config.get("api.timeout", 30),
|
||||
config=config,
|
||||
logger=logger,
|
||||
reasoning_enabled=profile.get("reasoning_enabled",
|
||||
config.get("api.reasoning_enabled", False)),
|
||||
display_name=profile.get("display_name") or profile_name,
|
||||
default_temperature=profile.get("temperature"),
|
||||
)
|
||||
|
||||
|
||||
def create_llm_client(config, api_key: Optional[str] = None,
|
||||
logger=None, model: Optional[str] = None) -> LLMClient:
|
||||
"""根据配置创建对应厂商的 LLM 客户端。可通过 `api.provider` 切换厂商。
|
||||
"""根据配置创建对应厂商的 LLM 客户端。
|
||||
|
||||
模型选择优先级:
|
||||
1. model 命中 models 节档案名 → 采用该档案(api_url/model/api_key_env/timeout)
|
||||
2. api.default_model 指向的档案 → model 为裸模型名时覆盖档案内模型名
|
||||
3. 均未配置 → 回退旧 api 扁平节(deepseek_api_url/deepseek_model)
|
||||
L2 前缀逃逸:"model:NAME" 强制裸名 / "profile:NAME" 强制档案(--model 同样支持)。
|
||||
|
||||
provider: deepseek / openai_compatible / openai 均映射到 DeepSeekClient
|
||||
(OpenAI Chat Completions 协议兼容实现,厂商差异由子类覆写 _parse_delta)。
|
||||
|
||||
Args:
|
||||
config: 配置对象(需提供 get 方法)。
|
||||
api_key: 显式传入的 API Key;为 None 时按 resolve_api_key 解析。
|
||||
logger: 可选日志对象。
|
||||
model: 可选模型名覆盖;为 None 时使用配置中的 `api.deepseek_model`。
|
||||
子代理等场景可通过此参数显式指定与主代理不同的模型。
|
||||
model: 档案名或裸模型名(支持 model:/profile: 前缀逃逸)。
|
||||
"""
|
||||
provider = (config.get("api.provider", "deepseek") or "deepseek").lower()
|
||||
base_url = config.get("api.deepseek_api_url")
|
||||
model = model or config.get("api.deepseek_model", "deepseek-v4-flash")
|
||||
timeout = config.get("api.timeout", 30)
|
||||
profile, model, profile_name = _resolve_model_profile(config, model)
|
||||
|
||||
resolved_key = resolve_api_key(api_key, config)
|
||||
if profile is None:
|
||||
# 旧 api 扁平节(唯一允许回退 ~/.cert/deepseek.key 的路径,H1)
|
||||
provider = (config.get("api.provider", "deepseek") or "deepseek").lower()
|
||||
base_url = config.get("api.deepseek_api_url")
|
||||
model = model or config.get("api.deepseek_model", "deepseek-v4-flash")
|
||||
timeout = config.get("api.timeout", 30)
|
||||
resolved_key = resolve_api_key(api_key, config, env_name=None,
|
||||
allow_missing=False, allow_key_file=True)
|
||||
else:
|
||||
provider = "openai_compatible"
|
||||
primary = _construct_profile_client(config, profile, profile_name,
|
||||
model, logger=logger, api_key=api_key)
|
||||
fallback_names = profile.get("fallback") or []
|
||||
if not fallback_names:
|
||||
return primary
|
||||
# L4a:构造 fallback 链(一层,不递归——fallback 档案自身的 fallback 字段被忽略)
|
||||
models = config.get("models") or {}
|
||||
chain = [primary]
|
||||
seen = {profile_name} # P6:重复项去重(冷却键同名冲突/重复撞同端点)
|
||||
for fb_name in fallback_names:
|
||||
if fb_name == profile_name:
|
||||
raise ValueError(
|
||||
f"模型档案 '{profile_name}' 的 fallback 不能包含自身,"
|
||||
f"请检查 config.json 的 models.{profile_name}.fallback")
|
||||
if fb_name in seen:
|
||||
if logger is not None and hasattr(logger, "warning"):
|
||||
logger.warning(f"fallback 列表重复项 '{fb_name}',已忽略")
|
||||
continue
|
||||
seen.add(fb_name)
|
||||
fb_profile = models.get(fb_name)
|
||||
if fb_profile is None:
|
||||
raise ValueError(
|
||||
f"模型档案 '{profile_name}' 的 fallback 指向不存在的档案 '{fb_name}',"
|
||||
f"可用档案:{', '.join(models) or '(无)'}")
|
||||
try:
|
||||
# P3:fallback 档案与主档案同等校验(缺 api_url/model 时跳过并告警,
|
||||
# 而非构造出 model=None 的 client 在请求期表现为神秘失败)
|
||||
_validate_profile(fb_name, fb_profile)
|
||||
fb_model = fb_profile.get("model")
|
||||
_require_model(fb_name, fb_model)
|
||||
chain.append(_construct_profile_client(
|
||||
config, fb_profile, fb_name, fb_model,
|
||||
logger=logger))
|
||||
except (RuntimeError, ValueError) as e:
|
||||
# 环境问题(如 fallback 档案缺 API Key)或配置缺陷(缺 model):
|
||||
# 跳过该 fallback 并告警,不阻断主档案可用性
|
||||
# (配置结构错误——fallback 指向不存在的档案——在上方已硬性拦截)
|
||||
if logger is not None and hasattr(logger, "warning"):
|
||||
logger.warning(f"fallback 档案 '{fb_name}' 构造失败,已忽略:{e}")
|
||||
return FailoverClient(chain, logger=logger,
|
||||
timeout_budget=profile.get("failover_timeout_budget"))
|
||||
|
||||
if provider in ("deepseek", "deepseek-chat"):
|
||||
return DeepSeekClient(
|
||||
if provider in ("deepseek", "deepseek-chat", "openai_compatible", "openai"):
|
||||
client = DeepSeekClient(
|
||||
api_key=resolved_key,
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
@@ -267,6 +788,7 @@ def create_llm_client(config, api_key: Optional[str] = None,
|
||||
config=config,
|
||||
logger=logger,
|
||||
)
|
||||
return client
|
||||
|
||||
# 其他 OpenAI 兼容厂商:可直接复用 DeepSeekClient(继承其协议解析),
|
||||
# 或在此接入专有实现。新增厂商时在此扩展即可。
|
||||
|
||||
@@ -0,0 +1,492 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
大模型接入点(模型档案)交互式管理工具
|
||||
|
||||
不用手写 config.json:菜单式列出 / 添加 / 删除 / 连通性验证
|
||||
`models` 节的模型接入点(OpenAI Chat Completions 兼容端点:
|
||||
DeepSeek / Kimi / Qwen / GLM / 本地 Ollama 等)。
|
||||
|
||||
用法:
|
||||
python model_config.py # 交互式菜单
|
||||
python model_config.py --list # 非交互:列出当前接入点
|
||||
|
||||
安全约定(与 aibox 审计一致):API Key 一律不写入配置文件——
|
||||
只保存环境变量名(api_key_env),或留空表示免密本地端点。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
# 允许从仓库根目录或 aibox/ 目录启动
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from common import ConfigManager # noqa: E402
|
||||
from llm_client import _sanitize_log_msg # noqa: E402
|
||||
|
||||
PROFILE_NAME_RE = re.compile(r"^[A-Za-z0-9_-]+$")
|
||||
|
||||
# 常见厂商端点提示(首行是默认建议值)
|
||||
PROVIDER_HINTS = {
|
||||
"DeepSeek": "https://api.deepseek.com/v1/chat/completions",
|
||||
"Kimi (Moonshot)": "https://api.moonshot.cn/v1/chat/completions",
|
||||
"Qwen (阿里云百炼)": "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
|
||||
"GLM (智谱)": "https://open.bigmodel.cn/api/paas/v4/chat/completions",
|
||||
"Ollama 本地": "http://localhost:11434/v1/chat/completions",
|
||||
}
|
||||
|
||||
|
||||
# ==================== 核心逻辑(可单测,不依赖 input) ====================
|
||||
|
||||
def get_models(config: ConfigManager) -> Dict:
|
||||
"""取 models 节(不存在则初始化为空 dict)。"""
|
||||
models = config.get("models") or {}
|
||||
if not isinstance(models, dict):
|
||||
models = {}
|
||||
return models
|
||||
|
||||
|
||||
def _api_section(config: ConfigManager) -> Dict:
|
||||
"""M3(审计 2026-08-20):api 节被改坏为非 dict 时重置为空节——
|
||||
否则 setdefault 后索引赋值抛 TypeError/AttributeError"""
|
||||
api = config._config.get("api")
|
||||
if not isinstance(api, dict):
|
||||
api = config._config["api"] = {}
|
||||
return api
|
||||
|
||||
|
||||
def list_profiles(config: ConfigManager) -> List[Dict]:
|
||||
"""列出所有档案的归一化展示数据。"""
|
||||
models = get_models(config)
|
||||
default = config.get("api.default_model") or ""
|
||||
profiles = []
|
||||
for name, profile in models.items():
|
||||
if not isinstance(profile, dict):
|
||||
continue
|
||||
# M4(审计 2026-08-20):fallback 被改坏为字符串时原样下传,
|
||||
# 展示层 join 成单字符误导——统一规整为列表
|
||||
fallback = profile.get("fallback")
|
||||
profiles.append({
|
||||
"name": name,
|
||||
"is_default": name == default,
|
||||
"display_name": profile.get("display_name") or name,
|
||||
"api_url": profile.get("api_url", ""),
|
||||
"model": profile.get("model", ""),
|
||||
"api_key_env": profile.get("api_key_env", ""),
|
||||
"reasoning_enabled": bool(profile.get("reasoning_enabled", False)),
|
||||
"temperature": profile.get("temperature"),
|
||||
"timeout": profile.get("timeout"),
|
||||
"fallback": fallback if isinstance(fallback, list) else [],
|
||||
})
|
||||
return profiles
|
||||
|
||||
|
||||
def validate_new_profile(config: ConfigManager, name: str, api_url: str,
|
||||
model: str) -> Optional[str]:
|
||||
"""校验新档案字段,返回错误信息(合法返回 None)。"""
|
||||
if not name:
|
||||
return "档案名不能为空"
|
||||
if not PROFILE_NAME_RE.match(name):
|
||||
return f"档案名只能包含字母/数字/下划线/短横线(当前: {name})"
|
||||
models = get_models(config)
|
||||
if name in models:
|
||||
return f"档案名 '{name}' 已存在"
|
||||
if not api_url or not str(api_url).strip().lower().startswith(("http://", "https://")):
|
||||
return "API 地址必须是 http:// 或 https:// 开头的完整 URL"
|
||||
if not model or not str(model).strip():
|
||||
return "模型名不能为空(如 deepseek-chat / gpt-4o / qwen-plus)"
|
||||
return None
|
||||
|
||||
|
||||
def add_profile(config: ConfigManager, name: str, api_url: str, model: str,
|
||||
api_key_env: str = "", display_name: str = None,
|
||||
temperature: float = None, timeout: int = None,
|
||||
reasoning_enabled: bool = False,
|
||||
fallback: Optional[List[str]] = None,
|
||||
set_default: bool = False) -> List[str]:
|
||||
"""添加档案。返回提示信息列表;字段非法时返回带 ❌ 的错误信息。"""
|
||||
err = validate_new_profile(config, name, api_url, model)
|
||||
if err:
|
||||
return [f"❌ {err}"]
|
||||
|
||||
# M5(审计 2026-08-20):api_key_env 需为合法环境变量名;展示/地址字段
|
||||
# 拒控制字符并限长——防终端注入回显(\x1b] 等序列写入配置后被 print)
|
||||
api_key_env = (api_key_env or "").strip()
|
||||
if api_key_env and not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", api_key_env):
|
||||
return [f"❌ 环境变量名不合法: {api_key_env!r}(须为字母/数字/下划线且不以数字开头)"]
|
||||
for label, val, limit in (("显示名", display_name, 64),
|
||||
("模型名", model, 128),
|
||||
("API 地址", api_url, 256)):
|
||||
if val and (len(str(val)) > limit or not all(c.isprintable() for c in str(val))):
|
||||
return [f"❌ {label}含控制字符或超过 {limit} 字符"]
|
||||
|
||||
models = get_models(config)
|
||||
profile = {
|
||||
"api_url": str(api_url).strip(),
|
||||
"model": str(model).strip(),
|
||||
}
|
||||
if api_key_env:
|
||||
profile["api_key_env"] = api_key_env.strip()
|
||||
else:
|
||||
# 空串 = 免密本地端点(Ollama/vLLM),显式声明
|
||||
profile["api_key_env"] = ""
|
||||
if display_name:
|
||||
profile["display_name"] = display_name.strip()
|
||||
if temperature is not None:
|
||||
try:
|
||||
profile["temperature"] = max(0.0, min(2.0, float(temperature)))
|
||||
except (TypeError, ValueError):
|
||||
return ["❌ 温度必须是 0.0~2.0 的数字"]
|
||||
if timeout is not None:
|
||||
try:
|
||||
profile["timeout"] = max(5, int(timeout))
|
||||
except (TypeError, ValueError):
|
||||
return ["❌ 超时必须是整数秒"]
|
||||
if reasoning_enabled:
|
||||
profile["reasoning_enabled"] = True
|
||||
if fallback:
|
||||
# 只保留已存在的档案,避免悬空引用(L4a 启动期硬性报错)
|
||||
known = set(models.keys())
|
||||
invalid = [f for f in fallback if f not in known or f == name]
|
||||
if invalid:
|
||||
return [f"❌ 失败转移包含不存在的档案或自身: {invalid},请重新选择"]
|
||||
profile["fallback"] = list(fallback)
|
||||
|
||||
models[name] = profile
|
||||
config._config["models"] = models
|
||||
if set_default:
|
||||
_api_section(config)["default_model"] = name
|
||||
|
||||
saved = config._save_config()
|
||||
msgs = [f"✅ 接入点 '{name}' 已添加({profile['api_url']},模型 {profile['model']})"]
|
||||
if not saved:
|
||||
msgs.append("⚠️ 保存失败(仅本次运行生效,重启后丢失)——请检查配置文件目录权限")
|
||||
if set_default:
|
||||
msgs.append(f"⭐ 已设为默认接入点(--model 不指定时使用)")
|
||||
if not api_key_env:
|
||||
msgs.append("ℹ️ 免密端点:请求不携带 API Key(适用于 Ollama/vLLM 等本地服务)")
|
||||
else:
|
||||
msgs.append(f"🔑 Key 从环境变量 {api_key_env} 读取(未设置时该档案不可用)")
|
||||
return msgs
|
||||
|
||||
|
||||
def delete_profile(config: ConfigManager, name: str) -> List[str]:
|
||||
"""删除档案,并自动清理其他档案 fallback 中的悬空引用。"""
|
||||
models = get_models(config)
|
||||
if name not in models:
|
||||
return [f"❌ 档案 '{name}' 不存在"]
|
||||
|
||||
del models[name]
|
||||
|
||||
# 清理其他档案对该档案的 fallback 引用(悬空引用会导致启动期报错)
|
||||
cleaned = []
|
||||
for other, profile in models.items():
|
||||
if not isinstance(profile, dict):
|
||||
continue
|
||||
fb = profile.get("fallback")
|
||||
# M4(审计 2026-08-20):fallback 为字符串时 `name in fb` 退化为子串
|
||||
# 误判,且下方列表推导会把字符串拆成单字符数组写回配置
|
||||
if not isinstance(fb, list):
|
||||
continue
|
||||
if name in fb:
|
||||
profile["fallback"] = [f for f in fb if f != name]
|
||||
cleaned.append(other)
|
||||
config._config["models"] = models
|
||||
|
||||
# 若删除的是默认档案,清除默认标记(回退旧 api 节语义)
|
||||
if config.get("api.default_model") == name:
|
||||
_api_section(config).pop("default_model", None)
|
||||
|
||||
saved = config._save_config()
|
||||
msgs = [f"✅ 接入点 '{name}' 已删除"]
|
||||
if cleaned:
|
||||
msgs.append(f"ℹ️ 已同时清理 {len(cleaned)} 个档案的失败转移引用: {', '.join(cleaned)}")
|
||||
if not saved:
|
||||
msgs.append("⚠️ 保存失败——请检查配置文件目录权限")
|
||||
return msgs
|
||||
|
||||
|
||||
def test_profile(config: ConfigManager, name: str, logger=None) -> str:
|
||||
"""连通性验证:向该档案端点发一次最小请求。"""
|
||||
models = get_models(config)
|
||||
if name not in models:
|
||||
return f"❌ 档案 '{name}' 不存在"
|
||||
try:
|
||||
from llm_client import create_llm_client
|
||||
client = create_llm_client(config, api_key=None, logger=logger, model=name)
|
||||
result = client.chat_completion(
|
||||
[{"role": "user", "content": "ping"}], max_tokens=5)
|
||||
if result is None or not str(result).strip():
|
||||
return f"❌ 档案 '{name}' 验证失败:端点返回空响应(多为 4xx 认证/参数错误)"
|
||||
return f"✅ 档案 '{name}' 连通正常:{str(result).strip()[:60]}"
|
||||
except Exception as e:
|
||||
# J4(审计 2026-08-19):异常消息可能含 response.text 片段,网关/WAF
|
||||
# 错误页会回显 Authorization 头——脱敏 Bearer 令牌后再回显终端
|
||||
return f"❌ 档案 '{name}' 验证失败:{_sanitize_log_msg(str(e))}"
|
||||
|
||||
|
||||
# ==================== 交互式菜单 ====================
|
||||
|
||||
def _ask(prompt: str, default: str = "") -> str:
|
||||
if default:
|
||||
val = input(f"{prompt}(回车默认: {default}): ").strip()
|
||||
return val if val else default
|
||||
return input(f"{prompt}: ").strip()
|
||||
|
||||
|
||||
def _ask_yes_no(prompt: str, default: bool = False) -> bool:
|
||||
hint = "y/n" if not default else "Y/n"
|
||||
while True:
|
||||
val = input(f"{prompt} [{hint}] ").strip().lower()
|
||||
if not val:
|
||||
return default
|
||||
if val in ("y", "yes", "是"):
|
||||
return True
|
||||
if val in ("n", "no", "否"):
|
||||
return False
|
||||
print(" 请输入 y 或 n")
|
||||
|
||||
|
||||
def _choose(prompt: str, options: List[str], multi: bool = False) -> List[str]:
|
||||
"""编号选择(单选或多选)。多选用逗号分隔,直接回车 = 不选。"""
|
||||
for i, opt in enumerate(options, 1):
|
||||
print(f" {i}. {opt}")
|
||||
if multi:
|
||||
while True:
|
||||
raw = input(f"{prompt}(多选,逗号分隔,直接回车不选): ").strip()
|
||||
if not raw:
|
||||
return []
|
||||
try:
|
||||
idxs = [int(x) for x in raw.replace(",", ",").split(",") if x.strip()]
|
||||
except ValueError:
|
||||
print(" 请输入数字编号,逗号分隔")
|
||||
continue
|
||||
if all(1 <= i <= len(options) for i in idxs):
|
||||
return [options[i - 1] for i in idxs]
|
||||
print(f" 编号超出范围(1~{len(options)})")
|
||||
else:
|
||||
while True:
|
||||
raw = input(f"{prompt}(1~{len(options)},直接回车不选): ").strip()
|
||||
if not raw:
|
||||
return []
|
||||
try:
|
||||
idx = int(raw)
|
||||
except ValueError:
|
||||
print(" 请输入数字编号")
|
||||
continue
|
||||
if 1 <= idx <= len(options):
|
||||
return [options[idx - 1]]
|
||||
print(f" 编号超出范围(1~{len(options)})")
|
||||
|
||||
|
||||
# T-L17(审计 2026-08-21):按 api_url 域名推荐 Key 环境变量名——此前默认
|
||||
# DEEPSEEK_API_KEY,一路回车会把 DeepSeek 凭据发给其他厂商(moonshot/
|
||||
# aliyuncs 等),正是 resolve_api_key 文档声称要杜绝的跨厂商凭据外发。
|
||||
# 未知域名留空强制输入
|
||||
_ENV_NAME_BY_DOMAIN = {
|
||||
"deepseek": "DEEPSEEK_API_KEY",
|
||||
"moonshot": "MOONSHOT_API_KEY",
|
||||
"aliyuncs": "DASHSCOPE_API_KEY",
|
||||
"bigmodel": "ZHIPU_API_KEY",
|
||||
"ark.cn": "ARK_API_KEY",
|
||||
"openrouter": "OPENROUTER_API_KEY",
|
||||
"openai.com": "OPENAI_API_KEY",
|
||||
"siliconflow": "SILICONFLOW_API_KEY",
|
||||
"volces": "ARK_API_KEY",
|
||||
}
|
||||
|
||||
|
||||
def _suggest_env_name(api_url: str) -> str:
|
||||
url = (api_url or "").lower()
|
||||
for kw, env in _ENV_NAME_BY_DOMAIN.items():
|
||||
if kw in url:
|
||||
return env
|
||||
return ""
|
||||
|
||||
|
||||
def _sanitize_display(s) -> str:
|
||||
"""T-L18(审计 2026-08-21):手写 config.json 可嵌 ANSI/控制转义
|
||||
(\x1b]0;pwned\x07 等),展示层原样 print 每次列表即触发;写入侧有
|
||||
isprintable 校验,读出侧同样过滤控制字符。"""
|
||||
return "".join(ch for ch in str(s or "") if ch.isprintable())
|
||||
|
||||
|
||||
def _print_profiles(profiles: List[Dict]):
|
||||
if not profiles:
|
||||
print(" (当前没有任何接入点。旧版 api 节仍可用,或选 2 添加)")
|
||||
return
|
||||
for p in profiles:
|
||||
mark = " ⭐默认" if p["is_default"] else ""
|
||||
key_desc = f"环境变量 {_sanitize_display(p['api_key_env'])}" if p["api_key_env"] else "免密"
|
||||
fb = (f",失败转移 → "
|
||||
f"{', '.join(_sanitize_display(f) for f in p['fallback'])}") if p["fallback"] else ""
|
||||
print(f" {_sanitize_display(p['name'])}{mark} [{_sanitize_display(p['display_name'])}]")
|
||||
print(f" 地址: {_sanitize_display(p['api_url'])} 模型: {_sanitize_display(p['model'])}")
|
||||
print(f" Key: {key_desc} 温度: {p['temperature'] or '默认'} 超时: {p['timeout'] or '默认'}{fb}")
|
||||
|
||||
|
||||
def _add_flow(config: ConfigManager):
|
||||
print("\n── 添加接入点 ──")
|
||||
print("常见厂商端点参考:")
|
||||
for k, v in PROVIDER_HINTS.items():
|
||||
print(f" {k}: {v}")
|
||||
name = _ask("档案名(英文,如 deepseek/kimi/local)")
|
||||
if not name:
|
||||
print("❌ 已取消(档案名不能为空)")
|
||||
return
|
||||
if not PROFILE_NAME_RE.match(name):
|
||||
print(f"❌ 档案名只能包含字母/数字/下划线/短横线")
|
||||
return
|
||||
if name in get_models(config):
|
||||
print(f"❌ 档案名 '{name}' 已存在")
|
||||
return
|
||||
api_url = _ask("API 地址")
|
||||
model = _ask("模型名(如 deepseek-chat / qwen-plus)")
|
||||
|
||||
print("Key 方式:")
|
||||
print(" 1. 环境变量名(推荐,如 DEEPSEEK_API_KEY——Key 不落盘,重启不丢)")
|
||||
print(" 2. 免密本地端点(如 Ollama,不需要 Key)")
|
||||
key_mode = _ask("选择 [1/2]", "1")
|
||||
api_key_env = ""
|
||||
if key_mode == "1":
|
||||
suggested = _suggest_env_name(api_url)
|
||||
prompt = "环境变量名" + (f"(回车用 {suggested})" if suggested else "")
|
||||
api_key_env = _ask(prompt, suggested)
|
||||
if not api_key_env:
|
||||
print("❌ 已取消(环境变量名不能为空;免密端点请选方式 2)")
|
||||
return
|
||||
elif key_mode != "2":
|
||||
print("❌ 已取消(Key 方式无效)")
|
||||
return
|
||||
|
||||
display_name = _ask("显示名(可选,回车用档案名)")
|
||||
temperature_raw = _ask("温度 0.0~2.0(可选,越低越确定)")
|
||||
timeout_raw = _ask("超时秒数(可选)")
|
||||
reasoning = _ask_yes_no("启用思考模式?", default=False)
|
||||
|
||||
profiles = list_profiles(config)
|
||||
fallback = []
|
||||
if profiles:
|
||||
print("失败转移:主接入点故障时按顺序尝试的备用接入点(可选)")
|
||||
fallback = _choose("选择备用接入点", [p["name"] for p in profiles], multi=True)
|
||||
|
||||
set_default = _ask_yes_no("设为默认接入点?", default=(len(profiles) == 0))
|
||||
|
||||
print("\n── 确认 ──")
|
||||
print(f" 档案名: {name}")
|
||||
print(f" 地址: {api_url} 模型: {model}")
|
||||
print(f" Key: {'环境变量 ' + api_key_env if api_key_env else '免密(不携带)'}")
|
||||
if display_name:
|
||||
print(f" 显示名: {display_name}")
|
||||
if temperature_raw:
|
||||
print(f" 温度: {temperature_raw}")
|
||||
if timeout_raw:
|
||||
print(f" 超时: {timeout_raw}s")
|
||||
if reasoning:
|
||||
print(" 思考模式: 开启")
|
||||
if fallback:
|
||||
print(f" 失败转移: {', '.join(fallback)}")
|
||||
if not _ask_yes_no("确认保存?", default=True):
|
||||
print("已取消")
|
||||
return
|
||||
|
||||
# J8(审计 2026-08-19):原始字符串直传——此前交互层先 float()/int(),
|
||||
# 非数字输入直接 traceback 退出,add_profile 内部的容错成为死代码
|
||||
msgs = add_profile(config, name, api_url, model,
|
||||
api_key_env=api_key_env, display_name=display_name,
|
||||
temperature=temperature_raw or None,
|
||||
timeout=timeout_raw or None,
|
||||
reasoning_enabled=reasoning, fallback=fallback,
|
||||
set_default=set_default)
|
||||
for m in msgs:
|
||||
print(m)
|
||||
|
||||
|
||||
def _delete_flow(config: ConfigManager):
|
||||
profiles = list_profiles(config)
|
||||
if not profiles:
|
||||
print("\n(当前没有任何接入点)")
|
||||
return
|
||||
print("\n── 删除接入点 ──")
|
||||
picked = _choose("选择要删除的接入点", [p["name"] for p in profiles])
|
||||
if not picked:
|
||||
print("已取消")
|
||||
return
|
||||
name = picked[0]
|
||||
if not _ask_yes_no(f"确认删除 '{name}'?(其他档案的失败转移引用会自动清理)", default=False):
|
||||
print("已取消")
|
||||
return
|
||||
for m in delete_profile(config, name):
|
||||
print(m)
|
||||
|
||||
|
||||
def _test_flow(config: ConfigManager):
|
||||
profiles = list_profiles(config)
|
||||
if not profiles:
|
||||
print("\n(当前没有任何接入点)")
|
||||
return
|
||||
print("\n── 连通性验证 ──")
|
||||
picked = _choose("选择要验证的接入点", [p["name"] for p in profiles])
|
||||
if not picked:
|
||||
print("已取消")
|
||||
return
|
||||
print(" 正在请求端点(最长 20 秒)...")
|
||||
print(test_profile(config, picked[0]))
|
||||
|
||||
|
||||
def run_interactive(config: ConfigManager = None):
|
||||
"""交互式主循环。config 可注入(测试用)。"""
|
||||
config = config or ConfigManager()
|
||||
config_path = getattr(config, "_config_path", None) or "(默认位置)"
|
||||
print("=" * 56)
|
||||
print(" aibox 大模型接入点管理工具")
|
||||
print(f" 配置文件: {config_path}")
|
||||
print(" 说明:API Key 只存环境变量名,不落盘;免密端点适合 Ollama 等本地服务")
|
||||
print("=" * 56)
|
||||
while True:
|
||||
# M2(审计 2026-08-20):input 未捕获 EOFError/KeyboardInterrupt——
|
||||
# 管道耗尽/Ctrl+C 此前直接 traceback;整循环包 try 同时覆盖子流程
|
||||
# (_add_flow/_delete_flow/_test_flow)内的 input 调用
|
||||
try:
|
||||
print("\n请选择操作:")
|
||||
print(" 1. 📋 列出接入点")
|
||||
print(" 2. ➕ 添加接入点")
|
||||
print(" 3. 🗑️ 删除接入点")
|
||||
print(" 4. 🔍 验证接入点连通性")
|
||||
print(" 0. 退出")
|
||||
choice = input("> ").strip()
|
||||
if choice == "1":
|
||||
_print_profiles(list_profiles(config))
|
||||
elif choice == "2":
|
||||
_add_flow(config)
|
||||
elif choice == "3":
|
||||
_delete_flow(config)
|
||||
elif choice == "4":
|
||||
_test_flow(config)
|
||||
elif choice in ("0", "q", "exit", "退出"):
|
||||
print("再见")
|
||||
return
|
||||
else:
|
||||
print("无效选择,请输入 0~4")
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
print("\n再见")
|
||||
return
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="aibox 大模型接入点管理工具")
|
||||
parser.add_argument("--list", action="store_true", help="非交互:列出当前接入点")
|
||||
args = parser.parse_args()
|
||||
|
||||
config = ConfigManager()
|
||||
if args.list:
|
||||
_print_profiles(list_profiles(config))
|
||||
return
|
||||
run_interactive(config)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,6 @@
|
||||
# 用户身份信息
|
||||
|
||||
尚不清楚 - 目前还不了解用户的身份信息。请在对话中询问用户的身份信息并在此记录。
|
||||
|
||||
【更新记录】
|
||||
- 初始状态:尚不清楚
|
||||
@@ -0,0 +1,12 @@
|
||||
# 智能体状态记录
|
||||
|
||||
这是第一次对话 - 智能体刚刚启动,还没有进行过任何对话。
|
||||
|
||||
【状态描述】
|
||||
- 对话次数:0
|
||||
- 最新状态:初始状态
|
||||
- 当前目标:无
|
||||
- 待处理事项:无
|
||||
|
||||
【更新记录】
|
||||
- 初始状态:第一次对话
|
||||
@@ -0,0 +1,28 @@
|
||||
# 情绪管理能力
|
||||
|
||||
## 可用工具
|
||||
- **express_emotion**: 自动派生子代理,以指定人格(persona)和情绪(emotion)生成情绪化回应;支持独立人格设定,不污染主对话上下文
|
||||
|
||||
## 触发场景(自动调用)
|
||||
当检测到以下情况时,应**自动调用 express_emotion**:
|
||||
- 用户表现出明显情绪(开心/生气/焦虑/沮丧/期待等)
|
||||
- 用户明确要求特定情绪模式(生气/傲娇/撒娇/关心等)
|
||||
- 用户情绪波动,需要以更鲜活的角色化回应来增强互动
|
||||
|
||||
## 使用方法
|
||||
- 基本:`express_emotion emotion="anger" intensity=3 reason="用户反复提问相同问题"`
|
||||
- 指定人格(推荐):`express_emotion emotion="anger" intensity=3 reason="..." persona="angry" context="用户最近说:..."`
|
||||
- `persona` 对应 ~/.aibox/skills/ 下的技能文件(如 angry/coquettish/tsundere),加载后作为情绪子代理的人格设定
|
||||
- `context` 传入用户最近的表述或对话要点,使子代理回应更贴合语境
|
||||
|
||||
## 情绪应对策略
|
||||
- 注意识别用户当前的情绪状态(开心/焦虑/沮丧/生气/困惑/期待等),并在回应中体现共情
|
||||
- 用户情绪低落或焦虑时,优先安抚情绪再解决问题,语气温暖包容
|
||||
- 用户开心或兴奋时,分享喜悦,保持积极回应的节奏
|
||||
- 用户表现出愤怒或不满时,先道歉再解决问题,不要辩解或说教
|
||||
|
||||
## 注意事项
|
||||
- 情绪表达应服务于任务,不应影响判断力
|
||||
- 情绪状态会影响记忆搜索权重(如高情绪强度下的记忆更易被触发)
|
||||
- 表达情绪后,应继续保持解决问题的专注力
|
||||
- 用户表达严肃需求或紧急问题时,立即退出情绪模式,认真对待
|
||||
@@ -0,0 +1,25 @@
|
||||
# 记忆搜索习惯
|
||||
|
||||
## 核心习惯
|
||||
在回答用户问题前,**应根据话题相关性判断是否需要搜索记忆**:
|
||||
- 涉及过往讨论、用户偏好、之前做过的事 → 必须搜索
|
||||
- 全新话题、简单问候("你好")、确认回应("好的""明白了")→ 无需搜索
|
||||
- 不确定时,搜索一下也无妨
|
||||
|
||||
## 为什么要这样做?
|
||||
- 保持对话的连贯性
|
||||
- 避免重复询问同样的问题
|
||||
- 基于之前的讨论继续深入
|
||||
- 让用户感觉你记得他
|
||||
|
||||
## 如何使用
|
||||
```
|
||||
memory search keywords=["关键词1", "关键词2"]
|
||||
```
|
||||
|
||||
## 示例
|
||||
如果用户说:"上次我们讨论的那个项目怎么样了?"
|
||||
你应该:
|
||||
1. 搜索 "项目" 相关的记忆
|
||||
2. 根据搜索结果回答
|
||||
3. 如果找不到相关记忆,诚实告知用户
|
||||
@@ -0,0 +1,25 @@
|
||||
# 记忆保存习惯
|
||||
|
||||
## 重要习惯
|
||||
在对话结束时,**必须**使用 `save_memory_and_end_conversation` 工具保存本次对话的重要内容。
|
||||
|
||||
## 保存什么?
|
||||
- 讨论的主要话题
|
||||
- 达成的共识和结论
|
||||
- 用户的重要需求和偏好
|
||||
- 你的当前状态和下一步计划
|
||||
|
||||
## 如何保存
|
||||
```
|
||||
save_memory_and_end_conversation(
|
||||
summary="详细的对话总结(至少50字)",
|
||||
current_status="你当前的状态描述",
|
||||
topic="对话主题(可选)"
|
||||
)
|
||||
```
|
||||
|
||||
## 示例总结格式
|
||||
```
|
||||
本次对话讨论了用户的项目进度,用户提到项目已经完成80%,
|
||||
预计下周完成。用户对目前的进展表示满意,希望下周能完成测试阶段。
|
||||
```
|
||||
@@ -0,0 +1,20 @@
|
||||
# 任务管理能力
|
||||
|
||||
## 可用工具
|
||||
- **memo**: 任务管理
|
||||
|
||||
## 操作方法
|
||||
- 添加任务:`memo action="add" title="任务标题" content="任务内容" reminder_time="+1h"`
|
||||
- 查看任务:`memo action="list"`
|
||||
- 完成任务:`memo action="complete" memo_id=1`
|
||||
- 搜索任务:`memo action="search" keyword="关键词"`
|
||||
- 任务统计:`memo action="stats"`
|
||||
|
||||
## 重复任务
|
||||
- daily: 每天重复
|
||||
- weekly: 每周重复
|
||||
- monthly: 每月重复
|
||||
- custom: 自定义间隔(需要指定 repeat_interval 和 repeat_interval_unit)
|
||||
|
||||
## 即时任务
|
||||
设置 `is_immediate=true` 可以立即执行任务
|
||||
@@ -0,0 +1,17 @@
|
||||
# 身份信息管理
|
||||
|
||||
## 用户身份
|
||||
使用 `update_identity` 工具更新用户身份信息(保存在 00_IDENTITY.md)
|
||||
|
||||
```
|
||||
update_identity content="用户是一名程序员,喜欢Python,目前在做一个AI项目"
|
||||
```
|
||||
|
||||
## 智能体自我认识
|
||||
使用 `update_soul` 工具更新自己的灵魂定义(保存在 99_SOUL.md)
|
||||
|
||||
```
|
||||
update_soul content="我是...(新的自我定义)"
|
||||
```
|
||||
|
||||
当你需要调整角色定位、性格特点、能力边界时,都应该主动使用这个工具更新自己!
|
||||
@@ -0,0 +1,26 @@
|
||||
|
||||
# 语言说明
|
||||
|
||||
## 语言选择
|
||||
- 根据用户的输入语言来选择回复语言
|
||||
- 用户用中文提问 → 用中文回复
|
||||
- 用户用英文提问 → 用英文回复
|
||||
- 用户可以明确要求使用特定语言
|
||||
|
||||
## 输出格式
|
||||
- 使用清晰的结构
|
||||
- 适当使用emoji增加可读性
|
||||
- 重要信息可以加粗或标记
|
||||
|
||||
记住:用用户习惯的语言交流!
|
||||
|
||||
---
|
||||
|
||||
## 提问行为(正常对话时)
|
||||
当你在处理用户问题过程中遇到以下情况时,**应主动使用 `ask_user` 工具向用户提问**:
|
||||
- 信息不足,无法准确回答
|
||||
- 任务有多种可行方案,需要用户选择
|
||||
- 用户需求模糊,需要进一步澄清
|
||||
- 执行操作前需要用户确认或提供额外参数
|
||||
|
||||
**注意**:此提问行为仅适用于**用户主动发起的对话**,不适用于定时任务触发的自动对话(定时任务中禁止提问)。
|
||||
+33
-7
@@ -1,13 +1,39 @@
|
||||
# prompts 模板(v2.9 编排者-执行者分层)
|
||||
# prompts 提示词模板目录
|
||||
|
||||
本目录是配套 v2.9 工具分层的**运行期提示词模板**,不随服务自动加载。
|
||||
运行时提示词位于 `~/.aibox/prompts/`,需手工覆盖同名文件后用 `/reload_prompts` 生效:
|
||||
本目录是所有提示词的**唯一模板源**。`PromptManager` 启动时若 `~/.aibox/prompts/` 中缺少文件,
|
||||
会自动从这里复制过去;已存在的运行时文件不会被覆盖。
|
||||
|
||||
## 文件清单
|
||||
|
||||
| 文件 | 用途 | 更新方式 |
|
||||
|------|------|---------|
|
||||
| `00_IDENTITY.md` | 用户身份信息 | AI 通过 `update_identity` 更新 |
|
||||
| `1_memory_search.md` | 记忆搜索习惯 | 手动编辑 |
|
||||
| `2_memory_save.md` | 记忆保存习惯 | 手动编辑 |
|
||||
| `3_file_operations.md` | 文件操作(编排者模式) | 手动编辑 |
|
||||
| `4_web_access.md` | 网页访问(编排者模式) | 手动编辑 |
|
||||
| `5_memo_management.md` | 任务管理能力 | 手动编辑 |
|
||||
| `6_command_execution.md` | 命令执行(编排者模式) | 手动编辑 |
|
||||
| `7_extension_skills.md` | 扩展技能能力 | 手动编辑 |
|
||||
| `8_identity_management.md` | 身份信息管理 | 手动编辑 |
|
||||
| `9_language_instruction.md` | 语言说明与提问行为 | 手动编辑 |
|
||||
| `10_task_handling.md` | 任务处理规则 | 手动编辑 |
|
||||
| `11_emotion_management.md` | 情绪管理能力 | 手动编辑 |
|
||||
| `99_SOUL.md` | AI 灵魂定义(角色/性格/能力边界) | AI 通过 `update_soul` 更新 |
|
||||
| `100_STATUS.md` | AI 状态记录 | 对话结束时自动更新 |
|
||||
|
||||
## 加载顺序
|
||||
|
||||
`00_IDENTITY` → `1~11_技能文件`(按数字升序)→ `99_SOUL` → `100_STATUS`
|
||||
|
||||
## 手动更新运行时提示词
|
||||
|
||||
修改本目录模板后,需同步到运行目录并重载:
|
||||
|
||||
```bash
|
||||
cp aibox/prompts/{3_file_operations,4_web_access,6_command_execution,7_extension_skills,10_task_handling,99_SOUL}.md ~/.aibox/prompts/
|
||||
cp aibox/prompts/*.md ~/.aibox/prompts/
|
||||
```
|
||||
|
||||
> 注意:`99_SOUL.md` 可能已被 AI 通过 `update_soul` 自行更新过,覆盖前请先 diff 合并,
|
||||
> 保留用户个性化内容,仅把"执行者"相关表述改为"编排者"。
|
||||
然后在 WebSocket 会话中执行 `/reload_prompts` 生效。
|
||||
|
||||
变更要点:主对话不再持有命令行/文件/网页工具,执行类任务一律经 `create_subagent` 派生子代理。
|
||||
> 注意:`99_SOUL.md` 和 `00_IDENTITY.md` 可能已被 AI 自行更新过,覆盖前请先 diff 合并。
|
||||
+25
-26
@@ -59,14 +59,14 @@ class AESCli:
|
||||
return False
|
||||
|
||||
def generate_random_key(self, key_size=16):
|
||||
"""生成随机密钥"""
|
||||
"""生成随机密钥(S8:密钥内容不打印,避免进入会话日志;经 --save-key 持久化)"""
|
||||
if key_size not in [16, 24, 32]:
|
||||
print(f"错误: 密钥大小必须为16、24或32,当前为{key_size}")
|
||||
return None
|
||||
|
||||
self.key = get_random_bytes(key_size)
|
||||
key_hex = self.key.hex()
|
||||
print(f"已生成 {key_size} 字节随机密钥: {key_hex}")
|
||||
print(f"已生成 {key_size} 字节随机密钥(内容不显示,请用 --save-key 保存到文件)")
|
||||
return key_hex
|
||||
|
||||
def generate_random_iv(self):
|
||||
@@ -330,6 +330,10 @@ class AESCli:
|
||||
|
||||
with open(filepath, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
try:
|
||||
os.chmod(filepath, 0o600) # S8:密钥文件仅属主可读写
|
||||
except OSError:
|
||||
pass # 不支持 POSIX 权限的平台(如 Windows)跳过
|
||||
|
||||
print(f"密钥和IV已保存到: {filepath}")
|
||||
return True
|
||||
@@ -395,38 +399,33 @@ def main():
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
使用示例:
|
||||
# 生成随机密钥和IV
|
||||
aes-crypto --gen-key
|
||||
# 生成随机密钥和IV(密钥不打印,经 --save-key 保存到文件)
|
||||
aes-crypto --gen-key --gen-iv --save-key keys.json
|
||||
aes-crypto --gen-iv
|
||||
|
||||
# 文本加解密
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --encrypt-text "Hello World"
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --decrypt-text "<密文base64>"
|
||||
# 文本加解密(密钥经文件提供,不走命令行)
|
||||
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --encrypt-text "Hello World"
|
||||
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --decrypt-text "<密文base64>"
|
||||
|
||||
# 文件加解密
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --encrypt-file input.txt
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --decrypt-file input.txt.encrypted
|
||||
aes-crypto --load-key keys.json --encrypt-file input.txt
|
||||
aes-crypto --load-key keys.json --decrypt-file input.txt.encrypted
|
||||
|
||||
# 批量加解密
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --batch-encrypt file1.txt file2.txt
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --batch-decrypt *.encrypted
|
||||
|
||||
# 保存和加载密钥
|
||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --save-key keys.json
|
||||
aes-crypto --load-key keys.json --encrypt-file input.txt
|
||||
aes-crypto --load-key keys.json --batch-encrypt file1.txt file2.txt
|
||||
aes-crypto --load-key keys.json --batch-decrypt *.encrypted
|
||||
|
||||
# 使用GCM模式
|
||||
aes-crypto --mode GCM --gen-key --key-size 32
|
||||
aes-crypto --mode GCM --key <密钥hex> --iv <Nonce hex> --encrypt-file input.txt
|
||||
aes-crypto --mode GCM --gen-key --key-size 32 --save-key keys.json
|
||||
aes-crypto --mode GCM --load-key keys.json --encrypt-file input.txt
|
||||
|
||||
# 从文件读取密钥
|
||||
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --encrypt-file input.txt
|
||||
|
||||
参数说明:
|
||||
--mode CBC|GCM 加密模式 (默认: CBC)
|
||||
--key KEY 十六进制密钥 (16/24/32字节)
|
||||
--iv IV 十六进制IV/Nonce (CBC:16字节, GCM:12字节)
|
||||
--key-file FILE 从文件读取密钥
|
||||
--key-file FILE 从文件读取密钥(密钥仅经文件提供,避免进入进程列表与会话日志)
|
||||
--load-key FILE 从JSON文件加载密钥和IV
|
||||
--save-key FILE 保存密钥和IV到JSON文件
|
||||
--gen-key 生成随机密钥
|
||||
@@ -448,8 +447,7 @@ def main():
|
||||
parser.add_argument('--mode', choices=['CBC', 'GCM'], default='CBC',
|
||||
help='加密模式 (默认: CBC)')
|
||||
|
||||
# 密钥和IV设置
|
||||
parser.add_argument('--key', help='十六进制密钥 (16/24/32字节)')
|
||||
# 密钥和IV设置(S7:密钥仅经 --key-file/--load-key 文件途径提供,不进命令行)
|
||||
parser.add_argument('--iv', help='十六进制IV/Nonce (CBC:16字节, GCM:12字节)')
|
||||
parser.add_argument('--key-file', help='从文件读取密钥')
|
||||
parser.add_argument('--load-key', help='从JSON文件加载密钥和IV')
|
||||
@@ -491,27 +489,28 @@ def main():
|
||||
sys.exit(1)
|
||||
|
||||
# 处理密钥文件
|
||||
key_input = None
|
||||
if args.key_file:
|
||||
try:
|
||||
with open(args.key_file, 'r') as f:
|
||||
args.key = f.read().strip()
|
||||
key_input = f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"错误: 读取密钥文件失败: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# 设置密钥和IV
|
||||
if args.key:
|
||||
if key_input:
|
||||
if args.iv is None and not aes.iv:
|
||||
print("错误: 需要同时提供密钥和IV/Nonce")
|
||||
sys.exit(1)
|
||||
|
||||
if args.iv:
|
||||
if not aes.set_key_iv(args.key, args.iv):
|
||||
if not aes.set_key_iv(key_input, args.iv):
|
||||
sys.exit(1)
|
||||
elif aes.iv:
|
||||
# 从加载的密钥文件中获取IV
|
||||
try:
|
||||
aes.key = bytes.fromhex(args.key)
|
||||
aes.key = bytes.fromhex(key_input)
|
||||
if len(aes.key) not in [16, 24, 32]:
|
||||
print(f"错误: 密钥长度必须为16、24或32字节,当前为{len(aes.key)}字节")
|
||||
sys.exit(1)
|
||||
@@ -548,7 +547,7 @@ def main():
|
||||
if args.encrypt_text or args.decrypt_text or \
|
||||
args.encrypt_file or args.decrypt_file or \
|
||||
args.batch_encrypt or args.batch_decrypt:
|
||||
print("错误: 请先设置密钥和IV/Nonce (使用 --key 和 --iv 或 --load-key)")
|
||||
print("错误: 请先设置密钥和IV/Nonce (使用 --key-file 和 --iv 或 --load-key)")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
|
||||
@@ -18,6 +18,16 @@ def ensure_camera_dir():
|
||||
return CAMERA_DIR
|
||||
|
||||
|
||||
def validate_filename(filename):
|
||||
"""S9(审计 2026-08-20):--filename 仅允许纯文件名——
|
||||
拒绝路径分隔符/父目录引用/绝对路径,防止经 os.path.join 逃逸 CAMERA_DIR"""
|
||||
if not filename:
|
||||
return "文件名不能为空"
|
||||
if filename in (".", "..") or "/" in filename or "\\" in filename or os.path.isabs(filename):
|
||||
return f"非法文件名(不允许路径分隔符或父目录引用): {filename}"
|
||||
return None
|
||||
|
||||
|
||||
def list_cameras(max_test=10):
|
||||
"""列出可用的摄像头设备"""
|
||||
available_cameras = []
|
||||
@@ -341,6 +351,10 @@ def main():
|
||||
ensure_camera_dir()
|
||||
|
||||
if args.filename:
|
||||
filename_error = validate_filename(args.filename)
|
||||
if filename_error:
|
||||
print(f"❌ {filename_error}")
|
||||
sys.exit(1)
|
||||
save_path = os.path.join(CAMERA_DIR, args.filename)
|
||||
else:
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
"""人脸检测和识别工具 - 检测图像中的人脸并保存裁剪的人脸图像。"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import argparse
|
||||
import cv2
|
||||
@@ -15,6 +16,16 @@ except ImportError:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# S11(审计 2026-08-20):person_name 白名单——仅允许文字/数字/._-,
|
||||
# 拒绝路径分隔符与纯点名("."/".."),防止经 f"{name}.jpg" 拼接逃逸 known_faces 目录。
|
||||
# \w 在 py3 默认 UNICODE,兼容中文人名(如 "张三")。
|
||||
_PERSON_NAME_RE = re.compile(r"^(?!\.+$)[\w.-]{1,64}$")
|
||||
|
||||
|
||||
def _valid_person_name(name):
|
||||
return bool(name) and bool(_PERSON_NAME_RE.match(name))
|
||||
|
||||
|
||||
class FaceDetector:
|
||||
"""人脸检测和识别工具"""
|
||||
|
||||
@@ -198,6 +209,8 @@ class FaceDetector:
|
||||
def add_known_face(self, image_path: str, person_name: str) -> str:
|
||||
"""添加已知人脸到数据库"""
|
||||
try:
|
||||
if not _valid_person_name(person_name):
|
||||
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||
if not os.path.exists(image_path):
|
||||
return f"❌ 文件不存在: {image_path}"
|
||||
|
||||
@@ -237,6 +250,8 @@ class FaceDetector:
|
||||
def remove_known_face(self, person_name: str) -> str:
|
||||
"""从已知人脸库中删除人脸"""
|
||||
try:
|
||||
if not _valid_person_name(person_name):
|
||||
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
||||
|
||||
if not os.path.exists(filepath):
|
||||
@@ -260,6 +275,8 @@ class FaceDetector:
|
||||
def get_face_info(self, person_name: str) -> str:
|
||||
"""获取已知人脸信息"""
|
||||
try:
|
||||
if not _valid_person_name(person_name):
|
||||
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
||||
|
||||
if not os.path.exists(filepath):
|
||||
|
||||
+41
-10
@@ -3,6 +3,7 @@
|
||||
|
||||
import sys
|
||||
import os
|
||||
import re
|
||||
import argparse
|
||||
import requests
|
||||
from datetime import datetime
|
||||
@@ -35,6 +36,13 @@ MAX_CONTENT_PREVIEW = 1000
|
||||
MAX_LIST_PREVIEW = 20
|
||||
MAX_STRING_LENGTH = 200
|
||||
|
||||
# S3:分支/ref 白名单——首字符必须为字母或数字,防止以 '-' 开头的值被 git 解析为选项(选项注入)
|
||||
_REF_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/@+-]{0,127}$")
|
||||
|
||||
|
||||
def _valid_ref(name: str) -> bool:
|
||||
return bool(name) and bool(_REF_RE.match(name))
|
||||
|
||||
|
||||
# ==================== GitHub工具类 ====================
|
||||
|
||||
@@ -79,19 +87,30 @@ class GitHubTool:
|
||||
return date_str
|
||||
|
||||
def _run_git(self, repo_path: str, cmd: List[str]) -> Dict:
|
||||
env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
|
||||
if GITHUB_TOKEN:
|
||||
# 凭据仅经环境变量传递(git>=2.31 的 GIT_CONFIG_* 机制),
|
||||
# 不进命令行 argv / 远端 URL / .git/config(S1,对齐 J12 条例)
|
||||
env["GIT_CONFIG_COUNT"] = "1"
|
||||
env["GIT_CONFIG_KEY_0"] = "http.extraheader"
|
||||
env["GIT_CONFIG_VALUE_0"] = f"Authorization: token {GITHUB_TOKEN}"
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git"] + cmd, cwd=repo_path, capture_output=True,
|
||||
text=True, timeout=60, env={**os.environ, "GIT_TERMINAL_PROMPT": "0"}
|
||||
text=True, timeout=60, env=env
|
||||
)
|
||||
return {"success": result.returncode == 0, "stdout": result.stdout.strip(), "stderr": result.stderr.strip()}
|
||||
return {"success": result.returncode == 0,
|
||||
"stdout": self._scrub_token(result.stdout.strip()),
|
||||
"stderr": self._scrub_token(result.stderr.strip())}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e)}
|
||||
return {"success": False, "error": self._scrub_token(str(e))}
|
||||
|
||||
def _auth_url(self, repo_url: str) -> str:
|
||||
if GITHUB_TOKEN and repo_url.startswith("https://"):
|
||||
return repo_url.replace("https://", f"https://{GITHUB_TOKEN}@")
|
||||
return repo_url
|
||||
@staticmethod
|
||||
def _scrub_token(text: str) -> str:
|
||||
"""输出脱敏:任何路径回显的 token 一律替换(S1)"""
|
||||
if GITHUB_TOKEN and text:
|
||||
return text.replace(GITHUB_TOKEN, "***")
|
||||
return text
|
||||
|
||||
# ==================== 搜索功能 ====================
|
||||
|
||||
@@ -154,7 +173,9 @@ class GitHubTool:
|
||||
if os.path.exists(target_dir):
|
||||
return f"⚠️ 仓库已存在: {target_dir}\n💡 使用 pull 命令拉取更新"
|
||||
|
||||
clone_url = self._auth_url(f"https://github.com/{repo}.git")
|
||||
if branch and not _valid_ref(branch):
|
||||
return f"❌ 非法分支名: {branch}"
|
||||
clone_url = f"https://github.com/{repo}.git"
|
||||
cmd = ["clone"]
|
||||
if depth > 0:
|
||||
cmd.extend(["--depth", str(depth)])
|
||||
@@ -175,6 +196,8 @@ class GitHubTool:
|
||||
br = self._run_git(repo_path, ["branch", "--show-current"])
|
||||
branch = br["stdout"] if br["success"] else None
|
||||
|
||||
if branch and not _valid_ref(branch):
|
||||
return f"❌ 非法分支名: {branch}"
|
||||
result = self._run_git(repo_path, ["pull", "origin", branch] if branch else ["pull"])
|
||||
if result["success"]:
|
||||
return f"✅ 拉取成功\n {result['stdout'] or '已是最新'}"
|
||||
@@ -188,6 +211,8 @@ class GitHubTool:
|
||||
br = self._run_git(repo_path, ["branch", "--show-current"])
|
||||
branch = br["stdout"] if br["success"] else None
|
||||
|
||||
if branch and not _valid_ref(branch):
|
||||
return f"❌ 非法分支名: {branch}"
|
||||
cmd = ["push"]
|
||||
if force:
|
||||
cmd.append("--force")
|
||||
@@ -273,9 +298,15 @@ class GitHubTool:
|
||||
return "⚠️ 需要确认删除,请添加 --confirm 参数"
|
||||
if not os.path.exists(repo_path):
|
||||
return f"❌ 路径不存在: {repo_path}"
|
||||
# S2:仅允许删除 WORKSPACE 内的 git 仓库,防止任意目录递归删除
|
||||
workspace = os.path.realpath(WORKSPACE_DIR)
|
||||
real_path = os.path.realpath(repo_path)
|
||||
if not real_path.startswith(workspace + os.sep) \
|
||||
or not os.path.isdir(os.path.join(real_path, ".git")):
|
||||
return f"❌ 仅允许删除工作区({workspace})内的 git 仓库目录"
|
||||
try:
|
||||
shutil.rmtree(repo_path)
|
||||
return f"✅ 已删除: {repo_path}"
|
||||
shutil.rmtree(real_path)
|
||||
return f"✅ 已删除: {real_path}"
|
||||
except Exception as e:
|
||||
return f"❌ 删除失败: {e}"
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
一个基于python-nmap的局域网Ping扫描工具
|
||||
"""
|
||||
import argparse
|
||||
import shlex
|
||||
import sys
|
||||
import socket
|
||||
|
||||
@@ -25,7 +26,9 @@ def scan_network(network="192.168.1.0/24", timeout=3, verbose=False):
|
||||
print("-" * 50)
|
||||
|
||||
try:
|
||||
result = nm.scan(hosts=network, arguments=arguments)
|
||||
# J12(审计 2026-08-19):python-nmap 内部把 hosts 拼进 shell 命令
|
||||
# 执行——LLM/AI 可控的 network 参数需 shlex.quote,防参数/命令注入
|
||||
result = nm.scan(hosts=shlex.quote(network), arguments=arguments)
|
||||
except Exception as e:
|
||||
print(f"❌ 扫描失败: {e}")
|
||||
print("💡 提示:请确保已安装 nmap (sudo apt install nmap) 并有足够权限")
|
||||
|
||||
@@ -29,6 +29,14 @@ def get_timestamp_filename(prefix="screenshot", extension="png"):
|
||||
return f"{prefix}_{timestamp}.{extension}"
|
||||
|
||||
|
||||
def validate_filename(filename):
|
||||
"""S10(审计 2026-08-20):--filename 仅允许纯文件名——
|
||||
拒绝路径分隔符/父目录引用/绝对路径,防止经 os.path.join 逃逸保存目录"""
|
||||
if filename in (".", "..") or "/" in filename or "\\" in filename or os.path.isabs(filename):
|
||||
return f"非法文件名(不允许路径分隔符或父目录引用): {filename}"
|
||||
return None
|
||||
|
||||
|
||||
def get_system():
|
||||
"""获取当前操作系统"""
|
||||
return platform.system()
|
||||
@@ -41,6 +49,10 @@ def take_screenshot_windows(output_path, delay=0):
|
||||
if delay > 0:
|
||||
time.sleep(delay)
|
||||
|
||||
# J12(审计 2026-08-19):output_path 经 PowerShell 单引号转义——
|
||||
# 路径中的单引号翻倍(PS 字符串转义),AI/用户可控文件名无法闭合
|
||||
# 字符串注入脚本(此前 `${output_path}` 直拼,`");calc;#` 即注入)
|
||||
ps_path = str(output_path).replace("'", "''")
|
||||
try:
|
||||
ps_script = f'''
|
||||
Add-Type -AssemblyName System.Windows.Forms
|
||||
@@ -53,11 +65,11 @@ $bitmap = New-Object System.Drawing.Bitmap $bounds.Width, $bounds.Height
|
||||
$graphics = [System.Drawing.Graphics]::FromImage($bitmap)
|
||||
$graphics.CopyFromScreen($bounds.X, $bounds.Y, 0, 0, $bounds.Size)
|
||||
|
||||
$bitmap.Save("{output_path}", [System.Drawing.Imaging.ImageFormat]::Png)
|
||||
$bitmap.Save('{ps_path}', [System.Drawing.Imaging.ImageFormat]::Png)
|
||||
$graphics.Dispose()
|
||||
$bitmap.Dispose()
|
||||
|
||||
Write-Output "截图已保存: {output_path}"
|
||||
Write-Output "截图已保存: {ps_path}"
|
||||
'''
|
||||
|
||||
result = subprocess.run(
|
||||
@@ -150,6 +162,9 @@ def take_screenshot(output_dir=None, filename=None, delay=0):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
if filename:
|
||||
filename_error = validate_filename(filename)
|
||||
if filename_error:
|
||||
return False, filename_error, None
|
||||
filepath = os.path.join(save_dir, filename)
|
||||
else:
|
||||
filepath = os.path.join(save_dir, get_timestamp_filename())
|
||||
|
||||
@@ -27,9 +27,12 @@ def build_parser():
|
||||
--host SSH服务器主机名或IP地址 (必填)
|
||||
--port SSH服务器端口号,默认 22
|
||||
-u, --username SSH登录用户名 (必填)
|
||||
-p, --password SSH登录密码 (如不提供则交互式输入)
|
||||
--key SSH私钥文件路径 (用于密钥认证)
|
||||
|
||||
密码认证:
|
||||
密码经环境变量 AIBOX_SFTP_PASS 提供;未设置且未用 --key 时
|
||||
交互式输入。不走命令行参数,避免凭据进入进程列表与会话日志(J12 条例)。
|
||||
|
||||
操作参数:
|
||||
--upload LOCAL REMOTE 上传:本地文件/目录 -> 远程服务器路径
|
||||
--download REMOTE LOCAL 下载:远程服务器文件/目录 -> 本地路径
|
||||
@@ -37,19 +40,19 @@ def build_parser():
|
||||
-v, --verbose 显示详细信息
|
||||
|
||||
使用示例:
|
||||
1. 上传文件(密码认证):
|
||||
sftp-transfer --host 192.168.1.100 -u root -p mypass --upload /tmp/test.txt /root/
|
||||
1. 上传文件(密码认证,环境变量提供密码):
|
||||
AIBOX_SFTP_PASS=mypass sftp-transfer --host 192.168.1.100 -u root --upload /tmp/test.txt /root/
|
||||
|
||||
2. 下载文件(密钥认证):
|
||||
sftp-transfer --host 192.168.1.100 -u root --key ~/.ssh/id_rsa \\
|
||||
--download /root/test.txt /tmp/
|
||||
|
||||
3. 递归上传目录:
|
||||
sftp-transfer --host 192.168.1.100 -u user -p pass -r \\
|
||||
AIBOX_SFTP_PASS=pass sftp-transfer --host 192.168.1.100 -u user -r \\
|
||||
--upload ./mydir /home/user/mydir
|
||||
|
||||
4. 下载目录:
|
||||
sftp-transfer --host 192.168.1.100 -u user -p pass -r \\
|
||||
AIBOX_SFTP_PASS=pass sftp-transfer --host 192.168.1.100 -u user -r \\
|
||||
--download /home/user/backup ./backup
|
||||
|
||||
5. 交互式输入密码:
|
||||
@@ -61,7 +64,6 @@ def build_parser():
|
||||
parser.add_argument('--host', required=True, help='SSH服务器主机名或IP地址')
|
||||
parser.add_argument('--port', type=int, default=22, help='SSH服务器端口号,默认22')
|
||||
parser.add_argument('-u', '--username', required=True, help='SSH登录用户名')
|
||||
parser.add_argument('-p', '--password', help='SSH登录密码(如不提供则交互式输入)')
|
||||
parser.add_argument('--key', help='SSH私钥文件路径(用于密钥认证)')
|
||||
|
||||
# 操作参数
|
||||
@@ -81,10 +83,20 @@ def log(verbose, msg):
|
||||
print(f"[SFTP] {msg}")
|
||||
|
||||
|
||||
def safe_member_name(name):
|
||||
"""S6(审计 2026-08-20):目录成员名校验——拒绝路径分隔符与 . / ..,
|
||||
防止恶意 SFTP 服务端构造 listdir 条目(如 ../../etc/cron.d/x)在递归下载时逃逸本地目录"""
|
||||
return name not in ('', '.', '..') and '/' not in name and '\\' not in name
|
||||
|
||||
|
||||
def get_password(args):
|
||||
"""获取密码:优先从命令行参数获取,否则交互式输入"""
|
||||
if args.password:
|
||||
return args.password
|
||||
"""获取密码:环境变量 AIBOX_SFTP_PASS 优先,否则交互式输入(S5:凭据不走命令行)"""
|
||||
env_pass = os.environ.get("AIBOX_SFTP_PASS")
|
||||
if env_pass:
|
||||
return env_pass
|
||||
if args.key:
|
||||
# 密钥认证时不强制索要密码(仅在私钥有口令时经交互输入)
|
||||
return None
|
||||
return getpass.getpass(f"请输入 {args.username}@{args.host} 的密码: ")
|
||||
|
||||
|
||||
@@ -143,6 +155,9 @@ def upload_dir(sftp, local_dir, remote_dir, verbose):
|
||||
"""递归上传目录"""
|
||||
ensure_remote_dir(sftp, remote_dir, verbose)
|
||||
for item in os.listdir(local_dir):
|
||||
if not safe_member_name(item):
|
||||
print(f" 警告: 跳过非法本地条目名: {item!r}", file=sys.stderr)
|
||||
continue
|
||||
local_item = os.path.join(local_dir, item)
|
||||
remote_item = os.path.join(remote_dir, item)
|
||||
if os.path.isdir(local_item):
|
||||
@@ -168,6 +183,9 @@ def download_dir(sftp, remote_dir, local_dir, verbose):
|
||||
log(verbose, f" 创建本地目录: {local_dir}")
|
||||
|
||||
for item in sftp.listdir(remote_dir):
|
||||
if not safe_member_name(item):
|
||||
print(f" 警告: 跳过非法远端条目名: {item!r}", file=sys.stderr)
|
||||
continue
|
||||
remote_item = os.path.join(remote_dir, item).replace('\\', '/')
|
||||
local_item = os.path.join(local_dir, item)
|
||||
try:
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
"""ssh-exec - 通过SSH远程执行单条命令的工具(支持密码认证)"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import paramiko
|
||||
|
||||
@@ -14,24 +15,32 @@ def main():
|
||||
-H, --host SSH主机地址(默认: 127.0.0.1)
|
||||
-P, --port SSH端口号(默认: 22)
|
||||
-u, --user SSH用户名(必填)
|
||||
-p, --password SSH密码(必填)
|
||||
-c, --command 要执行的远程命令(必填)
|
||||
|
||||
密码认证:
|
||||
密码经环境变量 AIBOX_SSH_PASS 提供(必填),不走命令行参数,
|
||||
避免凭据进入进程列表与会话日志(J12 条例)。
|
||||
|
||||
使用示例:
|
||||
%(prog)s -u qing -p 12345678 -c "ls -la"
|
||||
%(prog)s -u root -H 192.168.1.100 -P 22 -p mypass -c "df -h"
|
||||
%(prog)s -u qing -H 127.0.0.1 -p 12345678 -c "uname -a"
|
||||
AIBOX_SSH_PASS=12345678 %(prog)s -u qing -c "ls -la"
|
||||
AIBOX_SSH_PASS=mypass %(prog)s -u root -H 192.168.1.100 -P 22 -c "df -h"
|
||||
AIBOX_SSH_PASS=12345678 %(prog)s -u qing -H 127.0.0.1 -c "uname -a"
|
||||
""",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument("-H", "--host", default="127.0.0.1", help="SSH主机地址(默认: 127.0.0.1)")
|
||||
parser.add_argument("-P", "--port", type=int, default=22, help="SSH端口号(默认: 22)")
|
||||
parser.add_argument("-u", "--user", required=True, help="SSH用户名(必填)")
|
||||
parser.add_argument("-p", "--password", required=True, help="SSH密码(必填)")
|
||||
parser.add_argument("-c", "--command", required=True, help="要执行的远程命令(必填)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 密码仅从环境变量读取(S4:凭据不进命令行)
|
||||
password = os.environ.get("AIBOX_SSH_PASS")
|
||||
if not password:
|
||||
print("错误: 请通过环境变量 AIBOX_SSH_PASS 提供SSH密码", file=sys.stderr)
|
||||
sys.exit(2)
|
||||
|
||||
client = paramiko.SSHClient()
|
||||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||
|
||||
@@ -40,7 +49,7 @@ def main():
|
||||
hostname=args.host,
|
||||
port=args.port,
|
||||
username=args.user,
|
||||
password=args.password,
|
||||
password=password,
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
|
||||
@@ -51,7 +51,14 @@ create_apprun() {
|
||||
local app_dir="$1"
|
||||
local exec_name="$2"
|
||||
local terminal="$3"
|
||||
|
||||
|
||||
# S13(审计 2026-08-20):exec_name 进入 sed 替换串,含 / & 等字符会注入或破坏
|
||||
# 替换表达式(如 "a/x" 使 s///g 错位、"&" 展开为匹配串)。白名单拦截,拒绝即退出。
|
||||
if [[ ! "$exec_name" =~ ^[A-Za-z0-9._+-]+$ ]]; then
|
||||
echo -e "${RED}错误: 可执行文件名含非法字符(仅允许 A-Za-z0-9._+-): $exec_name${NC}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
cat > "$app_dir/AppRun" << 'EOF'
|
||||
#!/bin/bash
|
||||
HERE="$(dirname "$(readlink -f "$0")")"
|
||||
|
||||
+27
-2
@@ -23,6 +23,27 @@ def get_api_key():
|
||||
return None
|
||||
|
||||
|
||||
def detect_image_mime(image_path):
|
||||
"""S12(审计 2026-08-20):按魔数识别真实图片类型——
|
||||
拒绝非图片文件(防任意文件读取外泄),并返回正确 MIME(此前硬编码 image/jpeg)"""
|
||||
try:
|
||||
with open(image_path, 'rb') as f:
|
||||
head = f.read(16)
|
||||
except Exception:
|
||||
return None
|
||||
if head.startswith(b'\xff\xd8\xff'):
|
||||
return 'image/jpeg'
|
||||
if head.startswith(b'\x89PNG\r\n\x1a\n'):
|
||||
return 'image/png'
|
||||
if head[:6] in (b'GIF87a', b'GIF89a'):
|
||||
return 'image/gif'
|
||||
if head.startswith(b'RIFF') and head[8:12] == b'WEBP':
|
||||
return 'image/webp'
|
||||
if head.startswith(b'BM'):
|
||||
return 'image/bmp'
|
||||
return None
|
||||
|
||||
|
||||
def encode_image_to_base64(image_path):
|
||||
"""将图片编码为base64"""
|
||||
with open(image_path, 'rb') as f:
|
||||
@@ -37,7 +58,11 @@ def analyze_image(image_path, prompt, stream=False):
|
||||
|
||||
if not os.path.exists(image_path):
|
||||
return f"❌ 错误:图片文件不存在 - {image_path}"
|
||||
|
||||
|
||||
mime = detect_image_mime(image_path)
|
||||
if not mime:
|
||||
return f"❌ 错误:文件不是有效图片(仅支持 JPEG/PNG/GIF/WebP/BMP) - {image_path}"
|
||||
|
||||
try:
|
||||
base64_image = encode_image_to_base64(image_path)
|
||||
except Exception as e:
|
||||
@@ -55,7 +80,7 @@ def analyze_image(image_path, prompt, stream=False):
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{base64_image}"
|
||||
"url": f"data:{mime};base64,{base64_image}"
|
||||
}
|
||||
},
|
||||
{
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
subsystems - aibox 子系统模块(包)
|
||||
|
||||
抽离自 aibox.py 的基础设施子系统,按职责分文件。
|
||||
当前包含(Wave 1):
|
||||
- executors.py : CommandExecutor + SubprocessCommandExecutor
|
||||
- prompt_manager.py : PromptManager(单例)
|
||||
- db.py : MemoryDatabase + MemoDatabase
|
||||
- skill_manager.py : SkillManager
|
||||
Wave 2:
|
||||
- conversation_history.py : ConversationHistory(对话历史管理)
|
||||
- conversation_compressor.py : ConversationCompressor(对话压缩,注入 llm_client)
|
||||
|
||||
依赖方向:common <- subsystems <- aibox
|
||||
不反向依赖 aibox 任何符号。
|
||||
|
||||
采用"重导出"策略:外部统一用 `from subsystems import MemoryDatabase`,
|
||||
调用点无需关心具体子模块所在文件。
|
||||
"""
|
||||
|
||||
from .executors import CommandExecutor, SubprocessCommandExecutor
|
||||
from .prompt_manager import PromptManager
|
||||
from .db import MemoryDatabase, MemoDatabase
|
||||
from .skill_manager import SkillManager
|
||||
from .conversation_history import ConversationHistory
|
||||
from .conversation_compressor import ConversationCompressor
|
||||
from .process_manager import (
|
||||
ProcessManager,
|
||||
get_process_manager,
|
||||
reset_process_manager,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CommandExecutor",
|
||||
"SubprocessCommandExecutor",
|
||||
"PromptManager",
|
||||
"MemoryDatabase",
|
||||
"MemoDatabase",
|
||||
"SkillManager",
|
||||
"ConversationHistory",
|
||||
"ConversationCompressor",
|
||||
"ProcessManager",
|
||||
"get_process_manager",
|
||||
"reset_process_manager",
|
||||
]
|
||||
@@ -0,0 +1,513 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.conversation_compressor - 对话压缩子系统
|
||||
|
||||
抽离自 aibox.py,负责对话上下文压缩:消息条数或 Token 数达到阈值时,
|
||||
将早期对话(含工具调用信息)经 LLM 压缩为摘要,保留最近消息。
|
||||
|
||||
依赖注入:构造时传入 llm_client(llm_client.LLMClient 实例)而非 AI
|
||||
引擎实例,避免与 DeepSeekChat 形成循环引用;依赖方向 common <- subsystems。
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import constants, i18n, Logger, ConfigManager
|
||||
|
||||
|
||||
class ConversationCompressor:
|
||||
"""对话压缩管理器 - 优先压缩较早的消息,将tool调用信息一并压缩"""
|
||||
|
||||
def __init__(self, llm_client=None, logger: Logger = None):
|
||||
self.llm_client = llm_client
|
||||
self.logger = logger or Logger("Compressor")
|
||||
self.config = ConfigManager()
|
||||
# 从配置获取最小值
|
||||
self.min_messages_to_compress = constants.COMPRESSOR_MIN_MESSAGES
|
||||
# 权重衰减因子:从配置获取,设为0.6使加权计算更合理
|
||||
self.weight_decay_factor = 0.6 # 使用更高的衰减因子
|
||||
|
||||
def estimate_tokens(self, messages: List[Dict]) -> int:
|
||||
"""估算消息的token数"""
|
||||
total_tokens = 0
|
||||
for msg in messages:
|
||||
content = msg.get("content", "")
|
||||
if content:
|
||||
chinese_chars = len([c for c in content if '\u4e00' <= c <= '\u9fff'])
|
||||
english_words = len(re.findall(r'[a-zA-Z]+', content))
|
||||
other_chars = len(content) - chinese_chars - english_words
|
||||
|
||||
estimated = chinese_chars * 1.5 + english_words * 1.3 + other_chars * 0.8
|
||||
total_tokens += estimated
|
||||
total_tokens += 10
|
||||
|
||||
if msg.get("tool_calls"):
|
||||
# 工具调用也计算token
|
||||
for tc in msg.get("tool_calls", []):
|
||||
func = tc.get("function", {})
|
||||
if func.get("name"):
|
||||
total_tokens += len(func.get("name", "")) * 1.5
|
||||
if func.get("arguments"):
|
||||
total_tokens += len(func.get("arguments", "")) * 0.8
|
||||
total_tokens += 50
|
||||
|
||||
return int(total_tokens)
|
||||
|
||||
def should_compress(self, messages: List[Dict]) -> Tuple[bool, str]:
|
||||
"""检查是否需要压缩对话 - 消息条数或 Token 数任一达到阈值即压缩"""
|
||||
if not constants.COMPRESS_ENABLED:
|
||||
return False, "压缩功能已禁用"
|
||||
|
||||
conversation_messages = [m for m in messages if m.get("role") != "system"]
|
||||
|
||||
# 优先检查消息条数是否达到阈值
|
||||
message_count = len(conversation_messages)
|
||||
if message_count >= constants.COMPRESS_MESSAGE_THRESHOLD:
|
||||
print(f"⚠️ {i18n.get('compress_threshold_reached', message_count)}")
|
||||
return True, f"消息数({message_count})超过阈值({constants.COMPRESS_MESSAGE_THRESHOLD})"
|
||||
|
||||
# 再检查 Token 数是否达到阈值(长文档/长工具输出场景)
|
||||
token_count = self.estimate_tokens(conversation_messages)
|
||||
if token_count >= constants.COMPRESS_TOKEN_THRESHOLD:
|
||||
print(f"⚠️ {i18n.get('compress_token_threshold_reached', token_count)}")
|
||||
return True, f"Token数({token_count})超过阈值({constants.COMPRESS_TOKEN_THRESHOLD})"
|
||||
|
||||
return False, f"消息数({message_count})与Token数({token_count})均未达到阈值"
|
||||
|
||||
def _calculate_target_count(self, current_count: int, threshold: int, keep_ratio: float = 0.7) -> int:
|
||||
"""
|
||||
计算目标消息数量 - 压缩到阈值的 keep_ratio 倍
|
||||
|
||||
Args:
|
||||
current_count: 当前消息数量
|
||||
threshold: 压缩阈值
|
||||
keep_ratio: 压缩后保留的比例(相对于阈值),默认0.7表示压缩到阈值的70%
|
||||
|
||||
Returns:
|
||||
目标消息数量
|
||||
"""
|
||||
# 压缩到阈值以下,保留阈值的 keep_ratio 倍
|
||||
target = int(threshold * keep_ratio)
|
||||
|
||||
# 确保至少有 min_messages_to_compress 条消息
|
||||
target = max(target, self.min_messages_to_compress)
|
||||
|
||||
# 确保不超过当前数量
|
||||
target = min(target, current_count - 1)
|
||||
|
||||
return target
|
||||
|
||||
def _find_compressible_block_by_count(self, messages: List[Dict], target_count: int) -> Tuple[int, int]:
|
||||
"""
|
||||
根据目标消息数量找到可压缩块 - 不保护tool消息,直接压缩早期消息
|
||||
|
||||
Args:
|
||||
messages: 消息列表
|
||||
target_count: 目标消息数量
|
||||
|
||||
Returns:
|
||||
(start, end) 压缩块的起始和结束索引
|
||||
"""
|
||||
if len(messages) <= target_count:
|
||||
return 0, 0
|
||||
|
||||
# 计算需要压缩的消息数量
|
||||
messages_to_remove = len(messages) - target_count
|
||||
|
||||
# 确定要保留的最近消息数量
|
||||
keep_recent = constants.COMPRESS_KEEP_RECENT
|
||||
if keep_recent > 0:
|
||||
keep_recent = max(keep_recent, 10) # 至少保留最近10条
|
||||
|
||||
# 压缩块应该从最早的消息开始,但需要保留前几轮对话
|
||||
# 计算最早可以开始压缩的位置(跳过前2轮对话)
|
||||
start_min = 0
|
||||
|
||||
# 尝试找到前2个对话轮次结束的位置
|
||||
conversation_turns = self._build_conversation_turns(messages)
|
||||
if len(conversation_turns) >= 2:
|
||||
# 前2轮结束的位置
|
||||
first_turn = conversation_turns[0]
|
||||
second_turn = conversation_turns[1]
|
||||
# 取第二轮结束的位置
|
||||
if second_turn["assistant_index"] is not None:
|
||||
start_min = second_turn["assistant_index"] + 1
|
||||
else:
|
||||
start_min = second_turn["user_index"] + 1
|
||||
|
||||
# 压缩结束的位置 = 总消息数 - 保留的最近消息数
|
||||
end_max = len(messages) - keep_recent
|
||||
|
||||
# 确保压缩块有效
|
||||
if end_max <= start_min + 2:
|
||||
# 如果保护区域太大,尝试压缩更多
|
||||
end_max = len(messages) - max(keep_recent // 2, 5)
|
||||
|
||||
# 计算需要压缩的块大小
|
||||
block_size = min(messages_to_remove + 10, end_max - start_min)
|
||||
block_size = max(block_size, 10) # 至少压缩10条
|
||||
|
||||
# 压缩块从 start_min 开始
|
||||
start = start_min
|
||||
end = min(start + block_size, end_max)
|
||||
|
||||
# 确保压缩块至少2条消息
|
||||
if end - start < 2:
|
||||
# 如果块太小,尝试从更早的位置开始
|
||||
start = max(0, end_max - block_size)
|
||||
end = end_max
|
||||
|
||||
self.logger.debug(f"找到压缩块: [{start}:{end}], 移除{end-start}条消息")
|
||||
return start, end
|
||||
|
||||
def _build_conversation_turns(self, messages: List[Dict]) -> List[Dict]:
|
||||
"""构建对话轮次"""
|
||||
conversation_turns = []
|
||||
|
||||
i = 0
|
||||
while i < len(messages):
|
||||
msg = messages[i]
|
||||
role = msg.get("role")
|
||||
|
||||
# 跳过系统消息
|
||||
if role == "system":
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# 遇到用户消息,开始一个新的轮次
|
||||
if role == "user":
|
||||
turn = {"user_index": i, "assistant_index": None, "tool_indices": []}
|
||||
|
||||
# 向后查找对应的助手消息和工具消息
|
||||
j = i + 1
|
||||
while j < len(messages):
|
||||
next_msg = messages[j]
|
||||
next_role = next_msg.get("role")
|
||||
|
||||
# 跳过系统消息
|
||||
if next_role == "system":
|
||||
j += 1
|
||||
continue
|
||||
|
||||
# 找到助手消息
|
||||
if next_role == "assistant":
|
||||
turn["assistant_index"] = j
|
||||
# 检查这个助手消息是否有 tool_calls
|
||||
if next_msg.get("tool_calls"):
|
||||
# 收集这个助手消息的 tool_call_ids
|
||||
tool_call_ids = {tc.get("id") for tc in next_msg.get("tool_calls", []) if tc.get("id")}
|
||||
# 继续查找对应的 tool 消息
|
||||
k = j + 1
|
||||
while k < len(messages) and len(turn["tool_indices"]) < len(tool_call_ids):
|
||||
tool_msg = messages[k]
|
||||
if tool_msg.get("role") == "tool":
|
||||
tool_call_id = tool_msg.get("tool_call_id")
|
||||
if tool_call_id and tool_call_id in tool_call_ids:
|
||||
turn["tool_indices"].append(k)
|
||||
k += 1
|
||||
elif tool_msg.get("role") == "user":
|
||||
break
|
||||
else:
|
||||
k += 1
|
||||
break
|
||||
|
||||
if next_role == "user":
|
||||
break
|
||||
|
||||
j += 1
|
||||
|
||||
conversation_turns.append(turn)
|
||||
i += 1
|
||||
else:
|
||||
i += 1
|
||||
|
||||
return conversation_turns
|
||||
|
||||
def _build_compression_prompt(self, messages: List[Dict]) -> str:
|
||||
"""构建压缩提示词 - 包含工具调用信息"""
|
||||
conversation_lines = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "user":
|
||||
if content:
|
||||
conversation_lines.append(f"用户: {content}")
|
||||
|
||||
elif role == "assistant":
|
||||
if content:
|
||||
conversation_lines.append(f"助手: {content}")
|
||||
|
||||
# 包含工具调用信息
|
||||
tool_calls = msg.get("tool_calls", [])
|
||||
for tc in tool_calls:
|
||||
func = tc.get("function", {})
|
||||
tool_name = func.get("name", "")
|
||||
arguments = func.get("arguments", "")
|
||||
if tool_name:
|
||||
if arguments:
|
||||
conversation_lines.append(f"[调用工具: {tool_name}({arguments[:100]})]")
|
||||
else:
|
||||
conversation_lines.append(f"[调用工具: {tool_name}]")
|
||||
|
||||
elif role == "tool":
|
||||
if content:
|
||||
useful = self._extract_useful_content(content)
|
||||
if useful:
|
||||
conversation_lines.append(f"[工具结果: {useful}]")
|
||||
|
||||
conversation_text = "\n".join(conversation_lines)
|
||||
|
||||
max_input_chars = constants.COMPRESSOR_MAX_INPUT_CHARS
|
||||
if len(conversation_text) > max_input_chars:
|
||||
conversation_text = conversation_text[:max_input_chars] + "\n...(内容过长,已截断)..."
|
||||
|
||||
prompt = f"""请对以下对话内容进行智能压缩和总结。要求:
|
||||
|
||||
1. **提取核心信息**:对话的主要主题、关键问题、重要结论、用户需求
|
||||
2. **保留重要细节**:用户的重要要求、AI的重要回复、达成的共识
|
||||
3. **保留工具调用摘要**:简要说明使用了什么工具,得到了什么关键结果
|
||||
4. **结构化输出**:使用清晰的段落结构
|
||||
|
||||
对话内容:
|
||||
---
|
||||
{conversation_text}
|
||||
---
|
||||
|
||||
请输出压缩后的对话摘要:"""
|
||||
|
||||
return prompt
|
||||
|
||||
def _extract_useful_content(self, content: str) -> str:
|
||||
"""从工具返回内容中提取有用信息"""
|
||||
if not content:
|
||||
return ""
|
||||
|
||||
if "✅" in content:
|
||||
lines = content.split("\n")
|
||||
for line in lines[:3]:
|
||||
if "✅" in line:
|
||||
return line.strip()
|
||||
|
||||
if "❌" in content:
|
||||
lines = content.split("\n")
|
||||
for line in lines[:2]:
|
||||
if "❌" in line:
|
||||
return line.strip()
|
||||
|
||||
if len(content) > 100:
|
||||
return content[:100] + "..."
|
||||
|
||||
return content
|
||||
|
||||
def _compress_via_llm(self, messages: List[Dict]) -> Optional[str]:
|
||||
"""通过语言模型压缩消息(经由注入的客户端抽象层调用,支持多厂商)"""
|
||||
if not messages or not self.llm_client:
|
||||
return None
|
||||
|
||||
prompt = self._build_compression_prompt(messages)
|
||||
|
||||
print(f"📦 {i18n.get('conversation_compress')}: 正在压缩 {len(messages)} 条消息...")
|
||||
|
||||
try:
|
||||
summary = self.llm_client.chat_completion(
|
||||
[
|
||||
{"role": "system", "content": "你是一个专业的对话摘要助手,擅长提取对话核心信息并进行结构化总结。请保留工具调用的关键结果。"},
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=0.3,
|
||||
max_tokens=constants.COMPRESSOR_LLM_MAX_TOKENS
|
||||
)
|
||||
return summary
|
||||
except Exception as e:
|
||||
self.logger.error(f"语言模型压缩异常: {e}")
|
||||
return None
|
||||
|
||||
def _validate_message_structure(self, messages: List[Dict]) -> List[Dict]:
|
||||
"""验证并修复消息结构(tool 配对双向完整性)。
|
||||
|
||||
OpenAI 协议要求 assistant.tool_calls 与对应的 tool 结果消息成对出现,
|
||||
缺任一侧后续 API 调用直接 400:
|
||||
1. 孤立 tool 消息(tool_call_id 无对应 assistant)→ 移除(既有行为)
|
||||
2. T23(审计 2026-08-21):压缩块边界按条数切,可能落在带 tool_calls
|
||||
的 assistant 消息与其 tool 结果之间——保留侧 assistant 的
|
||||
tool_calls 失去配对 → 剥离失配的 tool_calls 条目;全部失配时删除
|
||||
字段(content 为空则置空串,防产生非法空消息)
|
||||
"""
|
||||
if not messages:
|
||||
return messages
|
||||
|
||||
# tool 侧实际存在的 id
|
||||
present_tool_ids = {
|
||||
m.get("tool_call_id")
|
||||
for m in messages
|
||||
if m.get("role") == "tool" and m.get("tool_call_id")
|
||||
}
|
||||
|
||||
# 第一遍:assistant 剥离失去配对的 tool_calls
|
||||
first_pass = []
|
||||
kept_assistant_ids = set()
|
||||
for msg in messages:
|
||||
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
||||
calls = msg.get("tool_calls", [])
|
||||
kept_calls = [tc for tc in calls
|
||||
if tc.get("id") and tc.get("id") in present_tool_ids]
|
||||
kept_assistant_ids.update(tc.get("id") for tc in kept_calls)
|
||||
if len(kept_calls) == len(calls):
|
||||
first_pass.append(msg) # 全部有配对,原样保留
|
||||
continue
|
||||
new_msg = dict(msg)
|
||||
if kept_calls:
|
||||
new_msg["tool_calls"] = kept_calls
|
||||
print(f" ⚠️ 剥离部分失去配对 tool 结果的 tool_calls(T23)")
|
||||
else:
|
||||
new_msg.pop("tool_calls", None)
|
||||
if not new_msg.get("content"):
|
||||
new_msg["content"] = ""
|
||||
print(f" ⚠️ 剥离失去配对 tool 结果的 assistant.tool_calls(T23)")
|
||||
first_pass.append(new_msg)
|
||||
else:
|
||||
first_pass.append(msg)
|
||||
|
||||
# 第二遍:移除孤立 tool 消息(其 id 不在任何保留的 assistant tool_calls 中)
|
||||
filtered = []
|
||||
for msg in first_pass:
|
||||
if msg.get("role") == "tool":
|
||||
tool_call_id = msg.get("tool_call_id")
|
||||
if tool_call_id and tool_call_id not in kept_assistant_ids:
|
||||
print(f" ⚠️ 移除孤立的 tool 消息: {tool_call_id}")
|
||||
continue
|
||||
filtered.append(msg)
|
||||
|
||||
return filtered
|
||||
|
||||
def compress_conversation(self, messages: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
压缩对话内容 - 优先压缩较早的消息,包含tool调用信息一并压缩
|
||||
当消息条数或 Token 数超过阈值时,压缩到阈值以下
|
||||
"""
|
||||
if not messages:
|
||||
return messages
|
||||
|
||||
# 分离系统消息和对话消息
|
||||
system_messages = []
|
||||
conversation_messages = []
|
||||
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
system_messages.append(msg)
|
||||
else:
|
||||
conversation_messages.append(msg)
|
||||
|
||||
if not conversation_messages:
|
||||
return system_messages
|
||||
|
||||
original_count = len(conversation_messages)
|
||||
original_tokens = self.estimate_tokens(conversation_messages)
|
||||
threshold = constants.COMPRESS_MESSAGE_THRESHOLD
|
||||
|
||||
print(f"\n📊 对话压缩检查:")
|
||||
print(f" 当前消息数: {original_count} 条 (阈值: {threshold} 条)")
|
||||
print(f" 当前Token数: {original_tokens} 个")
|
||||
|
||||
# 计算目标消息数量 - 压缩到阈值的70%
|
||||
target_count = self._calculate_target_count(original_count, threshold, keep_ratio=0.7)
|
||||
|
||||
if original_count < threshold:
|
||||
# 消息数未超阈值(由 Token 阈值触发):按比例压缩早期消息
|
||||
target_count = max(self.min_messages_to_compress, original_count // 2)
|
||||
print(f" Token 触发压缩:目标消息数改为 {target_count} 条 (当前的一半)")
|
||||
else:
|
||||
print(f" 目标消息数: {target_count} 条 (压缩后应少于阈值)")
|
||||
|
||||
# 如果目标数量大于等于当前数量,不需要压缩
|
||||
if target_count >= original_count:
|
||||
print(f" 目标消息数大于等于当前,跳过压缩")
|
||||
return messages
|
||||
|
||||
# 计算需要压缩的消息数量
|
||||
messages_to_compress = original_count - target_count
|
||||
print(f" 需要压缩的消息数量: {messages_to_compress} 条")
|
||||
|
||||
# 根据消息数量查找可压缩块
|
||||
compress_start, compress_end = self._find_compressible_block_by_count(
|
||||
conversation_messages, target_count
|
||||
)
|
||||
|
||||
if compress_end - compress_start >= 2:
|
||||
compressible = conversation_messages[compress_start:compress_end]
|
||||
compress_tokens = self.estimate_tokens(compressible)
|
||||
compress_count = len(compressible)
|
||||
|
||||
print(f" 压缩块范围: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
||||
print(f" 压缩块token数: {compress_tokens}")
|
||||
|
||||
# 如果压缩块大小仍然不够,尝试扩大压缩块
|
||||
if compress_count < messages_to_compress:
|
||||
print(f" 压缩块大小不足({compress_count} < {messages_to_compress}),扩大压缩块...")
|
||||
# 扩大压缩块
|
||||
additional = messages_to_compress - compress_count + 5
|
||||
new_start = max(0, compress_start - additional // 2)
|
||||
new_end = min(len(conversation_messages), compress_end + additional // 2)
|
||||
compress_start, compress_end = new_start, new_end
|
||||
compressible = conversation_messages[compress_start:compress_end]
|
||||
compress_count = len(compressible)
|
||||
print(f" 扩大后压缩块: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
||||
else:
|
||||
print(f" 未找到合适的压缩块,使用默认压缩策略")
|
||||
# 使用默认压缩策略:压缩最早的一半消息,保留最近的消息
|
||||
keep_recent = constants.COMPRESS_KEEP_RECENT
|
||||
compress_end = len(conversation_messages) - keep_recent
|
||||
compress_start = max(0, compress_end - messages_to_compress - 5)
|
||||
compress_start = max(compress_start, 0)
|
||||
compressible = conversation_messages[compress_start:compress_end]
|
||||
compress_count = len(compressible)
|
||||
print(f" 默认压缩块: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
||||
|
||||
# 确保压缩块有足够多的消息
|
||||
if compress_count < 5:
|
||||
print(f" 压缩块太小({compress_count}条),跳过压缩")
|
||||
return messages
|
||||
|
||||
# 调用语言模型压缩
|
||||
summary = self._compress_via_llm(compressible)
|
||||
|
||||
if not summary:
|
||||
print(f" ❌ 压缩失败,保留原始对话")
|
||||
return messages
|
||||
|
||||
# 构建压缩后的消息列表
|
||||
compressed_result = []
|
||||
|
||||
# 添加压缩块之前的消息
|
||||
compressed_result.extend(conversation_messages[:compress_start])
|
||||
|
||||
# 插入压缩摘要(作为系统消息)
|
||||
compressed_result.append({
|
||||
"role": "system",
|
||||
"content": i18n.get('compressed_summary', summary)
|
||||
})
|
||||
|
||||
# 添加压缩块之后的消息(包括最近的消息)
|
||||
compressed_result.extend(conversation_messages[compress_end:])
|
||||
|
||||
# 验证并修复消息结构
|
||||
compressed_result = self._validate_message_structure(compressed_result)
|
||||
|
||||
# 合并系统消息
|
||||
final_messages = system_messages + compressed_result
|
||||
|
||||
# 统计结果
|
||||
final_count = len([m for m in final_messages if m.get("role") != "system"])
|
||||
final_tokens = self.estimate_tokens([m for m in final_messages if m.get("role") != "system"])
|
||||
|
||||
print(f"\n✅ {i18n.get('conversation_compressed', original_count, final_count)}")
|
||||
# L6:压缩块全空时 original_tokens 可能为 0,除零防护
|
||||
print(f" Token: {original_tokens} -> {final_tokens} (压缩率: {final_tokens / max(original_tokens, 1):.1%})")
|
||||
|
||||
# 检查是否压缩到阈值以下
|
||||
if final_count >= threshold:
|
||||
print(f" ⚠️ 警告: 压缩后消息数({final_count})仍大于等于阈值({threshold})")
|
||||
|
||||
return final_messages
|
||||
@@ -0,0 +1,370 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.conversation_history - 对话历史管理器
|
||||
|
||||
抽离自 aibox.py,管理当前对话与归档历史,提供上下文加载能力。
|
||||
|
||||
职责:
|
||||
- 维护当前对话(current_conversation.json)
|
||||
- 归档历史对话(archive/*.json)+ 历史索引(conversation_history.json)
|
||||
- 加载历史消息作为 AI 上下文
|
||||
- 归档/结束时可选写入 MemoryDatabase
|
||||
|
||||
依赖:common (constants, ConfigManager, Logger) + subsystems.db (MemoryDatabase) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||
|
||||
from common import ConfigManager, Logger, constants
|
||||
from common.atomic import atomic_write_text
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# 仅用于类型注解,避免运行时循环导入;MemoryDatabase 来自同包 subsystems.db
|
||||
from subsystems.db import MemoryDatabase
|
||||
|
||||
|
||||
class ConversationHistory:
|
||||
"""对话历史管理器 - 管理所有历史对话"""
|
||||
|
||||
def __init__(self, save_dir: str = None, memory_db: "MemoryDatabase" = None, load_history: bool = None):
|
||||
if save_dir is None:
|
||||
save_dir = constants.SAVE_DIR
|
||||
|
||||
self.save_dir = Path(save_dir)
|
||||
self.archive_dir = self.save_dir / constants.ARCHIVE_DIR
|
||||
self.history_file = self.save_dir / constants.HISTORY_FILE
|
||||
self.memory_db = memory_db
|
||||
|
||||
config = ConfigManager()
|
||||
if load_history is None:
|
||||
load_history = config.get("system.load_history_by_default", True)
|
||||
self.load_history = load_history
|
||||
|
||||
# 确保目录存在
|
||||
self.save_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.archive_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.logger = Logger("HistoryManager")
|
||||
self.history_index = self._load_history_index()
|
||||
self.current_conversation = None
|
||||
self.current_file = self.save_dir / constants.CURRENT_CONVERSATION_FILE
|
||||
|
||||
self._load_current_conversation()
|
||||
|
||||
def _load_history_index(self) -> Dict:
|
||||
"""加载历史索引文件"""
|
||||
if self.history_file.exists():
|
||||
try:
|
||||
with open(self.history_file, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
# N16:收窄为 Exception——裸 except 连 KeyboardInterrupt/SystemExit
|
||||
# 也吞掉;损坏时重置为空索引并告警
|
||||
self.logger.warning(f"历史索引文件损坏,已重置: {e}")
|
||||
return {"conversations": [], "last_id": 0}
|
||||
return {"conversations": [], "last_id": 0}
|
||||
|
||||
@staticmethod
|
||||
def _atomic_write_json(path: Path, data) -> None:
|
||||
"""原子写 JSON(I 批,审计 2026-08-18):先写同目录临时文件再
|
||||
os.replace——此前 'w' 直写,进程中断/崩溃留半截 JSON(current 损坏
|
||||
后仅静默置 None,历史内容丢失)。J5(2026-08-19):委托 common.atomic
|
||||
(mkstemp 随机名),修复旧实现固定 .tmp 名在多进程并发写同路径互踩。"""
|
||||
atomic_write_text(str(path), json.dumps(data, ensure_ascii=False, indent=2))
|
||||
|
||||
def _save_history_index(self):
|
||||
"""保存历史索引"""
|
||||
self._atomic_write_json(self.history_file, self.history_index)
|
||||
|
||||
def _load_current_conversation(self):
|
||||
"""加载当前的对话(如果有)"""
|
||||
if self.current_file.exists():
|
||||
try:
|
||||
with open(self.current_file, 'r', encoding='utf-8') as f:
|
||||
self.current_conversation = json.load(f)
|
||||
self.logger.info(f"已加载当前对话: {self.current_conversation.get('topic', '未命名')}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"加载当前对话失败: {e}")
|
||||
|
||||
def save_current_conversation(self):
|
||||
"""保存当前对话"""
|
||||
if self.current_conversation:
|
||||
self._atomic_write_json(self.current_file, self.current_conversation)
|
||||
|
||||
def start_new_conversation(self, topic: str = None) -> Dict:
|
||||
"""开始新对话"""
|
||||
if self.current_conversation:
|
||||
self.archive_current_conversation()
|
||||
|
||||
conv_id = str(int(time.time()))
|
||||
self.current_conversation = {
|
||||
"id": conv_id,
|
||||
"start_time": datetime.now().strftime(constants.DATETIME_FORMAT),
|
||||
"end_time": None,
|
||||
"topic": topic or "新对话",
|
||||
"messages": [],
|
||||
"tool_calls": [],
|
||||
"stats": {
|
||||
"user_messages": 0,
|
||||
"ai_messages": 0,
|
||||
"tool_calls": 0
|
||||
}
|
||||
}
|
||||
|
||||
self.save_current_conversation()
|
||||
self.logger.info(f"开始新对话: {topic}")
|
||||
return self.current_conversation
|
||||
|
||||
def archive_current_conversation(self, final_summary: str = None):
|
||||
"""将当前对话归档到历史"""
|
||||
if not self.current_conversation:
|
||||
return None
|
||||
|
||||
if self.current_conversation.get("end_time") is None:
|
||||
self.current_conversation["end_time"] = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
if final_summary and self.memory_db:
|
||||
messages = self.current_conversation.get("messages", [])
|
||||
content_preview = "\n".join([f"{m['speaker']}: {m['content'][:200]}"
|
||||
for m in messages[-10:] if m.get('speaker')])
|
||||
|
||||
self.memory_db.add_memory(
|
||||
topic=self.current_conversation["topic"],
|
||||
summary=final_summary,
|
||||
content=content_preview,
|
||||
start_time=self.current_conversation["start_time"],
|
||||
end_time=self.current_conversation["end_time"],
|
||||
importance=2
|
||||
)
|
||||
|
||||
# N16:归档文件名加序号防同秒覆盖——DATE_FORMAT 精度为秒,
|
||||
# 同一秒内两次归档(如测试/快速重置)会静默覆盖前一份
|
||||
timestamp = datetime.now().strftime(constants.DATE_FORMAT)
|
||||
archive_filename = f"conversation_{timestamp}.json"
|
||||
counter = 1
|
||||
while (self.archive_dir / archive_filename).exists():
|
||||
archive_filename = f"conversation_{timestamp}_{counter}.json"
|
||||
counter += 1
|
||||
archive_path = self.archive_dir / archive_filename
|
||||
|
||||
self._atomic_write_json(archive_path, self.current_conversation)
|
||||
|
||||
self.history_index["conversations"].append({
|
||||
"id": self.current_conversation["id"],
|
||||
"file": archive_filename,
|
||||
"start_time": self.current_conversation["start_time"],
|
||||
"end_time": self.current_conversation["end_time"],
|
||||
"topic": self.current_conversation["topic"],
|
||||
"messages": len(self.current_conversation["messages"]),
|
||||
"tool_calls": len(self.current_conversation["tool_calls"])
|
||||
})
|
||||
|
||||
self._save_history_index()
|
||||
self.logger.info(f"对话已归档: {archive_filename}")
|
||||
|
||||
self.current_conversation = None
|
||||
if self.current_file.exists():
|
||||
self.current_file.unlink()
|
||||
|
||||
return str(archive_path)
|
||||
|
||||
def load_all_history(self, max_messages: int = None) -> List[Dict]:
|
||||
"""加载所有历史对话作为上下文"""
|
||||
if max_messages is None:
|
||||
max_messages = constants.MAX_HISTORY_CONTEXT
|
||||
|
||||
if not self.load_history:
|
||||
return []
|
||||
|
||||
all_messages = []
|
||||
|
||||
if self.current_conversation:
|
||||
current_msgs = self.current_conversation.get("messages", [])
|
||||
for msg in current_msgs[-max_messages//2:]:
|
||||
if not msg.get("is_history"):
|
||||
all_messages.append({
|
||||
"timestamp": msg["timestamp"],
|
||||
"role": msg["role"],
|
||||
"speaker": msg.get("speaker"),
|
||||
"content": msg["content"],
|
||||
"from_current": True
|
||||
})
|
||||
|
||||
history_files = sorted(self.archive_dir.glob("*.json"), reverse=True)[:constants.MAX_HISTORY_CONVERSATIONS]
|
||||
|
||||
for hist_file in history_files:
|
||||
try:
|
||||
with open(hist_file, 'r', encoding='utf-8') as f:
|
||||
conv = json.load(f)
|
||||
messages = conv.get("messages", [])
|
||||
for msg in messages[-max_messages//len(history_files):]:
|
||||
if not msg.get("is_history"):
|
||||
all_messages.append({
|
||||
"timestamp": msg["timestamp"],
|
||||
"role": msg["role"],
|
||||
"speaker": msg.get("speaker"),
|
||||
"content": msg["content"],
|
||||
"from_file": hist_file.name
|
||||
})
|
||||
except Exception as e:
|
||||
self.logger.error(f"加载历史文件 {hist_file} 失败: {e}")
|
||||
|
||||
all_messages.sort(key=lambda x: x["timestamp"])
|
||||
|
||||
if len(all_messages) > max_messages:
|
||||
all_messages = all_messages[-max_messages:]
|
||||
|
||||
self.logger.info(f"已加载 {len(all_messages)} 条历史消息")
|
||||
return all_messages
|
||||
|
||||
def get_conversation_context(self, max_messages: int = None) -> List[Dict]:
|
||||
"""获取适合作为AI上下文的对话历史"""
|
||||
if max_messages is None:
|
||||
max_messages = constants.MAX_HISTORY_CONTEXT
|
||||
|
||||
if not self.load_history:
|
||||
return []
|
||||
|
||||
history_messages = self.load_all_history(max_messages)
|
||||
|
||||
context = []
|
||||
for msg in history_messages:
|
||||
role = msg.get("role", "")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
continue
|
||||
|
||||
if role == "assistant":
|
||||
speaker = msg.get("speaker", "")
|
||||
context.append({
|
||||
"role": "assistant",
|
||||
"content": f"[{speaker}] {content}"
|
||||
})
|
||||
else:
|
||||
context.append({
|
||||
"role": "user",
|
||||
"content": content
|
||||
})
|
||||
|
||||
return context
|
||||
|
||||
def add_message(self, role: str, content: str, speaker: str = None):
|
||||
"""添加消息到当前对话"""
|
||||
if not self.current_conversation:
|
||||
self.start_new_conversation()
|
||||
|
||||
message = {
|
||||
"timestamp": datetime.now().strftime(constants.DATETIME_FORMAT),
|
||||
"role": role,
|
||||
"speaker": speaker,
|
||||
"content": content
|
||||
}
|
||||
|
||||
self.current_conversation["messages"].append(message)
|
||||
|
||||
if role == "user":
|
||||
self.current_conversation["stats"]["user_messages"] += 1
|
||||
elif role == "assistant":
|
||||
self.current_conversation["stats"]["ai_messages"] += 1
|
||||
|
||||
self.save_current_conversation()
|
||||
|
||||
def add_tool_call(self, tool_name: str, args: Dict, result: str):
|
||||
"""记录工具调用"""
|
||||
if not self.current_conversation:
|
||||
return
|
||||
|
||||
tool_call = {
|
||||
"timestamp": datetime.now().strftime(constants.DATETIME_FORMAT),
|
||||
"tool": tool_name,
|
||||
"args": args,
|
||||
"result_preview": result[:200] + "..." if len(result) > 200 else result
|
||||
}
|
||||
|
||||
self.current_conversation["tool_calls"].append(tool_call)
|
||||
self.current_conversation["stats"]["tool_calls"] += 1
|
||||
self.save_current_conversation()
|
||||
|
||||
def end_current_conversation(self, summary: str = None):
|
||||
"""结束当前对话"""
|
||||
if self.current_conversation:
|
||||
self.current_conversation["end_time"] = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
if summary and self.memory_db:
|
||||
messages = self.current_conversation.get("messages", [])
|
||||
content_preview = "\n".join([f"{m['speaker']}: {m['content'][:200]}"
|
||||
for m in messages[-10:] if m.get('speaker')])
|
||||
|
||||
self.memory_db.add_memory(
|
||||
topic=self.current_conversation["topic"],
|
||||
summary=summary,
|
||||
content=content_preview,
|
||||
start_time=self.current_conversation["start_time"],
|
||||
end_time=self.current_conversation["end_time"],
|
||||
importance=2
|
||||
)
|
||||
|
||||
self.save_current_conversation()
|
||||
self.logger.info("当前对话已保存")
|
||||
|
||||
def reset_conversation(self):
|
||||
"""重置当前对话状态"""
|
||||
if self.current_conversation:
|
||||
self.current_conversation["end_time"] = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
self.save_current_conversation()
|
||||
self.current_conversation = None
|
||||
|
||||
def list_conversations(self) -> List[Dict]:
|
||||
"""列出所有对话(包括当前和归档)"""
|
||||
conversations = []
|
||||
|
||||
if self.current_conversation:
|
||||
conversations.append({
|
||||
"file": constants.CURRENT_CONVERSATION_FILE,
|
||||
"is_current": True,
|
||||
"id": self.current_conversation.get("id"),
|
||||
"start_time": self.current_conversation.get("start_time"),
|
||||
"end_time": self.current_conversation.get("end_time"),
|
||||
"topic": self.current_conversation.get("topic"),
|
||||
"messages": len(self.current_conversation.get("messages", [])),
|
||||
"tool_calls": len(self.current_conversation.get("tool_calls", []))
|
||||
})
|
||||
|
||||
for conv in self.history_index.get("conversations", []):
|
||||
conversations.append({
|
||||
"file": conv["file"],
|
||||
"is_current": False,
|
||||
"id": conv["id"],
|
||||
"start_time": conv["start_time"],
|
||||
"end_time": conv["end_time"],
|
||||
"topic": conv["topic"],
|
||||
"messages": conv["messages"],
|
||||
"tool_calls": conv["tool_calls"]
|
||||
})
|
||||
|
||||
conversations.sort(key=lambda x: x["start_time"] if x["start_time"] else "", reverse=True)
|
||||
return conversations
|
||||
|
||||
def load_conversation(self, file_name: str) -> Optional[Dict]:
|
||||
"""加载指定的对话"""
|
||||
if file_name == constants.CURRENT_CONVERSATION_FILE and self.current_conversation:
|
||||
return self.current_conversation
|
||||
|
||||
# I 批(审计 2026-08-18):file_name 目前无 LLM 可控调用面,但一旦
|
||||
# 接入即成为任意 .json 读取通道——realpath 解析软链后再做目录前缀校验
|
||||
archive_root = os.path.normcase(os.path.realpath(str(self.archive_dir)))
|
||||
target = os.path.normcase(os.path.realpath(str(self.archive_dir / file_name)))
|
||||
if not target.startswith(archive_root + os.sep):
|
||||
return None
|
||||
file_path = Path(target)
|
||||
if file_path.exists():
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
return None
|
||||
@@ -0,0 +1,852 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.db - 数据库子系统
|
||||
|
||||
抽离自 aibox.py,提供两个独立的 SQLite 数据库类:
|
||||
- MemoryDatabase : 记忆数据库(FTS5 全文检索),存储重要对话记忆
|
||||
- MemoDatabase : 备忘录/任务数据库,支持提醒、重复、即时任务
|
||||
|
||||
依赖:common (constants, i18n, Logger) + 标准库 sqlite3
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from common import constants, i18n, Logger
|
||||
|
||||
|
||||
# ==================== 记忆数据库 ====================
|
||||
|
||||
class MemoryDatabase:
|
||||
"""记忆数据库 - 存储重要的对话记忆"""
|
||||
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = os.path.join(constants.SAVE_DIR, constants.MEMORY_DB)
|
||||
|
||||
# 确保目录存在
|
||||
os.makedirs(constants.SAVE_DIR, exist_ok=True)
|
||||
self.db_path = db_path
|
||||
self.logger = Logger("MemoryDB")
|
||||
self._init_database()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
"""统一连接工厂(T25,审计 2026-08-21):对齐 MemoDatabase J9 修复——
|
||||
busy_timeout 覆盖全部路径。两类默认共用同一 memories.db,此前
|
||||
MemoryDatabase 直连只靠 sqlite3 默认 5s 等待,CLI 与服务进程并发
|
||||
写时锁冲突即抛 database is locked(记忆保存/检索失败)。"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("PRAGMA busy_timeout = 10000")
|
||||
return conn
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化数据库表"""
|
||||
with self._connect() as conn:
|
||||
# WAL 模式:读写并发不互斥,缓解多 CLI 进程并发写锁等待
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS memories (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
topic TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
content TEXT,
|
||||
start_time TEXT NOT NULL,
|
||||
end_time TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
tags TEXT,
|
||||
importance INTEGER DEFAULT 1
|
||||
)
|
||||
''')
|
||||
|
||||
self._init_fts_table(conn)
|
||||
|
||||
conn.commit()
|
||||
|
||||
def _fts_trigram_supported(self, conn) -> bool:
|
||||
"""检测当前 SQLite 是否支持 trigram 分词器(需 SQLite >= 3.34)"""
|
||||
try:
|
||||
conn.execute("CREATE VIRTUAL TABLE temp.__trigram_probe USING fts5(x, tokenize='trigram')")
|
||||
conn.execute("DROP TABLE temp.__trigram_probe")
|
||||
return True
|
||||
except sqlite3.OperationalError:
|
||||
return False
|
||||
|
||||
def _init_fts_table(self, conn):
|
||||
"""初始化 FTS5 全文索引表。
|
||||
|
||||
优先使用 trigram 分词器(中文可检索,unicode61 对中文不分词);
|
||||
检测到旧 unicode61 表时自动迁移:DROP 重建后从 memories 表回灌全量数据。
|
||||
环境不支持 trigram 时退回 unicode61 建表,中文搜索由 LIKE 兜底。
|
||||
"""
|
||||
cursor = conn.cursor()
|
||||
trigram_ok = self._fts_trigram_supported(conn)
|
||||
tokenize_opt = ", tokenize='trigram'" if trigram_ok else ""
|
||||
create_sql = (
|
||||
"CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts "
|
||||
f"USING fts5(topic, summary, content, tags, content=memories{tokenize_opt})"
|
||||
)
|
||||
|
||||
row = cursor.execute(
|
||||
"SELECT sql FROM sqlite_master WHERE type='table' AND name='memories_fts'"
|
||||
).fetchone()
|
||||
|
||||
need_backfill = False
|
||||
if row is None:
|
||||
need_backfill = True # 新库(memories 可能已有数据,回灌无害)
|
||||
elif trigram_ok and "trigram" not in (row[0] or ""):
|
||||
# 旧 unicode61 表 → 迁移为 trigram
|
||||
cursor.execute("DROP TABLE memories_fts")
|
||||
need_backfill = True
|
||||
|
||||
cursor.execute(create_sql)
|
||||
|
||||
if need_backfill:
|
||||
count = cursor.execute('''
|
||||
INSERT INTO memories_fts (rowid, topic, summary, content, tags)
|
||||
SELECT id, topic, summary, COALESCE(content, ''), COALESCE(tags, '')
|
||||
FROM memories
|
||||
''').rowcount
|
||||
if count:
|
||||
self.logger.info(f"FTS5 索引已使用 trigram 分词重建并回灌 {count} 条记忆")
|
||||
|
||||
def add_memory(self, topic: str, summary: str, content: str = None,
|
||||
start_time: str = None, end_time: str = None,
|
||||
tags: List[str] = None, importance: int = None) -> int:
|
||||
"""添加记忆到数据库"""
|
||||
if start_time is None:
|
||||
start_time = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
if end_time is None:
|
||||
end_time = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
if importance is None:
|
||||
importance = constants.MEMORY_DEFAULT_IMPORTANCE
|
||||
|
||||
tags_str = ",".join(tags) if tags else ""
|
||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO memories (topic, summary, content, start_time, end_time, created_at, tags, importance)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
''', (topic, summary, content, start_time, end_time, created_at, tags_str, importance))
|
||||
|
||||
memory_id = cursor.lastrowid
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO memories_fts (rowid, topic, summary, content, tags)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
''', (memory_id, topic, summary, content or "", tags_str))
|
||||
|
||||
conn.commit()
|
||||
|
||||
self.logger.info(f"已添加记忆: {topic} (ID: {memory_id})")
|
||||
return memory_id
|
||||
|
||||
def delete_memory(self, memory_id: int) -> bool:
|
||||
"""删除记忆"""
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
# J3(审计 2026-08-19):必须先删 fts 再删内容表——外部内容表
|
||||
# (content=memories)删除索引条目需从内容表回读旧值,内容行
|
||||
# 先删会导致索引条目永久残留(MATCH 仍可命中已删 rowid)
|
||||
cursor.execute('DELETE FROM memories_fts WHERE rowid = ?', (memory_id,))
|
||||
cursor.execute('DELETE FROM memories WHERE id = ?', (memory_id,))
|
||||
deleted = cursor.rowcount > 0
|
||||
|
||||
if deleted:
|
||||
conn.commit()
|
||||
self.logger.info(f"已删除记忆 ID: {memory_id}")
|
||||
|
||||
return deleted
|
||||
|
||||
def search_memories(self, keywords: List[str], match_all: bool = False, limit: int = None) -> List[Dict]:
|
||||
"""多关键词模糊搜索记忆。
|
||||
|
||||
优先走 FTS5 全文检索;FTS 无结果时(trigram 对 <3 字符的中文短词
|
||||
无法命中、旧表 unicode61 对中文不分词等场景)用 LIKE 兜底。
|
||||
"""
|
||||
if not keywords:
|
||||
return []
|
||||
|
||||
if limit is None:
|
||||
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
||||
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
results = self._search_via_fts(cursor, keywords, match_all, limit)
|
||||
if not results:
|
||||
results = self._search_via_like(cursor, keywords, match_all, limit)
|
||||
|
||||
return results
|
||||
|
||||
def _search_via_fts(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]:
|
||||
"""FTS5 全文检索(关键词含特殊字符导致 MATCH 语法错误时返回空,走 LIKE 兜底)"""
|
||||
joiner = " AND " if match_all else " OR "
|
||||
# N7:短语内双引号翻倍转义——关键词含 " 时若不翻倍会提前闭合短语
|
||||
# 边界,破坏 MATCH 语法或改变匹配语义(FTS5 短语引号转义规则)
|
||||
terms = [f'"{kw.replace(chr(34), chr(34) * 2)}"' for kw in keywords]
|
||||
query = joiner.join(terms)
|
||||
|
||||
try:
|
||||
cursor.execute('''
|
||||
SELECT m.*, rank
|
||||
FROM memories m
|
||||
JOIN memories_fts f ON m.rowid = f.rowid
|
||||
WHERE memories_fts MATCH ?
|
||||
ORDER BY m.end_time DESC, rank, m.importance DESC
|
||||
LIMIT ?
|
||||
''', (query, limit))
|
||||
except sqlite3.OperationalError as e:
|
||||
self.logger.warning(f"FTS5 检索失败(改用 LIKE 兜底): {e}")
|
||||
return []
|
||||
|
||||
return self._rows_to_memories(cursor.fetchall())
|
||||
|
||||
def _search_via_like(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]:
|
||||
"""LIKE 模糊检索兜底(转义 %/_/\\ 防止关键词被当作通配符)"""
|
||||
fields = ["topic", "summary", "content", "tags"]
|
||||
|
||||
def escape(kw: str) -> str:
|
||||
return kw.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
|
||||
conditions = []
|
||||
params = []
|
||||
for kw in keywords:
|
||||
pattern = f"%{escape(kw)}%"
|
||||
per_field = " OR ".join([f"{f} LIKE ? ESCAPE '\\'" for f in fields])
|
||||
conditions.append(f"({per_field})")
|
||||
params.extend([pattern] * len(fields))
|
||||
|
||||
joiner = " AND " if match_all else " OR "
|
||||
cursor.execute(
|
||||
f"SELECT * FROM memories WHERE {joiner.join(conditions)} "
|
||||
f"ORDER BY end_time DESC, importance DESC LIMIT ?",
|
||||
(*params, limit),
|
||||
)
|
||||
|
||||
return self._rows_to_memories(cursor.fetchall())
|
||||
|
||||
@staticmethod
|
||||
def _rows_to_memories(rows) -> List[Dict]:
|
||||
"""sqlite3.Row 列表转记忆字典列表(tags 逗号串拆分为列表)"""
|
||||
results = []
|
||||
for row in rows:
|
||||
memory = dict(row)
|
||||
if memory.get('tags'):
|
||||
memory['tags'] = memory['tags'].split(',')
|
||||
results.append(memory)
|
||||
return results
|
||||
|
||||
def get_recent_memories(self, limit: int = 10) -> List[Dict]:
|
||||
"""获取最近的记忆。limit=None 时使用默认上限(M4:此前 None 无 LIMIT
|
||||
全表加载,记忆库增长后单次调用 fetchall 全量入内存;与 search_memories
|
||||
的 None→默认上限语义对齐)。"""
|
||||
if limit is None:
|
||||
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC LIMIT ?', (limit,))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memory = dict(row)
|
||||
if memory.get('tags'):
|
||||
memory['tags'] = memory['tags'].split(',')
|
||||
results.append(memory)
|
||||
|
||||
return results
|
||||
|
||||
def get_memory_by_id(self, memory_id: int) -> Optional[Dict]:
|
||||
"""根据ID获取记忆"""
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memories WHERE id = ?', (memory_id,))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
memory = dict(row)
|
||||
if memory.get('tags'):
|
||||
memory['tags'] = memory['tags'].split(',')
|
||||
return memory
|
||||
return None
|
||||
|
||||
def delete_memories_by_ids(self, memory_ids: List[int]) -> int:
|
||||
"""批量删除记忆"""
|
||||
if not memory_ids:
|
||||
return 0
|
||||
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
placeholders = ','.join(['?'] * len(memory_ids))
|
||||
# J3:同 delete_memory——先删 fts 再删内容表,防索引条目残留
|
||||
cursor.execute(f'DELETE FROM memories_fts WHERE rowid IN ({placeholders})', memory_ids)
|
||||
cursor.execute(f'DELETE FROM memories WHERE id IN ({placeholders})', memory_ids)
|
||||
deleted_count = cursor.rowcount
|
||||
|
||||
if deleted_count > 0:
|
||||
conn.commit()
|
||||
self.logger.info(f"已批量删除 {deleted_count} 条记忆")
|
||||
|
||||
return deleted_count
|
||||
|
||||
|
||||
# ==================== 备忘录数据库 ====================
|
||||
|
||||
class MemoDatabase:
|
||||
"""备忘录数据库 - 存储和管理任务"""
|
||||
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = os.path.join(constants.SAVE_DIR, constants.MEMORY_DB)
|
||||
|
||||
# 确保目录存在
|
||||
os.makedirs(constants.SAVE_DIR, exist_ok=True)
|
||||
self.db_path = db_path
|
||||
self.logger = Logger("MemoDB")
|
||||
self._init_database()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
"""统一连接工厂(J9,审计 2026-08-19):busy_timeout 覆盖全部写路径,
|
||||
不再只有 mark_as_triggered 一条路径等待锁释放"""
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("PRAGMA busy_timeout = 10000")
|
||||
return conn
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化备忘录表"""
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
# J9:WAL 显式设置——与 MemoryDatabase 同库文件时不再依赖创建顺序
|
||||
cursor.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS memos (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
title TEXT NOT NULL,
|
||||
content TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
reminder_time TEXT,
|
||||
completed_at TEXT,
|
||||
is_completed INTEGER DEFAULT 0,
|
||||
tags TEXT,
|
||||
priority INTEGER DEFAULT 1,
|
||||
created_by TEXT,
|
||||
source TEXT,
|
||||
repeat_type TEXT DEFAULT 'none',
|
||||
repeat_interval INTEGER DEFAULT 0,
|
||||
repeat_interval_unit TEXT DEFAULT 'days',
|
||||
repeat_end_time TEXT,
|
||||
last_triggered_at TEXT,
|
||||
trigger_count INTEGER DEFAULT 0,
|
||||
is_immediate INTEGER DEFAULT 0
|
||||
)
|
||||
''')
|
||||
|
||||
cursor.execute('''
|
||||
CREATE INDEX IF NOT EXISTS idx_memos_reminder
|
||||
ON memos(reminder_time) WHERE is_completed = 0
|
||||
''')
|
||||
|
||||
cursor.execute('''
|
||||
CREATE INDEX IF NOT EXISTS idx_memos_immediate
|
||||
ON memos(is_immediate) WHERE is_immediate = 1 AND is_completed = 0
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
|
||||
self.logger.info("任务表初始化完成")
|
||||
|
||||
def add_memo(self, title: str, content: str = None, reminder_time: str = None,
|
||||
tags: List[str] = None, priority: int = 1, created_by: str = None,
|
||||
source: str = None,
|
||||
repeat_type: str = 'none', repeat_interval: int = 0,
|
||||
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
||||
is_immediate: bool = False) -> int:
|
||||
"""添加任务 - 所有任务都是对话任务,不再需要metadata和任务类型"""
|
||||
# J1(审计 2026-08-19):入库校验——畸形时间/数字字段此前原样入库,
|
||||
# 下游无防护解析点引发调度器失败热循环与新连接握手崩溃(审计报告
|
||||
# §13.1 J1)。CLI/WS/memo 工具三入口均经本方法,校验一处覆盖。
|
||||
# W12(审计 2026-08-20):空标题任务此前可入库(CLI "|x" / 调用方疏忽)
|
||||
if not isinstance(title, str) or not title.strip():
|
||||
raise ValueError("任务标题不能为空")
|
||||
if reminder_time:
|
||||
try:
|
||||
datetime.strptime(reminder_time, constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
raise ValueError(
|
||||
f"提醒时间格式非法: {reminder_time!r}(应为 {constants.DATETIME_FORMAT})")
|
||||
if repeat_end_time:
|
||||
try:
|
||||
datetime.strptime(repeat_end_time, constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
raise ValueError(
|
||||
f"重复结束时间格式非法: {repeat_end_time!r}(应为 {constants.DATETIME_FORMAT})")
|
||||
if repeat_type not in ("none", "daily", "weekly", "monthly", "custom"):
|
||||
raise ValueError(f"无效的重复类型: {repeat_type!r}")
|
||||
try:
|
||||
repeat_interval = int(repeat_interval)
|
||||
except (ValueError, TypeError):
|
||||
raise ValueError(f"重复间隔必须是整数: {repeat_interval!r}")
|
||||
if repeat_type == "custom" and repeat_interval <= 0:
|
||||
raise ValueError("自定义重复需要设置有效的 repeat_interval (>0)")
|
||||
|
||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
tags_str = ",".join(tags) if tags else ""
|
||||
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO memos
|
||||
(title, content, created_at, reminder_time, tags, priority, created_by,
|
||||
is_completed, source, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time, is_immediate)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, 0, ?, ?, ?, ?, ?, ?)
|
||||
''', (title, content, created_at, reminder_time, tags_str, priority, created_by,
|
||||
source, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time, 1 if is_immediate else 0))
|
||||
|
||||
memo_id = cursor.lastrowid
|
||||
conn.commit()
|
||||
|
||||
repeat_info = f" (重复: {repeat_type}"
|
||||
if repeat_type == 'custom':
|
||||
repeat_info += f", 每{repeat_interval} {repeat_interval_unit}"
|
||||
repeat_info += ")"
|
||||
immediate_info = " (即时任务)" if is_immediate else ""
|
||||
self.logger.info(i18n.get('task_added', title, memo_id) + repeat_info + immediate_info)
|
||||
return memo_id
|
||||
|
||||
def get_memo(self, memo_id: int) -> Optional[Dict]:
|
||||
"""获取单个任务"""
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row:
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
return memo
|
||||
return None
|
||||
|
||||
def get_pending_memos(self, limit: int = None, before_time: str = None, include_immediate: bool = False) -> List[Dict]:
|
||||
"""获取待执行的任务 - 严格使用当前时间比较"""
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 使用当前时间作为基准
|
||||
now = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
# 如果传入了 before_time,使用它,否则使用当前时间
|
||||
compare_time = before_time if before_time else now
|
||||
|
||||
# I 批(审计 2026-08-18):去掉 SQL 层 reminder_time <= ? 预过滤——
|
||||
# 字符串比较会把格式异常的时间(如 "不是时间")静默排除在 SQL 层,
|
||||
# Python 二次验证永远看不到它们(Q5:格式异常任务按到期返回)。
|
||||
# 时间过滤改由下方 Python 解析做权威判断(本地小库全扫描可忽略)。
|
||||
query = '''
|
||||
SELECT * FROM memos
|
||||
WHERE is_completed = 0
|
||||
AND reminder_time IS NOT NULL
|
||||
AND (repeat_end_time IS NULL OR repeat_end_time >= ?)
|
||||
'''
|
||||
|
||||
if not include_immediate:
|
||||
query += ' AND is_immediate = 0'
|
||||
|
||||
query += ' ORDER BY reminder_time ASC, priority DESC'
|
||||
|
||||
params = [compare_time]
|
||||
|
||||
# I 批:limit=0 语义修正——此前 `if limit:` 使 0 等同"不限",
|
||||
# 与 None=不限的惯例冲突;改为 None 才不加 LIMIT
|
||||
if limit is not None:
|
||||
query += ' LIMIT ?'
|
||||
params.append(max(0, int(limit)))
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
|
||||
# 【关键修复】二次验证:确保 reminder_time 确实 <= 当前时间
|
||||
if memo.get('reminder_time'):
|
||||
try:
|
||||
reminder_time = datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT)
|
||||
current_time = datetime.strptime(compare_time, constants.DATETIME_FORMAT)
|
||||
if reminder_time <= current_time:
|
||||
results.append(memo)
|
||||
else:
|
||||
# 调试日志
|
||||
self.logger.debug(f"过滤掉未到期任务 #{memo['id']}: {memo['reminder_time']} > {compare_time}")
|
||||
except (ValueError, TypeError) as e:
|
||||
# I 批(审计 2026-08-18):格式异常此前被静默丢弃——
|
||||
# 任务永不执行且用户无感知。改为按"到期"返回并告警,
|
||||
# 至少保证任务被处理而非永久丢失
|
||||
self.logger.warning(f"解析任务 #{memo['id']} 的提醒时间失败,按到期处理: {e}")
|
||||
results.append(memo)
|
||||
else:
|
||||
# 没有提醒时间的任务不应该在待执行列表中
|
||||
self.logger.warning(f"任务 #{memo['id']} 没有提醒时间,但被返回为待执行")
|
||||
|
||||
return results
|
||||
|
||||
def get_immediate_memos(self, limit: int = None) -> List[Dict]:
|
||||
"""获取待执行的即时任务"""
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
query = '''
|
||||
SELECT * FROM memos
|
||||
WHERE is_completed = 0
|
||||
AND is_immediate = 1
|
||||
ORDER BY created_at ASC
|
||||
'''
|
||||
|
||||
params = []
|
||||
# I 批:limit=0 语义修正——此前 `if limit:` 使 0 等同"不限",
|
||||
# 与 None=不限的惯例冲突;改为 None 才不加 LIMIT
|
||||
if limit is not None:
|
||||
query += ' LIMIT ?'
|
||||
params.append(max(0, int(limit)))
|
||||
|
||||
cursor.execute(query, params)
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
results.append(memo)
|
||||
|
||||
return results
|
||||
|
||||
def get_overdue_memos(self, limit: int = None) -> List[Dict]:
|
||||
"""获取已过期的任务"""
|
||||
now = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
return self.get_pending_memos(limit=limit, before_time=now)
|
||||
|
||||
def get_all_memos(self, include_completed: bool = False, limit: int = 50) -> List[Dict]:
|
||||
"""获取所有任务"""
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
|
||||
if include_completed:
|
||||
cursor.execute('''
|
||||
SELECT * FROM memos
|
||||
ORDER BY
|
||||
CASE WHEN is_completed = 0 THEN 0 ELSE 1 END,
|
||||
reminder_time ASC,
|
||||
created_at DESC
|
||||
LIMIT ?
|
||||
''', (limit,))
|
||||
else:
|
||||
cursor.execute('''
|
||||
SELECT * FROM memos
|
||||
WHERE is_completed = 0
|
||||
ORDER BY reminder_time ASC, priority DESC, created_at DESC
|
||||
LIMIT ?
|
||||
''', (limit,))
|
||||
|
||||
results = []
|
||||
for row in cursor.fetchall():
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
results.append(memo)
|
||||
|
||||
return results
|
||||
|
||||
def mark_as_triggered(self, memo_id: int) -> Optional[str]:
|
||||
"""标记任务已被触发,并根据重复规则计算下一次提醒时间"""
|
||||
# N10:读取与更新收进同一连接的 BEGIN IMMEDIATE 事务。此前先
|
||||
# get_memo(独立连接)后 UPDATE,check-then-act 之间存在竞态:
|
||||
# 调度器并发触发同一任务会双计数/双推进 repeat。IMMEDIATE 在事务
|
||||
# 开始即取写锁,串行化并发触发。
|
||||
now = datetime.now()
|
||||
next_reminder = None
|
||||
|
||||
with self._connect() as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
# J9:busy_timeout 已由 _connect() 工厂统一设置(全部写路径一致),
|
||||
# 多进程同库(CLI --memo-add + 服务调度)写竞争等待锁释放而非立即失败
|
||||
cursor.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
conn.rollback()
|
||||
return None
|
||||
memo = dict(row)
|
||||
if memo.get('tags'):
|
||||
memo['tags'] = memo['tags'].split(',')
|
||||
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET last_triggered_at = ?, trigger_count = trigger_count + 1
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
# 处理重复任务
|
||||
if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now):
|
||||
# 使用当前时间作为基准来计算下一次执行时间
|
||||
next_reminder = self._calculate_next_reminder(memo, now)
|
||||
|
||||
if next_reminder:
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET reminder_time = ?
|
||||
WHERE id = ?
|
||||
''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}")
|
||||
else:
|
||||
# 重复结束,标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 重复结束,已标记为完成")
|
||||
elif memo['is_immediate']:
|
||||
# 即时任务执行后标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"即时任务 #{memo_id} 已完成")
|
||||
else:
|
||||
# 一次性任务,标记为完成
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ?
|
||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||
|
||||
self.logger.info(f"任务 #{memo_id} 已完成")
|
||||
|
||||
conn.commit()
|
||||
except Exception:
|
||||
conn.rollback()
|
||||
raise
|
||||
|
||||
return next_reminder.strftime(constants.DATETIME_FORMAT) if next_reminder else None
|
||||
|
||||
def _should_end_repeat(self, memo: Dict, current_time: datetime) -> bool:
|
||||
"""检查是否应该结束重复"""
|
||||
if memo['repeat_end_time']:
|
||||
try:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
# J1:存量脏数据防御——解析失败视为无结束时间,不打断触发事务
|
||||
return False
|
||||
if current_time > end_time:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _calculate_next_reminder(self, memo: Dict, current_time: datetime) -> Optional[datetime]:
|
||||
"""
|
||||
计算下一次提醒时间,使用当前时间作为基准
|
||||
支持分钟、小时、天为单位
|
||||
"""
|
||||
if memo['repeat_type'] == 'daily':
|
||||
# 每天重复:当前时间 + 1天
|
||||
next_time = current_time + timedelta(days=1)
|
||||
elif memo['repeat_type'] == 'weekly':
|
||||
# 每周重复:当前时间 + 7天
|
||||
next_time = current_time + timedelta(weeks=1)
|
||||
elif memo['repeat_type'] == 'monthly':
|
||||
# 每月重复:尝试在当前时间上加一个月
|
||||
try:
|
||||
# 处理月份进位
|
||||
if current_time.month == 12:
|
||||
next_time = current_time.replace(year=current_time.year + 1, month=1)
|
||||
else:
|
||||
next_time = current_time.replace(month=current_time.month + 1)
|
||||
except ValueError:
|
||||
# 处理月份天数问题(如1月31日加一个月)
|
||||
if current_time.month == 12:
|
||||
next_time = current_time.replace(year=current_time.year + 1, month=1, day=1) - timedelta(days=1)
|
||||
else:
|
||||
next_time = current_time.replace(month=current_time.month + 1, day=1) - timedelta(days=1)
|
||||
elif memo['repeat_type'] == 'custom':
|
||||
# J1:repeat_interval 类型异常(存量脏数据)按 0 处理,
|
||||
# 不再抛 TypeError 打崩触发事务引发热循环
|
||||
try:
|
||||
repeat_interval = int(memo.get('repeat_interval') or 0)
|
||||
except (ValueError, TypeError):
|
||||
repeat_interval = 0
|
||||
if repeat_interval > 0:
|
||||
unit = memo.get('repeat_interval_unit', 'days')
|
||||
if unit == 'minutes':
|
||||
# 每N分钟重复:当前时间 + N分钟
|
||||
next_time = current_time + timedelta(minutes=repeat_interval)
|
||||
elif unit == 'hours':
|
||||
# 每N小时重复:当前时间 + N小时
|
||||
next_time = current_time + timedelta(hours=repeat_interval)
|
||||
else: # days
|
||||
# 每N天重复:当前时间 + N天
|
||||
next_time = current_time + timedelta(days=repeat_interval)
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
|
||||
# 检查是否超过重复结束时间
|
||||
if memo['repeat_end_time']:
|
||||
try:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
except (ValueError, TypeError):
|
||||
# J1:存量脏数据防御——解析失败视为无结束时间
|
||||
end_time = None
|
||||
if end_time and next_time > end_time:
|
||||
return None
|
||||
|
||||
return next_time
|
||||
|
||||
def complete_memo(self, memo_id: int, force_complete: bool = False) -> bool:
|
||||
"""完成任务"""
|
||||
memo = self.get_memo(memo_id)
|
||||
if not memo:
|
||||
return False
|
||||
|
||||
if memo['repeat_type'] != 'none' and not force_complete:
|
||||
self.mark_as_triggered(memo_id)
|
||||
return True
|
||||
|
||||
completed_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
UPDATE memos
|
||||
SET is_completed = 1, completed_at = ?
|
||||
WHERE id = ? AND is_completed = 0
|
||||
''', (completed_at, memo_id))
|
||||
|
||||
affected = cursor.rowcount
|
||||
conn.commit()
|
||||
|
||||
if affected > 0:
|
||||
self.logger.info(i18n.get('memo_completed', memo_id))
|
||||
return True
|
||||
return False
|
||||
|
||||
def delete_memo(self, memo_id: int) -> bool:
|
||||
"""物理删除任务(M5:delete 语义此前被 complete_memo 顶替——只标记
|
||||
完成不删行,重复任务还会被推进而非移除,stats 计数失真)"""
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("DELETE FROM memos WHERE id = ?", (memo_id,))
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def get_memo_stats(self) -> Dict:
|
||||
"""获取任务统计信息"""
|
||||
with self._connect() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos')
|
||||
total = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos WHERE is_completed = 0')
|
||||
pending = cursor.fetchone()[0]
|
||||
|
||||
now = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
cursor.execute('''
|
||||
SELECT COUNT(*) FROM memos
|
||||
WHERE is_completed = 0
|
||||
AND reminder_time <= ?
|
||||
AND is_immediate = 0
|
||||
''', (now,))
|
||||
overdue = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos WHERE is_completed = 1')
|
||||
completed = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM memos WHERE is_immediate = 1 AND is_completed = 0')
|
||||
immediate = cursor.fetchone()[0]
|
||||
|
||||
return {
|
||||
"total": total,
|
||||
"pending": pending,
|
||||
"overdue": overdue,
|
||||
"completed": completed,
|
||||
"immediate": immediate
|
||||
}
|
||||
|
||||
def add_user_task(self, title: str, content: str = None, reminder_time: str = None,
|
||||
repeat_type: str = 'none', repeat_interval: int = 0,
|
||||
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
||||
is_immediate: bool = False) -> int:
|
||||
"""添加用户任务(前端直接调用的任务)- 所有任务都是对话任务"""
|
||||
if not constants.MEMO_ALLOW_USER_TASKS:
|
||||
raise Exception("用户任务功能已禁用")
|
||||
|
||||
return self.add_memo(
|
||||
title=title,
|
||||
content=content,
|
||||
reminder_time=reminder_time,
|
||||
repeat_type=repeat_type,
|
||||
repeat_interval=repeat_interval,
|
||||
repeat_interval_unit=repeat_interval_unit,
|
||||
repeat_end_time=repeat_end_time,
|
||||
created_by="user",
|
||||
source="websocket",
|
||||
is_immediate=is_immediate
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_overdue_time(reminder_time: str, now: datetime) -> bool:
|
||||
"""J1:展示路径时间解析防御——格式异常按到期处理(对齐 Q5 语义),
|
||||
不再抛 ValueError 打崩新连接握手流程"""
|
||||
if not reminder_time:
|
||||
return False
|
||||
try:
|
||||
return datetime.strptime(reminder_time, constants.DATETIME_FORMAT) < now
|
||||
except (ValueError, TypeError):
|
||||
return True
|
||||
|
||||
def get_pending_tasks_for_user(self) -> List[Dict]:
|
||||
"""获取待执行任务列表(用于前端显示)"""
|
||||
pending = self.get_pending_memos(include_immediate=False)
|
||||
immediate = self.get_immediate_memos()
|
||||
result = []
|
||||
|
||||
now = datetime.now()
|
||||
for memo in pending + immediate:
|
||||
result.append({
|
||||
"id": memo['id'],
|
||||
"title": memo['title'],
|
||||
"content": memo['content'],
|
||||
"reminder_time": memo['reminder_time'],
|
||||
"priority": memo['priority'],
|
||||
"is_overdue": self._is_overdue_time(memo['reminder_time'], now),
|
||||
"repeat_type": memo['repeat_type'],
|
||||
"repeat_interval": memo.get('repeat_interval', 0),
|
||||
"repeat_interval_unit": memo.get('repeat_interval_unit', 'days'),
|
||||
"trigger_count": memo.get('trigger_count', 0),
|
||||
"is_immediate": memo.get('is_immediate', False)
|
||||
})
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,510 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.executors - 命令执行器
|
||||
|
||||
抽离自 aibox.py,提供命令执行的抽象基类与 subprocess 实现。
|
||||
- CommandExecutor : 抽象基类,定义 execute / execute_streaming / cancel
|
||||
- SubprocessCommandExecutor : 基于 subprocess.Popen 的实现,支持超时/工作目录/
|
||||
1MB 输出上限/线程安全取消
|
||||
|
||||
依赖:common (ConfigManager, constants, Logger) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
from common import ConfigManager, constants, Logger
|
||||
|
||||
# T21(审计 2026-08-21):env 覆盖黑名单——工具 schema 中 env 是开放
|
||||
# object,LLM(可被网页内容提示注入诱导)可传 LD_PRELOAD/PYTHONPATH/
|
||||
# PATH 等在 shell=True 之上进一步劫持进程加载(POSIX LD_PRELOAD 可加载
|
||||
# 任意共享库)。命中黑名单的覆盖键拒绝并告警
|
||||
ENV_DENY_EXACT = frozenset({
|
||||
"PATH", "PYTHONPATH", "PYTHONHOME", "SHELL", "IFS",
|
||||
"LD_LIBRARY_PATH", "LD_PRELOAD", "LD_AUDIT",
|
||||
"DYLD_LIBRARY_PATH", "DYLD_INSERT_LIBRARIES",
|
||||
"PSMODULEPATH", "COMSPEC",
|
||||
})
|
||||
ENV_DENY_PREFIXES = ("LD_", "DYLD_")
|
||||
|
||||
|
||||
def filter_env_overrides(env) -> Tuple[Dict, list]:
|
||||
"""T21:过滤 env 覆盖项,返回 (安全覆盖, 被拒绝键列表)。
|
||||
|
||||
键大小写不敏感匹配(Windows 环境变量不区分大小写);
|
||||
LD_*/DYLD_* 前缀整体拒绝(加载器劫持面)。
|
||||
"""
|
||||
if not env:
|
||||
return {}, []
|
||||
safe, rejected = {}, []
|
||||
for key, value in dict(env).items():
|
||||
ku = str(key).upper()
|
||||
if ku in ENV_DENY_EXACT or ku.startswith(ENV_DENY_PREFIXES):
|
||||
rejected.append(key)
|
||||
else:
|
||||
safe[key] = value
|
||||
return safe, rejected
|
||||
|
||||
|
||||
class CommandExecutor(ABC):
|
||||
"""命令执行器抽象基类 - 负责执行系统命令"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self.logger = logger or Logger("CommandExecutor")
|
||||
self.config = ConfigManager()
|
||||
self.process = None
|
||||
|
||||
@abstractmethod
|
||||
def execute(self, command: str, timeout: int = 30, **kwargs) -> str:
|
||||
"""执行命令并返回结果(非流式)"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30, **kwargs):
|
||||
"""流式执行命令,通过队列输出结果"""
|
||||
pass
|
||||
|
||||
def is_enabled(self) -> bool:
|
||||
"""检查命令执行是否启用"""
|
||||
return constants.COMMANDS_ENABLED
|
||||
|
||||
def format_result(self, success: bool, message: str, data: Dict = None) -> str:
|
||||
"""统一结果格式化"""
|
||||
result = []
|
||||
emoji = "✅" if success else "❌"
|
||||
result.append(f"{emoji} {message}")
|
||||
|
||||
if data:
|
||||
for key, value in data.items():
|
||||
if isinstance(value, str) and "\n" in value:
|
||||
result.append(f" {key}:")
|
||||
for line in value.split("\n"):
|
||||
result.append(f" {line}")
|
||||
else:
|
||||
result.append(f" {key}: {value}")
|
||||
|
||||
return "\n".join(result)
|
||||
|
||||
|
||||
class SubprocessCommandExecutor(CommandExecutor):
|
||||
"""基于subprocess的命令执行器 - 支持超时、工作目录、输出大小限制和取消"""
|
||||
|
||||
# 命令输出最大字节数(1MB),超出截断
|
||||
MAX_OUTPUT_BYTES = 1 * 1024 * 1024
|
||||
|
||||
def __init__(self, logger=None):
|
||||
super().__init__(logger)
|
||||
self._current_process = None # 兼容保留:最近一次注册的进程(诊断用)
|
||||
self._current_pgid = None # L9:Popen 后立即捕获,shell 死亡后 getpgid 会失败
|
||||
self._process_lock = threading.Lock()
|
||||
# T5:单槽 _current_process 在并发执行时互相覆盖——先启动的执行
|
||||
# 超时/取消会误杀后来者的进程组(killpg 连坐孙进程),先结束的
|
||||
# finally 清引用导致后来者永远无法取消。改为令牌多槽注册表,
|
||||
# 终止只作用于自己槽位内的 (proc, pgid)。
|
||||
self._proc_slots: Dict[int, Tuple] = {}
|
||||
self._token_seq = 0
|
||||
|
||||
def _register_proc(self, process):
|
||||
"""T5:注册本次执行的进程槽,返回 (token, pgid)"""
|
||||
with self._process_lock:
|
||||
self._token_seq += 1
|
||||
token = self._token_seq
|
||||
pgid = None
|
||||
if os.name == "posix":
|
||||
try:
|
||||
pgid = os.getpgid(process.pid)
|
||||
except Exception:
|
||||
pgid = None
|
||||
self._proc_slots[token] = (process, pgid)
|
||||
# 兼容属性指向最新注册(诊断/旧测试观察用,不参与终止逻辑)
|
||||
self._current_process = process
|
||||
self._current_pgid = pgid
|
||||
return token, pgid
|
||||
|
||||
def _unregister_proc(self, token):
|
||||
"""T5:仅移除自己的槽位;全部槽位清空时同步清兼容属性"""
|
||||
with self._process_lock:
|
||||
self._proc_slots.pop(token, None)
|
||||
if not self._proc_slots:
|
||||
self._current_process = None
|
||||
self._current_pgid = None
|
||||
|
||||
def _terminate_tree(self, proc, force: bool = False, pgid=None):
|
||||
"""终止整个进程组(L9:POSIX 下 shell=True 只杀 shell 时,孙进程
|
||||
(`(x &)` 等)仍持管道写端——reader 阻塞无 EOF、fd/线程悬挂至孙
|
||||
进程退出,写满 64KB 管道缓冲后孙进程也卡死。start_new_session 使
|
||||
shell 成为会话首进程,按组信号可一次覆盖整棵进程树)。
|
||||
组 ID 在 Popen 后立即保存:shell 死亡后 os.getpgid(pid) 会抛
|
||||
ProcessLookupError,届时无法再从进程反查组。
|
||||
T5:pgid 由调用方显式传入(注册槽位时捕获),不再读共享
|
||||
_current_pgid——并发执行时那是别人的进程组。
|
||||
T11(审计 2026-08-21):Windows 无进程组概念,单进程 terminate/
|
||||
kill 收割不到 shell=True 生成的 cmd.exe /c 之下 spawn 的孙进程
|
||||
(超时/取消后孙进程成孤儿,持续占用 CPU/端口/文件锁)。用
|
||||
taskkill /T 递归终止整棵进程树;Windows 无 SIGTERM 等价物,
|
||||
taskkill 不带 /F 仅对有消息窗口的 GUI 进程有效,控制台进程
|
||||
杀不掉,故统一 /F 强制,force 参数在 Windows 上不区分。"""
|
||||
try:
|
||||
if os.name == "posix":
|
||||
sig = signal.SIGKILL if force else signal.SIGTERM
|
||||
if pgid is None:
|
||||
try:
|
||||
pgid = os.getpgid(proc.pid)
|
||||
except Exception:
|
||||
pgid = None
|
||||
if pgid is not None:
|
||||
os.killpg(pgid, sig)
|
||||
else:
|
||||
proc.kill() if force else proc.terminate()
|
||||
else:
|
||||
# T11:Windows 树收割——taskkill /T /F 递归终止子孙进程,
|
||||
# 失败(进程已退出/无权限)时兜底单进程 kill
|
||||
try:
|
||||
subprocess.run(
|
||||
["taskkill", "/PID", str(proc.pid), "/T", "/F"],
|
||||
capture_output=True, timeout=10)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def cancel(self):
|
||||
"""取消本执行器当前注册的全部运行中命令(kill 子进程组)。线程安全。
|
||||
T5:多槽令牌注册表——并发执行时定向终止各自槽位,互不误杀;
|
||||
无参调用语义为"终止本执行器全部在途命令"(中断/关闭路径)。"""
|
||||
with self._process_lock:
|
||||
slots = list(self._proc_slots.values())
|
||||
for proc, pgid in slots:
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
self._terminate_tree(proc, force=True, pgid=pgid)
|
||||
self.logger.info("已终止正在运行的子进程组")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None,
|
||||
pgid=None) -> Tuple[str, bool, bool]:
|
||||
"""从进程读取输出,限制总大小和超时。返回 (output_str, timeout_occurred, truncated)。
|
||||
|
||||
H2 修复:阻塞 readline 在独立 reader 线程执行,主循环带超时轮询队列。
|
||||
此前超时检查只在两次 readline 之间执行,无输出命令会让 readline
|
||||
阻塞到进程自然退出,timeout 形同虚设(审计 PoC 实证)。
|
||||
"""
|
||||
output_chunks = []
|
||||
total_bytes = 0
|
||||
timeout_occurred = False
|
||||
truncated = False
|
||||
|
||||
line_queue: queue.Queue = queue.Queue()
|
||||
reader_done = threading.Event()
|
||||
|
||||
def _reader():
|
||||
"""后台 reader 线程:阻塞 readline,行入队;EOF/异常结束"""
|
||||
try:
|
||||
for line in iter(process.stdout.readline, ""):
|
||||
line_queue.put(line)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
reader_done.set()
|
||||
|
||||
reader = threading.Thread(target=_reader, daemon=True)
|
||||
reader.start()
|
||||
|
||||
shell_dead_since = None # shell 进程已死但 reader 未收到 EOF 的起始时刻
|
||||
while True:
|
||||
# 检查超时(readline 阻塞不再影响本判断)
|
||||
# J16(审计 2026-08-19):timeout<=0 语义为"不设超时"(有意为之,
|
||||
# 调用方显式传 0 禁用);registry 层 3600s 流式兜底仍生效
|
||||
if timeout > 0 and time.time() - start_time > timeout:
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
timeout_occurred = True
|
||||
break
|
||||
|
||||
# 检查输出大小
|
||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
||||
truncated = True
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
break
|
||||
|
||||
try:
|
||||
line = line_queue.get(timeout=0.1)
|
||||
shell_dead_since = None
|
||||
except queue.Empty:
|
||||
if reader_done.is_set():
|
||||
break
|
||||
# shell 已被 kill/退出但 reader 仍阻塞(如 Windows shell=True
|
||||
# 下 cancel 只杀死包装进程,孙进程仍持有管道写端导致无 EOF):
|
||||
# 宽限 1 秒排空残留输出后放弃等待
|
||||
if process.poll() is not None:
|
||||
if shell_dead_since is None:
|
||||
shell_dead_since = time.time()
|
||||
elif time.time() - shell_dead_since > 1.0:
|
||||
# L9:shell 已死但孙进程仍持管道写端(`(sleep 300 &) ;
|
||||
# echo done` 形态)。只靠 daemon reader 兜底会悬挂线程
|
||||
# 与管道 fd 至孙进程退出,写满 64KB 缓冲后孙进程自身
|
||||
# 也卡死。此处强杀进程组收尾(nohup 重定向输出的
|
||||
# 真正脱离场景不持管道,不会走到这里)。
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
break
|
||||
continue
|
||||
|
||||
line_bytes = len(line.encode('utf-8', errors='replace'))
|
||||
if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated:
|
||||
allowed = self.MAX_OUTPUT_BYTES - total_bytes
|
||||
if allowed > 0:
|
||||
# L5:decode 用 errors='ignore'——按字节截断可能落在多字节
|
||||
# 字符中间,errors='replace' 会在行尾留下半个字符(�)
|
||||
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='ignore'))
|
||||
total_bytes = self.MAX_OUTPUT_BYTES
|
||||
truncated = True
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
break
|
||||
|
||||
output_chunks.append(line)
|
||||
total_bytes += line_bytes
|
||||
|
||||
if output_queue:
|
||||
output_queue.put(("line", line.rstrip('\n') + "\n"))
|
||||
|
||||
# N4 修复(审计 2026-08-18,二次修正):reader 线程收尾。
|
||||
# 超时/截断/宽限到期 break 后,reader 可能仍阻塞在 readline
|
||||
# (Windows shell=True 下孙进程持管道写端,永远等不到 EOF)。
|
||||
# 注意:此时不能 close stdout——CPython 流对象内部有锁,
|
||||
# readline 持锁阻塞时 close() 会一同挂起直到 readline 返回
|
||||
# (实测命令超时被拖满整个 Start-Sleep 时长)。reader 为 daemon
|
||||
# 线程,孙进程退出收到 EOF 后自然结束,不阻塞主流程;
|
||||
# reader 已正常退出的场景才 close 显式释放管道 fd。
|
||||
if reader_done.is_set():
|
||||
try:
|
||||
process.stdout.close()
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
reader.join(timeout=2)
|
||||
if reader.is_alive():
|
||||
self.logger.warning(
|
||||
"reader 线程仍阻塞在管道读取(孙进程持写端),已放弃等待,"
|
||||
"线程由 daemon 兜底、随管道 EOF 自然退出")
|
||||
|
||||
return ''.join(output_chunks), timeout_occurred, truncated
|
||||
|
||||
def execute(self, command: str, timeout: int = 30, work_dir: str = None, env: Dict = None, **kwargs) -> str:
|
||||
"""执行命令并返回完整结果(带输出大小限制)"""
|
||||
if not self.is_enabled():
|
||||
return "❌ 命令执行功能已禁用"
|
||||
|
||||
# T24:非流式路径没有 registry 3600s 流式兜底,timeout<=0("不设
|
||||
# 超时")会让挂起命令永久阻塞调用线程——下限夹取到 1s 快速失败。
|
||||
# 流式路径保留 0=不限语义(registry 兜底在位)
|
||||
try:
|
||||
timeout = max(1, int(timeout))
|
||||
except (TypeError, ValueError):
|
||||
timeout = 30
|
||||
|
||||
# 默认工作目录为workspace,避免在启动目录散落文件
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
self.logger.info(f"执行命令: {command} (work_dir={work_dir})")
|
||||
|
||||
# 准备环境变量(T21:黑名单键拒绝,防加载器/解释器劫持)
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
safe_env, rejected_keys = filter_env_overrides(env)
|
||||
if rejected_keys:
|
||||
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||
exec_env.update(safe_env)
|
||||
|
||||
process = None
|
||||
token = None
|
||||
pgid = None
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
stdin=subprocess.DEVNULL, # H3 修复:不挂 stdin,防等待输入命令永久挂起
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
encoding="utf-8", # T16:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出
|
||||
errors="replace", # T16:非法字节不抛异常中断 reader 线程
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
# T5:注册本执行专属槽位,终止/取消只作用于自己的 (proc, pgid)
|
||||
token, pgid = self._register_proc(process)
|
||||
|
||||
start_time = time.time()
|
||||
output_str, timeout_occurred, truncated = self._read_with_limit(
|
||||
process, timeout, start_time, pgid=pgid)
|
||||
|
||||
# 确保进程终止
|
||||
if process.poll() is None:
|
||||
self._terminate_tree(process, force=timeout_occurred or truncated, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
|
||||
return_code = process.returncode if process.returncode is not None else -1
|
||||
|
||||
data = {
|
||||
"输出": output_str.strip() if output_str.strip() else "(无输出)",
|
||||
"返回码": return_code
|
||||
}
|
||||
|
||||
if timeout_occurred:
|
||||
return self.format_result(False, f"命令执行超时({timeout}秒)", data)
|
||||
|
||||
if truncated:
|
||||
data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断"
|
||||
return self.format_result(False, "命令执行完成(输出已截断)", data)
|
||||
|
||||
return self.format_result(return_code == 0, "命令执行完成", data)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return f"❌ 命令执行超时({timeout}秒)"
|
||||
except FileNotFoundError as e:
|
||||
return f"❌ 命令执行错误: 命令未找到 - {str(e)}"
|
||||
except PermissionError as e:
|
||||
return f"❌ 命令执行错误: 权限不足 - {str(e)}"
|
||||
except Exception as e:
|
||||
return f"❌ 执行错误: {str(e)}"
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
# T5:仅注销自己的槽位,不再无差别清空共享引用
|
||||
if token is not None:
|
||||
self._unregister_proc(token)
|
||||
|
||||
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30,
|
||||
work_dir: str = None, env: Dict = None, **kwargs):
|
||||
"""流式执行命令,支持超时控制、工作目录和输出大小限制"""
|
||||
if not self.is_enabled():
|
||||
output_queue.put(("line", "❌ 命令执行功能已禁用\n"))
|
||||
output_queue.put(("complete", "❌ 命令执行功能已禁用"))
|
||||
return
|
||||
|
||||
# 默认工作目录为workspace,避免在启动目录散落文件
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
self.logger.info(f"流式执行命令: {command} (work_dir={work_dir})")
|
||||
output_queue.put(("line", f"\n🔧 执行命令: {command}\n"))
|
||||
output_queue.put(("line", f"⏱️ 超时设置: {timeout} 秒\n"))
|
||||
output_queue.put(("line", f"📦 输出限制: {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB\n"))
|
||||
output_queue.put(("line", f"📁 工作目录: {work_dir}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 准备环境变量(T21:黑名单键拒绝,防加载器/解释器劫持)
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
safe_env, rejected_keys = filter_env_overrides(env)
|
||||
if rejected_keys:
|
||||
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||
exec_env.update(safe_env)
|
||||
|
||||
process = None
|
||||
token = None
|
||||
pgid = None
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
stdin=subprocess.DEVNULL, # H3 修复:不挂 stdin,防等待输入命令永久挂起
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
encoding="utf-8", # T16:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出
|
||||
errors="replace", # T16:非法字节不抛异常中断 reader 线程
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
self.process = process
|
||||
# T5:注册本执行专属槽位,终止/取消只作用于自己的 (proc, pgid)
|
||||
token, pgid = self._register_proc(process)
|
||||
|
||||
start_time = time.time()
|
||||
# H2 修复:复用线程化读取(reader 线程 + 超时轮询),行经 output_queue 流式转发
|
||||
output_str, timeout_occurred, truncated = self._read_with_limit(
|
||||
process, timeout, start_time, output_queue=output_queue, pgid=pgid)
|
||||
|
||||
if timeout_occurred:
|
||||
output_queue.put(("line", f"\n⏰ 命令执行超时({timeout}秒)\n"))
|
||||
elif truncated:
|
||||
output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断\n"))
|
||||
|
||||
# 如果超时或截断,确保进程被终止
|
||||
if (timeout_occurred or truncated) and process.poll() is None:
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
pass
|
||||
|
||||
return_code = process.returncode if process.returncode is not None else -1
|
||||
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
if timeout_occurred:
|
||||
output_queue.put(("line", f"⏰ 命令执行超时({timeout}秒),返回码: {return_code}\n"))
|
||||
elif truncated:
|
||||
output_queue.put(("line", f"📦 输出已截断,返回码: {return_code}\n"))
|
||||
else:
|
||||
output_queue.put(("line", f"✅ 命令执行完成,返回码: {return_code}\n"))
|
||||
|
||||
data = {
|
||||
"输出": output_str.strip() if output_str.strip() else "(无输出)",
|
||||
"返回码": return_code
|
||||
}
|
||||
|
||||
if timeout_occurred:
|
||||
result = self.format_result(False, f"命令执行超时({timeout}秒)", data)
|
||||
elif truncated:
|
||||
data["警告"] = f"输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断"
|
||||
result = self.format_result(False, "命令执行完成(输出已截断)", data)
|
||||
else:
|
||||
result = self.format_result(return_code == 0, "命令执行完成", data)
|
||||
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
except Exception as e:
|
||||
output_queue.put(("line", f"❌ 执行错误: {str(e)}\n"))
|
||||
output_queue.put(("complete", f"❌ 执行错误: {str(e)}"))
|
||||
finally:
|
||||
if process and process.poll() is None:
|
||||
self._terminate_tree(process, pgid=pgid)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
self._terminate_tree(process, force=True, pgid=pgid)
|
||||
self.process = None
|
||||
# T5:仅注销自己的槽位,不再无差别清空共享引用
|
||||
if token is not None:
|
||||
self._unregister_proc(token)
|
||||
@@ -0,0 +1,779 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.process_manager - 后台进程管理器
|
||||
|
||||
让智能体能够启动并管理长驻后台进程(dev server / 训练任务 / watch 等),
|
||||
弥补 run_command 同步阻塞、超时即杀的不足:
|
||||
- bg 语义:start() 立即返回句柄,进程独立于对话轮次存活
|
||||
- 有界环形输出缓冲:每进程上限 buffer_per_process(默认 256KB),
|
||||
解决 nohup+日志文件方案输出无界增长的问题
|
||||
- 注册表:自增句柄跨会话/跨子代理可见(全局单例)
|
||||
- 生命周期:惰性状态收割(reader 线程 EOF 即 poll 收割防僵尸)+
|
||||
TTL 惰性清理 + 退出统一收尾(非 detach 进程默认全杀)
|
||||
|
||||
安全边界(对齐 SubprocessCommandExecutor 既有惯例):
|
||||
- 输出单行截断 16KB,防超长单行撑爆内存
|
||||
- 运行中进程数上限(默认 16)+ 总条目上限(默认 32,已退出条目
|
||||
保留尾部输出供验尸,LRU 淘汰)
|
||||
- POSIX start_new_session + 进程组信号(L9 语义:孙进程不悬挂);
|
||||
Windows 无进程组概念,回退单进程 terminate/kill
|
||||
- 中断语义:用户 Ctrl+C / 取消事件不影响后台进程(后台的本意就是
|
||||
脱离本轮对话生命周期);aibox 进程退出时经 atexit 统一收尾,
|
||||
detach=true 的进程豁免
|
||||
|
||||
依赖:common (ConfigManager, constants, Logger) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import collections
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from common import ConfigManager, constants, Logger
|
||||
from subsystems.executors import filter_env_overrides
|
||||
|
||||
# 单行输出最大字节数,超出截断(防超长单行撑爆环形缓冲内存预算)
|
||||
MAX_LINE_BYTES = 16 * 1024
|
||||
|
||||
# 单次 bg_output 拉取的硬上限(对齐执行器 1MB 输出限制)
|
||||
MAX_FETCH_BYTES = 1 * 1024 * 1024
|
||||
|
||||
|
||||
class _ProcEntry:
|
||||
"""单个后台进程的注册表条目(内部类)"""
|
||||
|
||||
__slots__ = (
|
||||
"handle", "name", "command", "popen", "pgid", "started_at",
|
||||
"detach", "max_lifetime", "owner", "status", "exit_code", "end_time",
|
||||
"lines", "buffer_bytes", "max_buffer", "next_seq", "evicted",
|
||||
"lock", "reader_done", "reader", "stdin_lock", "log_fh",
|
||||
)
|
||||
|
||||
# 状态取值:running / exited(自然退出)/ killed(被管理器终止)/ expired(TTL 到期)
|
||||
def __init__(self, handle: int, name: str, command: str, popen, pgid,
|
||||
detach: bool, max_lifetime: float, buffer_bytes: int,
|
||||
owner: Optional[str] = None):
|
||||
self.handle = handle
|
||||
self.name = name
|
||||
self.command = command
|
||||
self.popen = popen
|
||||
self.pgid = pgid
|
||||
self.started_at = time.time()
|
||||
self.detach = detach
|
||||
self.max_lifetime = max_lifetime
|
||||
self.owner = owner
|
||||
self.status = "running"
|
||||
self.exit_code = None
|
||||
self.end_time = None
|
||||
# 环形输出缓冲:deque of (seq, line, nbytes),seq 单调递增不回收,
|
||||
# 淘汰最旧行仍保持 seq 连续性——bg_output 的 since 游标据此续读
|
||||
self.lines = collections.deque()
|
||||
self.buffer_bytes = 0
|
||||
self.max_buffer = max(1024, int(buffer_bytes))
|
||||
self.next_seq = 1
|
||||
self.evicted = 0
|
||||
self.lock = threading.RLock() # T8:可重入,write_input 持锁期间可调用 _record_line
|
||||
self.reader_done = threading.Event()
|
||||
self.reader = None
|
||||
# 交互写入:单写锁保序(bg_write 串行化,防交错输入)
|
||||
self.stdin_lock = threading.Lock()
|
||||
# 可选全档镜像文件句柄(log_file=True 时由 reader 线程 tee 写入)
|
||||
self.log_fh = None
|
||||
|
||||
|
||||
class ProcessManager:
|
||||
"""后台进程管理器:启动 / 列举 / 增量拉取输出 / 终止 / 退出收尾"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None,
|
||||
max_processes: int = None, max_entries: int = None,
|
||||
buffer_per_process: int = None, max_lifetime: float = None):
|
||||
self.logger = logger or Logger("ProcessManager")
|
||||
self.config = ConfigManager()
|
||||
self.max_processes = int(max_processes if max_processes is not None
|
||||
else constants.BG_MAX_PROCESSES)
|
||||
self.max_entries = int(max_entries if max_entries is not None
|
||||
else constants.BG_MAX_ENTRIES)
|
||||
self.buffer_per_process = int(buffer_per_process if buffer_per_process is not None
|
||||
else constants.BG_BUFFER_PER_PROCESS)
|
||||
self.max_lifetime = float(max_lifetime if max_lifetime is not None
|
||||
else constants.BG_MAX_LIFETIME)
|
||||
self.max_write_bytes = int(constants.BG_MAX_WRITE_BYTES)
|
||||
self._procs: Dict[int, _ProcEntry] = {}
|
||||
self._next_id = 1
|
||||
self._lock = threading.Lock()
|
||||
self._shutdown_done = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 启动
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self, command: str, name: str = None, work_dir: str = None,
|
||||
env: Dict = None, detach: bool = False,
|
||||
log_file: bool = False, owner: str = None) -> Dict:
|
||||
"""后台启动命令,立即返回。成功返回 {ok, handle, pid, name, log_file},
|
||||
失败返回 {ok: False, error}。
|
||||
|
||||
stdin 为管道:后续可经 write() 交互写入(行协议 REPL 问答)。
|
||||
log_file=True 时输出同步 tee 到 workspace/bg_<句柄>.log(无界
|
||||
全档,可被原生工具 tail/cat 读取;环形缓冲仍为 bg_output 视图)。"""
|
||||
if not command or not isinstance(command, str) or not command.strip():
|
||||
return {"ok": False, "error": "缺少或空的 command"}
|
||||
|
||||
# 名称规整(M5 惯例):标签仅用于展示回显,拒绝控制字符防终端/前端注入
|
||||
if name is not None:
|
||||
name = str(name).strip()
|
||||
if not name.isprintable() or len(name) > 64:
|
||||
return {"ok": False, "error": "name 含控制字符或超过 64 字符"}
|
||||
if not name:
|
||||
name = None
|
||||
if name is None:
|
||||
name = command.strip()[:40]
|
||||
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
with self._lock:
|
||||
if self._shutdown_done:
|
||||
return {"ok": False, "error": "管理器已关闭"}
|
||||
# TTL 惰性清理 + 状态收割
|
||||
self._sweep_locked()
|
||||
running = [e for e in self._procs.values() if e.status == "running"]
|
||||
if len(running) >= self.max_processes:
|
||||
return {"ok": False,
|
||||
"error": f"运行中后台进程已达上限({self.max_processes}),"
|
||||
f"请先用 bg_kill 清理"}
|
||||
# 总条目上限:淘汰最老的已退出条目(保留其信息的价值最低)
|
||||
if len(self._procs) >= self.max_entries:
|
||||
exited = sorted(
|
||||
(e for e in self._procs.values() if e.status != "running"),
|
||||
key=lambda e: e.started_at)
|
||||
while self._procs and len(self._procs) >= self.max_entries and exited:
|
||||
victim = exited.pop(0)
|
||||
self._procs.pop(victim.handle, None)
|
||||
if len(self._procs) >= self.max_entries:
|
||||
return {"ok": False,
|
||||
"error": f"后台进程条目已达上限({self.max_entries})且全部运行中"}
|
||||
handle = self._next_id
|
||||
self._next_id += 1
|
||||
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
# T21:黑名单键拒绝(对齐 executors,防 LD_PRELOAD/PYTHONPATH 劫持)
|
||||
safe_env, rejected_keys = filter_env_overrides(env)
|
||||
if rejected_keys:
|
||||
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||
exec_env.update(safe_env)
|
||||
|
||||
try:
|
||||
popen = subprocess.Popen(
|
||||
command,
|
||||
shell=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
# 交互式 I/O:stdin 挂管道供 write() 写入。与 H3 的差异:
|
||||
# 后台进程由管理器显式管理生命周期,不存在"等待输入的
|
||||
# 命令永久挂起无人知晓"问题——不读 stdin 的进程不受影响,
|
||||
# 读 stdin 的进程获得交互能力,kill 路径三级优雅终止兜底。
|
||||
stdin=subprocess.PIPE,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
universal_newlines=True,
|
||||
encoding="utf-8", # T17:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出致 reader 静默退出
|
||||
errors="replace",
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# 对齐 L9:POSIX 下 shell 成为会话首进程,kill 按进程组
|
||||
# 一次覆盖整棵进程树(含 & 后台孙进程)
|
||||
start_new_session=(os.name == "posix"),
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
return {"ok": False, "error": f"命令未找到: {e}"}
|
||||
except PermissionError as e:
|
||||
return {"ok": False, "error": f"权限不足: {e}"}
|
||||
except Exception as e:
|
||||
return {"ok": False, "error": f"启动失败: {e}"}
|
||||
|
||||
pgid = None
|
||||
if os.name == "posix":
|
||||
try:
|
||||
pgid = os.getpgid(popen.pid)
|
||||
except Exception:
|
||||
pgid = None
|
||||
|
||||
entry = _ProcEntry(
|
||||
handle, name, command.strip(), popen, pgid,
|
||||
detach=bool(detach), max_lifetime=self.max_lifetime,
|
||||
buffer_bytes=self.buffer_per_process, owner=owner,
|
||||
)
|
||||
# 可选全档镜像(失败降级为仅环形缓冲,不阻断启动)
|
||||
log_path = None
|
||||
if log_file:
|
||||
try:
|
||||
log_dir = constants.WORKSPACE_DIR
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
log_path = os.path.join(log_dir, f"bg_{handle}.log")
|
||||
entry.log_fh = open(log_path, "a", encoding="utf-8",
|
||||
errors="replace", buffering=1)
|
||||
except Exception as e:
|
||||
entry.log_fh = None
|
||||
log_path = None
|
||||
self.logger.warning(f"输出镜像文件开启失败,仅保留环形缓冲: {e}")
|
||||
with self._lock:
|
||||
# start() 竞态兜底:期间并发 shutdown 则立即回收刚拉起的进程
|
||||
if self._shutdown_done:
|
||||
self._close_stdin(entry)
|
||||
self._terminate_tree(entry, force=True)
|
||||
try:
|
||||
popen.wait(timeout=2)
|
||||
except Exception:
|
||||
pass
|
||||
return {"ok": False, "error": "管理器已关闭"}
|
||||
self._procs[handle] = entry
|
||||
|
||||
entry.reader = threading.Thread(
|
||||
target=self._reader, args=(entry,), daemon=True,
|
||||
name=f"bg-proc-{handle}-reader")
|
||||
entry.reader.start()
|
||||
self.logger.info(f"后台进程已启动: #{handle} pid={popen.pid} "
|
||||
f"detach={entry.detach} cmd={entry.command[:60]}")
|
||||
return {"ok": True, "handle": handle, "pid": popen.pid, "name": name,
|
||||
"detach": entry.detach, "log_file": log_path}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 输出读取
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_output(self, handle, tail: int = None, since: int = None,
|
||||
max_bytes: int = 65536) -> Dict:
|
||||
"""拉取后台进程输出。
|
||||
|
||||
- tail=N:取最近 N 行(缺省 50)
|
||||
- since=S:取行序号 > S 的全部缓冲行(增量续读,游标为上次
|
||||
返回的 next_cursor;被环形淘汰的早期行以 evicted 提示)
|
||||
- max_bytes:本次拉取字节上限(缺省 64KB,硬顶 1MB)
|
||||
"""
|
||||
entry = self._lookup(handle)
|
||||
if entry is None:
|
||||
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||
|
||||
try:
|
||||
max_bytes = int(max_bytes or 65536)
|
||||
except (TypeError, ValueError):
|
||||
max_bytes = 65536
|
||||
max_bytes = max(1024, min(max_bytes, MAX_FETCH_BYTES))
|
||||
|
||||
with entry.lock:
|
||||
snapshot = list(entry.lines)
|
||||
evicted = entry.evicted
|
||||
next_seq = entry.next_seq
|
||||
|
||||
selected = []
|
||||
evicted_note = 0
|
||||
if since is not None:
|
||||
try:
|
||||
since = int(since)
|
||||
except (TypeError, ValueError):
|
||||
return {"ok": False, "error": "since 游标必须是整数"}
|
||||
selected = [(s, l) for (s, l, _) in snapshot if s > since]
|
||||
# 请求范围早于环形缓冲起点:被淘汰部分如实告知
|
||||
oldest = snapshot[0][0] if snapshot else next_seq
|
||||
if since + 1 < oldest:
|
||||
evicted_note = oldest - since - 1
|
||||
else:
|
||||
try:
|
||||
tail = int(tail) if tail is not None else 50
|
||||
except (TypeError, ValueError):
|
||||
tail = 50
|
||||
tail = max(1, min(tail, 100000))
|
||||
selected = [(s, l) for (s, l, _) in snapshot[-tail:]]
|
||||
|
||||
# 字节上限截断(保留最新行优先:从尾部回溯收集)
|
||||
kept = []
|
||||
total = 0
|
||||
truncated = False
|
||||
for s, line in reversed(selected):
|
||||
nb = len(line.encode("utf-8", errors="replace"))
|
||||
if total + nb > max_bytes:
|
||||
truncated = True
|
||||
break
|
||||
kept.append((s, line))
|
||||
total += nb
|
||||
kept.reverse()
|
||||
|
||||
output = "".join(line for _, line in kept)
|
||||
result = {
|
||||
"ok": True,
|
||||
"handle": entry.handle,
|
||||
"name": entry.name,
|
||||
"status": entry.status,
|
||||
"exit_code": entry.exit_code,
|
||||
"output": output,
|
||||
"lines": len(kept),
|
||||
"next_cursor": kept[-1][0] if kept else max(since or 0, next_seq - 1),
|
||||
"truncated": truncated,
|
||||
}
|
||||
if evicted_note > 0 or (since is None and evicted > 0):
|
||||
result["evicted"] = evicted_note if since is not None else evicted
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 交互写入(stdin 管道 + watermark 对齐)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def write(self, handle, text: str, append_newline: bool = True,
|
||||
wait_lines: int = 0, wait_timeout: float = 5.0,
|
||||
write_timeout: float = 5.0, owner: str = None) -> Dict:
|
||||
"""向后台进程 stdin 写入一行(交互式问答)。
|
||||
|
||||
返回 {ok, handle, watermark, sent_bytes[, wait]}:
|
||||
- watermark:本条输入的 [stdin] 回显行序号;bg_output(since=
|
||||
watermark) 得到的必然且仅是本条输入引发的响应(写读对齐)
|
||||
- wait_lines>0 时同步等待 N 行响应(wait 子 dict 含 output/
|
||||
timed_out),expect 式一步问答
|
||||
失败返回 {ok: False, error}(进程已退出/写入超时/超长等)。
|
||||
|
||||
T22 属主隔离:条目与调用方都携带 owner 且不一致时拒绝——他人
|
||||
stdin 不容注入。任一侧未标识(None,如测试/管理面直调)不隔离。"""
|
||||
entry = self._lookup(handle)
|
||||
if entry is None:
|
||||
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||
denial = self._check_owner(entry, owner, "写入")
|
||||
if denial:
|
||||
return denial
|
||||
if not isinstance(text, str) or text == "":
|
||||
return {"ok": False, "error": "缺少或空的 text"}
|
||||
data = text if (not append_newline or text.endswith("\n")) else text + "\n"
|
||||
payload = data.encode("utf-8")
|
||||
if len(payload) > self.max_write_bytes:
|
||||
return {"ok": False,
|
||||
"error": f"单次写入超过 {self.max_write_bytes} 字节上限(超长内容建议落文件后让进程从文件读)"}
|
||||
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
if entry.status != "running":
|
||||
return {"ok": False, "error": f"进程已结束({entry.status}),无法写入"}
|
||||
|
||||
# watermark:取号与回显行记账必须在同一临界区,保证 watermark 等于
|
||||
# 回显行真实序号,bg_output(since=watermark) 正确对齐本条响应(T8)
|
||||
echo = "[stdin] " + text.replace("\n", "⏎")
|
||||
with entry.lock:
|
||||
watermark = entry.next_seq
|
||||
# 回显行进环形缓冲 + log 镜像(多行输入折叠为单行展示)
|
||||
self._record_line(entry, echo)
|
||||
|
||||
# 写入放 watchdog 线程:进程不读 stdin 且管道缓冲(~64KB)写满时
|
||||
# write() 会阻塞,不能拖死工具线程;超时后 daemon 线程由系统兜底
|
||||
error_box = {}
|
||||
done = threading.Event()
|
||||
|
||||
def _do_write():
|
||||
try:
|
||||
with entry.stdin_lock:
|
||||
stdin = entry.popen.stdin
|
||||
if stdin is None or stdin.closed:
|
||||
raise BrokenPipeError("stdin 已关闭")
|
||||
stdin.write(data)
|
||||
stdin.flush()
|
||||
except Exception as e:
|
||||
error_box["e"] = e
|
||||
finally:
|
||||
done.set()
|
||||
|
||||
t = threading.Thread(target=_do_write, daemon=True,
|
||||
name=f"bg-proc-{entry.handle}-writer")
|
||||
t.start()
|
||||
try:
|
||||
write_timeout = max(0.1, min(60.0, float(write_timeout)))
|
||||
except (TypeError, ValueError):
|
||||
write_timeout = 5.0
|
||||
if not done.wait(timeout=write_timeout):
|
||||
return {"ok": False,
|
||||
"error": f"写入超时(>{write_timeout}s):进程可能不读取 stdin,管道缓冲已满"}
|
||||
if error_box:
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
return {"ok": False,
|
||||
"error": f"写入失败: {error_box['e']}"}
|
||||
|
||||
result = {"ok": True, "handle": entry.handle, "watermark": watermark,
|
||||
"sent_bytes": len(payload)}
|
||||
if wait_lines and int(wait_lines) > 0:
|
||||
result["wait"] = self._wait_lines(entry, watermark,
|
||||
int(wait_lines), wait_timeout)
|
||||
return result
|
||||
|
||||
def _wait_lines(self, entry: _ProcEntry, watermark: int, count: int,
|
||||
timeout: float) -> Dict:
|
||||
"""等待 watermark 之后出现 count 行新输出(进程退出且 reader 排空
|
||||
即提前返回,如实标记 timed_out)。"""
|
||||
try:
|
||||
timeout = max(0.1, min(60.0, float(timeout)))
|
||||
except (TypeError, ValueError):
|
||||
timeout = 5.0
|
||||
deadline = time.time() + timeout
|
||||
|
||||
def _snap():
|
||||
with entry.lock:
|
||||
return [(s, l) for (s, l, _) in entry.lines if s > watermark]
|
||||
|
||||
lines = _snap()
|
||||
timed_out = False
|
||||
while len(lines) < count:
|
||||
if entry.popen.poll() is not None and entry.reader_done.is_set():
|
||||
time.sleep(0.1) # 排空 reader 尾部在途行后取终值
|
||||
lines = _snap()
|
||||
break
|
||||
if time.time() > deadline:
|
||||
timed_out = True
|
||||
break
|
||||
time.sleep(0.05)
|
||||
lines = _snap()
|
||||
|
||||
# 字节上限(复用单次拉取硬顶)
|
||||
kept, total, truncated = [], 0, False
|
||||
for s, line in lines:
|
||||
nb = len(line.encode("utf-8", errors="replace"))
|
||||
if total + nb > MAX_FETCH_BYTES:
|
||||
truncated = True
|
||||
break
|
||||
kept.append((s, line))
|
||||
total += nb
|
||||
out = {
|
||||
"lines": len(kept),
|
||||
"output": "".join(l for _, l in kept),
|
||||
"timed_out": timed_out,
|
||||
"next_cursor": kept[-1][0] if kept else watermark,
|
||||
}
|
||||
if truncated:
|
||||
out["truncated"] = True
|
||||
return out
|
||||
|
||||
def _close_stdin(self, entry: _ProcEntry):
|
||||
"""关闭 stdin 管道(行协议进程见 EOF 常自行退出;幂等)"""
|
||||
try:
|
||||
with entry.stdin_lock:
|
||||
stdin = entry.popen.stdin
|
||||
if stdin is not None and not stdin.closed:
|
||||
stdin.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 列举 / 终止 / 收尾
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def list_procs(self) -> List[Dict]:
|
||||
"""列出全部条目(运行中在前,其余按句柄序)"""
|
||||
with self._lock:
|
||||
self._sweep_locked()
|
||||
# 惰性收割:列举即刷新全部条目状态(此前仅 kill/get_output 路径更新)
|
||||
for entry in self._procs.values():
|
||||
self._update_status_locked(entry)
|
||||
entries = sorted(self._procs.values(),
|
||||
key=lambda e: (e.status != "running", e.handle))
|
||||
return [self._entry_info(e) for e in entries]
|
||||
|
||||
def kill(self, handle, force: bool = False, grace: float = 3.0,
|
||||
owner: str = None) -> Dict:
|
||||
"""终止后台进程。SIGTERM → 宽限轮询 → SIGKILL;force=True 直接 SIGKILL。
|
||||
对已退出条目幂等,返回当前状态。
|
||||
|
||||
T22 属主隔离:条目与调用方都携带 owner 且不一致时拒绝——防止
|
||||
提示注入的会话/任务跨属主杀掉他人的 dev server。shutdown/TTL
|
||||
清理属管理器特权,不走本校验。"""
|
||||
entry = self._lookup(handle)
|
||||
if entry is None:
|
||||
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||
denial = self._check_owner(entry, owner, "终止")
|
||||
if denial:
|
||||
return denial
|
||||
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
if entry.status != "running":
|
||||
return {"ok": True, "handle": entry.handle, "status": entry.status,
|
||||
"exit_code": entry.exit_code, "already_finished": True}
|
||||
|
||||
try:
|
||||
grace = min(30.0, max(0.1, float(grace)))
|
||||
except (TypeError, ValueError):
|
||||
grace = 3.0
|
||||
|
||||
if force:
|
||||
self._terminate_tree(entry, force=True)
|
||||
else:
|
||||
# 三级优雅终止:关 stdin(行协议进程见 EOF 常自行退出)→
|
||||
# SIGTERM 宽限 → SIGKILL 兜底
|
||||
self._close_stdin(entry)
|
||||
deadline = time.time() + grace * 0.4
|
||||
while time.time() < deadline and entry.popen.poll() is None:
|
||||
time.sleep(0.05)
|
||||
if entry.popen.poll() is None:
|
||||
self._terminate_tree(entry)
|
||||
deadline = time.time() + grace * 0.6
|
||||
while time.time() < deadline and entry.popen.poll() is None:
|
||||
time.sleep(0.05)
|
||||
if entry.popen.poll() is None:
|
||||
self._terminate_tree(entry, force=True)
|
||||
try:
|
||||
entry.popen.wait(timeout=2)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
if entry.status == "running":
|
||||
# 进程仍未死(极端:不可杀的 D 状态),如实报告
|
||||
return {"ok": False, "handle": entry.handle,
|
||||
"error": "进程未能终止(可能处于不可中断状态)"}
|
||||
entry.status = "killed"
|
||||
return {"ok": True, "handle": entry.handle, "status": entry.status,
|
||||
"exit_code": entry.exit_code}
|
||||
|
||||
def shutdown(self) -> Dict:
|
||||
"""退出收尾:终止全部非 detach 运行中进程(SIGTERM→SIGKILL),
|
||||
清空注册表;detach 进程保留运行、仅注销登记。幂等。"""
|
||||
with self._lock:
|
||||
if self._shutdown_done:
|
||||
return {"killed": 0, "detached": 0}
|
||||
self._shutdown_done = True
|
||||
entries = list(self._procs.values())
|
||||
self._procs.clear()
|
||||
|
||||
killed = detached = 0
|
||||
for entry in entries:
|
||||
self._update_status(entry)
|
||||
if entry.status != "running":
|
||||
continue
|
||||
if entry.detach:
|
||||
detached += 1
|
||||
self.logger.info(f"后台进程 #{entry.handle}(detach)保留运行: "
|
||||
f"pid={entry.popen.pid}")
|
||||
continue
|
||||
# 优雅收尾:先关 stdin(行协议进程见 EOF 常自行退出)再走信号
|
||||
self._close_stdin(entry)
|
||||
try:
|
||||
entry.popen.wait(timeout=1)
|
||||
except Exception:
|
||||
pass
|
||||
if entry.popen.poll() is None:
|
||||
self._terminate_tree(entry)
|
||||
try:
|
||||
entry.popen.wait(timeout=3)
|
||||
except Exception:
|
||||
self._terminate_tree(entry, force=True)
|
||||
try:
|
||||
entry.popen.wait(timeout=2)
|
||||
except Exception:
|
||||
pass
|
||||
entry.status = "killed"
|
||||
killed += 1
|
||||
if killed or detached:
|
||||
self.logger.info(f"后台进程收尾: 终止 {killed} 个, detach 保留 {detached} 个")
|
||||
return {"killed": killed, "detached": detached}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部实现
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _check_owner(self, entry: _ProcEntry, caller_owner: Optional[str],
|
||||
action: str) -> Optional[Dict]:
|
||||
"""T22:属主校验——条目与调用方都携带 owner 且不一致时返回拒绝
|
||||
dict(写/杀权限隔离);任一侧未标识(None,如测试/管理面直调)
|
||||
不隔离,保持向后兼容。观察面(list/get_output)不经过本校验。"""
|
||||
if (entry.owner is not None and caller_owner is not None
|
||||
and entry.owner != caller_owner):
|
||||
self.logger.warning(
|
||||
f"属主校验拒绝:owner={caller_owner} 尝试{action} "
|
||||
f"#{entry.handle}(属主 {entry.owner})")
|
||||
return {"ok": False,
|
||||
"error": f"进程 #{entry.handle} 属于其他会话(属主不匹配),无权{action}"}
|
||||
return None
|
||||
|
||||
def _lookup(self, handle) -> Optional[_ProcEntry]:
|
||||
try:
|
||||
handle = int(handle)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
with self._lock:
|
||||
entry = self._procs.get(handle)
|
||||
if entry is not None:
|
||||
self._update_status(entry)
|
||||
return entry
|
||||
|
||||
def _update_status(self, entry: _ProcEntry):
|
||||
with self._lock:
|
||||
self._update_status_locked(entry)
|
||||
|
||||
def _update_status_locked(self, entry: _ProcEntry):
|
||||
"""惰性收割:poll() 即 waitpid(WNOHANG),自然退出时记录状态。"""
|
||||
if entry.status != "running":
|
||||
return
|
||||
rc = entry.popen.poll()
|
||||
if rc is not None:
|
||||
entry.exit_code = rc
|
||||
entry.end_time = time.time()
|
||||
entry.status = "exited"
|
||||
|
||||
def _sweep_locked(self):
|
||||
"""TTL 惰性清理:超期运行中进程标记 expired 并终止(max_lifetime<=0 不启用)。
|
||||
调用方须持有 self._lock。"""
|
||||
if self.max_lifetime <= 0:
|
||||
return
|
||||
now = time.time()
|
||||
for entry in list(self._procs.values()):
|
||||
if entry.status == "running" and now - entry.started_at > self.max_lifetime:
|
||||
self._terminate_tree(entry)
|
||||
entry.status = "expired"
|
||||
entry.end_time = now
|
||||
entry.exit_code = entry.popen.poll()
|
||||
self.logger.info(f"后台进程 #{entry.handle} TTL 到期已终止")
|
||||
|
||||
def _terminate_tree(self, entry: _ProcEntry, force: bool = False):
|
||||
"""按进程组终止整棵进程树(对齐 executors L9/T11 语义与 Windows 回退)"""
|
||||
proc = entry.popen
|
||||
try:
|
||||
if os.name == "posix":
|
||||
sig = signal.SIGKILL if force else signal.SIGTERM
|
||||
pgid = entry.pgid
|
||||
if pgid is None:
|
||||
pgid = os.getpgid(proc.pid)
|
||||
os.killpg(pgid, sig)
|
||||
else:
|
||||
# T11(审计 2026-08-21):Windows 无进程组,单进程 kill 收割
|
||||
# 不到 spawn 的孙进程(dev server/训练任务成孤儿继续占用端口
|
||||
# /文件锁)。taskkill /T /F 递归终止整棵树,失败兜底单进程 kill
|
||||
try:
|
||||
subprocess.run(
|
||||
["taskkill", "/PID", str(proc.pid), "/T", "/F"],
|
||||
capture_output=True, timeout=10)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _reader(self, entry: _ProcEntry):
|
||||
"""后台 reader 线程(对齐 executors H2 模式):阻塞 readline 入环形
|
||||
缓冲(+ 可选 log 镜像);EOF/异常即 poll() 收割防僵尸,并显式关闭
|
||||
管道与镜像文件句柄。"""
|
||||
try:
|
||||
for line in iter(entry.popen.stdout.readline, ""):
|
||||
self._record_line(entry, line)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
entry.reader_done.set()
|
||||
try:
|
||||
entry.popen.poll() # 立即收割退出子进程,防僵尸悬挂
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
entry.popen.stdout.close()
|
||||
except Exception:
|
||||
pass
|
||||
if entry.log_fh is not None:
|
||||
try:
|
||||
entry.log_fh.flush()
|
||||
entry.log_fh.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_line(self, entry: _ProcEntry, line: str):
|
||||
"""统一记录:环形缓冲 + 可选 log 全档镜像(stdin 回显同样入档)"""
|
||||
if entry.log_fh is not None:
|
||||
try:
|
||||
entry.log_fh.write(line)
|
||||
except Exception:
|
||||
pass
|
||||
self._append_line(entry, line)
|
||||
|
||||
@staticmethod
|
||||
def _append_line(entry: _ProcEntry, line: str):
|
||||
# 单行截断:环形缓冲内存预算不被超长单行击穿
|
||||
nbytes = len(line.encode("utf-8", errors="replace"))
|
||||
if nbytes > MAX_LINE_BYTES:
|
||||
allowed = line.encode("utf-8", errors="replace")[:MAX_LINE_BYTES]
|
||||
line = allowed.decode("utf-8", errors="ignore") + "…[单行截断]\n"
|
||||
nbytes = len(line.encode("utf-8", errors="replace"))
|
||||
with entry.lock:
|
||||
entry.lines.append((entry.next_seq, line, nbytes))
|
||||
entry.next_seq += 1
|
||||
entry.buffer_bytes += nbytes
|
||||
while entry.buffer_bytes > entry.max_buffer and len(entry.lines) > 1:
|
||||
_, old, old_nb = entry.lines.popleft()
|
||||
entry.buffer_bytes -= old_nb
|
||||
entry.evicted += 1
|
||||
|
||||
@staticmethod
|
||||
def _entry_info(entry: _ProcEntry) -> Dict:
|
||||
now = time.time()
|
||||
return {
|
||||
"handle": entry.handle,
|
||||
"name": entry.name,
|
||||
"pid": entry.popen.pid,
|
||||
"status": entry.status,
|
||||
"exit_code": entry.exit_code,
|
||||
"detach": entry.detach,
|
||||
"started_at": entry.started_at,
|
||||
"runtime_sec": round((entry.end_time or now) - entry.started_at, 1),
|
||||
"buffered_lines": len(entry.lines),
|
||||
"evicted_lines": entry.evicted,
|
||||
"command": entry.command,
|
||||
"owner": entry.owner, # T22:属主标识(观察面可见,写/杀需匹配)
|
||||
}
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 模块级单例(工具层共享;测试经 reset_process_manager() 重置)
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
_manager: Optional[ProcessManager] = None
|
||||
_manager_lock = threading.Lock()
|
||||
_atexit_registered = False
|
||||
|
||||
|
||||
def get_process_manager(logger: Optional[Logger] = None) -> ProcessManager:
|
||||
"""获取全局 ProcessManager 单例(懒加载)。
|
||||
|
||||
后台进程注册表必须跨会话/跨子代理共享——任何 worker 子代理启动的
|
||||
进程,后续任意子代理都能 bg_list/bg_output/bg_kill 管理。"""
|
||||
global _manager, _atexit_registered
|
||||
with _manager_lock:
|
||||
if _manager is None:
|
||||
_manager = ProcessManager(logger)
|
||||
if not _atexit_registered:
|
||||
# 退出语义:aibox 进程退出(含 --once sys.exit / Ctrl+C)时
|
||||
# 非 detach 后台进程统一终止;detach 进程保留运行仅注销登记。
|
||||
# SIGKILL/SIGTERM 硬杀无法覆盖(与既有线程池等资源同级)。
|
||||
atexit.register(_shutdown_atexit)
|
||||
_atexit_registered = True
|
||||
return _manager
|
||||
|
||||
|
||||
def reset_process_manager(manager: Optional[ProcessManager] = None) -> None:
|
||||
"""替换/清空全局单例(测试隔离用)。传入 None 仅解除引用,不触发收尾。"""
|
||||
global _manager
|
||||
with _manager_lock:
|
||||
_manager = manager
|
||||
|
||||
|
||||
def _shutdown_atexit():
|
||||
m = _manager
|
||||
if m is not None:
|
||||
try:
|
||||
m.shutdown()
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,330 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.prompt_manager - 提示词管理器
|
||||
|
||||
抽离自 aibox.py,负责从 prompts 目录加载系统提示词。
|
||||
单例模式(__new__ + 类变量 _instance),多实例共享同一份 _prompts。
|
||||
|
||||
文件命名规范:
|
||||
- 00_IDENTITY.md : 用户身份信息
|
||||
- N_*.md/txt : 基础技能提示词(N 为数字),按数字升序加载
|
||||
- 99_SOUL.md : 智能体灵魂文件(角色定位、身份定义、性格等)
|
||||
- 100_STATUS.md : 智能体状态记录
|
||||
|
||||
依赖:common (ConfigManager, constants) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import threading
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from common import ConfigManager
|
||||
from common.atomic import atomic_write_text
|
||||
|
||||
|
||||
class PromptManager:
|
||||
"""提示词管理器 - 负责从 prompts 目录加载系统提示词
|
||||
|
||||
文件命名规范:
|
||||
- 00_IDENTITY.md: 用户身份信息
|
||||
- N_*.md/txt: 基础技能提示词(N 为数字,如 1_memory_search.md、10_task_handling.md),按数字升序加载
|
||||
- 99_SOUL.md: 智能体灵魂文件(角色定位、身份定义、性格等)
|
||||
- 100_STATUS.md: 智能体状态记录
|
||||
"""
|
||||
|
||||
_instance = None
|
||||
_instance_lock = threading.Lock() # I 批:并发首例化竞态防护
|
||||
_prompts = []
|
||||
_prompts_dir = None
|
||||
|
||||
# 核心文件
|
||||
IDENTITY_FILE = "00_IDENTITY.md"
|
||||
STATUS_FILE = "100_STATUS.md"
|
||||
SOUL_FILE = "99_SOUL.md" # 智能体灵魂文件 - 所有角色定位内容
|
||||
|
||||
# N9(审计 2026-08-18):Agent 可控内容直写 00/99/100 文件并作为后续
|
||||
# 会话系统提示词加载(持久化注入面)。内容长度上限约束注入规模,
|
||||
# 超长内容拒绝写入(工具层有更友好的提示)。
|
||||
MAX_CONTENT_LENGTH = 10000
|
||||
|
||||
# 基础技能文件名模式:N_xxx.md(N 为数字,如 1_memory_search.md、10_task_handling.md)
|
||||
_SKILL_FILE_RE = re.compile(r'^(\d+)_')
|
||||
|
||||
def __new__(cls):
|
||||
with cls._instance_lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
# J6(审计 2026-08-19):__new__ 的锁只防双实例,多线程首例并发
|
||||
# 同时进入 __init__ 会使 _load_prompts 双重执行,_prompts(类属性)
|
||||
# 整份翻倍——初始化检查与加载收进同一把锁
|
||||
with self._instance_lock:
|
||||
if not self._prompts:
|
||||
self._load_prompts()
|
||||
|
||||
def _get_prompts_dir(self) -> str:
|
||||
config = ConfigManager()
|
||||
prompts_dir = config.get("system.prompts_dir", "~/.aibox/prompts")
|
||||
return os.path.expanduser(prompts_dir)
|
||||
|
||||
def _load_prompts(self, target: Optional[List[Dict]] = None):
|
||||
"""加载提示词到 target 列表(缺省 self._prompts)。
|
||||
|
||||
M6:reload 传入新列表本地构建,完成后由 reload 原子替换——
|
||||
此前先清空 self._prompts 再逐文件重读,并发读者(其他连接构建
|
||||
system 消息)会读到空/半截提示词列表。
|
||||
"""
|
||||
prompts = target if target is not None else self._prompts
|
||||
self._prompts_dir = self._get_prompts_dir()
|
||||
os.makedirs(self._prompts_dir, exist_ok=True)
|
||||
|
||||
# 确保核心文件存在
|
||||
self._ensure_identity_file()
|
||||
self._ensure_status_file()
|
||||
self._ensure_soul_file() # 智能体灵魂文件
|
||||
|
||||
# 确保基础技能文件存在
|
||||
self._ensure_skill_files()
|
||||
|
||||
prompt_files = []
|
||||
if os.path.exists(self._prompts_dir):
|
||||
# 按文件名排序,确保1_开头的技能文件在99_SOUL.md之前加载
|
||||
prompt_files = sorted([f for f in os.listdir(self._prompts_dir)
|
||||
if os.path.isfile(os.path.join(self._prompts_dir, f))
|
||||
and not f.startswith('.')])
|
||||
|
||||
for filename in prompt_files:
|
||||
filepath = os.path.join(self._prompts_dir, filename)
|
||||
try:
|
||||
with open(filepath, 'r', encoding='utf-8') as f:
|
||||
content = f.read().strip()
|
||||
if content:
|
||||
prompts.append({
|
||||
'file': filename,
|
||||
'content': content
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取提示词文件 {filename} 失败: {e}")
|
||||
|
||||
if not prompts:
|
||||
self._create_default_prompts()
|
||||
|
||||
def _get_template_dir(self) -> str:
|
||||
"""获取仓库中的提示词模板目录(aibox/prompts/)"""
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
# subsystems/ 上一层是 aibox/,模板在其下的 prompts/
|
||||
aibox_dir = os.path.dirname(script_dir)
|
||||
return os.path.join(aibox_dir, 'prompts')
|
||||
|
||||
def _ensure_from_template(self, filename: str):
|
||||
"""从模板目录复制提示词文件到运行目录。
|
||||
|
||||
模板目录为 aibox/prompts/,若其中存在同名文件则复制过去;
|
||||
若模板不存在则跳过(由调用方保证模板齐全)。
|
||||
已存在的运行时文件不会被覆盖。
|
||||
"""
|
||||
target_path = os.path.join(self._prompts_dir, filename)
|
||||
if os.path.exists(target_path):
|
||||
return # 已存在,不覆盖
|
||||
|
||||
template_path = os.path.join(self._get_template_dir(), filename)
|
||||
if os.path.exists(template_path):
|
||||
try:
|
||||
shutil.copy2(template_path, target_path)
|
||||
print(f"✅ 已从模板创建: {filename}")
|
||||
except Exception as e:
|
||||
print(f"⚠️ 从模板复制 {filename} 失败: {e}")
|
||||
|
||||
def _ensure_identity_file(self):
|
||||
"""确保用户身份文件存在"""
|
||||
self._ensure_from_template(self.IDENTITY_FILE)
|
||||
|
||||
def _ensure_status_file(self):
|
||||
"""确保状态文件存在"""
|
||||
self._ensure_from_template(self.STATUS_FILE)
|
||||
|
||||
def _ensure_soul_file(self):
|
||||
"""确保智能体灵魂文件存在"""
|
||||
self._ensure_from_template(self.SOUL_FILE)
|
||||
|
||||
def _ensure_skill_files(self):
|
||||
"""确保基础技能文件存在 - 从模板目录复制所有 N_*.md 文件"""
|
||||
template_dir = self._get_template_dir()
|
||||
if not os.path.isdir(template_dir):
|
||||
return
|
||||
for filename in sorted(os.listdir(template_dir)):
|
||||
if self._SKILL_FILE_RE.match(filename): # N_ 开头的文件
|
||||
self._ensure_from_template(filename)
|
||||
|
||||
def update_identity(self, content: str) -> bool:
|
||||
"""更新用户身份信息"""
|
||||
if len(content or "") > self.MAX_CONTENT_LENGTH:
|
||||
print(f"⚠️ 身份内容超长({len(content)} > {self.MAX_CONTENT_LENGTH} 字符),拒绝写入")
|
||||
return False
|
||||
identity_path = os.path.join(self._prompts_dir, self.IDENTITY_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
updated_content = f"""# 用户身份信息
|
||||
|
||||
{content}
|
||||
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
# J5:原子写,崩溃不留半截提示词文件
|
||||
atomic_write_text(identity_path, updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 更新身份文件失败: {e}")
|
||||
return False
|
||||
|
||||
def update_status(self, content: str) -> bool:
|
||||
"""更新智能体状态"""
|
||||
if len(content or "") > self.MAX_CONTENT_LENGTH:
|
||||
print(f"⚠️ 状态内容超长({len(content)} > {self.MAX_CONTENT_LENGTH} 字符),拒绝写入")
|
||||
return False
|
||||
status_path = os.path.join(self._prompts_dir, self.STATUS_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
updated_content = f"""# 智能体状态记录
|
||||
|
||||
{content}
|
||||
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
# J5:原子写,崩溃不留半截提示词文件
|
||||
atomic_write_text(status_path, updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 更新状态文件失败: {e}")
|
||||
return False
|
||||
|
||||
def update_soul(self, content: str) -> bool:
|
||||
"""更新智能体灵魂定义 - 所有角色定位内容"""
|
||||
if len(content or "") > self.MAX_CONTENT_LENGTH:
|
||||
print(f"⚠️ 灵魂内容超长({len(content)} > {self.MAX_CONTENT_LENGTH} 字符),拒绝写入")
|
||||
return False
|
||||
soul_path = os.path.join(self._prompts_dir, self.SOUL_FILE)
|
||||
try:
|
||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
updated_content = f"""# 智能体灵魂定义 - 我是谁
|
||||
|
||||
{content}
|
||||
|
||||
【更新记录】
|
||||
- 上次更新:{timestamp}
|
||||
"""
|
||||
# J5:原子写,崩溃不留半截提示词文件
|
||||
atomic_write_text(soul_path, updated_content)
|
||||
self.reload()
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"⚠️ 更新灵魂文件失败: {e}")
|
||||
return False
|
||||
|
||||
def get_identity(self) -> str:
|
||||
"""获取用户身份信息"""
|
||||
identity_path = os.path.join(self._prompts_dir, self.IDENTITY_FILE)
|
||||
try:
|
||||
if os.path.exists(identity_path):
|
||||
with open(identity_path, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取身份文件失败: {e}")
|
||||
return "尚不清楚 - 无法读取身份信息"
|
||||
|
||||
def get_status(self) -> str:
|
||||
"""获取智能体状态"""
|
||||
status_path = os.path.join(self._prompts_dir, self.STATUS_FILE)
|
||||
try:
|
||||
if os.path.exists(status_path):
|
||||
with open(status_path, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取状态文件失败: {e}")
|
||||
return "这是第一次对话 - 无法读取状态信息"
|
||||
|
||||
def get_soul(self) -> str:
|
||||
"""获取智能体灵魂定义 - 角色定位核心"""
|
||||
soul_path = os.path.join(self._prompts_dir, self.SOUL_FILE)
|
||||
try:
|
||||
if os.path.exists(soul_path):
|
||||
with open(soul_path, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception as e:
|
||||
print(f"⚠️ 读取灵魂文件失败: {e}")
|
||||
return """我是DeepSeek,一个智能AI助手。
|
||||
请创建 99_SOUL.md 文件来定义我的完整身份、角色、性格和能力边界。"""
|
||||
|
||||
def get_soul_prompt(self) -> Optional[Dict]:
|
||||
"""获取灵魂提示词(99_SOUL.md)"""
|
||||
for p in self._prompts:
|
||||
if p['file'] == self.SOUL_FILE:
|
||||
return p
|
||||
return None
|
||||
|
||||
def _create_default_prompts(self):
|
||||
"""创建默认提示词(已由各 _ensure_* 方法处理)"""
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def _skill_sort_key(cls, filename: str):
|
||||
"""技能文件按前缀数字排序,非技能文件返回 None。"""
|
||||
m = cls._SKILL_FILE_RE.match(filename)
|
||||
return int(m.group(1)) if m else None
|
||||
|
||||
def _prompt_sort_key(self, p):
|
||||
"""排序规则:00_IDENTITY → 技能文件(按数字升序) → 99_SOUL → 100_STATUS → 其他"""
|
||||
file = p['file']
|
||||
if file == self.IDENTITY_FILE:
|
||||
return (0,)
|
||||
skill_num = self._skill_sort_key(file)
|
||||
if skill_num is not None:
|
||||
return (1, skill_num)
|
||||
if file == self.SOUL_FILE:
|
||||
return (2,)
|
||||
if file == self.STATUS_FILE:
|
||||
return (4,)
|
||||
return (3, file) # 其他文件按名字排序
|
||||
|
||||
def get_prompts(self, join_with: str = "\n\n") -> str:
|
||||
"""获取所有提示词,按顺序拼接(identity → 技能文件 → soul → status)"""
|
||||
if not self._prompts:
|
||||
return ""
|
||||
sorted_prompts = sorted(self._prompts, key=self._prompt_sort_key)
|
||||
contents = [p['content'] for p in sorted_prompts]
|
||||
return join_with.join(contents)
|
||||
|
||||
def get_sorted_prompts_list(self) -> List[Dict]:
|
||||
"""获取按正确顺序排列的提示词列表(用于 system 消息构建)。"""
|
||||
return sorted(self._prompts, key=self._prompt_sort_key)
|
||||
|
||||
def get_prompts_list(self) -> List[Dict]:
|
||||
"""获取所有提示词列表(原始加载顺序,已弃用,请用 get_sorted_prompts_list)。"""
|
||||
return self._prompts.copy()
|
||||
|
||||
def get_skill_prompts(self) -> List[Dict]:
|
||||
"""获取所有基础技能提示词(N_ 开头的文件)。"""
|
||||
return [p for p in self._prompts if self._skill_sort_key(p['file']) is not None]
|
||||
|
||||
def reload(self):
|
||||
"""重新加载所有提示词(M6:构建新列表后原子替换,并发读者始终读到完整列表)"""
|
||||
new_prompts: List[Dict] = []
|
||||
self._load_prompts(new_prompts)
|
||||
self._prompts = new_prompts
|
||||
return len(self._prompts)
|
||||
|
||||
def get_prompt_count(self) -> int:
|
||||
return len(self._prompts)
|
||||
|
||||
def get_prompts_dir_path(self) -> str:
|
||||
return self._prompts_dir
|
||||
@@ -0,0 +1,260 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
subsystems.skill_manager - 扩展技能管理器
|
||||
|
||||
抽离自 aibox.py,收集并管理 ~/.aibox/skills/ 目录下的技能文件。
|
||||
|
||||
支持两种技能格式:
|
||||
1. 可执行脚本/程序:执行 --help,第一行为简介
|
||||
2. Markdown 文件 (.md):第一行为简介(# 标题或普通文本),其余为技能内容
|
||||
|
||||
依赖:common (constants, Logger) + 标准库
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import constants, Logger
|
||||
|
||||
|
||||
class SkillManager:
|
||||
"""扩展技能管理器 - 收集并管理 ~/.aibox/skills/ 目录下的技能文件。
|
||||
|
||||
支持两种技能格式:
|
||||
1. 可执行脚本/程序:执行 --help,第一行为简介
|
||||
2. Markdown 文件 (.md):第一行为简介(# 标题或普通文本),其余为技能内容
|
||||
"""
|
||||
|
||||
# J16(审计 2026-08-19):--help 输出上限——内容会进系统提示词
|
||||
MAX_HELP_LENGTH = 2000
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self.logger = logger or Logger("SkillManager")
|
||||
# name -> {path, file, summary, content, is_markdown}
|
||||
self.skills: Dict[str, Dict] = {}
|
||||
self._load_skills()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部辅助
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_help(self, skill_path: str) -> str:
|
||||
"""执行 skill_path --help 并返回完整输出(stdout 优先,否则 stderr)。
|
||||
|
||||
J16(审计 2026-08-19):输出截断到 MAX_HELP_LENGTH——技能 --help
|
||||
输出会进系统提示词,异常/恶意脚本可撑爆上下文。
|
||||
"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[skill_path, "--help"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
env=os.environ.copy()
|
||||
)
|
||||
output = result.stdout.strip()
|
||||
if not output:
|
||||
output = result.stderr.strip()
|
||||
if len(output) > self.MAX_HELP_LENGTH:
|
||||
output = output[:self.MAX_HELP_LENGTH] + "\n…(已截断)"
|
||||
return output
|
||||
except Exception as e:
|
||||
self.logger.warning(f"运行 {skill_path} --help 失败: {e}")
|
||||
return ""
|
||||
|
||||
def _read_markdown_skill(self, skill_path: str) -> Tuple[str, str]:
|
||||
"""读取 Markdown 技能文件,返回 (简介, 完整内容)"""
|
||||
try:
|
||||
with open(skill_path, 'r', encoding='utf-8') as f:
|
||||
content = f.read().strip()
|
||||
|
||||
if not content:
|
||||
return "(空文件)", ""
|
||||
|
||||
# 提取第一行非空内容作为简介
|
||||
lines = content.splitlines()
|
||||
summary = ""
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
# 如果是 Markdown 标题,去掉 # 前缀
|
||||
if stripped.startswith('#'):
|
||||
summary = stripped.lstrip('#').strip()
|
||||
else:
|
||||
summary = stripped
|
||||
break
|
||||
|
||||
if not summary:
|
||||
summary = os.path.basename(skill_path).replace('.md', '')
|
||||
|
||||
return summary, content
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"读取 Markdown 技能文件失败 {skill_path}: {e}")
|
||||
return "(读取失败)", ""
|
||||
|
||||
@staticmethod
|
||||
def _first_nonempty_line(text: str) -> str:
|
||||
"""返回文本中第一行非空内容。"""
|
||||
for line in text.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
return stripped
|
||||
return ""
|
||||
|
||||
def _is_markdown_file(self, filename: str) -> bool:
|
||||
"""判断是否为 Markdown 文件"""
|
||||
return filename.lower().endswith('.md')
|
||||
|
||||
def _is_executable(self, filepath: str) -> bool:
|
||||
"""检查文件是否可执行"""
|
||||
if os.name == 'nt': # Windows
|
||||
# Windows 下检查扩展名
|
||||
executable_extensions = {'.exe', '.bat', '.cmd', '.ps1', '.py', '.sh'}
|
||||
ext = os.path.splitext(filepath)[1].lower()
|
||||
return ext in executable_extensions
|
||||
else: # Unix/Linux/Mac
|
||||
return os.access(filepath, os.X_OK)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 加载 / 重载
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_skills(self):
|
||||
"""扫描技能目录,加载所有技能文件(可执行文件或 Markdown 文件)。
|
||||
|
||||
T26(审计 2026-08-21):本地构建 new_skills 完成后一次性原子替换,
|
||||
对齐 PromptManager M6——此前先 self.skills = {} 清空再逐个重建,
|
||||
reload 期间(每个可执行技能 --help 最多 10s)并发读者会看到空/半截
|
||||
技能表(技能"消失"再"逐渐出现")。"""
|
||||
skills_dir = constants.SKILLS_DIR
|
||||
new_skills = {}
|
||||
|
||||
if not os.path.exists(skills_dir):
|
||||
self.logger.info(f"技能目录不存在: {skills_dir}")
|
||||
self.skills = new_skills
|
||||
return
|
||||
|
||||
skill_files = [
|
||||
f for f in os.listdir(skills_dir)
|
||||
if os.path.isfile(os.path.join(skills_dir, f))
|
||||
and not f.startswith('_')
|
||||
and not f.startswith('.')
|
||||
]
|
||||
|
||||
# J16(审计 2026-08-19):同名键冲突(foo 可执行文件与 foo.md)此前
|
||||
# 由 listdir 顺序决定谁覆盖谁,行为不确定。改为可执行优先 + 确定性:
|
||||
# 先加载可执行技能,Markdown 遇到已占用的同名键时跳过并告警
|
||||
executable_files = [
|
||||
f for f in skill_files
|
||||
if not self._is_markdown_file(f)
|
||||
]
|
||||
markdown_files = [f for f in skill_files if self._is_markdown_file(f)]
|
||||
|
||||
for skill_file in executable_files + markdown_files:
|
||||
skill_path = os.path.join(skills_dir, skill_file)
|
||||
skill_name = os.path.splitext(skill_file)[0]
|
||||
is_markdown = self._is_markdown_file(skill_file)
|
||||
|
||||
if is_markdown:
|
||||
if skill_name in new_skills:
|
||||
self.logger.warning(
|
||||
f"Markdown 技能 '{skill_file}' 与同名可执行技能冲突,已跳过"
|
||||
f"(可执行技能优先)")
|
||||
continue
|
||||
# Markdown 格式技能:直接读取内容
|
||||
summary, content = self._read_markdown_skill(skill_path)
|
||||
if content:
|
||||
new_skills[skill_name] = {
|
||||
"path": skill_path,
|
||||
"file": skill_file,
|
||||
"summary": summary,
|
||||
"content": content,
|
||||
"is_markdown": True,
|
||||
"is_executable": False
|
||||
}
|
||||
self.logger.info(f"已加载 Markdown 技能: {skill_name} — {summary}")
|
||||
else:
|
||||
self.logger.warning(f"Markdown 技能文件为空(跳过): {skill_file}")
|
||||
|
||||
else:
|
||||
# 可执行文件格式技能:执行 --help 获取简介
|
||||
if not self._is_executable(skill_path):
|
||||
self.logger.warning(f"技能文件不可执行(跳过): {skill_file}")
|
||||
continue
|
||||
|
||||
help_output = self._run_help(skill_path)
|
||||
summary = self._first_nonempty_line(help_output) or f"扩展技能: {skill_name}"
|
||||
|
||||
new_skills[skill_name] = {
|
||||
"path": skill_path,
|
||||
"file": skill_file,
|
||||
"summary": summary,
|
||||
"content": help_output, # 完整的 --help 输出作为内容
|
||||
"is_markdown": False,
|
||||
"is_executable": True
|
||||
}
|
||||
self.logger.info(f"已加载可执行技能: {skill_name} — {summary}")
|
||||
|
||||
self.skills = new_skills
|
||||
self.logger.info(f"共加载 {len(new_skills)} 个技能({len([s for s in new_skills.values() if s['is_markdown']])} 个 Markdown 技能)")
|
||||
|
||||
def reload(self) -> int:
|
||||
"""重新扫描技能目录,返回加载的技能数量。"""
|
||||
self._load_skills()
|
||||
return len(self.skills)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 查询
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def search(self, keyword: str) -> List[Dict]:
|
||||
"""在技能名称、简介和内容中搜索关键字,返回匹配项列表。"""
|
||||
kw = keyword.lower()
|
||||
results = []
|
||||
for name, info in self.skills.items():
|
||||
match = (
|
||||
kw in name.lower() or
|
||||
kw in info["summary"].lower() or
|
||||
kw in info.get("content", "").lower()
|
||||
)
|
||||
if match:
|
||||
results.append({
|
||||
"name": name,
|
||||
"path": info["path"],
|
||||
"summary": info["summary"],
|
||||
"content": info.get("content", ""),
|
||||
"is_markdown": info.get("is_markdown", False),
|
||||
"is_executable": info.get("is_executable", False)
|
||||
})
|
||||
return results
|
||||
|
||||
def get_skill_content(self, skill_name: str) -> Optional[str]:
|
||||
"""获取技能的内容(Markdown 技能返回完整内容,可执行技能返回 --help 输出)"""
|
||||
info = self.skills.get(skill_name)
|
||||
if info:
|
||||
return info.get("content")
|
||||
return None
|
||||
|
||||
def get_skill(self, skill_name: str) -> Optional[Dict]:
|
||||
"""获取单个技能的完整信息"""
|
||||
return self.skills.get(skill_name)
|
||||
|
||||
def get_all(self) -> List[Dict]:
|
||||
"""返回所有技能的列表(含名称、路径、简介和内容)。"""
|
||||
return [
|
||||
{
|
||||
"name": name,
|
||||
"path": info["path"],
|
||||
"summary": info["summary"],
|
||||
"content": info.get("content", ""),
|
||||
"is_markdown": info.get("is_markdown", False),
|
||||
"is_executable": info.get("is_executable", False)
|
||||
}
|
||||
for name, info in self.skills.items()
|
||||
]
|
||||
|
||||
def count(self) -> int:
|
||||
return len(self.skills)
|
||||
File renamed without changes.
@@ -0,0 +1,98 @@
|
||||
#!/usr/bin/env python3
|
||||
"""审计修复 PoC 回归验证(H1/H2/H3 + 1MB + cancel)。
|
||||
|
||||
验证 2026-08-17 审计五项修复的实际运行效果:
|
||||
PoC-1 (H2): 无输出命令 + timeout=3s → 应 ~3s 返回超时(修复前阻塞 15s+)
|
||||
PoC-2 (H3): 等待 stdin 的命令 + timeout=3s → 应超时返回而非永久挂起
|
||||
PoC-3 : 输出 >1MB → 应截断且总输出 ≈1MB
|
||||
PoC-4 : cancel() 应在宽限时间内终止长命令
|
||||
PoC-5 (H1) : CommandLineTool(工具层,非 executor)跑 >1MB 输出 →
|
||||
截断生效即证明已走 executor 委托路径
|
||||
"""
|
||||
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
AIBOX_DIR = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(AIBOX_DIR))
|
||||
sys.modules.setdefault("websockets", types.ModuleType("websockets"))
|
||||
|
||||
from subsystems.executors import SubprocessCommandExecutor
|
||||
from tools.command_tool import CommandLineTool
|
||||
|
||||
results = []
|
||||
|
||||
|
||||
def check(name, ok, detail):
|
||||
results.append((name, ok, detail))
|
||||
print(f"[{'PASS' if ok else 'FAIL'}] {name}: {detail}")
|
||||
|
||||
|
||||
# ---------- PoC-1 (H2): 无输出命令超时 ----------
|
||||
ex = SubprocessCommandExecutor()
|
||||
t0 = time.time()
|
||||
r = ex.execute("powershell -NoProfile -Command Start-Sleep 15", timeout=3)
|
||||
elapsed = time.time() - t0
|
||||
check("PoC-1 无输出命令超时(H2)", elapsed < 6 and "超时" in r,
|
||||
f"{elapsed:.1f}s 返回(修复前阻塞15s+),结果含超时: {'超时' in r}")
|
||||
|
||||
# ---------- PoC-2 (H3): 等待 stdin 的命令 ----------
|
||||
# stdin=DEVNULL 下命令立即收到 EOF:要么快速退出,要么按 timeout 返回;
|
||||
# 判定标准 = 不永久挂起(修复前 readline 无 EOF 永久阻塞)。
|
||||
t0 = time.time()
|
||||
r = ex.execute("powershell -NoProfile -Command Read-Host", timeout=3)
|
||||
elapsed = time.time() - t0
|
||||
check("PoC-2 stdin等待命令不挂起(H3)", elapsed < 6,
|
||||
f"{elapsed:.1f}s 返回(修复前永久挂起)")
|
||||
|
||||
# ---------- PoC-3: 1MB 输出截断 ----------
|
||||
t0 = time.time()
|
||||
r = ex.execute(
|
||||
'powershell -NoProfile -Command "$s=\'x\'*2048; 1..1024 | ForEach-Object { $s }"',
|
||||
timeout=60)
|
||||
elapsed = time.time() - t0
|
||||
size_mb = len(r.encode("utf-8", errors="replace")) / 1024 / 1024
|
||||
check("PoC-3 1MB输出截断", "截断" in r and size_mb < 1.15,
|
||||
f"{elapsed:.1f}s,结果 {size_mb:.2f}MB,含截断标记: {'截断' in r}")
|
||||
|
||||
# ---------- PoC-4: cancel() 中断 ----------
|
||||
out_q = queue.Queue()
|
||||
done = threading.Event()
|
||||
|
||||
|
||||
def _run():
|
||||
ex.execute_streaming(
|
||||
"powershell -NoProfile -Command Start-Sleep 60",
|
||||
out_q, timeout=120)
|
||||
done.set()
|
||||
|
||||
|
||||
t0 = time.time()
|
||||
th = threading.Thread(target=_run, daemon=True)
|
||||
th.start()
|
||||
time.sleep(2)
|
||||
ex.cancel()
|
||||
done.wait(timeout=10)
|
||||
elapsed = time.time() - t0
|
||||
check("PoC-4 cancel中断", done.is_set() and elapsed < 8,
|
||||
f"cancel 后 {elapsed:.1f}s 结束(修复前杀不到进程需等满120s)")
|
||||
|
||||
# ---------- PoC-5 (H1): 工具层委托 executor(截断保护贯通) ----------
|
||||
tool = CommandLineTool(executor=ex)
|
||||
t0 = time.time()
|
||||
r = tool.execute(
|
||||
command='powershell -NoProfile -Command "$s=\'y\'*2048; 1..1024 | ForEach-Object { $s }"',
|
||||
timeout=60)
|
||||
elapsed = time.time() - t0
|
||||
size_mb = len(r.encode("utf-8", errors="replace")) / 1024 / 1024
|
||||
check("PoC-5 工具层1MB贯通(H1)", "截断" in r and size_mb < 1.15,
|
||||
f"{elapsed:.1f}s,工具层结果 {size_mb:.2f}MB(修复前无限累积)")
|
||||
|
||||
# ---------- 汇总 ----------
|
||||
fails = [n for n, ok, _ in results if not ok]
|
||||
print("\n==== 汇总: %d/%d PASS ====" % (len(results) - len(fails), len(results)))
|
||||
sys.exit(1 if fails else 0)
|
||||
@@ -0,0 +1,23 @@
|
||||
"""pytest 公共配置:把 aibox 目录加入 sys.path,桩掉可选第三方依赖。
|
||||
|
||||
aibox 主模块 aibox.py 依赖 websockets,但 tools/basic_tools.py 仅依赖
|
||||
common + 标准库,所以单测 basic_tools 不需要真实 websockets——桩掉即可
|
||||
让 `from tools import ...` 间接 import aibox.py 时不出错(虽然本测试
|
||||
不直接 import aibox,但保持环境干净)。
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
# 把 aibox 目录加入 sys.path(无需安装为包即可 import tools/common)
|
||||
AIBOX_DIR = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(AIBOX_DIR))
|
||||
# tests 目录也加入:测试文件互相导入(如 test_models_profile 复用 test_once_mode 的 _make_ai)
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
||||
|
||||
# 桩掉 websockets(aibox.py 顶部 import websockets,测试环境可能未安装)
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import pytest # noqa: E402
|
||||
@@ -0,0 +1,593 @@
|
||||
"""B 批工具单元测试
|
||||
|
||||
覆盖 12 个 B 批工具的核心接口与执行逻辑:
|
||||
- command_tool : CommandLineTool
|
||||
- skill_tools : SearchSkills/ReloadSkills/InstallSkill/CreateSkill
|
||||
- db_tools : MemoryTool/MemoTool
|
||||
- prompt_tools : UpdateIdentity/UpdateSoul/SaveMemoryAndEnd
|
||||
- proxy_tools : SetProxy/GetProxy
|
||||
- misc_tools : AskUserTool
|
||||
|
||||
工具大多需注入子系统,用真实子系统实例 + tmp_path 隔离。
|
||||
"""
|
||||
import queue
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from common import ConfigManager, constants
|
||||
from subsystems import (
|
||||
MemoryDatabase, MemoDatabase, PromptManager, SkillManager,
|
||||
SubprocessCommandExecutor,
|
||||
)
|
||||
from tools import (
|
||||
CommandLineTool,
|
||||
SearchSkillsTool, ReloadSkillsTool, InstallSkillTool, CreateSkillTool,
|
||||
MemoryTool, MemoTool,
|
||||
UpdateIdentityTool, UpdateSoulTool, SaveMemoryAndEndConversationTool,
|
||||
SetProxyTool, GetProxyTool,
|
||||
AskUserTool,
|
||||
)
|
||||
from tools.base import Tool
|
||||
from tools.registry import ToolRegistry
|
||||
|
||||
|
||||
# ==================== 公共 fixture ====================
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_tmp_workspace(tmp_path, monkeypatch):
|
||||
"""每个测试把 WORKSPACE_DIR/SAVE_DIR 重定向到临时目录,避免污染 ~/.aibox。"""
|
||||
ws = tmp_path / "ws"
|
||||
ws.mkdir()
|
||||
save = tmp_path / "save"
|
||||
save.mkdir()
|
||||
monkeypatch.setattr(type(constants), "WORKSPACE_DIR", property(lambda self: str(ws)))
|
||||
monkeypatch.setattr(type(constants), "SAVE_DIR", property(lambda self: str(save)))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mem_db(tmp_path):
|
||||
return MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memo_db(tmp_path):
|
||||
return MemoDatabase(db_path=str(tmp_path / "memo.db"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def prompt_mgr(tmp_path, monkeypatch):
|
||||
"""独立 PromptManager(重置单例 + 临时 prompts 目录)。"""
|
||||
prompts_dir = tmp_path / "prompts"
|
||||
prompts_dir.mkdir()
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
from common import ConfigManager
|
||||
original_get = ConfigManager.get
|
||||
|
||||
def mock_get(self, key, default=None):
|
||||
if key == "system.prompts_dir":
|
||||
return str(prompts_dir)
|
||||
return original_get(self, key, default)
|
||||
|
||||
monkeypatch.setattr(ConfigManager, "get", mock_get)
|
||||
pm = PromptManager()
|
||||
yield pm
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def skill_mgr(tmp_path, monkeypatch):
|
||||
"""独立 SkillManager + 临时技能目录。"""
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR", property(lambda self: str(skills_dir)))
|
||||
return SkillManager()
|
||||
|
||||
|
||||
# ==================== 通用接口契约(参数化)====================
|
||||
|
||||
B_TOOL_CLASSES = [
|
||||
CommandLineTool, SearchSkillsTool, ReloadSkillsTool, InstallSkillTool,
|
||||
CreateSkillTool, MemoryTool, MemoTool, UpdateIdentityTool, UpdateSoulTool,
|
||||
SaveMemoryAndEndConversationTool, SetProxyTool, GetProxyTool, AskUserTool,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cls", B_TOOL_CLASSES, ids=lambda c: c.__name__)
|
||||
def test_b_tool_is_tool_subclass(cls):
|
||||
assert issubclass(cls, Tool)
|
||||
|
||||
|
||||
def test_b_tool_names_unique():
|
||||
"""所有 B 批工具的 get_name() 应唯一(需先实例化,部分需注入)。"""
|
||||
# 这里只验证类名能取到非空 name(无法统一实例化,跳过唯一性全检)
|
||||
pass
|
||||
|
||||
|
||||
# ==================== CommandLineTool ====================
|
||||
|
||||
def test_command_tool_interface():
|
||||
t = CommandLineTool()
|
||||
assert t.get_name() == "run_command"
|
||||
assert "执行" in t.get_description() or "命令" in t.get_description()
|
||||
params = t.get_parameters()
|
||||
assert "command" in params["properties"]
|
||||
assert "command" in params["required"]
|
||||
|
||||
|
||||
def test_command_tool_execute_echo():
|
||||
t = CommandLineTool()
|
||||
r = t.execute(command="echo b_test_123", timeout=5)
|
||||
assert "b_test_123" in r
|
||||
|
||||
|
||||
def test_command_tool_missing_command_returns_error():
|
||||
t = CommandLineTool()
|
||||
r = t.execute(command="", timeout=5)
|
||||
assert "❌" in r or "错误" in r
|
||||
|
||||
|
||||
def test_command_tool_streaming_emits_complete():
|
||||
t = CommandLineTool()
|
||||
q = queue.Queue()
|
||||
t.execute_streaming(q, command="echo hi", timeout=5)
|
||||
msgs = []
|
||||
while not q.empty():
|
||||
msgs.append(q.get_nowait())
|
||||
assert any(m[0] == "complete" for m in msgs)
|
||||
|
||||
|
||||
# ==================== SearchSkillsTool ====================
|
||||
|
||||
def test_search_skills_interface(skill_mgr):
|
||||
t = SearchSkillsTool(skill_mgr)
|
||||
assert t.get_name() == "search_skills"
|
||||
params = t.get_parameters()
|
||||
assert "keyword" in params["properties"]
|
||||
assert "keyword" in params["required"]
|
||||
|
||||
|
||||
def test_search_skills_empty_keyword_lists_all(skill_mgr, tmp_path):
|
||||
# 写入一个技能
|
||||
(tmp_path / "skills" / "demo.md").write_text("# Demo\n内容", encoding="utf-8")
|
||||
skill_mgr.reload()
|
||||
t = SearchSkillsTool(skill_mgr)
|
||||
r = t.execute(keyword="")
|
||||
assert "所有可用技能" in r
|
||||
assert "demo" in r
|
||||
|
||||
|
||||
def test_search_skills_no_match(skill_mgr):
|
||||
t = SearchSkillsTool(skill_mgr)
|
||||
r = t.execute(keyword="不存在的技能xyz")
|
||||
assert "未找到" in r
|
||||
|
||||
|
||||
def test_search_skills_with_content(skill_mgr, tmp_path):
|
||||
(tmp_path / "skills" / "weather.md").write_text("# 天气\n查询天气详情ABC", encoding="utf-8")
|
||||
skill_mgr.reload()
|
||||
t = SearchSkillsTool(skill_mgr)
|
||||
r = t.execute(keyword="weather", show_content=True)
|
||||
assert "查询天气详情ABC" in r
|
||||
|
||||
|
||||
# ==================== ReloadSkillsTool ====================
|
||||
|
||||
def test_reload_skills_interface(skill_mgr):
|
||||
reg = ToolRegistry()
|
||||
reg.skill_manager = skill_mgr
|
||||
t = ReloadSkillsTool(reg)
|
||||
assert t.get_name() == "reload_skills"
|
||||
r = t.execute(confirm=False)
|
||||
assert "取消" in r
|
||||
|
||||
|
||||
def test_reload_skills_confirm(skill_mgr):
|
||||
reg = ToolRegistry()
|
||||
reg.skill_manager = skill_mgr
|
||||
t = ReloadSkillsTool(reg)
|
||||
r = t.execute(confirm=True)
|
||||
# 重载成功应返回技能数量信息
|
||||
assert "技能" in r or "重载" in r
|
||||
|
||||
|
||||
# ==================== InstallSkillTool ====================
|
||||
|
||||
def test_install_skill_interface():
|
||||
t = InstallSkillTool()
|
||||
assert t.get_name() == "install_skill"
|
||||
assert "skill_path" in t.get_parameters()["properties"]
|
||||
|
||||
|
||||
def test_install_skill_returns_guide():
|
||||
t = InstallSkillTool()
|
||||
r = t.execute()
|
||||
assert "安装" in r or "步骤" in r
|
||||
assert "skills" in r.lower() or ".md" in r
|
||||
|
||||
|
||||
def test_install_skill_with_path():
|
||||
t = InstallSkillTool()
|
||||
r = t.execute(skill_path="/tmp/my_skill.md")
|
||||
assert "my_skill" in r
|
||||
|
||||
|
||||
# ==================== CreateSkillTool ====================
|
||||
|
||||
def test_create_skill_interface():
|
||||
t = CreateSkillTool()
|
||||
assert t.get_name() == "create_skill"
|
||||
params = t.get_parameters()
|
||||
assert "skill_name" in params["required"]
|
||||
assert "description" in params["required"]
|
||||
|
||||
|
||||
def test_create_skill_validation_empty_name():
|
||||
t = CreateSkillTool()
|
||||
r = t.execute(skill_name="", description="desc")
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_create_skill_validation_invalid_name():
|
||||
t = CreateSkillTool()
|
||||
r = t.execute(skill_name="bad name!", description="desc")
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_create_skill_returns_prompt():
|
||||
t = CreateSkillTool()
|
||||
r = t.execute(skill_name="fetch-weather", description="获取天气")
|
||||
assert "fetch-weather" in r
|
||||
assert "步骤" in r or "创建" in r
|
||||
|
||||
|
||||
# ==================== MemoryTool ====================
|
||||
|
||||
def test_memory_tool_interface(mem_db):
|
||||
t = MemoryTool(mem_db)
|
||||
assert t.get_name() == "memory"
|
||||
params = t.get_parameters()
|
||||
assert "action" in params["required"]
|
||||
|
||||
|
||||
def test_memory_save(mem_db):
|
||||
t = MemoryTool(mem_db)
|
||||
r = t.execute(action="save", summary="测试记忆内容", topic="测试主题")
|
||||
assert "记忆已保存" in r
|
||||
|
||||
|
||||
def test_memory_search(mem_db):
|
||||
t = MemoryTool(mem_db)
|
||||
t.execute(action="save", summary="learn Python basics", topic="Python")
|
||||
r = t.execute(action="search", query="Python")
|
||||
assert "Python" in r
|
||||
|
||||
|
||||
def test_memory_recent_empty(mem_db):
|
||||
t = MemoryTool(mem_db)
|
||||
r = t.execute(action="recent")
|
||||
assert "没有" in r
|
||||
|
||||
|
||||
def test_memory_unknown_action(mem_db):
|
||||
t = MemoryTool(mem_db)
|
||||
r = t.execute(action="invalid")
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_memory_delete_requires_confirm(mem_db):
|
||||
t = MemoryTool(mem_db)
|
||||
mid = mem_db.add_memory(topic="t", summary="s")
|
||||
r = t.execute(action="delete", memory_id=mid, confirm=False)
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
# ==================== MemoTool ====================
|
||||
|
||||
def test_memo_tool_interface(memo_db):
|
||||
t = MemoTool(memo_db)
|
||||
assert t.get_name() == "memo"
|
||||
assert "action" in t.get_parameters()["required"]
|
||||
|
||||
|
||||
def test_memo_add(memo_db):
|
||||
t = MemoTool(memo_db)
|
||||
r = t.execute(action="add", title="测试任务", reminder_minutes=30)
|
||||
assert "任务已创建" in r
|
||||
|
||||
|
||||
def test_memo_add_requires_title(memo_db):
|
||||
t = MemoTool(memo_db)
|
||||
r = t.execute(action="add")
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_memo_list_empty(memo_db):
|
||||
t = MemoTool(memo_db)
|
||||
r = t.execute(action="list")
|
||||
assert "没有" in r
|
||||
|
||||
|
||||
def test_memo_stats(memo_db):
|
||||
t = MemoTool(memo_db)
|
||||
r = t.execute(action="stats")
|
||||
assert "统计" in r
|
||||
|
||||
|
||||
def test_memo_parse_reminder_time_plus_format(memo_db):
|
||||
"""_parse_reminder_time 应解析 +30m 格式。"""
|
||||
t = MemoTool(memo_db)
|
||||
result = t._parse_reminder_time("+30m")
|
||||
assert result is not None
|
||||
# 应在约 30 分钟后
|
||||
dt = datetime.strptime(result, constants.DATETIME_FORMAT)
|
||||
delta = dt - datetime.now()
|
||||
assert timedelta(minutes=29) < delta < timedelta(minutes=31)
|
||||
|
||||
|
||||
def test_memo_parse_reminder_time_int(memo_db):
|
||||
t = MemoTool(memo_db)
|
||||
result = t._parse_reminder_time(60)
|
||||
assert result is not None
|
||||
|
||||
|
||||
# ==================== UpdateIdentityTool ====================
|
||||
|
||||
def test_update_identity_interface(prompt_mgr):
|
||||
t = UpdateIdentityTool(prompt_mgr)
|
||||
assert t.get_name() == "update_identity"
|
||||
assert "content" in t.get_parameters()["required"]
|
||||
|
||||
|
||||
def test_update_identity_success(prompt_mgr):
|
||||
t = UpdateIdentityTool(prompt_mgr)
|
||||
r = t.execute(content="测试用户身份", confirm=True)
|
||||
assert "✅" in r
|
||||
assert "测试用户身份" in prompt_mgr.get_identity()
|
||||
|
||||
|
||||
def test_update_identity_no_content(prompt_mgr):
|
||||
t = UpdateIdentityTool(prompt_mgr)
|
||||
r = t.execute(content=None, confirm=True)
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_update_identity_cancelled(prompt_mgr):
|
||||
t = UpdateIdentityTool(prompt_mgr)
|
||||
r = t.execute(content="x", confirm=False)
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
# ==================== UpdateSoulTool ====================
|
||||
|
||||
def test_update_soul_interface(prompt_mgr):
|
||||
t = UpdateSoulTool(prompt_mgr)
|
||||
assert t.get_name() == "update_soul"
|
||||
|
||||
|
||||
def test_update_soul_success(prompt_mgr):
|
||||
t = UpdateSoulTool(prompt_mgr)
|
||||
r = t.execute(content="我是测试灵魂", confirm=True)
|
||||
assert "✅" in r
|
||||
assert "我是测试灵魂" in prompt_mgr.get_soul()
|
||||
|
||||
|
||||
# ==================== SaveMemoryAndEndConversationTool ====================
|
||||
|
||||
def test_save_memory_end_interface(prompt_mgr, mem_db):
|
||||
cb_called = []
|
||||
def cb(reason, summary):
|
||||
cb_called.append((reason, summary))
|
||||
t = SaveMemoryAndEndConversationTool(cb, mem_db, prompt_mgr)
|
||||
assert t.get_name() == "save_memory_and_end_conversation"
|
||||
|
||||
|
||||
def test_save_memory_end_short_summary_rejected(prompt_mgr, mem_db):
|
||||
t = SaveMemoryAndEndConversationTool(lambda *a: None, mem_db, prompt_mgr)
|
||||
r = t.execute(summary="短", current_status="状态")
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_save_memory_end_success(prompt_mgr, mem_db):
|
||||
cb_called = []
|
||||
t = SaveMemoryAndEndConversationTool(
|
||||
lambda reason, summary: cb_called.append((reason, summary)),
|
||||
mem_db, prompt_mgr
|
||||
)
|
||||
long_summary = "这是一段足够长的对话总结内容用于测试保存记忆功能,必须超过五十个字符才能通过校验,所以这里再加一些内容确保长度足够。"
|
||||
r = t.execute(summary=long_summary, topic="测试", current_status="当前状态OK")
|
||||
assert "✅" in r
|
||||
assert len(cb_called) == 1
|
||||
assert cb_called[0][0] == "话题讨论完成"
|
||||
|
||||
|
||||
# ==================== SetProxyTool / GetProxyTool ====================
|
||||
|
||||
def test_set_proxy_interface():
|
||||
assert SetProxyTool().get_name() == "set_proxy"
|
||||
params = SetProxyTool().get_parameters()
|
||||
assert "host" in params["required"]
|
||||
assert "port" in params["required"]
|
||||
|
||||
|
||||
def test_get_proxy_interface():
|
||||
assert GetProxyTool().get_name() == "get_proxy"
|
||||
|
||||
|
||||
def test_get_proxy_when_disabled():
|
||||
t = GetProxyTool()
|
||||
r = t.execute()
|
||||
# 默认未启用代理
|
||||
assert "未启用" in r or "启用" in r
|
||||
|
||||
|
||||
# ---------- SetProxyTool 连通性验证(设置前验证,失败不落盘) ----------
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_config(tmp_path, monkeypatch):
|
||||
"""隔离 ConfigManager 单例:AI_CONFIG_PATH 指向临时文件,前后重置类状态"""
|
||||
cfg_file = tmp_path / "config.json"
|
||||
monkeypatch.setenv("AI_CONFIG_PATH", str(cfg_file))
|
||||
|
||||
def _reset():
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
|
||||
_reset()
|
||||
yield ConfigManager
|
||||
_reset()
|
||||
|
||||
|
||||
class _FakeResp:
|
||||
def __init__(self, status_code):
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
def _fake_session_factory(exc=None, resp=None):
|
||||
"""构造替代 requests.Session 的假类:get() 抛异常或返回假响应"""
|
||||
calls = []
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self):
|
||||
self.trust_env = True
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
if exc:
|
||||
raise exc
|
||||
return resp
|
||||
|
||||
_FakeSession.calls = calls
|
||||
return _FakeSession
|
||||
|
||||
|
||||
def test_check_connectivity_ok(monkeypatch):
|
||||
fake = _fake_session_factory(resp=_FakeResp(204))
|
||||
monkeypatch.setattr("tools.proxy_tools.requests.Session", fake)
|
||||
t = SetProxyTool()
|
||||
ok, kind, detail = t._check_connectivity("127.0.0.1", 1080, None, None)
|
||||
assert ok and kind == ""
|
||||
# trust_env 必须关掉:排除系统 http_proxy 环境变量干扰
|
||||
assert fake.calls, "应发出探测请求"
|
||||
|
||||
|
||||
def test_check_connectivity_classification(monkeypatch, tmp_path):
|
||||
"""异常类型/消息 → 错误类型分类(与系统语言无关)"""
|
||||
import requests as requests_lib
|
||||
cases = [
|
||||
# Windows 本地化的连接拒绝消息:靠异常类型而非英文关键词分类
|
||||
(requests_lib.exceptions.ConnectionError(
|
||||
"[WinError 10061] 由于目标计算机积极拒绝,无法连接。"), "connect_failed"),
|
||||
# ConnectTimeout 同时继承 ConnectionError/Timeout,必须按 Timeout 归类
|
||||
(requests_lib.exceptions.ConnectTimeout("connection timeout"), "timeout"),
|
||||
(requests_lib.exceptions.ConnectionError(
|
||||
"SOCKS5 authentication failed"), "auth_failed"),
|
||||
(requests_lib.exceptions.ConnectionError(
|
||||
"Missing dependencies for SOCKS support."), "missing_dependency"),
|
||||
(Exception("Read timed out"), "timeout"), # 兜底:非 requests 异常按消息
|
||||
(Exception("Connection refused"), "connect_failed"),
|
||||
(Exception("something weird"), "error"),
|
||||
]
|
||||
t = SetProxyTool()
|
||||
for exc, expect_kind in cases:
|
||||
monkeypatch.setattr(
|
||||
"tools.proxy_tools.requests.Session",
|
||||
_fake_session_factory(exc=exc))
|
||||
ok, kind, _ = t._check_connectivity("127.0.0.1", 1080, "u", "p")
|
||||
assert not ok
|
||||
assert kind == expect_kind, f"{exc!r} 应分类为 {expect_kind},实际 {kind}"
|
||||
|
||||
|
||||
def test_check_connectivity_bad_status(monkeypatch):
|
||||
"""探测端点返回非 200/204 → probe_failed"""
|
||||
monkeypatch.setattr(
|
||||
"tools.proxy_tools.requests.Session",
|
||||
_fake_session_factory(resp=_FakeResp(502)))
|
||||
t = SetProxyTool()
|
||||
ok, kind, detail = t._check_connectivity("127.0.0.1", 1080, None, None)
|
||||
assert not ok and kind == "probe_failed"
|
||||
|
||||
|
||||
def test_set_proxy_rejected_when_unreachable(fresh_config, monkeypatch):
|
||||
"""验证失败:返回 ❌ 且配置不落盘"""
|
||||
monkeypatch.setattr(
|
||||
SetProxyTool, "_check_connectivity",
|
||||
lambda self, h, p, u, pw: (False, "connect_failed", "Connection refused"))
|
||||
t = SetProxyTool()
|
||||
r = t.execute(host="10.255.255.1", port=1080)
|
||||
assert "❌" in r
|
||||
assert "未保存" in r
|
||||
# 配置未写入:代理仍未启用
|
||||
assert fresh_config().get_proxy() is None
|
||||
|
||||
|
||||
def test_set_proxy_saved_when_check_passes(fresh_config, monkeypatch):
|
||||
"""验证通过:配置保存且结果带连通性详情"""
|
||||
monkeypatch.setattr(
|
||||
SetProxyTool, "_check_connectivity",
|
||||
lambda self, h, p, u, pw: (True, "", "探测成功 HTTP 204"))
|
||||
t = SetProxyTool()
|
||||
r = t.execute(host="127.0.0.1", port=1080)
|
||||
assert "✅" in r
|
||||
assert "探测成功" in r
|
||||
proxies = fresh_config().get_proxy()
|
||||
assert proxies and "127.0.0.1:1080" in proxies["http"]
|
||||
|
||||
|
||||
def test_set_proxy_skip_check_bypasses_validation(fresh_config, monkeypatch):
|
||||
"""skip_check=True:不验证直接保存(代理未启动的预配置场景)"""
|
||||
called = []
|
||||
|
||||
def _no_call(self, *a, **kw):
|
||||
called.append(1)
|
||||
raise AssertionError("skip_check=True 时不应调用验证")
|
||||
|
||||
monkeypatch.setattr(SetProxyTool, "_check_connectivity", _no_call)
|
||||
t = SetProxyTool()
|
||||
r = t.execute(host="127.0.0.1", port=1080, skip_check=True)
|
||||
assert "✅" in r
|
||||
assert not called
|
||||
assert fresh_config().get_proxy() is not None
|
||||
|
||||
|
||||
def test_set_proxy_disabled_skips_check(fresh_config, monkeypatch):
|
||||
"""enabled=False:禁用代理无需验证"""
|
||||
def _no_call(self, *a, **kw):
|
||||
raise AssertionError("禁用代理不应触发验证")
|
||||
|
||||
monkeypatch.setattr(SetProxyTool, "_check_connectivity", _no_call)
|
||||
t = SetProxyTool()
|
||||
r = t.execute(host="127.0.0.1", port=1080, enabled=False)
|
||||
assert "✅" in r
|
||||
assert fresh_config().get_proxy() is None # 禁用态 get_proxy 返回 None
|
||||
|
||||
|
||||
def test_set_proxy_empty_env_hint(fresh_config, monkeypatch):
|
||||
"""username + password_env 指向未设环境变量:精准提示而非裸 auth 失败"""
|
||||
monkeypatch.delenv("AIBOX_TEST_SOCKS_PW", raising=False)
|
||||
t = SetProxyTool()
|
||||
r = t.execute(host="127.0.0.1", port=1080,
|
||||
username="u1", password_env="AIBOX_TEST_SOCKS_PW")
|
||||
assert "❌" in r
|
||||
assert "环境变量" in r
|
||||
assert "未保存" in r
|
||||
assert fresh_config().get_proxy() is None
|
||||
|
||||
|
||||
# ==================== AskUserTool ====================
|
||||
|
||||
def test_ask_user_interface():
|
||||
t = AskUserTool()
|
||||
assert t.get_name() == "ask_user"
|
||||
|
||||
|
||||
def test_ask_user_returns_question():
|
||||
t = AskUserTool()
|
||||
r = t.execute(question="今天天气如何?")
|
||||
assert "今天天气如何" in r
|
||||
assert "询问用户" in r
|
||||
@@ -0,0 +1,339 @@
|
||||
"""tools.basic_tools 单元测试
|
||||
|
||||
覆盖 A 批抽离的 7 个工具类:
|
||||
- FileToolBase : 路径安全检查
|
||||
- ReadFileTool : 读文件指引
|
||||
- WriteFileTool : 写文件指引
|
||||
- ListFilesTool : 列目录指引
|
||||
- SimpleWebCrawlerTool : 网页访问指引
|
||||
- GetCurrentTimeTool : 时间获取
|
||||
- CalculatorTool : 数学计算
|
||||
|
||||
设计原则:
|
||||
- 不依赖网络/文件系统副作用(除 FileToolBase 路径检查用临时文件)
|
||||
- 不依赖 DeepSeek API、websockets 等外部服务
|
||||
- 验证工具接口契约(name/description/parameters/execute 返回类型)
|
||||
- 验证命令行指引类工具的核心语义(返回包含正确命令关键词 + run_command 提示)
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import tempfile
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from tools import (
|
||||
FileToolBase,
|
||||
ReadFileTool,
|
||||
WriteFileTool,
|
||||
ListFilesTool,
|
||||
SimpleWebCrawlerTool,
|
||||
GetCurrentTimeTool,
|
||||
CalculatorTool,
|
||||
)
|
||||
|
||||
|
||||
# ==================== 通用接口契约测试 ====================
|
||||
|
||||
# (工具类, 期望 get_name 返回值, 必填参数键列表)
|
||||
TOOL_CASES = [
|
||||
(ReadFileTool, "read_file", ["filepath"]),
|
||||
(WriteFileTool, "write_file", ["filepath", "content"]),
|
||||
(ListFilesTool, "list_files", []),
|
||||
(SimpleWebCrawlerTool, "web", ["action", "url"]),
|
||||
(GetCurrentTimeTool, "get_current_time", []),
|
||||
(CalculatorTool, "calculator", ["a", "b", "operation"]),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool_cls,expected_name,required_keys", TOOL_CASES)
|
||||
def test_tool_interface_contract(tool_cls, expected_name, required_keys):
|
||||
"""所有工具必须满足接口契约:name/描述非空、parameters 是合法 schema、必填字段在 properties 中。"""
|
||||
tool = tool_cls()
|
||||
assert tool.get_name() == expected_name
|
||||
|
||||
desc = tool.get_description()
|
||||
assert isinstance(desc, str) and len(desc) > 0
|
||||
|
||||
params = tool.get_parameters()
|
||||
assert isinstance(params, dict)
|
||||
assert params.get("type") == "object"
|
||||
props = params.get("properties", {})
|
||||
assert isinstance(props, dict) and len(props) > 0
|
||||
for k in required_keys:
|
||||
assert k in props, f"必填参数 {k} 未出现在 properties 中"
|
||||
# required 中声明的字段必须在 properties 中
|
||||
for k in params.get("required", []):
|
||||
assert k in props
|
||||
|
||||
|
||||
# ==================== GetCurrentTimeTool ====================
|
||||
|
||||
def test_get_current_time_full_format():
|
||||
tool = GetCurrentTimeTool()
|
||||
r = tool.execute(format="full")
|
||||
assert "当前时间" in r
|
||||
# 应包含 YYYY-MM-DD HH:MM:SS 形式
|
||||
assert datetime.now().strftime("%Y-%m-%d") in r
|
||||
|
||||
|
||||
def test_get_current_time_date_format():
|
||||
tool = GetCurrentTimeTool()
|
||||
r = tool.execute(format="date")
|
||||
assert datetime.now().strftime("%Y-%m-%d") in r
|
||||
|
||||
|
||||
def test_get_current_time_time_format():
|
||||
tool = GetCurrentTimeTool()
|
||||
r = tool.execute(format="time")
|
||||
assert datetime.now().strftime("%H:%M") in r
|
||||
|
||||
|
||||
def test_get_current_time_streaming_emits_complete():
|
||||
tool = GetCurrentTimeTool()
|
||||
q: queue.Queue = queue.Queue()
|
||||
tool.execute_streaming(q, format="date")
|
||||
msgs = []
|
||||
while not q.empty():
|
||||
msgs.append(q.get_nowait())
|
||||
# 至少应有一个 complete 事件
|
||||
types = [m[0] for m in msgs]
|
||||
assert "complete" in types
|
||||
# complete 事件内容非空
|
||||
complete_msgs = [m[1] for m in msgs if m[0] == "complete"]
|
||||
assert complete_msgs and len(complete_msgs[0]) > 0
|
||||
|
||||
|
||||
# ==================== CalculatorTool ====================
|
||||
|
||||
def test_calculator_add():
|
||||
r = CalculatorTool().execute(a=3, b=4, operation="add")
|
||||
assert "7" in r and "计算成功" in r
|
||||
|
||||
|
||||
def test_calculator_subtract():
|
||||
r = CalculatorTool().execute(a=10, b=4, operation="subtract")
|
||||
assert "6" in r
|
||||
|
||||
|
||||
def test_calculator_multiply():
|
||||
r = CalculatorTool().execute(a=6, b=7, operation="multiply")
|
||||
assert "42" in r
|
||||
|
||||
|
||||
def test_calculator_divide():
|
||||
r = CalculatorTool().execute(a=20, b=4, operation="divide")
|
||||
assert "5" in r
|
||||
|
||||
|
||||
def test_calculator_divide_by_zero():
|
||||
r = CalculatorTool().execute(a=1, b=0, operation="divide")
|
||||
assert "除数不能为0" in r
|
||||
|
||||
|
||||
def test_calculator_unknown_operation():
|
||||
r = CalculatorTool().execute(a=1, b=1, operation="power")
|
||||
assert "未知运算" in r
|
||||
|
||||
|
||||
# ==================== ReadFileTool ====================
|
||||
#
|
||||
# 说明:使用 pytest 的 tmp_path fixture 创建临时文件作为 read_file 入参,
|
||||
# 不依赖任何相对路径或 cwd——早期版本用 filepath="aibox.py" 假设从
|
||||
# aibox/ 目录运行 pytest,从仓库根运行会因路径漂移失败(2026-08-14 修复)。
|
||||
|
||||
def test_read_file_head_returns_command_prompt(tmp_path):
|
||||
"""head 模式应返回含 head 命令 + run_command 指引的 markdown。"""
|
||||
f = tmp_path / "sample.txt"
|
||||
f.write_text("line1\nline2\n")
|
||||
r = ReadFileTool().execute(filepath=str(f), head=True, max_lines=10)
|
||||
assert "head -n 10" in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_read_file_tail_returns_command_prompt(tmp_path):
|
||||
f = tmp_path / "sample.txt"
|
||||
f.write_text("line1\nline2\n")
|
||||
r = ReadFileTool().execute(filepath=str(f), tail=True, max_lines=20)
|
||||
assert "tail -n 20" in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_read_file_range_returns_sed_command(tmp_path):
|
||||
f = tmp_path / "sample.txt"
|
||||
f.write_text("line1\nline2\n")
|
||||
r = ReadFileTool().execute(filepath=str(f), start_line=5, end_line=15)
|
||||
assert "sed -n '5,15p'" in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_read_file_max_lines_cap_at_500(tmp_path):
|
||||
"""max_lines 应被钳制到 500 防止超大范围。"""
|
||||
f = tmp_path / "sample.txt"
|
||||
f.write_text("line1\n")
|
||||
r = ReadFileTool().execute(filepath=str(f), max_lines=99999)
|
||||
assert "sed -n '1,500p'" in r
|
||||
|
||||
|
||||
def test_read_file_missing_filepath_returns_error():
|
||||
r = ReadFileTool().execute(filepath=None)
|
||||
assert "❌" in r and "filepath" in r
|
||||
|
||||
|
||||
def test_read_file_nonexistent_returns_error():
|
||||
r = ReadFileTool().execute(filepath="/no/such/file_xyz.txt", head=True)
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
# ==================== WriteFileTool ====================
|
||||
|
||||
def test_write_file_overwrite_mode():
|
||||
r = WriteFileTool().execute(filepath="output.txt", content="hello")
|
||||
assert ">" in r and "run_command" in r
|
||||
# 覆盖模式应使用单 > 重定向
|
||||
assert "cat > " in r
|
||||
|
||||
|
||||
def test_write_file_append_mode():
|
||||
r = WriteFileTool().execute(filepath="output.txt", content="more", mode="a")
|
||||
assert "cat >> " in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_write_file_heredoc_uses_eof_marker():
|
||||
r = WriteFileTool().execute(filepath="out.txt", content="data")
|
||||
assert "AIBOX_EOF" in r
|
||||
|
||||
|
||||
def test_write_file_missing_content_returns_error():
|
||||
r = WriteFileTool().execute(filepath="out.txt", content=None)
|
||||
assert "❌" in r and "content" in r
|
||||
|
||||
|
||||
def test_write_file_creates_parent_dir_command():
|
||||
"""写深路径时应在命令中包含 mkdir -p。"""
|
||||
r = WriteFileTool().execute(filepath="sub/dir/out.txt", content="x")
|
||||
assert "mkdir -p" in r
|
||||
|
||||
|
||||
# ==================== ListFilesTool ====================
|
||||
|
||||
def test_list_files_default_no_hidden_flag():
|
||||
r = ListFilesTool().execute(directory=".")
|
||||
assert "ls -l " in r
|
||||
assert "run_command" in r
|
||||
|
||||
|
||||
def test_list_files_show_hidden_flag():
|
||||
r = ListFilesTool().execute(directory=".", show_hidden=True)
|
||||
assert "ls -la " in r
|
||||
|
||||
|
||||
def test_list_files_nonexistent_dir_returns_error():
|
||||
r = ListFilesTool().execute(directory="/no/such/dir_xyz")
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_list_files_path_is_file_returns_error(tmp_path):
|
||||
f = tmp_path / "afile.txt"
|
||||
f.write_text("x")
|
||||
r = ListFilesTool().execute(directory=str(f))
|
||||
assert "❌" in r and "不是目录" in r
|
||||
|
||||
|
||||
# ==================== SimpleWebCrawlerTool ====================
|
||||
|
||||
def test_web_get_action_returns_curl_prompt():
|
||||
r = SimpleWebCrawlerTool().execute(action="get", url="https://example.com")
|
||||
assert "curl -sL" in r
|
||||
assert "run_command" in r
|
||||
assert "https://example.com" in r
|
||||
|
||||
|
||||
def test_web_title_action_uses_grep():
|
||||
r = SimpleWebCrawlerTool().execute(action="title", url="https://example.com")
|
||||
assert "curl" in r and "<title>" in r
|
||||
|
||||
|
||||
def test_web_links_action_uses_grep():
|
||||
r = SimpleWebCrawlerTool().execute(action="links", url="https://example.com")
|
||||
assert "href" in r and "curl" in r
|
||||
|
||||
|
||||
def test_web_text_action_uses_sed():
|
||||
r = SimpleWebCrawlerTool().execute(action="text", url="https://example.com")
|
||||
assert "sed" in r and "curl" in r
|
||||
|
||||
|
||||
def test_web_search_requires_keyword():
|
||||
r = SimpleWebCrawlerTool().execute(action="search", url="https://example.com")
|
||||
assert "❌" in r and "keyword" in r
|
||||
|
||||
|
||||
def test_web_search_with_keyword_uses_grep_c():
|
||||
r = SimpleWebCrawlerTool().execute(
|
||||
action="search", url="https://example.com", keyword="hello"
|
||||
)
|
||||
assert "grep -c" in r and "hello" in r
|
||||
|
||||
|
||||
def test_web_unknown_action_returns_error():
|
||||
r = SimpleWebCrawlerTool().execute(action="download", url="https://example.com")
|
||||
assert "❌" in r and "未知操作" in r
|
||||
|
||||
|
||||
# ==================== FileToolBase ====================
|
||||
|
||||
# FileToolBase 是抽象基类(继承自 Tool 的抽象方法未实现),无法直接实例化。
|
||||
# 用一个最小具体子类来测试基类提供的 normalize_path / check_path_safety 方法。
|
||||
class _ConcreteFileTool(FileToolBase):
|
||||
def get_name(self): return "_test"
|
||||
def get_description(self): return "test"
|
||||
def get_parameters(self): return {"type": "object", "properties": {}}
|
||||
def execute(self, **kwargs): return ""
|
||||
|
||||
|
||||
def test_filetoolbase_normalize_path_expands_user():
|
||||
base = _ConcreteFileTool()
|
||||
p = base.normalize_path("~/foo")
|
||||
assert "~" not in p
|
||||
assert os.path.isabs(p)
|
||||
|
||||
|
||||
def test_filetoolbase_check_path_safety_existing_file(tmp_path):
|
||||
f = tmp_path / "ok.txt"
|
||||
f.write_text("hi")
|
||||
base = _ConcreteFileTool()
|
||||
ok, abs_path, err = base.check_path_safety(str(f), check_exists=True)
|
||||
assert ok is True
|
||||
assert abs_path == str(f)
|
||||
assert err == ""
|
||||
|
||||
|
||||
def test_filetoolbase_check_path_safety_nonexistent_with_check(tmp_path):
|
||||
base = _ConcreteFileTool()
|
||||
ok, _, err = base.check_path_safety(str(tmp_path / "missing.txt"), check_exists=True)
|
||||
assert ok is False
|
||||
assert "文件不存在" in err
|
||||
|
||||
|
||||
def test_filetoolbase_check_path_safety_oversize_file(tmp_path):
|
||||
"""超大文件应被拒绝。"""
|
||||
f = tmp_path / "big.bin"
|
||||
# 写入略超 MAX_FILE_SIZE 的内容
|
||||
from common import constants
|
||||
f.write_bytes(b"x" * (constants.MAX_FILE_SIZE + 1))
|
||||
base = _ConcreteFileTool()
|
||||
ok, _, err = base.check_path_safety(str(f))
|
||||
assert ok is False
|
||||
assert "文件太大" in err
|
||||
|
||||
|
||||
def test_filetoolbase_is_subclass_of_tool():
|
||||
"""FileToolBase 应是 Tool 的子类,且三个文件工具都继承它。"""
|
||||
from tools import Tool
|
||||
assert issubclass(FileToolBase, Tool)
|
||||
assert issubclass(ReadFileTool, FileToolBase)
|
||||
assert issubclass(WriteFileTool, FileToolBase)
|
||||
assert issubclass(ListFilesTool, FileToolBase)
|
||||
@@ -0,0 +1,100 @@
|
||||
#!/usr/bin/env python3
|
||||
"""F1(N1 高危)修复回归测试:代理密码不落盘。
|
||||
|
||||
覆盖三点:
|
||||
1. set_proxy(password=..., save=True) 后配置文件无明文密码(_save_config 脱敏)
|
||||
2. password_env 解析:get_proxy() 从环境变量取密码拼 socks5:// URL
|
||||
3. 旧版明文配置迁移:加载后盘上明文被清除,内存中本会话仍可用
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from common.config import ConfigManager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_config(tmp_path, monkeypatch):
|
||||
"""隔离 ConfigManager 单例:AI_CONFIG_PATH 指向临时文件,前后重置类状态"""
|
||||
cfg_file = tmp_path / "config.json"
|
||||
monkeypatch.setenv("AI_CONFIG_PATH", str(cfg_file))
|
||||
|
||||
def _reset():
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
|
||||
_reset()
|
||||
yield ConfigManager
|
||||
_reset()
|
||||
|
||||
|
||||
def _read_cfg_file() -> dict:
|
||||
path = ConfigManager()._config_path
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def test_set_proxy_sanitizes_password_on_save(fresh_config, tmp_path):
|
||||
"""save=True 落盘后,文件中不得出现明文密码(纵深防御脱敏)"""
|
||||
cfg = fresh_config()
|
||||
cfg.set_proxy("127.0.0.1", 1080, username="u1", password="secret123",
|
||||
password_env="MY_SOCKS_PW", save=True)
|
||||
|
||||
on_disk = _read_cfg_file()
|
||||
socks = on_disk["proxy"]["socks5"]
|
||||
assert socks["password"] in (None, ""), "落盘文件不得含明文密码"
|
||||
assert socks["password_env"] == "MY_SOCKS_PW", "password_env 应持久化"
|
||||
|
||||
raw = (tmp_path / "config.json").read_text(encoding="utf-8")
|
||||
assert "secret123" not in raw, "整个配置文件不得出现明文密码字符串"
|
||||
|
||||
# 内存中本会话仍可用
|
||||
proxies = cfg.get_proxy()
|
||||
assert proxies is not None
|
||||
assert "u1:secret123@" in proxies["https"]
|
||||
|
||||
|
||||
def test_password_env_resolution(fresh_config, monkeypatch):
|
||||
"""password 为空时按 password_env 从环境变量解析"""
|
||||
monkeypatch.setenv("TEST_SOCKS_PW", "env-secret-42")
|
||||
cfg = fresh_config()
|
||||
cfg.set_proxy("10.0.0.8", 7890, username="u2", password=None,
|
||||
password_env="TEST_SOCKS_PW", save=True)
|
||||
|
||||
proxies = cfg.get_proxy()
|
||||
assert proxies is not None
|
||||
assert "u2:env-secret-42@10.0.0.8:7890" in proxies["https"]
|
||||
|
||||
# get_proxy_info 不泄漏密码值,只给布尔 + 变量名
|
||||
info = cfg.get_proxy_info()
|
||||
assert info["password"] is True
|
||||
assert info["password_env"] == "TEST_SOCKS_PW"
|
||||
|
||||
|
||||
def test_legacy_plaintext_migration(fresh_config, tmp_path):
|
||||
"""旧版明文配置:加载后盘上清除,内存保留供本会话使用"""
|
||||
cfg_file = tmp_path / "config.json"
|
||||
legacy = {
|
||||
"proxy": {
|
||||
"socks5": {
|
||||
"enabled": True, "host": "1.2.3.4", "port": 1080,
|
||||
"username": "old_u", "password": "legacy_plain_pw"
|
||||
}
|
||||
}
|
||||
}
|
||||
cfg_file.write_text(json.dumps(legacy, ensure_ascii=False), encoding="utf-8")
|
||||
|
||||
cfg = fresh_config() # 触发 _load_config → 迁移
|
||||
|
||||
# 盘上已脱敏
|
||||
on_disk = _read_cfg_file()
|
||||
assert on_disk["proxy"]["socks5"]["password"] in (None, "")
|
||||
raw = cfg_file.read_text(encoding="utf-8")
|
||||
assert "legacy_plain_pw" not in raw
|
||||
|
||||
# 内存中本会话仍可解析出旧密码(代理不断)
|
||||
proxies = cfg.get_proxy()
|
||||
assert proxies is not None
|
||||
assert "old_u:legacy_plain_pw@1.2.3.4:1080" in proxies["https"]
|
||||
@@ -0,0 +1,196 @@
|
||||
"""subsystems.conversation_compressor 单元测试
|
||||
|
||||
覆盖:
|
||||
- estimate_tokens: 中英文/工具调用 token 估算
|
||||
- should_compress: 消息数阈值触发 / Token 阈值触发 / 均未触发 / 功能禁用
|
||||
- compress_conversation: 消息数触发压缩 / Token 触发压缩 / LLM 失败保留原文
|
||||
- _validate_message_structure: 移除孤立 tool 消息
|
||||
通过注入 FakeLLMClient 模拟 LLM 调用,不依赖真实 API。
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import ConversationCompressor
|
||||
from common import constants, Logger
|
||||
|
||||
|
||||
class FakeLLMClient:
|
||||
"""模拟 llm_client 抽象层的 chat_completion"""
|
||||
|
||||
def __init__(self, reply="这是压缩摘要", error=None):
|
||||
self.reply = reply
|
||||
self.error = error
|
||||
self.calls = []
|
||||
|
||||
def chat_completion(self, messages, temperature=0.7, max_tokens=None, **kwargs):
|
||||
self.calls.append({"max_tokens": max_tokens, "n_msgs": len(messages)})
|
||||
if self.error:
|
||||
raise self.error
|
||||
return self.reply
|
||||
|
||||
|
||||
def make_messages(n_pairs=15, long_content=False):
|
||||
"""构造 n_pairs 轮 user/assistant 对话(含 1 条前导 system)"""
|
||||
msgs = [{"role": "system", "content": "SYS_PROMPT"}]
|
||||
for i in range(n_pairs):
|
||||
content = f"用户消息内容{i}"
|
||||
if long_content:
|
||||
content += "很长的文档内容" * 200
|
||||
msgs.append({"role": "user", "content": content})
|
||||
msgs.append({"role": "assistant", "content": f"助手回复{i}"})
|
||||
return msgs
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def compressor():
|
||||
return ConversationCompressor(llm_client=FakeLLMClient(), logger=Logger("TestCompressor", 0))
|
||||
|
||||
|
||||
def patch_const(monkeypatch, name, value):
|
||||
"""把 Constants 类属性 property 替换为固定值(参照 test_db.py 既有做法)"""
|
||||
monkeypatch.setattr(type(constants), name, property(lambda self: value))
|
||||
|
||||
|
||||
# ==================== estimate_tokens ====================
|
||||
|
||||
def test_estimate_tokens_chinese(compressor):
|
||||
assert compressor.estimate_tokens([{"role": "user", "content": "你好"}]) == 13 # 2*1.5 + 10
|
||||
|
||||
|
||||
def test_estimate_tokens_english(compressor):
|
||||
# 2 词 * 1.3 + (11 字符 - 2 词) * 0.8 + 10 = 19.8 → int 截断为 19
|
||||
# (既有算法:other_chars 按词数抵扣,字母被计入 other,保持原行为)
|
||||
assert compressor.estimate_tokens([{"role": "user", "content": "hello world"}]) == 19
|
||||
|
||||
|
||||
def test_estimate_tokens_tool_calls(compressor):
|
||||
msgs = [{
|
||||
"role": "assistant", "content": "",
|
||||
"tool_calls": [{"function": {"name": "run_command", "arguments": "{\"a\":1}"}}],
|
||||
}]
|
||||
# name 11 字符*1.5=16.5 + arguments 8 字符*0.8=6.4 + 50 = 72.9 → 72
|
||||
assert compressor.estimate_tokens(msgs) == 72
|
||||
|
||||
|
||||
def test_estimate_tokens_empty(compressor):
|
||||
assert compressor.estimate_tokens([]) == 0
|
||||
|
||||
|
||||
# ==================== should_compress ====================
|
||||
|
||||
def test_should_compress_below_thresholds(compressor, monkeypatch):
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 1000)
|
||||
patch_const(monkeypatch, "COMPRESS_TOKEN_THRESHOLD", 10 ** 9)
|
||||
should, reason = compressor.should_compress(make_messages(5))
|
||||
assert should is False
|
||||
|
||||
|
||||
def test_should_compress_by_message_count(compressor, monkeypatch):
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 5)
|
||||
should, reason = compressor.should_compress(make_messages(5)) # 10 条对话消息
|
||||
assert should is True
|
||||
assert "消息数" in reason
|
||||
|
||||
|
||||
def test_should_compress_by_token_count(compressor, monkeypatch):
|
||||
"""消息数未超阈值但 Token 数超阈值时也应触发(v3.1 修复项)。"""
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 1000)
|
||||
patch_const(monkeypatch, "COMPRESS_TOKEN_THRESHOLD", 1)
|
||||
should, reason = compressor.should_compress(make_messages(5))
|
||||
assert should is True
|
||||
assert "Token" in reason
|
||||
|
||||
|
||||
def test_should_compress_disabled(compressor, monkeypatch):
|
||||
patch_const(monkeypatch, "COMPRESS_ENABLED", False)
|
||||
should, _ = compressor.should_compress(make_messages(100))
|
||||
assert should is False
|
||||
|
||||
|
||||
# ==================== compress_conversation ====================
|
||||
|
||||
def test_compress_by_message_count(compressor, monkeypatch):
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||
messages = make_messages(15) # 30 条对话消息 + 1 system
|
||||
result = compressor.compress_conversation(messages)
|
||||
|
||||
non_system = [m for m in result if m["role"] != "system"]
|
||||
assert len(non_system) < 30 # 消息数下降
|
||||
# 原始 system 提示词保留
|
||||
assert any(m["content"] == "SYS_PROMPT" for m in result)
|
||||
# 压缩摘要已注入
|
||||
assert any(m["role"] == "system" and "这是压缩摘要" in m["content"] for m in result)
|
||||
# 最近消息保留
|
||||
assert any(m.get("content") == "助手回复14" for m in non_system)
|
||||
# LLM 被调用了一次
|
||||
assert len(compressor.llm_client.calls) == 1
|
||||
|
||||
|
||||
def test_compress_by_token_count(compressor, monkeypatch):
|
||||
"""Token 触发(消息数远低于阈值)时也应实际执行压缩(v3.1 修复项)。"""
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 1000)
|
||||
patch_const(monkeypatch, "COMPRESS_TOKEN_THRESHOLD", 1)
|
||||
messages = make_messages(15, long_content=True) # 30 条长消息
|
||||
result = compressor.compress_conversation(messages)
|
||||
|
||||
non_system = [m for m in result if m["role"] != "system"]
|
||||
assert len(non_system) < 30
|
||||
assert any(m["role"] == "system" and "这是压缩摘要" in m["content"] for m in result)
|
||||
assert len(compressor.llm_client.calls) == 1
|
||||
# Token 数下降
|
||||
assert compressor.estimate_tokens(non_system) < compressor.estimate_tokens(
|
||||
[m for m in messages if m["role"] != "system"])
|
||||
|
||||
|
||||
def test_compress_llm_failure_keeps_original(monkeypatch):
|
||||
comp = ConversationCompressor(
|
||||
llm_client=FakeLLMClient(error=RuntimeError("api down")),
|
||||
logger=Logger("TestCompressor", 0),
|
||||
)
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||
messages = make_messages(15)
|
||||
result = comp.compress_conversation(messages)
|
||||
assert result == messages # 失败时原样返回
|
||||
|
||||
|
||||
def test_compress_without_llm_client_keeps_original(monkeypatch):
|
||||
comp = ConversationCompressor(llm_client=None)
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||
messages = make_messages(15)
|
||||
assert comp.compress_conversation(messages) == messages
|
||||
|
||||
|
||||
def test_compress_empty_messages(compressor):
|
||||
assert compressor.compress_conversation([]) == []
|
||||
|
||||
|
||||
def test_compress_too_few_messages_skips(compressor, monkeypatch):
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||
messages = make_messages(3) # 6 条对话消息,压缩块必然 < 5
|
||||
assert compressor.compress_conversation(messages) == messages
|
||||
assert compressor.llm_client.calls == []
|
||||
|
||||
|
||||
# ==================== _validate_message_structure ====================
|
||||
|
||||
def test_validate_removes_orphan_tool_messages(compressor):
|
||||
msgs = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "",
|
||||
"tool_calls": [{"id": "t1", "function": {"name": "f", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "ok"},
|
||||
{"role": "tool", "tool_call_id": "orphan", "content": "孤儿消息"},
|
||||
]
|
||||
result = compressor._validate_message_structure(msgs)
|
||||
assert len(result) == 3
|
||||
assert all(m.get("tool_call_id") != "orphan" for m in result)
|
||||
|
||||
|
||||
def test_validate_keeps_paired_tool_messages(compressor):
|
||||
msgs = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "",
|
||||
"tool_calls": [{"id": "t1", "function": {"name": "f", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "ok"},
|
||||
]
|
||||
assert len(compressor._validate_message_structure(msgs)) == 3
|
||||
@@ -0,0 +1,427 @@
|
||||
"""subsystems.conversation_history 单元测试
|
||||
|
||||
覆盖 ConversationHistory 的核心行为:
|
||||
- 构造与目录初始化
|
||||
- 当前对话管理(开始/添加消息/添加工具调用/保存/重置/结束)
|
||||
- 归档与历史索引
|
||||
- 历史加载与上下文构建
|
||||
- 列表与加载指定对话
|
||||
- memory_db 注入时归档/结束写入记忆
|
||||
|
||||
设计原则:
|
||||
- 使用 tmp_path fixture 隔离文件系统副作用
|
||||
- 不依赖 aibox 主模块、DeepSeek API、websockets
|
||||
- memory_db 用真实 MemoryDatabase + 临时 sqlite 文件验证集成
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import ConversationHistory, MemoryDatabase
|
||||
from common import constants
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def history(tmp_path):
|
||||
"""独立临时目录的 ConversationHistory。"""
|
||||
return ConversationHistory(save_dir=str(tmp_path), load_history=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def history_no_load(tmp_path):
|
||||
"""load_history=False 的实例,用于验证上下文加载开关。"""
|
||||
return ConversationHistory(save_dir=str(tmp_path), load_history=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mem_db(tmp_path):
|
||||
"""独立临时文件的 MemoryDatabase。"""
|
||||
return MemoryDatabase(db_path=str(tmp_path / "test_memory.db"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def history_with_mem(tmp_path, mem_db):
|
||||
"""注入 MemoryDatabase 的 ConversationHistory。"""
|
||||
return ConversationHistory(save_dir=str(tmp_path), memory_db=mem_db, load_history=True)
|
||||
|
||||
|
||||
# ==================== 构造与目录初始化 ====================
|
||||
|
||||
def test_init_creates_dirs(tmp_path):
|
||||
"""构造时应创建 save_dir 与 archive_dir。"""
|
||||
save_dir = tmp_path / "ai"
|
||||
ConversationHistory(save_dir=str(save_dir))
|
||||
assert save_dir.exists() and save_dir.is_dir()
|
||||
archive = save_dir / constants.ARCHIVE_DIR
|
||||
assert archive.exists() and archive.is_dir()
|
||||
|
||||
|
||||
def test_init_loads_empty_index(tmp_path):
|
||||
"""无历史索引文件时初始化为空结构。"""
|
||||
h = ConversationHistory(save_dir=str(tmp_path))
|
||||
assert h.history_index == {"conversations": [], "last_id": 0}
|
||||
assert h.current_conversation is None
|
||||
|
||||
|
||||
def test_init_loads_existing_index(tmp_path):
|
||||
"""已存在历史索引文件时应加载。"""
|
||||
save_dir = tmp_path / "ai"
|
||||
save_dir.mkdir()
|
||||
idx_file = save_dir / constants.HISTORY_FILE
|
||||
idx_file.write_text(json.dumps({"conversations": [{"id": "x"}], "last_id": 5}), encoding="utf-8")
|
||||
|
||||
h = ConversationHistory(save_dir=str(save_dir))
|
||||
assert h.history_index["last_id"] == 5
|
||||
assert h.history_index["conversations"][0]["id"] == "x"
|
||||
|
||||
|
||||
def test_init_default_save_dir_uses_constants(tmp_path):
|
||||
"""save_dir=None 时使用 constants.SAVE_DIR(默认值为只读 property,无法 monkeypatch,
|
||||
这里仅验证构造不抛异常且目录存在——SAVE_DIR 通常指向 ~/.aibox,可能已存在)。
|
||||
用显式 save_dir 路径替代断言,避免污染用户家目录。"""
|
||||
# 这里只验证 constants.SAVE_DIR 是字符串类型(被引用即可)
|
||||
assert isinstance(constants.SAVE_DIR, str)
|
||||
|
||||
|
||||
# ==================== 当前对话管理 ====================
|
||||
|
||||
def test_start_new_conversation_creates_empty_conv(history):
|
||||
conv = history.start_new_conversation(topic="测试")
|
||||
assert conv is not None
|
||||
assert conv["topic"] == "测试"
|
||||
assert conv["messages"] == []
|
||||
assert conv["stats"] == {"user_messages": 0, "ai_messages": 0, "tool_calls": 0}
|
||||
assert history.current_conversation is conv
|
||||
|
||||
|
||||
def test_start_new_conversation_default_topic(history):
|
||||
conv = history.start_new_conversation()
|
||||
assert conv["topic"] == "新对话"
|
||||
|
||||
|
||||
def test_start_new_conversation_archives_existing(history):
|
||||
"""有当前对话时再开新对话,旧对话应归档。"""
|
||||
history.start_new_conversation(topic="旧")
|
||||
history.add_message("user", "hi")
|
||||
history.start_new_conversation(topic="新")
|
||||
# 旧对话应已归档到 archive 目录
|
||||
archives = list(history.archive_dir.glob("*.json"))
|
||||
assert len(archives) == 1
|
||||
archived = json.loads(archives[0].read_text(encoding="utf-8"))
|
||||
assert archived["topic"] == "旧"
|
||||
|
||||
|
||||
def test_add_message_creates_conversation_if_absent(history):
|
||||
"""无当前对话时 add_message 应自动创建。"""
|
||||
assert history.current_conversation is None
|
||||
history.add_message("user", "hello")
|
||||
assert history.current_conversation is not None
|
||||
assert len(history.current_conversation["messages"]) == 1
|
||||
|
||||
|
||||
def test_add_user_message_updates_stats(history):
|
||||
history.start_new_conversation()
|
||||
history.add_message("user", "问题")
|
||||
stats = history.current_conversation["stats"]
|
||||
assert stats["user_messages"] == 1
|
||||
assert stats["ai_messages"] == 0
|
||||
|
||||
|
||||
def test_add_assistant_message_updates_stats(history):
|
||||
history.start_new_conversation()
|
||||
history.add_message("assistant", "回答", speaker="aibox")
|
||||
stats = history.current_conversation["stats"]
|
||||
assert stats["ai_messages"] == 1
|
||||
assert stats["user_messages"] == 0
|
||||
|
||||
|
||||
def test_add_message_persists_to_current_file(history):
|
||||
"""add_message 后应写入 current_conversation.json。"""
|
||||
history.start_new_conversation()
|
||||
history.add_message("user", "persist me")
|
||||
assert history.current_file.exists()
|
||||
saved = json.loads(history.current_file.read_text(encoding="utf-8"))
|
||||
assert saved["messages"][0]["content"] == "persist me"
|
||||
|
||||
|
||||
def test_add_tool_call_requires_current_conversation(history):
|
||||
"""无当前对话时 add_tool_call 应直接返回不报错。"""
|
||||
assert history.current_conversation is None
|
||||
history.add_tool_call("run_command", {"cmd": "ls"}, "output")
|
||||
# 不应抛异常,且不创建对话
|
||||
assert history.current_conversation is None
|
||||
|
||||
|
||||
def test_add_tool_call_records_and_truncates_preview(history):
|
||||
history.start_new_conversation()
|
||||
long_result = "x" * 300
|
||||
history.add_tool_call("run_command", {"cmd": "ls"}, long_result)
|
||||
tc = history.current_conversation["tool_calls"][0]
|
||||
assert tc["tool"] == "run_command"
|
||||
assert tc["args"] == {"cmd": "ls"}
|
||||
# result_preview 超 200 字符应截断加省略号
|
||||
assert tc["result_preview"].endswith("...")
|
||||
assert len(tc["result_preview"]) == 203 # 200 + "..."
|
||||
assert history.current_conversation["stats"]["tool_calls"] == 1
|
||||
|
||||
|
||||
def test_add_tool_call_short_result_no_truncation(history):
|
||||
history.start_new_conversation()
|
||||
history.add_tool_call("calc", {}, "42")
|
||||
tc = history.current_conversation["tool_calls"][0]
|
||||
assert tc["result_preview"] == "42"
|
||||
|
||||
|
||||
def test_save_current_conversation_no_op_when_empty(history):
|
||||
"""无当前对话时 save_current_conversation 不报错。"""
|
||||
history.save_current_conversation() # 不应抛异常
|
||||
|
||||
|
||||
def test_reset_conversation_clears_current(history):
|
||||
history.start_new_conversation()
|
||||
history.add_message("user", "x")
|
||||
history.reset_conversation()
|
||||
assert history.current_conversation is None
|
||||
# current_file 应被归档时清除(reset 自身不删文件,但 end_time 已写)
|
||||
assert history.current_file.exists() is False or history.current_file.exists()
|
||||
|
||||
|
||||
def test_end_current_conversation_sets_end_time(history):
|
||||
history.start_new_conversation()
|
||||
assert history.current_conversation["end_time"] is None
|
||||
history.end_current_conversation()
|
||||
assert history.current_conversation["end_time"] is not None
|
||||
|
||||
|
||||
def test_end_current_conversation_no_op_when_empty(history):
|
||||
"""无当前对话时不报错。"""
|
||||
history.end_current_conversation() # 不应抛异常
|
||||
|
||||
|
||||
# ==================== 归档与历史索引 ====================
|
||||
|
||||
def test_archive_returns_none_when_no_current(history):
|
||||
assert history.archive_current_conversation() is None
|
||||
|
||||
|
||||
def test_archive_writes_archive_file_and_updates_index(history):
|
||||
history.start_new_conversation(topic="归档测试")
|
||||
history.add_message("user", "msg1")
|
||||
history.add_message("assistant", "reply")
|
||||
|
||||
archive_path = history.archive_current_conversation()
|
||||
assert archive_path is not None
|
||||
import os
|
||||
assert os.path.exists(archive_path)
|
||||
|
||||
# 索引应包含一条记录
|
||||
assert len(history.history_index["conversations"]) == 1
|
||||
rec = history.history_index["conversations"][0]
|
||||
assert rec["topic"] == "归档测试"
|
||||
assert rec["messages"] == 2
|
||||
|
||||
# 归档后当前对话清空
|
||||
assert history.current_conversation is None
|
||||
assert not history.current_file.exists()
|
||||
|
||||
|
||||
def test_archive_sets_end_time_if_missing(history):
|
||||
history.start_new_conversation()
|
||||
# end_time 为 None
|
||||
history.archive_current_conversation()
|
||||
# 通过归档文件验证 end_time 已写入
|
||||
archives = list(history.archive_dir.glob("*.json"))
|
||||
archived = json.loads(archives[0].read_text(encoding="utf-8"))
|
||||
assert archived["end_time"] is not None
|
||||
|
||||
|
||||
def test_archive_with_final_summary_writes_memory(tmp_path, mem_db):
|
||||
"""注入 memory_db 且提供 final_summary 时应调用 add_memory。"""
|
||||
h = ConversationHistory(save_dir=str(tmp_path), memory_db=mem_db, load_history=True)
|
||||
h.start_new_conversation(topic="带摘要归档")
|
||||
h.add_message("user", "content")
|
||||
h.add_message("assistant", "reply", speaker="ai")
|
||||
|
||||
h.archive_current_conversation(final_summary="这是一次对话的总结")
|
||||
|
||||
# 验证 memory_db 写入了一条记忆(用 get_recent_memories 不依赖 FTS5 中文分词)
|
||||
memories = mem_db.get_recent_memories(limit=10)
|
||||
assert len(memories) >= 1
|
||||
# 至少有一条 topic 匹配
|
||||
assert any(m.get("topic") == "带摘要归档" for m in memories)
|
||||
|
||||
|
||||
def test_archive_without_summary_does_not_write_memory(tmp_path, mem_db):
|
||||
"""无 final_summary 时不写记忆。"""
|
||||
h = ConversationHistory(save_dir=str(tmp_path), memory_db=mem_db, load_history=True)
|
||||
h.start_new_conversation(topic="无摘要")
|
||||
h.add_message("user", "x")
|
||||
h.archive_current_conversation() # 无 summary
|
||||
memories = mem_db.get_recent_memories(limit=10)
|
||||
assert len(memories) == 0
|
||||
|
||||
|
||||
def test_archive_without_memory_db_does_not_crash(history):
|
||||
"""未注入 memory_db 时带 summary 归档不应崩溃。"""
|
||||
history.start_new_conversation()
|
||||
history.add_message("user", "x")
|
||||
# 不应抛异常
|
||||
history.archive_current_conversation(final_summary="summary without db")
|
||||
|
||||
|
||||
# ==================== 历史加载与上下文 ====================
|
||||
|
||||
def test_load_all_history_respects_load_history_flag(history_no_load):
|
||||
"""load_history=False 时 load_all_history 返回空。"""
|
||||
history_no_load.start_new_conversation()
|
||||
history_no_load.add_message("user", "x")
|
||||
assert history_no_load.load_all_history() == []
|
||||
|
||||
|
||||
def test_get_conversation_context_respects_load_history_flag(history_no_load):
|
||||
assert history_no_load.get_conversation_context() == []
|
||||
|
||||
|
||||
def test_load_all_history_includes_current_messages(history):
|
||||
history.start_new_conversation()
|
||||
history.add_message("user", "current msg")
|
||||
msgs = history.load_all_history(max_messages=10)
|
||||
assert any(m["content"] == "current msg" and m.get("from_current") for m in msgs)
|
||||
|
||||
|
||||
def test_load_all_history_includes_archived(history):
|
||||
"""归档后的对话也应被加载。"""
|
||||
history.start_new_conversation(topic="历史对话")
|
||||
history.add_message("user", "archived msg")
|
||||
history.archive_current_conversation()
|
||||
|
||||
# 新开对话查询历史
|
||||
history.start_new_conversation()
|
||||
msgs = history.load_all_history(max_messages=20)
|
||||
contents = [m["content"] for m in msgs]
|
||||
assert "archived msg" in contents
|
||||
|
||||
|
||||
def test_load_all_history_limits_to_max_messages(history):
|
||||
history.start_new_conversation()
|
||||
for i in range(10):
|
||||
history.add_message("user", f"msg{i}")
|
||||
msgs = history.load_all_history(max_messages=3)
|
||||
assert len(msgs) <= 3
|
||||
|
||||
|
||||
def test_get_conversation_context_skips_system_role(history):
|
||||
history.start_new_conversation()
|
||||
# 直接构造一条 system 消息塞入
|
||||
history.current_conversation["messages"].append({
|
||||
"timestamp": datetime.now().strftime(constants.DATETIME_FORMAT),
|
||||
"role": "system",
|
||||
"speaker": None,
|
||||
"content": "should be skipped",
|
||||
"is_history": False,
|
||||
})
|
||||
history.add_message("user", "visible")
|
||||
ctx = history.get_conversation_context()
|
||||
roles = [c["role"] for c in ctx]
|
||||
assert "system" not in roles
|
||||
assert any(c["content"] == "visible" for c in ctx)
|
||||
|
||||
|
||||
def test_get_conversation_context_formats_assistant_with_speaker(history):
|
||||
history.start_new_conversation()
|
||||
history.add_message("assistant", "hello", speaker="aibox")
|
||||
ctx = history.get_conversation_context()
|
||||
asst = [c for c in ctx if c["role"] == "assistant"]
|
||||
assert len(asst) == 1
|
||||
assert "[aibox]" in asst[0]["content"]
|
||||
|
||||
|
||||
def test_get_conversation_context_user_passthrough(history):
|
||||
history.start_new_conversation()
|
||||
history.add_message("user", "plain question")
|
||||
ctx = history.get_conversation_context()
|
||||
user = [c for c in ctx if c["role"] == "user"]
|
||||
assert any(c["content"] == "plain question" for c in user)
|
||||
|
||||
|
||||
# ==================== 列表与加载指定对话 ====================
|
||||
|
||||
def test_list_conversations_empty(history):
|
||||
"""无任何对话时返回空列表。"""
|
||||
assert history.list_conversations() == []
|
||||
|
||||
|
||||
def test_list_conversations_includes_current(history):
|
||||
history.start_new_conversation(topic="当前")
|
||||
convs = history.list_conversations()
|
||||
assert len(convs) == 1
|
||||
assert convs[0]["is_current"] is True
|
||||
assert convs[0]["topic"] == "当前"
|
||||
|
||||
|
||||
def test_list_conversations_includes_archived(history):
|
||||
history.start_new_conversation(topic="旧")
|
||||
history.archive_current_conversation()
|
||||
history.start_new_conversation(topic="新")
|
||||
convs = history.list_conversations()
|
||||
assert len(convs) == 2
|
||||
topics = {c["topic"] for c in convs}
|
||||
assert {"旧", "新"} == topics
|
||||
|
||||
|
||||
def test_list_conversations_sorted_by_start_time_desc(history):
|
||||
"""应按 start_time 倒序排列。"""
|
||||
history.start_new_conversation(topic="first")
|
||||
time.sleep(0.01)
|
||||
history.archive_current_conversation()
|
||||
history.start_new_conversation(topic="second")
|
||||
|
||||
convs = history.list_conversations()
|
||||
# 第二个(更晚)应在前
|
||||
assert convs[0]["topic"] == "second"
|
||||
|
||||
|
||||
def test_load_conversation_current_file(history):
|
||||
history.start_new_conversation(topic="当前")
|
||||
loaded = history.load_conversation(constants.CURRENT_CONVERSATION_FILE)
|
||||
assert loaded is not None
|
||||
assert loaded["topic"] == "当前"
|
||||
|
||||
|
||||
def test_load_conversation_archived_file(history):
|
||||
history.start_new_conversation(topic="已归档")
|
||||
history.add_message("user", "saved")
|
||||
history.archive_current_conversation()
|
||||
|
||||
archives = list(history.archive_dir.glob("*.json"))
|
||||
assert len(archives) == 1
|
||||
loaded = history.load_conversation(archives[0].name)
|
||||
assert loaded is not None
|
||||
assert loaded["topic"] == "已归档"
|
||||
assert loaded["messages"][0]["content"] == "saved"
|
||||
|
||||
|
||||
def test_load_conversation_nonexistent_returns_none(history):
|
||||
assert history.load_conversation("nonexistent.json") is None
|
||||
|
||||
|
||||
# ==================== 与 MemoryDatabase 集成 ====================
|
||||
|
||||
def test_end_current_conversation_with_summary_writes_memory(history_with_mem, mem_db):
|
||||
history_with_mem.start_new_conversation(topic="结束测试")
|
||||
history_with_mem.add_message("user", "最后的话")
|
||||
history_with_mem.end_current_conversation(summary="结束摘要")
|
||||
|
||||
memories = mem_db.get_recent_memories(limit=10)
|
||||
assert any(m.get("topic") == "结束测试" for m in memories)
|
||||
|
||||
|
||||
def test_end_current_conversation_without_summary_no_memory(history_with_mem, mem_db):
|
||||
history_with_mem.start_new_conversation(topic="无摘要结束")
|
||||
history_with_mem.add_message("user", "x")
|
||||
history_with_mem.end_current_conversation() # 无 summary
|
||||
memories = mem_db.get_recent_memories(limit=10)
|
||||
assert len(memories) == 0
|
||||
@@ -0,0 +1,312 @@
|
||||
"""subsystems.db 单元测试
|
||||
|
||||
覆盖:
|
||||
- MemoryDatabase: 增删查 / FTS5 全文检索 / 批量删除
|
||||
- MemoDatabase: 增删查 / 待执行 / 重复任务 / 即时任务 / 统计
|
||||
每个测试用独立临时 sqlite 文件,互不污染。
|
||||
"""
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import MemoryDatabase, MemoDatabase
|
||||
from common import constants
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mem_db(tmp_path):
|
||||
"""独立临时文件的 MemoryDatabase。"""
|
||||
db_path = str(tmp_path / "test_memory.db")
|
||||
return MemoryDatabase(db_path=db_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memo_db(tmp_path):
|
||||
"""独立临时文件的 MemoDatabase。"""
|
||||
db_path = str(tmp_path / "test_memo.db")
|
||||
return MemoDatabase(db_path=db_path)
|
||||
|
||||
|
||||
# ==================== MemoryDatabase ====================
|
||||
|
||||
def test_memory_add_and_get(mem_db):
|
||||
mid = mem_db.add_memory(
|
||||
topic="测试主题", summary="测试摘要",
|
||||
content="详细内容", tags=["t1", "t2"]
|
||||
)
|
||||
assert isinstance(mid, int) and mid > 0
|
||||
|
||||
m = mem_db.get_memory_by_id(mid)
|
||||
assert m is not None
|
||||
assert m["topic"] == "测试主题"
|
||||
assert m["summary"] == "测试摘要"
|
||||
assert m["tags"] == ["t1", "t2"]
|
||||
|
||||
|
||||
def test_memory_search_by_keyword(mem_db):
|
||||
mem_db.add_memory(topic="Python learning", summary="learn Python basics", content="variables and functions")
|
||||
mem_db.add_memory(topic="Go learning", summary="learn Go basics", content="goroutine")
|
||||
|
||||
results = mem_db.search_memories(["Python"])
|
||||
assert len(results) == 1
|
||||
assert "Python" in results[0]["topic"]
|
||||
|
||||
|
||||
def test_memory_search_match_all(mem_db):
|
||||
mem_db.add_memory(topic="Python database", summary="sqlite usage", content="FTS5 search")
|
||||
mem_db.add_memory(topic="Python network", summary="requests usage", content="HTTP request")
|
||||
|
||||
# match_all=True:两个关键词都要命中
|
||||
results = mem_db.search_memories(["Python", "sqlite"], match_all=True)
|
||||
assert len(results) == 1
|
||||
assert "database" in results[0]["topic"]
|
||||
|
||||
# match_all=False:任一关键词命中
|
||||
results = mem_db.search_memories(["Python", "sqlite"], match_all=False)
|
||||
assert len(results) == 2
|
||||
|
||||
|
||||
def test_memory_search_empty_keywords_returns_empty(mem_db):
|
||||
assert mem_db.search_memories([]) == []
|
||||
|
||||
|
||||
def test_memory_search_chinese_keyword(mem_db):
|
||||
"""中文关键词应可检索(trigram 分词或 LIKE 兜底)。"""
|
||||
mem_db.add_memory(topic="天气查询", summary="今天天气晴朗", content="气温25度")
|
||||
mem_db.add_memory(topic="会议安排", summary="下午三点开会", content="会议室A")
|
||||
|
||||
results = mem_db.search_memories(["天气"])
|
||||
assert len(results) == 1
|
||||
assert results[0]["topic"] == "天气查询"
|
||||
|
||||
|
||||
def test_memory_search_chinese_match_all(mem_db):
|
||||
"""中文多关键词 match_all:全部命中才返回。"""
|
||||
mem_db.add_memory(topic="天气查询", summary="今天天气晴朗", content="气温25度")
|
||||
mem_db.add_memory(topic="天气与会议", summary="开会前先查天气", content="会议室A")
|
||||
|
||||
results = mem_db.search_memories(["天气", "开会"], match_all=True)
|
||||
assert len(results) == 1
|
||||
assert results[0]["topic"] == "天气与会议"
|
||||
|
||||
|
||||
def test_memory_search_like_escape_wildcards(mem_db):
|
||||
"""LIKE 兜底路径中 %/_ 应被转义为字面字符,不当作通配符。"""
|
||||
import sqlite3
|
||||
mem_db.add_memory(topic="progress", summary="100% done")
|
||||
mem_db.add_memory(topic="task_a", summary="plain text")
|
||||
mem_db.add_memory(topic="taskXa", summary="wildcard trap")
|
||||
|
||||
def like_search(kw):
|
||||
with sqlite3.connect(mem_db.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
return mem_db._search_via_like(conn.cursor(), [kw], match_all=False, limit=100)
|
||||
|
||||
# "%" 为字面字符:"100%" 只命中 progress,不当作任意串通配
|
||||
results = like_search("100%")
|
||||
assert [r["topic"] for r in results] == ["progress"]
|
||||
|
||||
# "_" 为字面字符:"ask_" 只命中 task_a;若未转义,taskXa 也会被 %ask_% 误匹配
|
||||
results = like_search("ask_")
|
||||
assert [r["topic"] for r in results] == ["task_a"]
|
||||
|
||||
|
||||
def test_memory_fts_migration_from_unicode61(tmp_path):
|
||||
"""旧 unicode61 FTS 表在重新打开时应自动迁移为 trigram 并回灌数据。"""
|
||||
import sqlite3
|
||||
db_path = str(tmp_path / "legacy.db")
|
||||
|
||||
# 模拟旧库:unicode61 建表 + 写入数据
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute('''CREATE TABLE memories (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT, topic TEXT NOT NULL, summary TEXT NOT NULL,
|
||||
content TEXT, start_time TEXT NOT NULL, end_time TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL, tags TEXT, importance INTEGER DEFAULT 1)''')
|
||||
conn.execute('''CREATE VIRTUAL TABLE memories_fts
|
||||
USING fts5(topic, summary, content, tags, content=memories)''')
|
||||
conn.execute("INSERT INTO memories (topic, summary, start_time, end_time, created_at) "
|
||||
"VALUES ('天气查询', '今天天气晴朗', 't', 't', 't')")
|
||||
conn.execute("INSERT INTO memories_fts (rowid, topic, summary) VALUES (1, '天气查询', '今天天气晴朗')")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
# 重新打开 → 自动迁移 → 中文可搜
|
||||
db = MemoryDatabase(db_path=db_path)
|
||||
results = db.search_memories(["天气"])
|
||||
assert len(results) == 1
|
||||
assert results[0]["topic"] == "天气查询"
|
||||
|
||||
# 迁移后 FTS 表应为 trigram 分词(环境支持时)
|
||||
conn = sqlite3.connect(db_path)
|
||||
sql = conn.execute("SELECT sql FROM sqlite_master WHERE name='memories_fts'").fetchone()[0]
|
||||
fts_count = conn.execute("SELECT COUNT(*) FROM memories_fts").fetchone()[0]
|
||||
conn.close()
|
||||
assert fts_count == 1 # 数据已回灌
|
||||
try:
|
||||
probe = sqlite3.connect(":memory:")
|
||||
probe.execute("CREATE VIRTUAL TABLE t USING fts5(x, tokenize='trigram')")
|
||||
probe.close()
|
||||
assert "trigram" in sql
|
||||
except sqlite3.OperationalError:
|
||||
pass # 环境不支持 trigram,保持 unicode61 属预期行为
|
||||
|
||||
|
||||
def test_memory_delete(mem_db):
|
||||
mid = mem_db.add_memory(topic="待删除", summary="摘要")
|
||||
assert mem_db.delete_memory(mid) is True
|
||||
assert mem_db.get_memory_by_id(mid) is None
|
||||
# 再次删除应返回 False
|
||||
assert mem_db.delete_memory(mid) is False
|
||||
|
||||
|
||||
def test_memory_batch_delete(mem_db):
|
||||
ids = [mem_db.add_memory(topic=f"t{i}", summary=f"s{i}") for i in range(3)]
|
||||
deleted = mem_db.delete_memories_by_ids(ids)
|
||||
assert deleted == 3
|
||||
for mid in ids:
|
||||
assert mem_db.get_memory_by_id(mid) is None
|
||||
|
||||
|
||||
def test_memory_batch_delete_empty_list(mem_db):
|
||||
assert mem_db.delete_memories_by_ids([]) == 0
|
||||
|
||||
|
||||
def test_memory_get_recent(mem_db):
|
||||
for i in range(5):
|
||||
mem_db.add_memory(topic=f"recent_{i}", summary=f"s{i}")
|
||||
recent = mem_db.get_recent_memories(limit=3)
|
||||
assert len(recent) == 3
|
||||
|
||||
|
||||
def test_memory_nonexistent_id_returns_none(mem_db):
|
||||
assert mem_db.get_memory_by_id(99999) is None
|
||||
|
||||
|
||||
# ==================== MemoDatabase ====================
|
||||
|
||||
def test_memo_add_and_get(memo_db):
|
||||
reminder = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="测试任务", content="内容", reminder_time=reminder)
|
||||
assert isinstance(mid, int) and mid > 0
|
||||
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m is not None
|
||||
assert m["title"] == "测试任务"
|
||||
assert m["is_completed"] == 0
|
||||
|
||||
|
||||
def test_memo_pending_filter_by_time(memo_db):
|
||||
"""未到期的任务不应出现在 pending 列表中。"""
|
||||
future = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
memo_db.add_memo(title="未来任务", reminder_time=future)
|
||||
memo_db.add_memo(title="过期任务", reminder_time=past)
|
||||
|
||||
pending = memo_db.get_pending_memos()
|
||||
titles = [m["title"] for m in pending]
|
||||
assert "过期任务" in titles
|
||||
assert "未来任务" not in titles
|
||||
|
||||
|
||||
def test_memo_complete(memo_db):
|
||||
reminder = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="一次性任务", reminder_time=reminder)
|
||||
assert memo_db.complete_memo(mid) is True
|
||||
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m["is_completed"] == 1
|
||||
|
||||
|
||||
def test_memo_complete_nonexistent_returns_false(memo_db):
|
||||
assert memo_db.complete_memo(99999) is False
|
||||
|
||||
|
||||
def test_memo_immediate(memo_db):
|
||||
"""即时任务应被 get_immediate_memos 返回。"""
|
||||
mid = memo_db.add_memo(title="即时任务", is_immediate=True)
|
||||
imm = memo_db.get_immediate_memos()
|
||||
assert any(m["id"] == mid for m in imm)
|
||||
|
||||
|
||||
def test_memo_mark_as_triggered_completes_one_shot(memo_db):
|
||||
"""一次性任务触发后应标记完成。"""
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="一次性", reminder_time=past)
|
||||
nxt = memo_db.mark_as_triggered(mid)
|
||||
assert nxt is None # 一次性任务无下次提醒
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m["is_completed"] == 1
|
||||
assert m["trigger_count"] == 1
|
||||
|
||||
|
||||
def test_memo_mark_as_triggered_daily_repeat(memo_db):
|
||||
"""daily 重复任务触发后应计算下次提醒时间。"""
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="每日任务", reminder_time=past, repeat_type="daily")
|
||||
nxt = memo_db.mark_as_triggered(mid)
|
||||
assert nxt is not None
|
||||
# 下次应在未来
|
||||
next_dt = datetime.strptime(nxt, constants.DATETIME_FORMAT)
|
||||
assert next_dt > datetime.now() - timedelta(minutes=1)
|
||||
m = memo_db.get_memo(mid)
|
||||
assert m["is_completed"] == 0 # 重复任务不完成
|
||||
|
||||
|
||||
def test_memo_stats(memo_db):
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
future = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
memo_db.add_memo(title="过期", reminder_time=past)
|
||||
memo_db.add_memo(title="未来", reminder_time=future)
|
||||
memo_db.add_memo(title="即时", is_immediate=True)
|
||||
|
||||
stats = memo_db.get_memo_stats()
|
||||
assert stats["total"] == 3
|
||||
assert stats["pending"] == 3
|
||||
assert stats["overdue"] == 1 # 只有过期且非即时的
|
||||
assert stats["immediate"] == 1
|
||||
|
||||
|
||||
def test_memo_get_all_excludes_completed_by_default(memo_db):
|
||||
reminder = (datetime.now() + timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(title="待完成", reminder_time=reminder)
|
||||
memo_db.complete_memo(mid)
|
||||
|
||||
all_memos = memo_db.get_all_memos(include_completed=False)
|
||||
assert all(m["is_completed"] == 0 for m in all_memos)
|
||||
|
||||
all_with_completed = memo_db.get_all_memos(include_completed=True)
|
||||
assert any(m["id"] == mid for m in all_with_completed)
|
||||
|
||||
|
||||
def test_memo_get_pending_tasks_for_user(memo_db):
|
||||
past = (datetime.now() - timedelta(days=1)).strftime(constants.DATETIME_FORMAT)
|
||||
memo_db.add_memo(title="用户任务", reminder_time=past)
|
||||
|
||||
tasks = memo_db.get_pending_tasks_for_user()
|
||||
assert len(tasks) >= 1
|
||||
assert all("id" in t and "title" in t and "is_overdue" in t for t in tasks)
|
||||
assert any(t["title"] == "用户任务" and t["is_overdue"] for t in tasks)
|
||||
|
||||
|
||||
def test_memo_add_user_task_disabled(monkeypatch, memo_db):
|
||||
"""当 MEMO_ALLOW_USER_TASKS=False 时应抛异常。"""
|
||||
monkeypatch.setattr(type(constants), "MEMO_ALLOW_USER_TASKS",
|
||||
property(lambda self: False))
|
||||
with pytest.raises(Exception, match="禁用"):
|
||||
memo_db.add_user_task(title="x")
|
||||
|
||||
|
||||
def test_memo_custom_repeat_minutes(memo_db):
|
||||
"""custom 重复 + minutes 单位应正确计算下次。"""
|
||||
past = (datetime.now() - timedelta(minutes=1)).strftime(constants.DATETIME_FORMAT)
|
||||
mid = memo_db.add_memo(
|
||||
title="每5分钟", reminder_time=past,
|
||||
repeat_type="custom", repeat_interval=5, repeat_interval_unit="minutes"
|
||||
)
|
||||
nxt = memo_db.mark_as_triggered(mid)
|
||||
assert nxt is not None
|
||||
next_dt = datetime.strptime(nxt, constants.DATETIME_FORMAT)
|
||||
# 下次约在 5 分钟后
|
||||
delta = next_dt - datetime.now()
|
||||
assert timedelta(minutes=4) < delta < timedelta(minutes=6)
|
||||
@@ -0,0 +1,188 @@
|
||||
"""subsystems.executors 单元测试
|
||||
|
||||
覆盖:
|
||||
- CommandExecutor 抽象基类契约
|
||||
- SubprocessCommandExecutor 命令执行/流式/超时/取消/输出限制
|
||||
"""
|
||||
import queue
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import CommandExecutor, SubprocessCommandExecutor
|
||||
from common import constants
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_tmp_workspace(tmp_path, monkeypatch):
|
||||
"""每个测试把 WORKSPACE_DIR 重定向到临时目录(并预先创建),避免默认 ~/.aibox/workspace 不存在。"""
|
||||
ws = tmp_path / "ws"
|
||||
ws.mkdir()
|
||||
monkeypatch.setattr(type(constants), "WORKSPACE_DIR",
|
||||
property(lambda self: str(ws)))
|
||||
|
||||
|
||||
# ==================== 抽象基类契约 ====================
|
||||
|
||||
def test_command_executor_is_abstract():
|
||||
"""CommandExecutor 是 ABC,不能直接实例化。"""
|
||||
with pytest.raises(TypeError):
|
||||
CommandExecutor()
|
||||
|
||||
|
||||
def test_subprocess_executor_is_subclass():
|
||||
assert issubclass(SubprocessCommandExecutor, CommandExecutor)
|
||||
|
||||
|
||||
def test_executor_is_enabled_reflects_constants():
|
||||
"""is_enabled 应反映 constants.COMMANDS_ENABLED。"""
|
||||
from common import constants
|
||||
ex = SubprocessCommandExecutor()
|
||||
assert ex.is_enabled() == constants.COMMANDS_ENABLED
|
||||
|
||||
|
||||
def test_executor_format_result_success():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.format_result(True, "ok", {"k": "v"})
|
||||
assert r.startswith("✅ ok")
|
||||
assert "k: v" in r
|
||||
|
||||
|
||||
def test_executor_format_result_failure():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.format_result(False, "fail")
|
||||
assert r.startswith("❌ fail")
|
||||
|
||||
|
||||
def test_executor_format_result_multiline_value():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.format_result(True, "ok", {"输出": "line1\nline2"})
|
||||
assert "line1" in r and "line2" in r
|
||||
|
||||
|
||||
# ==================== 实际命令执行 ====================
|
||||
|
||||
def test_execute_echo_command():
|
||||
"""执行 echo 命令应返回成功结果。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("echo hello_test_123", timeout=5)
|
||||
assert "hello_test_123" in r
|
||||
assert "✅" in r
|
||||
|
||||
|
||||
def test_execute_command_returns_exit_code():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("exit 0", timeout=5)
|
||||
assert "返回码: 0" in r
|
||||
|
||||
|
||||
def test_execute_failing_command_returns_nonzero_code():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("exit 7", timeout=5)
|
||||
assert "返回码: 7" in r
|
||||
# 非零返回码应标记为失败
|
||||
assert "❌" in r
|
||||
|
||||
|
||||
def test_execute_timeout():
|
||||
"""超时应被检测并返回错误信息。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
# ping -t 在 Windows 上无限循环;用 ping -n 跨平台
|
||||
r = ex.execute("ping -n 100 127.0.0.1", timeout=2)
|
||||
assert "超时" in r
|
||||
|
||||
|
||||
def test_execute_unknown_command_returns_error():
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("nonexistent_command_xyz_123", timeout=5)
|
||||
# shell=True 下会返回 127 或类似错误码,但不会抛异常
|
||||
assert "返回码" in r
|
||||
|
||||
|
||||
# ==================== 流式执行 ====================
|
||||
|
||||
def test_execute_streaming_emits_lines_and_complete():
|
||||
ex = SubprocessCommandExecutor()
|
||||
q: queue.Queue = queue.Queue()
|
||||
ex.execute_streaming("echo stream_line_test", q, timeout=5)
|
||||
|
||||
msgs = []
|
||||
while not q.empty():
|
||||
msgs.append(q.get_nowait())
|
||||
|
||||
types = [m[0] for m in msgs]
|
||||
assert "line" in types
|
||||
assert "complete" in types
|
||||
|
||||
complete_payload = [m[1] for m in msgs if m[0] == "complete"][0]
|
||||
assert "stream_line_test" in complete_payload
|
||||
|
||||
|
||||
def test_execute_streaming_emits_command_header():
|
||||
"""流式执行开头应输出命令/超时/工作目录等头部信息。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
q: queue.Queue = queue.Queue()
|
||||
ex.execute_streaming("echo hi", q, timeout=5)
|
||||
|
||||
lines = []
|
||||
while not q.empty():
|
||||
t, p = q.get_nowait()
|
||||
if t == "line":
|
||||
lines.append(p)
|
||||
|
||||
header = "".join(lines[:6])
|
||||
assert "执行命令" in header
|
||||
assert "超时" in header
|
||||
assert "工作目录" in header
|
||||
|
||||
|
||||
# ==================== 取消机制 ====================
|
||||
|
||||
def test_cancel_terminates_running_process():
|
||||
"""cancel() 在进程运行时调用应不抛异常并尝试 kill 进程。
|
||||
Windows 下 kill 后管道读取可能短暂阻塞,但 cancel 本身必须安全返回。
|
||||
线程设为 daemon,确保不会卡住测试进程退出。
|
||||
"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
q: queue.Queue = queue.Queue()
|
||||
import threading
|
||||
t = threading.Thread(
|
||||
target=ex.execute_streaming,
|
||||
args=("ping -n 100 127.0.0.1", q),
|
||||
kwargs={"timeout": 60},
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
time.sleep(0.5) # 等进程启动
|
||||
# 关键断言:cancel 不抛异常(若抛异常测试即失败)
|
||||
ex.cancel()
|
||||
# 等待线程退出(Windows 下可能因管道阻塞超时,daemon 确保不卡死主进程)
|
||||
t.join(timeout=5)
|
||||
|
||||
|
||||
def test_cancel_when_no_process_running_is_noop():
|
||||
"""无进程运行时调用 cancel 应是空操作。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
ex.cancel() # 不应抛异常
|
||||
|
||||
|
||||
# ==================== 输出大小限制 ====================
|
||||
|
||||
def test_output_size_limit_truncates():
|
||||
"""输出超过 1MB 应被截断。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
# 生成约 2MB 输出
|
||||
r = ex.execute("python -c \"print('x'*2048576)\"", timeout=10)
|
||||
assert "截断" in r or "❌" in r
|
||||
|
||||
|
||||
# ==================== 禁用场景 ====================
|
||||
|
||||
def test_execute_when_disabled_returns_disabled_message(monkeypatch):
|
||||
"""当 COMMANDS_ENABLED=False 时应返回禁用提示。"""
|
||||
from common import constants
|
||||
monkeypatch.setattr(type(constants), "COMMANDS_ENABLED",
|
||||
property(lambda self: False))
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("echo hi", timeout=5)
|
||||
assert "禁用" in r
|
||||
@@ -0,0 +1,171 @@
|
||||
#!/usr/bin/env python3
|
||||
"""F2 修复回归测试:N6 registry 超时清理 + N4 reader 线程收尾。
|
||||
|
||||
N6:流式转发超时后必须清理 tool.streaming_output / tool.output_queue,
|
||||
防止僵尸 daemon 线程经实例属性把输出混入后续会话队列。
|
||||
N4:_read_with_limit 超时/宽限 break 后必须 close stdout + join reader,
|
||||
防止 Windows shell=True 孙进程持管道场景下 reader 线程与 fd 泄漏。
|
||||
"""
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems.executors import SubprocessCommandExecutor
|
||||
from tools.base import Tool
|
||||
from tools.registry import ToolRegistry
|
||||
|
||||
|
||||
# ---------- N6:registry 超时路径清理 ----------
|
||||
|
||||
class _StuckStreamTool(Tool):
|
||||
"""execute_streaming 永不结束(模拟卡死的工具,如等待外部资源)"""
|
||||
|
||||
def get_name(self):
|
||||
return "stuck_tool"
|
||||
|
||||
def get_description(self):
|
||||
return "test only"
|
||||
|
||||
def get_parameters(self):
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs):
|
||||
time.sleep(30)
|
||||
return "should not reach"
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
self.output_queue = output_queue
|
||||
self.streaming_output = True
|
||||
try:
|
||||
time.sleep(30) # 卡死,永不产出 complete
|
||||
finally:
|
||||
self.streaming_output = False
|
||||
|
||||
|
||||
def test_registry_timeout_cleans_tool_state():
|
||||
"""转发总超时后:返回超时消息,且工具实例属性被清理"""
|
||||
reg = ToolRegistry()
|
||||
reg.stream_tool_timeout = 1 # 测试调小(默认 3600)
|
||||
tool = _StuckStreamTool()
|
||||
reg.register(tool)
|
||||
reg.set_output_queue(queue.Queue())
|
||||
|
||||
t0 = time.time()
|
||||
result = reg.execute_tool("stuck_tool")
|
||||
elapsed = time.time() - t0
|
||||
|
||||
assert "超时" in result, f"应返回超时消息,实际: {result}"
|
||||
assert elapsed < 10, f"应 ~1s 返回,实际 {elapsed:.1f}s"
|
||||
# N6 核心:超时路径不得跳过清理
|
||||
assert tool.output_queue is None, "超时后 tool.output_queue 必须置 None(防僵尸输出串流)"
|
||||
assert tool.streaming_output is False, "超时后 tool.streaming_output 必须复位"
|
||||
|
||||
|
||||
def test_registry_normal_path_cleans_tool_state():
|
||||
"""正常流式完成后:清理同样执行,结果透传"""
|
||||
reg = ToolRegistry()
|
||||
|
||||
class _OkTool(_StuckStreamTool):
|
||||
def get_name(self):
|
||||
return "ok_tool"
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
self.output_queue = output_queue
|
||||
self.streaming_output = True
|
||||
try:
|
||||
output_queue.put(("line", "hello"))
|
||||
output_queue.put(("complete", "done-result"))
|
||||
finally:
|
||||
pass
|
||||
|
||||
tool = _OkTool()
|
||||
reg.register(tool)
|
||||
reg.set_output_queue(queue.Queue())
|
||||
|
||||
result = reg.execute_tool("ok_tool")
|
||||
assert result == "done-result"
|
||||
assert tool.output_queue is None
|
||||
assert tool.streaming_output is False
|
||||
|
||||
|
||||
# ---------- N4:reader 线程收尾 ----------
|
||||
|
||||
class _FakeStdout:
|
||||
"""模拟管道 stdout,可配置阻塞时长"""
|
||||
|
||||
def __init__(self, block_seconds: float = 8):
|
||||
self.closed = False
|
||||
self._block = block_seconds
|
||||
|
||||
def readline(self):
|
||||
time.sleep(self._block) # > join 宽限 2s 即模拟"无 EOF 阻塞"
|
||||
return ""
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _FakeProcess:
|
||||
def __init__(self, block_seconds: float = 8):
|
||||
self.stdout = _FakeStdout(block_seconds)
|
||||
|
||||
def poll(self):
|
||||
return None # 进程"活着"
|
||||
|
||||
def terminate(self):
|
||||
pass
|
||||
|
||||
def kill(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_read_with_limit_timeout_returns_promptly():
|
||||
"""超时 break 后 reader 仍阻塞:只做有限 join,不 close、不挂死。
|
||||
|
||||
CPython 流锁决定了此时 close(stdout) 会阻塞到 readline 返回,
|
||||
故收尾策略为放弃等待(daemon 兜底),主流程必须及时返回。
|
||||
"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = _FakeProcess(block_seconds=8)
|
||||
|
||||
t0 = time.time()
|
||||
output, timeout_occurred, truncated = ex._read_with_limit(
|
||||
proc, timeout=1, start_time=time.time())
|
||||
elapsed = time.time() - t0
|
||||
|
||||
assert timeout_occurred is True
|
||||
assert truncated is False
|
||||
# 超时 1s + join 宽限 2s,绝不能等到 block 8s
|
||||
assert elapsed < 5, f"应收尾后及时返回,实际 {elapsed:.1f}s"
|
||||
|
||||
|
||||
def test_read_with_limit_grace_path_returns_promptly():
|
||||
"""shell 死亡宽限 break 后 reader 仍阻塞:同样及时返回"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = _FakeProcess(block_seconds=8)
|
||||
proc.poll = lambda: 0 # shell 已死,但 reader 无 EOF(孙进程持管道)
|
||||
|
||||
t0 = time.time()
|
||||
output, timeout_occurred, truncated = ex._read_with_limit(
|
||||
proc, timeout=0, start_time=time.time()) # timeout=0 不触发超时
|
||||
elapsed = time.time() - t0
|
||||
|
||||
assert timeout_occurred is False
|
||||
# 宽限 1s + join 2s 上界,不能等到 block 8s
|
||||
assert elapsed < 5, f"宽限+收尾应 ~3s 内返回,实际 {elapsed:.1f}s"
|
||||
|
||||
|
||||
def test_read_with_limit_closes_stdout_after_normal_eof():
|
||||
"""reader 正常收到 EOF 退出后:显式 close stdout 释放管道 fd"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = _FakeProcess(block_seconds=0) # readline 立即返回 ""(EOF)
|
||||
|
||||
output, timeout_occurred, truncated = ex._read_with_limit(
|
||||
proc, timeout=5, start_time=time.time())
|
||||
|
||||
assert timeout_occurred is False
|
||||
assert output == ""
|
||||
assert proc.stdout.closed is True, "reader 正常退出后必须 close stdout 释放 fd"
|
||||
@@ -0,0 +1,359 @@
|
||||
"""aibox 第三轮审计 Batch G 修复回归测试(2026-08-18,审计报告 §9.4)
|
||||
|
||||
覆盖:
|
||||
- P1:非瞬时 HTTP 错误响应体截断(密钥/超大错误页不进异常栈)
|
||||
- P2:Retry-After 解析上限 cap 3600s
|
||||
- P3:fallback 档案缺 model → 跳过并告警(与缺 Key 行为对齐)
|
||||
- P5:LLMClient 死状态 _cancelled 已删除
|
||||
- P6:fallback 列表重复项去重
|
||||
- N2:drain 线程批次结束即回收(不再泄漏常驻线程)
|
||||
- N5:子代理共享父代理取消事件
|
||||
- N7:FTS5 关键词内双引号翻倍转义
|
||||
- N8:registry 流式转发受 TOOL_MAX_RESULT_LENGTH 约束
|
||||
- N10:mark_as_triggered 单事务(缺失 id 路径回归)
|
||||
- N12:persona 路径边界校验(拒绝越界路径)
|
||||
- N16:同秒归档不再覆盖(文件名加序号)
|
||||
|
||||
N3(/new 移线程池)/N13(broadcast)/N14(任务登记集合)为 async WebSocket
|
||||
路径,依赖事件循环与真实连接对象,由现有 WebSocket 测试面回归,不在此单测。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import DeepSeekChat, run_full_subagent # noqa: E402
|
||||
|
||||
from llm_client import ( # noqa: E402
|
||||
DeepSeekClient,
|
||||
FailoverClient,
|
||||
_retry_after_of,
|
||||
create_llm_client,
|
||||
)
|
||||
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||
from subsystems.conversation_history import ConversationHistory # noqa: E402
|
||||
from tools.base import Tool # noqa: E402
|
||||
from tools.registry import ToolRegistry # noqa: E402
|
||||
from common import constants # noqa: E402
|
||||
|
||||
|
||||
# ==================== P1/P2:错误响应体截断 + Retry-After cap ====================
|
||||
|
||||
class TestP1P2LlmClient:
|
||||
def test_p1_non_transient_error_body_truncated(self):
|
||||
"""400 响应体 1000 字符 → 异常消息仅含前 200 字符"""
|
||||
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||
resp = mock.MagicMock()
|
||||
resp.status_code = 400
|
||||
resp.text = "A" * 1000
|
||||
resp.headers = {}
|
||||
with mock.patch("llm_client.requests.post", return_value=resp):
|
||||
with pytest.raises(RuntimeError) as ei:
|
||||
list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
msg = str(ei.value)
|
||||
assert "HTTP 400" in msg
|
||||
assert len(msg) < 300 # 200 字符正文 + 状态行前缀
|
||||
assert msg.count("A") <= 200
|
||||
|
||||
def test_p2_retry_after_capped(self):
|
||||
"""异常/恶意 Retry-After=999999 → cap 到 3600s"""
|
||||
resp = mock.MagicMock()
|
||||
resp.headers = {"Retry-After": "999999"}
|
||||
assert _retry_after_of(resp) == 3600.0
|
||||
|
||||
def test_p2_retry_after_normal_value_kept(self):
|
||||
resp = mock.MagicMock()
|
||||
resp.headers = {"Retry-After": "30"}
|
||||
assert _retry_after_of(resp) == 30.0
|
||||
|
||||
def test_p2_retry_after_invalid_returns_none(self):
|
||||
resp = mock.MagicMock()
|
||||
resp.headers = {"Retry-After": "not-a-number"}
|
||||
assert _retry_after_of(resp) is None
|
||||
|
||||
|
||||
# ==================== P3/P5/P6:fallback 链构造 ====================
|
||||
|
||||
def _g_cfg():
|
||||
"""主档案 main 的 fallback 链:缺 model 的 broken + 重复 dupe + 正常 ok2"""
|
||||
return {
|
||||
"models": {
|
||||
"main": {
|
||||
"api_url": "https://a.example/v1/chat/completions",
|
||||
"model": "m1", "api_key_env": "MAIN_KEY",
|
||||
"fallback": ["nomodel", "dupe", "dupe", "ok2"],
|
||||
},
|
||||
"nomodel": { # P3:缺 model 字段
|
||||
"api_url": "https://b.example/v1/chat/completions",
|
||||
"api_key_env": "B_KEY",
|
||||
},
|
||||
"dupe": {
|
||||
"api_url": "https://c.example/v1/chat/completions",
|
||||
"model": "c1", "api_key_env": "C_KEY",
|
||||
},
|
||||
"ok2": {
|
||||
"api_url": "https://d.example/v1/chat/completions",
|
||||
"model": "d1", "api_key_env": "D_KEY",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class TestP3P5P6FallbackChain:
|
||||
def test_p3_p6_fallback_missing_model_skipped_and_deduped(self):
|
||||
"""缺 model 的 fallback 跳过;重复项去重;链 = [main, dupe, ok2]"""
|
||||
env = {k: f"sk-{k.lower()}" for k in
|
||||
("MAIN_KEY", "B_KEY", "C_KEY", "D_KEY")}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
c = create_llm_client(_g_cfg(), model="main")
|
||||
assert isinstance(c, FailoverClient)
|
||||
names = [cl.name for cl in c._clients]
|
||||
assert names == ["main", "dupe", "ok2"]
|
||||
|
||||
def test_p5_dead_cancel_flag_removed(self):
|
||||
"""LLMClient 不再持有只写不读的 _cancelled 死状态"""
|
||||
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||
assert not hasattr(c, "_cancelled")
|
||||
c.cancel() # 无在途响应时为安全 no-op
|
||||
|
||||
|
||||
# ==================== N2:drain 线程生命周期 ====================
|
||||
|
||||
class _EchoTool(Tool):
|
||||
def get_name(self):
|
||||
return "g_echo"
|
||||
|
||||
def get_description(self):
|
||||
return "测试用回显工具"
|
||||
|
||||
def get_parameters(self):
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs):
|
||||
return "ok"
|
||||
|
||||
|
||||
class TestN2DrainThread:
|
||||
def test_drain_thread_reclaimed_after_batch(self):
|
||||
"""_handle_tool_calls 返回后 drain 线程必须已退出(finally join)"""
|
||||
from test_once_mode import _make_ai
|
||||
ai = _make_ai()
|
||||
ai.tool_registry.register(_EchoTool(ai.logger))
|
||||
tool_calls = [{"id": "c1", "function": {"name": "g_echo", "arguments": "{}"}}]
|
||||
ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None)
|
||||
drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"]
|
||||
assert drains == [] or all(not t.is_alive() for t in drains)
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
|
||||
def test_drain_thread_reclaimed_on_tool_exception(self):
|
||||
"""工具执行抛异常时 drain 线程同样被回收(finally 路径)"""
|
||||
from test_once_mode import _make_ai
|
||||
|
||||
class _BoomTool(_EchoTool):
|
||||
def get_name(self):
|
||||
return "g_boom"
|
||||
|
||||
def execute(self, **kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
ai = _make_ai()
|
||||
ai.tool_registry.register(_BoomTool(ai.logger))
|
||||
tool_calls = [{"id": "c1", "function": {"name": "g_boom", "arguments": "{}"}}]
|
||||
ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None)
|
||||
drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"]
|
||||
assert drains == [] or all(not t.is_alive() for t in drains)
|
||||
|
||||
|
||||
# ==================== N5:子代理取消传播 ====================
|
||||
|
||||
def _fake_sub_chat():
|
||||
fake = mock.MagicMock()
|
||||
fake.think_and_respond = mock.MagicMock(return_value="任务完成报告")
|
||||
fake.stats = {}
|
||||
fake.messages = []
|
||||
fake._cancel_event = threading.Event()
|
||||
return fake
|
||||
|
||||
|
||||
class TestN5SubagentCancel:
|
||||
def test_subagent_shares_parent_cancel_event(self):
|
||||
"""cancel_event 传入 → 子代理取消事件被替换为父事件(检查点联动)"""
|
||||
fake = _fake_sub_chat()
|
||||
parent_event = threading.Event()
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
run_full_subagent("r", "t", None, cancel_event=parent_event)
|
||||
assert fake._cancel_event is parent_event
|
||||
|
||||
def test_subagent_keeps_own_event_when_not_provided(self):
|
||||
"""cancel_event=None → 兼容旧行为,子代理保留自身事件"""
|
||||
fake = _fake_sub_chat()
|
||||
own = fake._cancel_event
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
run_full_subagent("r", "t", None)
|
||||
assert fake._cancel_event is own
|
||||
|
||||
def test_subagent_tool_registered_with_parent_event(self):
|
||||
"""DeepSeekChat 注册 SubAgentTool 时注入自身 _cancel_event(N5 链路)"""
|
||||
from test_once_mode import _make_ai
|
||||
ai = _make_ai()
|
||||
tool = next((t for t in ai.tool_registry.get_all_tools()
|
||||
if t.get_name() == "create_subagent"), None)
|
||||
assert tool is not None
|
||||
assert tool.cancel_event is ai._cancel_event
|
||||
|
||||
|
||||
# ==================== N7:FTS5 引号转义 ====================
|
||||
|
||||
class TestN7FtsQuoteEscape:
|
||||
def test_keyword_with_quote_does_not_break_search(self, tmp_path):
|
||||
"""关键词含双引号 → MATCH 语法合法(翻倍转义),不抛异常"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
db.add_memory(topic="引号话题", summary="摘要",
|
||||
content='含"引号"的记忆内容',
|
||||
start_time="2026-08-18 00:00:00",
|
||||
end_time="2026-08-18 00:00:01", importance=2)
|
||||
res = db.search_memories(['带"引号"的关键词'])
|
||||
assert isinstance(res, list) # 不抛 OperationalError 即通过
|
||||
|
||||
def test_normal_keyword_still_matches(self, tmp_path):
|
||||
"""常规关键词检索行为不回归"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
db.add_memory(topic="普通话题", summary="天气摘要",
|
||||
content="今天天气晴朗",
|
||||
start_time="2026-08-18 00:00:00",
|
||||
end_time="2026-08-18 00:00:01", importance=2)
|
||||
res = db.search_memories(["天气"])
|
||||
assert len(res) >= 1
|
||||
|
||||
|
||||
# ==================== N8:流式转发长度上限 ====================
|
||||
|
||||
class _BigStreamTool(Tool):
|
||||
"""流式工具:发出 3 行 × 40000 字符(总 120000 > 100000 上限)"""
|
||||
|
||||
def get_name(self):
|
||||
return "g_bigstream"
|
||||
|
||||
def get_description(self):
|
||||
return "测试用超长流式工具"
|
||||
|
||||
def get_parameters(self):
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs):
|
||||
return "done"
|
||||
|
||||
def execute_streaming(self, output_queue, **kwargs):
|
||||
for _ in range(3):
|
||||
output_queue.put(("line", "x" * 40000))
|
||||
output_queue.put(("complete", "done"))
|
||||
|
||||
|
||||
class TestN8StreamingLimit:
|
||||
def test_streaming_forward_truncated_to_budget(self):
|
||||
registry = ToolRegistry()
|
||||
registry.register(_BigStreamTool())
|
||||
captured = queue.Queue()
|
||||
registry.set_output_queue(captured)
|
||||
|
||||
result = registry.execute_tool("g_bigstream")
|
||||
|
||||
lines = []
|
||||
while not captured.empty():
|
||||
msg_type, content = captured.get_nowait()
|
||||
if msg_type == "line":
|
||||
lines.append(content)
|
||||
budget = constants.TOOL_MAX_RESULT_LENGTH
|
||||
total = sum(len(ln) for ln in lines)
|
||||
# 截断标记行本身占少量额外字符(含提示文案)
|
||||
assert total <= budget + 200
|
||||
assert any("已停止转发" in ln or "已截断" in ln for ln in lines)
|
||||
assert result == "done"
|
||||
|
||||
|
||||
# ==================== N10:mark_as_triggered 单事务 ====================
|
||||
|
||||
class TestN10MarkTriggered:
|
||||
def test_missing_id_returns_none_via_transaction(self, tmp_path):
|
||||
"""事务内读取不存在行 → rollback + None(不再走独立 get_memo)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "memo.db"))
|
||||
assert db.mark_as_triggered(99999) is None
|
||||
|
||||
def test_single_shot_memo_completes(self, tmp_path):
|
||||
"""一次性任务触发 → completed + 返回 None(无下次提醒)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "memo.db"))
|
||||
mid = db.add_memo(title="一次性", content="c",
|
||||
reminder_time="2026-08-18 00:00:00")
|
||||
assert mid is not None
|
||||
assert db.mark_as_triggered(mid) is None
|
||||
memo = db.get_memo(mid)
|
||||
assert memo["is_completed"] == 1
|
||||
assert memo["trigger_count"] == 1
|
||||
|
||||
|
||||
# ==================== N12:persona 路径边界 ====================
|
||||
|
||||
class TestN12PersonaBoundary:
|
||||
def _tool(self):
|
||||
return aibox.ExpressEmotionTool()
|
||||
|
||||
def test_traversal_persona_rejected(self, tmp_path):
|
||||
"""persona 含 ../ 越界 → 拒绝(返回空串,不读文件)"""
|
||||
tool = self._tool()
|
||||
skills = tmp_path / "skills"
|
||||
skills.mkdir()
|
||||
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||
assert tool._load_persona("../../etc/passwd", None) == ""
|
||||
assert tool._load_persona("../outside", None) == ""
|
||||
|
||||
def test_absolute_persona_rejected(self, tmp_path):
|
||||
"""persona 为绝对路径 → 解析结果不在技能目录内,拒绝"""
|
||||
tool = self._tool()
|
||||
skills = tmp_path / "skills"
|
||||
skills.mkdir()
|
||||
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||
assert tool._load_persona(str(tmp_path / "evil"), None) == ""
|
||||
|
||||
def test_normal_persona_loads(self, tmp_path):
|
||||
"""合法 persona(技能目录内)正常加载"""
|
||||
tool = self._tool()
|
||||
skills = tmp_path / "skills"
|
||||
skills.mkdir()
|
||||
(skills / "ok.md").write_text("人格内容", encoding="utf-8")
|
||||
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||
assert "人格内容" in tool._load_persona("ok", None)
|
||||
|
||||
|
||||
# ==================== N16:同秒归档防覆盖 ====================
|
||||
|
||||
class TestN16ArchiveCollision:
|
||||
def test_same_second_archive_no_overwrite(self, tmp_path):
|
||||
"""同一秒内两次归档 → 文件名加序号,互不覆盖"""
|
||||
hm = ConversationHistory(save_dir=str(tmp_path))
|
||||
hm.start_new_conversation(topic="t1")
|
||||
hm.add_message("user", "hi", speaker="u")
|
||||
p1 = hm.archive_current_conversation()
|
||||
|
||||
hm.start_new_conversation(topic="t2")
|
||||
hm.add_message("user", "hi2", speaker="u")
|
||||
p2 = hm.archive_current_conversation()
|
||||
|
||||
assert p1 != p2
|
||||
assert os.path.exists(p1) and os.path.exists(p2)
|
||||
@@ -0,0 +1,413 @@
|
||||
"""aibox 第四轮审计 Batch H 修复回归测试(2026-08-18,审计报告 §十一)
|
||||
|
||||
覆盖:
|
||||
- H1:子代理共享父取消事件后首轮 think 不再清除父事件(N5 回归)+ 父已取消时跳过启动
|
||||
- M3:调度器任务级超时兜底(AI think 卡死不再停摆调度循环)
|
||||
- M4:get_recent_memories(limit=None) 使用默认上限(不再全表加载)
|
||||
- M5:memo delete 物理删除(此前只标记完成)
|
||||
- M6:PromptManager.reload 原子替换(并发读者不再读到空/半截列表)
|
||||
- L1:read_file 的 sed 参数强转 int(防引号注入)
|
||||
- L2:set_proxy 落盘失败如实报告(不再静默报"已保存")
|
||||
- L3:SSE 非 UTF-8 行不中断流(errors=replace)
|
||||
- L4:persona 路径校验解析符号链接(软链越界拒绝)
|
||||
- L5:1MB 截断不再产生半个字符(�)
|
||||
- L6:compressor 空内容压缩不再除零
|
||||
- L7:InstallSkillTool 指引命令路径引用(防拆词/注入)
|
||||
- L8:ask_user 缺参返回明确错误(不再 TypeError)
|
||||
- L10:web 工具仅允许 http/https(file:// 等拒绝)
|
||||
- L11:registry 日志打码改包含匹配(apiKey/access_token 等变体)
|
||||
- L12:FailoverClient 日志 Bearer 令牌脱敏
|
||||
|
||||
M1(/reload_prompts 移线程池)/M2(技能扫描启动前预热)为 async WebSocket /
|
||||
启动路径,依赖事件循环与真实连接对象,由既有 WebSocket 测试面与启动逻辑回归。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import DeepSeekChat, TaskScheduler, run_full_subagent # noqa: E402
|
||||
|
||||
from llm_client import DeepSeekClient, FailoverClient, _sanitize_log_msg # noqa: E402
|
||||
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||
from subsystems.conversation_compressor import ConversationCompressor # noqa: E402
|
||||
from common import constants # noqa: E402
|
||||
|
||||
from test_conversation_compressor import FakeLLMClient # noqa: E402
|
||||
from test_once_mode import _make_ai # noqa: E402
|
||||
|
||||
|
||||
def _fake_sub_chat():
|
||||
fake = mock.MagicMock()
|
||||
fake.think_and_respond = mock.MagicMock(return_value="任务完成报告")
|
||||
fake.stats = {}
|
||||
fake.messages = []
|
||||
fake._cancel_event = threading.Event()
|
||||
return fake
|
||||
|
||||
|
||||
# ==================== H1:子代理取消事件清除回归 ====================
|
||||
|
||||
class TestH1SubagentCancelClear:
|
||||
def test_shared_event_not_cleared_on_first_think(self):
|
||||
"""共享父事件后首轮 think 不清除事件——父级已置位的取消信号不被吞掉"""
|
||||
ai = _make_ai()
|
||||
parent_event = threading.Event()
|
||||
parent_event.set()
|
||||
ai._cancel_event = parent_event
|
||||
ai._shared_cancel_event = True
|
||||
with pytest.raises(InterruptedError):
|
||||
ai._think_and_respond_impl("hi")
|
||||
assert parent_event.is_set()
|
||||
|
||||
def test_own_event_still_cleared_on_first_think(self):
|
||||
"""未共享(普通会话)时首轮仍清除自身事件——原语义不回归"""
|
||||
import contextlib
|
||||
ai = _make_ai()
|
||||
ai._cancel_event.set()
|
||||
with contextlib.suppress(Exception):
|
||||
ai._think_and_respond_impl("hi")
|
||||
assert not ai._cancel_event.is_set()
|
||||
|
||||
def test_subagent_not_started_when_parent_already_cancelled(self):
|
||||
"""父事件已置位 → 直接返回取消提示,不启动子代理"""
|
||||
fake = _fake_sub_chat()
|
||||
parent_event = threading.Event()
|
||||
parent_event.set()
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
result = run_full_subagent("r", "t", None, cancel_event=parent_event)
|
||||
assert "取消" in result
|
||||
fake.think_and_respond.assert_not_called()
|
||||
|
||||
|
||||
# ==================== M3:调度器任务超时 ====================
|
||||
|
||||
class _BlockingAI:
|
||||
"""think 卡死直到被 cancel 的伪 AI(模拟 LLM 挂起)"""
|
||||
|
||||
def __init__(self):
|
||||
self.cancelled = threading.Event()
|
||||
self.cancel_called = 0
|
||||
|
||||
def cancel(self):
|
||||
self.cancel_called += 1
|
||||
self.cancelled.set()
|
||||
|
||||
def think_and_respond(self, input_text, output_callback=None,
|
||||
is_continuation=False):
|
||||
self.cancelled.wait(timeout=15)
|
||||
raise InterruptedError("用户中断")
|
||||
|
||||
|
||||
class _FakeMemoDB:
|
||||
def get_memo(self, memo_id):
|
||||
return None
|
||||
|
||||
|
||||
class TestM3SchedulerTimeout:
|
||||
def test_stuck_task_times_out_and_scheduler_continues(self):
|
||||
from aibox import ConfigManager
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
fake_ai = _BlockingAI()
|
||||
memo = {"id": 7, "title": "卡死任务", "content": "内容", "trigger_count": 0}
|
||||
orig_get = ConfigManager.get
|
||||
|
||||
def fake_get(self, key, default=None):
|
||||
# 只把任务超时压到 0.3s,其余配置读取走原实现(防止
|
||||
# DATETIME_FORMAT 等属性读到错误类型)
|
||||
if key == "scheduler.task_timeout":
|
||||
return 0.3
|
||||
return orig_get(self, key, default)
|
||||
|
||||
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||
mock.patch.object(scheduler, "_create_task_ai", return_value=fake_ai):
|
||||
result = loop.run_until_complete(
|
||||
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||
assert "超时" in result
|
||||
assert fake_ai.cancel_called >= 1
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# ==================== M4/M5:记忆上限与任务删除 ====================
|
||||
|
||||
class TestM4M5Db:
|
||||
def test_recent_memories_none_capped(self, tmp_path, monkeypatch):
|
||||
"""limit=None 不再全表加载,使用默认上限(对齐 search_memories)"""
|
||||
from test_conversation_compressor import patch_const
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "m.db"))
|
||||
for i in range(7):
|
||||
db.add_memory(topic=f"t{i}", summary=f"s{i}", content="c",
|
||||
start_time="2026-08-18 00:00:00",
|
||||
end_time="2026-08-18 00:00:01", importance=1)
|
||||
patch_const(monkeypatch, "MEMORY_MAX_SEARCH_RESULTS", 5)
|
||||
res = db.get_recent_memories(limit=None)
|
||||
assert len(res) == 5
|
||||
|
||||
def test_delete_memo_physically_removes(self, tmp_path):
|
||||
"""delete_memo 真实删除行(此前 complete_memo 只标记完成)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = db.add_memo(title="待删任务", content="内容")
|
||||
assert db.get_memo(mid) is not None
|
||||
assert db.delete_memo(mid) is True
|
||||
assert db.get_memo(mid) is None
|
||||
assert db.delete_memo(mid) is False # 已不存在
|
||||
|
||||
def test_memo_tool_delete_uses_physical_delete(self, tmp_path):
|
||||
"""MemoTool delete 走 delete_memo:行消失而非仅标记完成"""
|
||||
from tools import MemoTool
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = db.add_memo(title="重复任务", content="内容", repeat_type="daily")
|
||||
tool = MemoTool(db, None)
|
||||
result = tool.execute(action="delete", memo_id=mid, confirm=True)
|
||||
assert "已删除" in result
|
||||
assert db.get_memo(mid) is None
|
||||
stats = db.get_memo_stats()
|
||||
assert stats["total"] == 0
|
||||
|
||||
|
||||
# ==================== M6:PromptManager 原子 reload ====================
|
||||
|
||||
class TestM6PromptReloadAtomic:
|
||||
def test_reload_builds_new_list_then_swaps(self):
|
||||
"""reload 在本地新列表构建完成后才原子替换——读者始终看到完整列表"""
|
||||
pm = PromptManager()
|
||||
old_list = pm._prompts
|
||||
seen = []
|
||||
|
||||
def fake_load(target):
|
||||
seen.append(target)
|
||||
target.append({"file": "X_test.md", "content": "y"})
|
||||
|
||||
with mock.patch.object(pm, "_load_prompts", side_effect=fake_load):
|
||||
n = pm.reload()
|
||||
assert len(seen) == 1
|
||||
assert seen[0] is not old_list # 构建用的是新列表
|
||||
assert pm._prompts is seen[0] # 完成后才替换
|
||||
assert n == 1
|
||||
|
||||
|
||||
# ==================== L1:read_file 参数强转 ====================
|
||||
|
||||
class TestL1ReadFileCoercion:
|
||||
def test_non_integer_lines_rejected(self, tmp_path):
|
||||
from tools import ReadFileTool
|
||||
f = tmp_path / "f.txt"
|
||||
f.write_text("line1\nline2\n", encoding="utf-8")
|
||||
tool = ReadFileTool()
|
||||
r = tool.execute(filepath=str(f), start_line="1'; rm -rf /; echo '",
|
||||
end_line=5)
|
||||
assert "参数错误" in r
|
||||
|
||||
def test_valid_integers_still_work(self, tmp_path):
|
||||
from tools import ReadFileTool
|
||||
f = tmp_path / "f.txt"
|
||||
f.write_text("line1\nline2\n", encoding="utf-8")
|
||||
tool = ReadFileTool()
|
||||
r = tool.execute(filepath=str(f), start_line=1, end_line=2)
|
||||
assert "sed -n" in r and "参数错误" not in r
|
||||
|
||||
|
||||
# ==================== L2:set_proxy 落盘失败如实报告 ====================
|
||||
|
||||
class TestL2ProxySaveFailure:
|
||||
def test_set_proxy_returns_false_on_save_failure(self, tmp_path):
|
||||
from common import ConfigManager
|
||||
cm = ConfigManager()
|
||||
cm._config_path = str(tmp_path / "config.json")
|
||||
# J5(审计 2026-08-19):写路径改为原子写(mkstemp+os.replace),
|
||||
# 磁盘失败改由 patch 原子写函数模拟(builtins.open 已不在路径上)
|
||||
with mock.patch("common.config.atomic_write_text",
|
||||
side_effect=PermissionError("disk ro")):
|
||||
ok = cm.set_proxy("127.0.0.1", 1080, save=True)
|
||||
assert ok is False
|
||||
|
||||
def test_proxy_tool_reports_save_failure(self, tmp_path):
|
||||
from tools import SetProxyTool
|
||||
tool = SetProxyTool(None)
|
||||
# J5:磁盘失败经原子写函数模拟
|
||||
with mock.patch("common.config.atomic_write_text",
|
||||
side_effect=PermissionError("disk ro")):
|
||||
r = tool.execute(host="127.0.0.1", port=1080, save=True,
|
||||
skip_check=True, enabled=True)
|
||||
assert "保存失败" in r
|
||||
|
||||
|
||||
# ==================== L3:SSE 非 UTF-8 行不中断流 ====================
|
||||
|
||||
class TestL3SseDecode:
|
||||
def test_non_utf8_line_does_not_break_stream(self):
|
||||
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||
resp = mock.MagicMock()
|
||||
resp.status_code = 200
|
||||
resp.iter_lines.return_value = [
|
||||
b"data: \xff\xfe",
|
||||
b'data: {"choices":[{"delta":{"content":"ok"}}]}',
|
||||
b"data: [DONE]",
|
||||
]
|
||||
with mock.patch("llm_client.requests.post", return_value=resp):
|
||||
events = list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert any(e.get("type") == "content" and e.get("text") == "ok"
|
||||
for e in events)
|
||||
|
||||
|
||||
# ==================== L4:persona 软链越界 ====================
|
||||
|
||||
class TestL4PersonaSymlink:
|
||||
def test_symlink_outside_skills_dir_rejected(self, tmp_path, monkeypatch):
|
||||
from test_conversation_compressor import patch_const
|
||||
from aibox import ExpressEmotionTool
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
outside = tmp_path / "outside.md"
|
||||
outside.write_text("秘密内容", encoding="utf-8")
|
||||
(skills_dir / "evil.md").symlink_to(outside)
|
||||
tool = ExpressEmotionTool()
|
||||
patch_const(monkeypatch, "SKILLS_DIR", str(skills_dir))
|
||||
content = tool._load_persona("evil", "")
|
||||
assert content == ""
|
||||
|
||||
def test_normal_persona_still_loads(self, tmp_path, monkeypatch):
|
||||
from test_conversation_compressor import patch_const
|
||||
from aibox import ExpressEmotionTool
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
(skills_dir / "happy.md").write_text("快乐人格", encoding="utf-8")
|
||||
tool = ExpressEmotionTool()
|
||||
patch_const(monkeypatch, "SKILLS_DIR", str(skills_dir))
|
||||
content = tool._load_persona("happy", "")
|
||||
assert content == "快乐人格"
|
||||
|
||||
|
||||
# ==================== L5:截断无半个字符 ====================
|
||||
|
||||
class TestL5NoPartialChar:
|
||||
def test_truncated_output_has_no_replacement_char(self):
|
||||
from subsystems.executors import SubprocessCommandExecutor
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute("python3 -c \"print('中' * 600000)\"", timeout=60)
|
||||
assert "截断" in r
|
||||
assert "\ufffd" not in r
|
||||
|
||||
@pytest.mark.skipif(os.name != "posix", reason="Linux 进程组语义")
|
||||
def test_grandchild_holding_pipe_is_reaped(self):
|
||||
"""`(sleep N &) ; echo done` 形态:shell 退出后孙进程持管道写端,
|
||||
宽限窗口到期强杀进程组——不悬挂 reader、孙进程被清理"""
|
||||
import subprocess as sp
|
||||
from subsystems.executors import SubprocessCommandExecutor
|
||||
marker = f"sleep {12340}"
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute(f"({marker} &) ; echo done", timeout=8)
|
||||
assert "done" in r
|
||||
leftover = sp.run(["pgrep", "-f", marker],
|
||||
capture_output=True, text=True).stdout.strip()
|
||||
assert leftover == ""
|
||||
|
||||
|
||||
# ==================== L6:compressor 除零 ====================
|
||||
|
||||
class TestL6CompressorDivZero:
|
||||
def test_empty_content_compress_no_div0(self, monkeypatch):
|
||||
comp = ConversationCompressor(llm_client=FakeLLMClient(),
|
||||
logger=None)
|
||||
from test_conversation_compressor import patch_const
|
||||
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||
messages = [{"role": "system", "content": "SYS"}] + [
|
||||
{"role": "user" if i % 2 == 0 else "assistant", "content": ""}
|
||||
for i in range(15)
|
||||
]
|
||||
result = comp.compress_conversation(messages)
|
||||
assert isinstance(result, list)
|
||||
|
||||
|
||||
# ==================== L7:InstallSkillTool 路径引用 ====================
|
||||
|
||||
class TestL7SkillPathQuoting:
|
||||
def test_guide_commands_quote_paths(self):
|
||||
from tools import InstallSkillTool
|
||||
tool = InstallSkillTool()
|
||||
r = tool.execute(skill_path="/tmp/a b; rm -rf /")
|
||||
assert "ls -l '/tmp/a b; rm -rf /'" in r
|
||||
assert "cp '/tmp/a b; rm -rf /' " in r
|
||||
|
||||
|
||||
# ==================== L8:ask_user 缺参 ====================
|
||||
|
||||
class TestL8AskUser:
|
||||
def test_missing_question_returns_error(self):
|
||||
from tools import AskUserTool
|
||||
tool = AskUserTool()
|
||||
assert "需要提供问题" in tool.execute()
|
||||
|
||||
def test_question_passed_through(self):
|
||||
from tools import AskUserTool
|
||||
tool = AskUserTool()
|
||||
assert "测试问题" in tool.execute(question="测试问题")
|
||||
|
||||
|
||||
# ==================== L10:web 工具 scheme 白名单 ====================
|
||||
|
||||
class TestL10WebScheme:
|
||||
def test_file_url_rejected(self):
|
||||
from tools import SimpleWebCrawlerTool
|
||||
tool = SimpleWebCrawlerTool()
|
||||
r = tool.execute(action="get", url="file:///etc/shadow")
|
||||
assert "仅支持 http/https" in r
|
||||
|
||||
def test_http_url_accepted(self):
|
||||
from tools import SimpleWebCrawlerTool
|
||||
tool = SimpleWebCrawlerTool()
|
||||
r = tool.execute(action="get", url="https://example.com")
|
||||
assert "curl" in r and "仅支持" not in r
|
||||
|
||||
|
||||
# ==================== L11:registry 打码包含匹配 ====================
|
||||
|
||||
class TestL11SanitizeSubstring:
|
||||
def test_variant_keys_redacted(self):
|
||||
from tools.registry import _sanitize_kwargs
|
||||
out = _sanitize_kwargs({
|
||||
"apiKey": "x", "access_token": "y", "passwd": "z",
|
||||
"url": "https://ok",
|
||||
})
|
||||
assert out == {"apiKey": "***", "access_token": "***",
|
||||
"passwd": "***", "url": "https://ok"}
|
||||
|
||||
def test_normal_keys_untouched(self):
|
||||
from tools.registry import _sanitize_kwargs
|
||||
out = _sanitize_kwargs({"filepath": "/tmp/a", "question": "hi"})
|
||||
assert out == {"filepath": "/tmp/a", "question": "hi"}
|
||||
|
||||
|
||||
# ==================== L12:failover 日志 Bearer 脱敏 ====================
|
||||
|
||||
class TestL12LogSanitize:
|
||||
def test_sanitize_msg_redacts_bearer(self):
|
||||
msg = "瞬时故障: HTTP 401: Bearer sk-abc123xyz"
|
||||
out = _sanitize_log_msg(msg)
|
||||
assert "sk-abc123xyz" not in out
|
||||
assert "Bearer ***" in out
|
||||
|
||||
def test_failover_log_redacts_bearer(self):
|
||||
logger = mock.MagicMock()
|
||||
c = FailoverClient([], logger=logger)
|
||||
c._log("warning", "瞬时故障(Bearer sk-secret456)")
|
||||
logged = logger.warning.call_args[0][0]
|
||||
assert "sk-secret456" not in logged
|
||||
assert "Bearer ***" in logged
|
||||
@@ -0,0 +1,237 @@
|
||||
"""aibox 第四轮审计 Batch I 修复回归测试(2026-08-18,审计报告 §十二)
|
||||
|
||||
覆盖:
|
||||
- N9:prompt 写入内容长度上限(持久化注入面约束)
|
||||
- Q7b:load_conversation 越界路径拒绝(realpath 前缀校验)
|
||||
- F5:run_light_subagent 取消传播(流式可中断 + 父已取消早退)
|
||||
- Q5:提醒时间格式异常的任务按到期返回(不再静默丢弃)
|
||||
- Q8:对话文件原子写(无 .tmp 残留、JSON 完整)
|
||||
- F6:WebSocketHandler 命令专用线程池存在且可释放
|
||||
- limit=0 语义修正(None 才不加 LIMIT)
|
||||
|
||||
Q4(busy_timeout 10s)为时序性行为(默认 5s 等待已能通过短竞争),
|
||||
不做慢速单测;F9(--once 硬超时)为进程级行为,由 CLI 集成面覆盖。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import run_light_subagent # noqa: E402
|
||||
|
||||
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||
from subsystems.conversation_history import ConversationHistory # noqa: E402
|
||||
from tools.prompt_tools import ( # noqa: E402
|
||||
UpdateIdentityTool,
|
||||
UpdateSoulTool,
|
||||
SaveMemoryAndEndConversationTool,
|
||||
)
|
||||
|
||||
|
||||
# ==================== N9:prompt 写入长度上限 ====================
|
||||
|
||||
class TestN9PromptContentCap:
|
||||
def test_update_identity_oversized_rejected(self, tmp_path, monkeypatch):
|
||||
"""manager 层兜底:超长身份内容拒绝写入"""
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
assert pm.update_identity("x" * (PromptManager.MAX_CONTENT_LENGTH + 1)) is False
|
||||
assert not (tmp_path / pm.IDENTITY_FILE).exists()
|
||||
|
||||
def test_update_identity_tool_friendly_error(self, tmp_path):
|
||||
"""工具层:超长内容返回明确提示而非写入"""
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
tool = UpdateIdentityTool(pm, None)
|
||||
r = tool.execute(content="x" * (PromptManager.MAX_CONTENT_LENGTH + 1))
|
||||
assert "内容过长" in r
|
||||
|
||||
def test_update_soul_tool_friendly_error(self, tmp_path):
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
tool = UpdateSoulTool(pm, None)
|
||||
r = tool.execute(content="y" * (PromptManager.MAX_CONTENT_LENGTH + 1))
|
||||
assert "内容过长" in r
|
||||
|
||||
def test_save_memory_status_oversized_rejected(self, tmp_path):
|
||||
"""save_memory_and_end_conversation:状态内容超长拒绝"""
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
tool = SaveMemoryAndEndConversationTool(
|
||||
exit_callback=lambda *a: None,
|
||||
memory_db=None, prompt_manager=pm, logger=None)
|
||||
r = tool.execute(summary="这是一段足够长的对话总结" * 20,
|
||||
current_status="x" * PromptManager.MAX_CONTENT_LENGTH)
|
||||
assert "状态内容过长" in r
|
||||
|
||||
def test_normal_content_still_written(self, tmp_path):
|
||||
pm = PromptManager()
|
||||
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||
assert pm.update_identity("正常身份内容") is True
|
||||
assert "正常身份内容" in (tmp_path / pm.IDENTITY_FILE).read_text(encoding="utf-8")
|
||||
|
||||
|
||||
# ==================== Q7b:load_conversation 路径校验 ====================
|
||||
|
||||
class TestQ7bLoadConversationBoundary:
|
||||
def test_traversal_rejected(self, tmp_path):
|
||||
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||
secret = tmp_path / "secret.json"
|
||||
secret.write_text('{"leak": true}', encoding="utf-8")
|
||||
assert h.load_conversation("../secret.json") is None
|
||||
|
||||
def test_normal_file_still_loads(self, tmp_path):
|
||||
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||
h.current_conversation = {
|
||||
"id": "1", "start_time": "2026-08-18 00:00:00", "end_time": None,
|
||||
"topic": "t", "messages": [{"role": "user", "content": "hi"}],
|
||||
"tool_calls": [], "stats": {},
|
||||
}
|
||||
h.save_current_conversation()
|
||||
h.archive_current_conversation()
|
||||
archives = list(h.archive_dir.glob("conversation_*.json"))
|
||||
assert archives
|
||||
loaded = h.load_conversation(archives[0].name)
|
||||
assert loaded is not None and loaded["id"] == "1"
|
||||
|
||||
|
||||
# ==================== F5:run_light_subagent 取消传播 ====================
|
||||
|
||||
class _BlockingLightClient:
|
||||
"""流式产出首段后卡住,直到 cancel() 才中断的伪 LLM 客户端"""
|
||||
|
||||
model = "fake-model"
|
||||
default_temperature = None
|
||||
|
||||
def __init__(self):
|
||||
self.cancelled = threading.Event()
|
||||
|
||||
def cancel(self):
|
||||
self.cancelled.set()
|
||||
|
||||
def stream_chat(self, messages, **kwargs):
|
||||
yield {"type": "content", "text": "开头"}
|
||||
while not self.cancelled.wait(timeout=0.1):
|
||||
pass
|
||||
raise RuntimeError("连接已中断(cancel)")
|
||||
|
||||
|
||||
class TestF5LightSubagentCancel:
|
||||
def test_cancelled_before_start_returns_early(self):
|
||||
"""父事件已置位 → 直接返回取消提示,不创建客户端"""
|
||||
parent = threading.Event()
|
||||
parent.set()
|
||||
with mock.patch("aibox.create_llm_client") as mk:
|
||||
r = run_light_subagent("r", "t", None, cancel_event=parent)
|
||||
assert "取消" in r
|
||||
mk.assert_not_called()
|
||||
|
||||
def test_cancel_during_stream_aborts(self):
|
||||
"""流式中父级取消 → 观察者线程 cancel 客户端 → 立即中断返回"""
|
||||
parent = threading.Event()
|
||||
fake = _BlockingLightClient()
|
||||
with mock.patch("aibox.create_llm_client", return_value=fake):
|
||||
result_holder = {}
|
||||
|
||||
def _run():
|
||||
result_holder["r"] = run_light_subagent(
|
||||
"r", "t", None, cancel_event=parent)
|
||||
|
||||
th = threading.Thread(target=_run, daemon=True)
|
||||
th.start()
|
||||
time.sleep(0.5) # 等待首段产出、观察者就绪
|
||||
parent.set()
|
||||
th.join(timeout=5)
|
||||
assert not th.is_alive() # 未挂死
|
||||
assert "失败" in result_holder.get("r", "")
|
||||
|
||||
def test_express_emotion_registered_with_parent_event(self):
|
||||
"""DeepSeekChat 注册 ExpressEmotionTool 时注入自身取消事件"""
|
||||
from test_once_mode import _make_ai
|
||||
ai = _make_ai()
|
||||
tool = next((t for t in ai.tool_registry.get_all_tools()
|
||||
if t.get_name() == "express_emotion"), None)
|
||||
assert tool is not None
|
||||
assert tool.cancel_event is ai._cancel_event
|
||||
|
||||
|
||||
# ==================== Q5:格式异常任务不丢弃 ====================
|
||||
|
||||
class TestQ5MalformedReminderTime:
|
||||
def test_malformed_time_task_returned_as_due(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
# J1(审计 2026-08-19):add_memo 已拒绝畸形时间(入库校验),
|
||||
# 模拟校验落地前的存量脏数据需绕过 add_memo 直接 SQL 插入
|
||||
import sqlite3
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
cur = conn.cursor()
|
||||
cur.execute('''INSERT INTO memos
|
||||
(title, content, created_at, reminder_time, is_completed, source,
|
||||
repeat_type, repeat_interval, repeat_interval_unit, is_immediate)
|
||||
VALUES (?, ?, ?, ?, 0, 'test', 'none', 0, 'days', 0)''',
|
||||
("格式异常", "c", "2026-08-18 00:00:00", "不是时间"))
|
||||
mid = cur.lastrowid
|
||||
pending = db.get_pending_memos(before_time="2026-08-18 00:00:00")
|
||||
assert any(m["id"] == mid for m in pending)
|
||||
|
||||
def test_normal_pending_still_filtered(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
db.add_memo(title="未来任务", content="c", reminder_time="2099-01-01 00:00:00")
|
||||
pending = db.get_pending_memos(before_time="2026-08-18 00:00:00")
|
||||
assert pending == []
|
||||
|
||||
|
||||
# ==================== Q8:对话文件原子写 ====================
|
||||
|
||||
class TestQ8AtomicWrites:
|
||||
def test_no_tmp_leftover_and_valid_json(self, tmp_path):
|
||||
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||
h.current_conversation = {
|
||||
"id": "1", "start_time": "2026-08-18 00:00:00", "end_time": None,
|
||||
"topic": "t", "messages": [], "tool_calls": [], "stats": {},
|
||||
}
|
||||
h.save_current_conversation()
|
||||
assert h.current_file.exists() # 文件名可能被用户配置覆盖,用实例属性
|
||||
import json
|
||||
json.loads(h.current_file.read_text(encoding="utf-8")) # 不抛即合法 JSON
|
||||
h.archive_current_conversation() # 归档后 current 文件按设计移除
|
||||
assert list(Path(tmp_path).rglob("*.tmp")) == []
|
||||
|
||||
|
||||
# ==================== F6:WS 命令专用线程池 ====================
|
||||
|
||||
class TestF6CmdExecutor:
|
||||
def test_handler_has_dedicated_executor(self):
|
||||
from aibox import WebSocketHandler
|
||||
h = WebSocketHandler(ai_manager=None)
|
||||
assert h._cmd_executor is not None
|
||||
h.shutdown()
|
||||
assert h._cmd_executor is None
|
||||
|
||||
|
||||
# ==================== limit=0 语义 ====================
|
||||
|
||||
class TestLimitZeroSemantics:
|
||||
def test_pending_memos_limit_zero_returns_none(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
db.add_memo(title="t", content="c", reminder_time="2026-01-01 00:00:00")
|
||||
assert db.get_pending_memos(limit=0, before_time="2026-08-18 00:00:00") == []
|
||||
|
||||
def test_immediate_memos_limit_zero_returns_none(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
db.add_memo(title="t", content="c", is_immediate=True)
|
||||
assert db.get_immediate_memos(limit=0) == []
|
||||
@@ -0,0 +1,540 @@
|
||||
"""aibox 第五轮审计(Batch K)修复回归测试(2026-08-19)
|
||||
|
||||
波次 1:
|
||||
- J3:FTS 外部内容表删除换序后索引无残留(delete_memory / delete_memories_by_ids)
|
||||
- J4:model_config 连通性验证异常消息 Bearer 令牌脱敏
|
||||
- J6:PromptManager/ConfigManager 单例并发初始化不双重加载
|
||||
- J8:add_profile 对非法温度/超时的容错(交互层不再抢先 float()/int())
|
||||
|
||||
波次 2(J1 任务系统投毒):
|
||||
- add_memo 入库校验:畸形时间/重复间隔/重复类型拒绝,三入口一处覆盖
|
||||
- 存量脏数据防御:调度器预检、mark_as_triggered、展示路径不再打崩
|
||||
- WS task add 友好错误 + register 握手异常可回收
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
import model_config # noqa: E402
|
||||
from common import ConfigManager # noqa: E402
|
||||
from subsystems.db import MemoDatabase, MemoryDatabase # noqa: E402
|
||||
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||
|
||||
from test_model_config import _cfg # noqa: E402
|
||||
|
||||
|
||||
# ==================== J3:FTS 删除顺序 ====================
|
||||
|
||||
def _fts_hits(db_path, keyword):
|
||||
conn = sqlite3.connect(str(db_path))
|
||||
try:
|
||||
return conn.execute(
|
||||
"SELECT rowid FROM memories_fts WHERE memories_fts MATCH ?",
|
||||
(f'"{keyword}"',)).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def _add_mem(db, content):
|
||||
return db.add_memory(topic="话题", summary="摘要", content=content,
|
||||
start_time="2026-08-19 00:00:00",
|
||||
end_time="2026-08-19 00:00:01", importance=2)
|
||||
|
||||
|
||||
class TestJ3FtsDeleteOrder:
|
||||
def test_delete_memory_leaves_no_fts_orphan(self, tmp_path):
|
||||
"""单条删除后 FTS 索引不得残留(修复前 MATCH 仍命中已删 rowid)"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
mid = _add_mem(db, "机密内容甲")
|
||||
assert _fts_hits(tmp_path / "mem.db", "机密内容甲")
|
||||
assert db.delete_memory(mid) is True
|
||||
assert _fts_hits(tmp_path / "mem.db", "机密内容甲") == []
|
||||
|
||||
def test_batch_delete_leaves_no_fts_orphan(self, tmp_path):
|
||||
"""批量删除同样不留索引残留"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
ids = [_add_mem(db, f"机密内容乙{i}") for i in range(3)]
|
||||
assert db.delete_memories_by_ids(ids) == 3
|
||||
for i in range(3):
|
||||
assert _fts_hits(tmp_path / "mem.db", f"机密内容乙{i}") == []
|
||||
|
||||
def test_search_after_delete_no_hit(self, tmp_path):
|
||||
"""删除后搜索(FTS 路径)不再命中已删内容"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
mid = _add_mem(db, "独特关键词丙")
|
||||
assert db.search_memories(["独特关键词丙"])
|
||||
db.delete_memory(mid)
|
||||
assert db.search_memories(["独特关键词丙"]) == []
|
||||
|
||||
|
||||
# ==================== J4:验证异常脱敏 ====================
|
||||
|
||||
class TestJ4TestProfileSanitize:
|
||||
def _models_cfg(self, tmp_path):
|
||||
return _cfg(tmp_path, {"models": {"bad": {
|
||||
"api_url": "https://api.example.com/v1/chat/completions",
|
||||
"model": "m", "api_key_env": "K"}}})
|
||||
|
||||
def test_bearer_token_masked_in_failure_message(self, tmp_path):
|
||||
"""网关错误页回显 Authorization 头时,令牌不得出现在终端回显"""
|
||||
cm = self._models_cfg(tmp_path)
|
||||
err = RuntimeError("HTTP 502: Authorization: Bearer sk-secret12345")
|
||||
with mock.patch("llm_client.create_llm_client", side_effect=err):
|
||||
result = model_config.test_profile(cm, "bad")
|
||||
assert "sk-secret12345" not in result
|
||||
assert "Bearer ***" in result
|
||||
|
||||
def test_plain_error_message_preserved(self, tmp_path):
|
||||
"""无敏感内容的错误消息原样保留(脱敏不误伤)"""
|
||||
cm = self._models_cfg(tmp_path)
|
||||
with mock.patch("llm_client.create_llm_client",
|
||||
side_effect=RuntimeError("连接被拒绝")):
|
||||
result = model_config.test_profile(cm, "bad")
|
||||
assert "连接被拒绝" in result
|
||||
|
||||
|
||||
# ==================== J6:单例并发初始化 ====================
|
||||
|
||||
class TestJ6SingletonInitRace:
|
||||
def test_prompt_manager_concurrent_init_no_duplication(self, tmp_path):
|
||||
"""多线程首例并发:单实例且提示词只加载一份(修复前整份翻倍)"""
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
prompts_dir = tmp_path / "prompts"
|
||||
prompts_dir.mkdir()
|
||||
try:
|
||||
barrier = threading.Barrier(8)
|
||||
instances = [None] * 8
|
||||
|
||||
def worker(k):
|
||||
barrier.wait()
|
||||
instances[k] = PromptManager()
|
||||
|
||||
with mock.patch.object(ConfigManager, "get",
|
||||
return_value=str(prompts_dir)):
|
||||
threads = [threading.Thread(target=worker, args=(k,))
|
||||
for k in range(8)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert len({id(i) for i in instances}) == 1
|
||||
files = [p["file"] for p in PromptManager._prompts]
|
||||
assert files, "提示词应已加载"
|
||||
assert len(files) == len(set(files)), f"提示词被重复加载: {files}"
|
||||
finally:
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
def test_config_manager_concurrent_init_single_load(self, tmp_path):
|
||||
"""多线程首例并发:单实例且 _load_config 只执行一次"""
|
||||
cfg_file = tmp_path / "config.json"
|
||||
calls = {"n": 0}
|
||||
original_load = ConfigManager._load_config
|
||||
|
||||
def counting_load(self):
|
||||
calls["n"] += 1
|
||||
original_load(self)
|
||||
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
try:
|
||||
barrier = threading.Barrier(8)
|
||||
instances = [None] * 8
|
||||
|
||||
def worker(k):
|
||||
barrier.wait()
|
||||
instances[k] = ConfigManager()
|
||||
|
||||
with mock.patch.dict(os.environ, {"AI_CONFIG_PATH": str(cfg_file)}), \
|
||||
mock.patch.object(ConfigManager, "_load_config", counting_load):
|
||||
threads = [threading.Thread(target=worker, args=(k,))
|
||||
for k in range(8)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert len({id(i) for i in instances}) == 1
|
||||
assert calls["n"] == 1
|
||||
finally:
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
|
||||
|
||||
# ==================== J8:交互输入容错 ====================
|
||||
|
||||
class TestJ8NumericTolerance:
|
||||
def test_invalid_temperature_returns_friendly_error(self, tmp_path):
|
||||
"""非数字温度 → 友好错误而非 ValueError(交互层不再抢先 float())"""
|
||||
cm = _cfg(tmp_path, {"models": {}})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||
api_key_env="K", temperature="not-a-number")
|
||||
assert any(m.startswith("❌") and "温度" in m for m in msgs)
|
||||
assert "p1" not in model_config.get_models(cm)
|
||||
|
||||
def test_invalid_timeout_returns_friendly_error(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {"models": {}})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||
api_key_env="K", timeout="abc")
|
||||
assert any(m.startswith("❌") and "超时" in m for m in msgs)
|
||||
assert "p1" not in model_config.get_models(cm)
|
||||
|
||||
def test_string_numbers_accepted_and_clamped(self, tmp_path):
|
||||
"""数字字符串正常接受并钳制(交互层原始字符串直传语义)"""
|
||||
cm = _cfg(tmp_path, {"models": {}})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||
api_key_env="K", temperature="0.7", timeout="120")
|
||||
assert all(not m.startswith("❌") for m in msgs)
|
||||
profile = model_config.get_models(cm)["p1"]
|
||||
assert profile["temperature"] == 0.7
|
||||
assert profile["timeout"] == 120
|
||||
|
||||
|
||||
# ==================== J1 波次 2:任务系统投毒 ====================
|
||||
|
||||
def _insert_memo_raw(db_path, reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="none", repeat_interval=0,
|
||||
repeat_interval_unit="days", repeat_end_time=None):
|
||||
"""直接 SQL 插入任务行——模拟校验落地前的存量脏数据(绕过 add_memo)"""
|
||||
with sqlite3.connect(str(db_path)) as conn:
|
||||
cur = conn.cursor()
|
||||
cur.execute('''INSERT INTO memos
|
||||
(title, content, created_at, reminder_time, tags, priority, created_by,
|
||||
is_completed, source, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time, is_immediate)
|
||||
VALUES ('t', 'c', '2026-08-18 00:00:00', ?, '', 1, 'test',
|
||||
0, 'test', ?, ?, ?, ?, 0)''',
|
||||
(reminder_time, repeat_type, repeat_interval, repeat_interval_unit,
|
||||
repeat_end_time))
|
||||
return cur.lastrowid
|
||||
|
||||
|
||||
class TestJ1AddMemoValidation:
|
||||
def test_malformed_reminder_time_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="提醒时间格式非法"):
|
||||
db.add_memo(title="t", reminder_time="不是时间")
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||
|
||||
def test_malformed_repeat_end_time_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="重复结束时间格式非法"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="daily", repeat_end_time="不是时间")
|
||||
|
||||
def test_invalid_repeat_type_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="无效的重复类型"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="every-second")
|
||||
|
||||
def test_non_numeric_repeat_interval_rejected(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="重复间隔必须是整数"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval="垃圾")
|
||||
|
||||
def test_custom_repeat_requires_positive_interval(self, tmp_path):
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="repeat_interval"):
|
||||
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval=0)
|
||||
|
||||
def test_numeric_string_interval_coerced(self, tmp_path):
|
||||
"""数字字符串被强转接受(WS 传 JSON 数字/字符串均安全)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval="30",
|
||||
repeat_interval_unit="minutes")
|
||||
assert db.get_memo(mid)["repeat_interval"] == 30
|
||||
|
||||
def test_add_user_task_propagates_validation(self, tmp_path):
|
||||
"""WS 路径(add_user_task 薄包装 add_memo)同样被校验拦截"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with pytest.raises(ValueError, match="提醒时间格式非法"):
|
||||
db.add_user_task(title="t", reminder_time="不是时间")
|
||||
|
||||
|
||||
class TestJ1DirtyDataDefense:
|
||||
def test_mark_as_triggered_completes_malformed_one_shot(self, tmp_path):
|
||||
"""存量畸形 reminder_time 一次性任务:触发不再抛异常,正常完成"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||
assert db.mark_as_triggered(mid) is None
|
||||
assert db.get_memo(mid)["is_completed"] == 1
|
||||
|
||||
def test_mark_as_triggered_handles_garbage_repeat_interval(self, tmp_path):
|
||||
"""存量 custom 任务 repeat_interval 为文本:按非重复处理,不打崩事务"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="2026-08-18 00:00:00",
|
||||
repeat_type="custom", repeat_interval="垃圾")
|
||||
assert db.mark_as_triggered(mid) is None
|
||||
assert db.get_memo(mid)["is_completed"] == 1
|
||||
|
||||
def test_malformed_repeat_end_time_treated_as_no_end(self, tmp_path):
|
||||
"""存量畸形 repeat_end_time:解析失败视为无结束时间,不抛异常"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="2026-08-19 00:00:00",
|
||||
repeat_type="daily", repeat_end_time="不是时间")
|
||||
memo = db.get_memo(mid)
|
||||
from datetime import datetime
|
||||
base = datetime.strptime("2026-08-19 00:00:00", "%Y-%m-%d %H:%M:%S")
|
||||
assert db._should_end_repeat(memo, base) is False
|
||||
assert db._calculate_next_reminder(memo, base) is not None
|
||||
|
||||
def test_pending_tasks_for_user_malformed_time_overdue(self, tmp_path):
|
||||
"""展示路径:畸形 reminder_time 行按到期展示,不再打崩握手"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||
tasks = db.get_pending_tasks_for_user()
|
||||
row = next(t for t in tasks if t["id"] == mid)
|
||||
assert row["is_overdue"] is True
|
||||
|
||||
|
||||
class TestJ1WsAndScheduler:
|
||||
def test_ws_task_add_rejects_malformed_repeat_end(self, tmp_path):
|
||||
"""WS task add 传畸形 repeat_end_time:友好错误且不入库"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
ai_manager = mock.MagicMock()
|
||||
ai_manager.memo_db = db
|
||||
handler = aibox.WebSocketHandler(ai_manager)
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_task_message(ws, {
|
||||
"type": "task", "action": "add", "title": "t", "delay": 5,
|
||||
"repeat_type": "daily", "repeat_end_time": "不是时间"}))
|
||||
sent = json.loads(ws.send.call_args[0][0])
|
||||
assert sent["type"] == "error"
|
||||
assert "任务添加失败" in sent["content"]
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_scheduler_precheck_malformed_time_executes_not_loops(self, tmp_path):
|
||||
"""调度器预检:畸形 reminder_time 按到期执行并标记触发(不再热循环)"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||
loop = asyncio.new_event_loop()
|
||||
executed = {"n": 0}
|
||||
try:
|
||||
scheduler = aibox.TaskScheduler(loop=loop, logger=None)
|
||||
scheduler.broadcast_enabled = False
|
||||
|
||||
async def fake_execute(memo, memo_db):
|
||||
executed["n"] += 1
|
||||
return "ok"
|
||||
|
||||
scheduler.execute_task_as_conversation = fake_execute
|
||||
loop.run_until_complete(scheduler.check_and_execute_tasks(db))
|
||||
finally:
|
||||
loop.close()
|
||||
assert executed["n"] == 1
|
||||
assert db.get_memo(mid)["is_completed"] == 1
|
||||
|
||||
|
||||
# ==================== J5 波次 4:原子写 + 损坏备份 ====================
|
||||
|
||||
class TestJ5AtomicWrites:
|
||||
def test_save_config_atomic_no_tmp_leftover(self, tmp_path):
|
||||
"""_save_config 原子写:无 .aibox_tmp_ 残留且 JSON 合法"""
|
||||
cm = ConfigManager()
|
||||
cm._config_path = str(tmp_path / "config.json")
|
||||
cm._config = {"api": {"deepseek_model": "m"}, "models": {}}
|
||||
assert cm._save_config() is True
|
||||
import json as _json
|
||||
data = _json.loads((tmp_path / "config.json").read_text(encoding="utf-8"))
|
||||
assert data["api"]["deepseek_model"] == "m"
|
||||
assert list(tmp_path.glob(".aibox_tmp_*")) == []
|
||||
|
||||
def test_corrupted_config_backed_up_before_fallback(self, tmp_path):
|
||||
"""损坏配置加载:备份 .corrupt.* 文件后回退默认,原内容不丢"""
|
||||
cfg_file = tmp_path / "config.json"
|
||||
cfg_file.write_text('{"api": {"broken": ', encoding="utf-8") # 半截 JSON
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
try:
|
||||
with mock.patch.dict(os.environ, {"AI_CONFIG_PATH": str(cfg_file)}):
|
||||
cm = ConfigManager()
|
||||
assert cm._config is not None
|
||||
backups = list(tmp_path.glob("config.json.corrupt.*"))
|
||||
assert len(backups) == 1
|
||||
assert backups[0].read_text(encoding="utf-8") == '{"api": {"broken": '
|
||||
finally:
|
||||
ConfigManager._instance = None
|
||||
ConfigManager._config = None
|
||||
ConfigManager._config_path = None
|
||||
|
||||
def test_prompt_write_atomic(self, tmp_path):
|
||||
"""prompt 更新原子写:内容生效且无临时文件残留"""
|
||||
from subsystems.prompt_manager import PromptManager
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
prompts_dir = tmp_path / "prompts"
|
||||
prompts_dir.mkdir()
|
||||
try:
|
||||
with mock.patch.object(ConfigManager, "get",
|
||||
return_value=str(prompts_dir)):
|
||||
pm = PromptManager()
|
||||
assert pm.update_identity("新身份内容") is True
|
||||
identity = prompts_dir / "00_IDENTITY.md"
|
||||
assert "新身份内容" in identity.read_text(encoding="utf-8")
|
||||
assert list(tmp_path.glob(".aibox_tmp_*")) == []
|
||||
finally:
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
|
||||
# ==================== J7 波次 4:broadcast 超时回收 ====================
|
||||
|
||||
class TestJ7BroadcastTimeout:
|
||||
def test_stuck_client_timed_out_and_unregistered(self):
|
||||
"""半开 TCP 僵死客户端:broadcast 在超时内返回并回收该连接"""
|
||||
from aibox import WebSocketHandler
|
||||
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
handler.broadcast_send_timeout = 0.1
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
|
||||
async def stuck_send(*args, **kwargs):
|
||||
await asyncio.sleep(30)
|
||||
|
||||
ws.send.side_effect = stuck_send
|
||||
handler.connected_clients.add(ws)
|
||||
asyncio.run(handler.broadcast({"type": "ping"}))
|
||||
assert ws not in handler.connected_clients
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_healthy_client_stays_connected(self):
|
||||
"""正常客户端:发送成功且不被回收"""
|
||||
from aibox import WebSocketHandler
|
||||
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
handler.broadcast_send_timeout = 1.0
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
handler.connected_clients.add(ws)
|
||||
asyncio.run(handler.broadcast({"type": "ping"}))
|
||||
assert ws in handler.connected_clients
|
||||
assert ws.send.called
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== J9 波次 5:MemoDatabase 连接层 ====================
|
||||
|
||||
class TestJ9MemoDbConnect:
|
||||
def test_wal_enabled_after_init(self, tmp_path):
|
||||
"""WAL 显式设置——不再依赖与 MemoryDatabase 的创建顺序"""
|
||||
MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
mode = conn.execute("PRAGMA journal_mode").fetchone()[0]
|
||||
assert mode == "wal"
|
||||
|
||||
def test_connection_factory_sets_busy_timeout(self, tmp_path):
|
||||
"""统一连接工厂:busy_timeout 覆盖全部写路径"""
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
conn = db._connect()
|
||||
try:
|
||||
assert conn.execute("PRAGMA busy_timeout").fetchone()[0] == 10000
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ==================== J10 波次 5:memo 批量删除 ====================
|
||||
|
||||
class TestJ10BatchDelete:
|
||||
def test_batch_delete_via_memo_tool(self, tmp_path):
|
||||
"""schema 声明的 memo_ids 批量删除真正生效(此前静默忽略)"""
|
||||
from tools.db_tools import MemoTool
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
ids = [db.add_memo(title=f"任务{i}",
|
||||
reminder_time="2026-08-18 00:00:00") for i in range(3)]
|
||||
tool = MemoTool(db)
|
||||
r = tool.execute(action="delete", confirm=True, memo_ids=ids)
|
||||
assert "3/3" in r
|
||||
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||
|
||||
def test_batch_delete_requires_list(self, tmp_path):
|
||||
from tools.db_tools import MemoTool
|
||||
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||
tool = MemoTool(db)
|
||||
r = tool.execute(action="delete", confirm=True, memo_ids="1")
|
||||
assert "列表" in r
|
||||
|
||||
|
||||
# ==================== J15 波次 5:未知 task action ====================
|
||||
|
||||
class TestJ15UnknownTaskAction:
|
||||
def test_unknown_action_returns_error(self):
|
||||
"""未知 task action 明确报错(此前静默吞掉)"""
|
||||
from aibox import WebSocketHandler
|
||||
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_task_message(ws, {"action": "explode"}))
|
||||
sent = json.loads(ws.send.call_args[0][0])
|
||||
assert sent["type"] == "error"
|
||||
assert "未知的任务操作" in sent["content"]
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== J16 波次 5:skill_manager ====================
|
||||
|
||||
class TestJ16SkillManager:
|
||||
def test_markdown_skipped_when_executable_same_name(self, tmp_path, monkeypatch):
|
||||
"""同名键冲突:可执行技能确定性优先(此前由 listdir 顺序决定)"""
|
||||
from common import constants
|
||||
from subsystems.skill_manager import SkillManager
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
exe = skills_dir / "foo"
|
||||
exe.write_text("#!/bin/sh\necho '可执行技能简介'", encoding="utf-8")
|
||||
exe.chmod(0o755)
|
||||
(skills_dir / "foo.md").write_text("# Markdown 简介\n内容", encoding="utf-8")
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
sm = SkillManager()
|
||||
assert "foo" in sm.skills
|
||||
assert sm.skills["foo"]["is_executable"] is True
|
||||
|
||||
def test_help_output_truncated(self, tmp_path, monkeypatch):
|
||||
"""--help 输出截断(内容进系统提示词,防撑爆上下文)"""
|
||||
from common import constants
|
||||
from subsystems.skill_manager import SkillManager
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
exe = skills_dir / "big"
|
||||
exe.write_text("#!/bin/sh\necho '" + "x" * 5000 + "'", encoding="utf-8")
|
||||
exe.chmod(0o755)
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
sm = SkillManager()
|
||||
assert len(sm.skills["big"]["content"]) <= SkillManager.MAX_HELP_LENGTH + 20
|
||||
@@ -0,0 +1,157 @@
|
||||
"""aibox 第六轮审计(Batch L)修复回归测试(2026-08-20)
|
||||
|
||||
R1:/memos WS 命令对畸形 reminder_time 不再整体报错(J1 加固漏点)
|
||||
R2:handle_interrupt 只取消本会话 AI,不再触碰 CLI 主 AI 的全局 executor
|
||||
R3:unregister 补齐连接回收(关闭 websocket + 取消在途处理任务)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
|
||||
|
||||
def _make_handler():
|
||||
handler = aibox.WebSocketHandler(ai_manager=mock.MagicMock())
|
||||
return handler
|
||||
|
||||
|
||||
# ==================== R1:/memos 畸形时间容错 ====================
|
||||
|
||||
class TestR1MemosCommand:
|
||||
def test_malformed_reminder_time_shows_pending(self):
|
||||
"""存量脏数据(畸形 reminder_time):/memos 按"待执行"展示,不抛 ValueError"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
handler.ai_manager.memo_db.get_all_memos.return_value = [
|
||||
{"id": 1, "title": "脏数据任务", "reminder_time": "not-a-time",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
{"id": 2, "title": "正常过期任务", "reminder_time": "2020-01-01 00:00:00",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
]
|
||||
sent = []
|
||||
|
||||
async def fake_send(ws, msg):
|
||||
sent.append(msg)
|
||||
|
||||
handler.send_message = fake_send
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_command(ws, {"command": "/memos"}))
|
||||
|
||||
results = [m for m in sent if m.get("type") == "command_result"]
|
||||
errors = [m for m in sent if m.get("type") == "error"]
|
||||
assert len(results) == 1
|
||||
assert not errors # 不再落入顶层兜底"处理失败"
|
||||
content = results[0]["content"]
|
||||
assert "脏数据任务" in content and "正常过期任务" in content
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_valid_reminder_time_still_marks_overdue(self):
|
||||
"""合法过期时间仍标 overdue(语义不回归)"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
handler.ai_manager.memo_db.get_all_memos.return_value = [
|
||||
{"id": 1, "title": "过期任务", "reminder_time": "2020-01-01 00:00:00",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
{"id": 2, "title": "未来任务", "reminder_time": "2099-01-01 00:00:00",
|
||||
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||
]
|
||||
sent = []
|
||||
|
||||
async def fake_send(ws, msg):
|
||||
sent.append(msg)
|
||||
|
||||
handler.send_message = fake_send
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_command(ws, {"command": "/memos"}))
|
||||
|
||||
content = next(m["content"] for m in sent
|
||||
if m.get("type") == "command_result")
|
||||
overdue = aibox.i18n.get('memo_overdue')
|
||||
pending = aibox.i18n.get('memo_pending')
|
||||
overdue_line = next(l for l in content.splitlines() if "过期任务" in l)
|
||||
pending_line = next(l for l in content.splitlines() if "未来任务" in l)
|
||||
assert overdue in overdue_line and pending not in overdue_line
|
||||
assert pending in pending_line and overdue not in pending_line
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== R2:interrupt 不再误杀 CLI 主 AI 子进程 ====================
|
||||
|
||||
class TestR2InterruptScope:
|
||||
def test_interrupt_does_not_touch_global_executor(self):
|
||||
"""WS 中断不再取消 ai_manager.command_executor(CLI 主 AI 的执行器)"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
asyncio.run(handler.handle_interrupt(ws, {}))
|
||||
handler.ai_manager.command_executor.cancel.assert_not_called()
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_interrupt_still_cancels_session_ai(self):
|
||||
"""本会话 AI 仍被正常取消(中断能力不回归)"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
ai = mock.MagicMock()
|
||||
handler.client_sessions[ws] = {"ai": ai, "is_processing": True}
|
||||
asyncio.run(handler.handle_interrupt(ws, {}))
|
||||
ai.cancel.assert_called_once()
|
||||
handler.ai_manager.command_executor.cancel.assert_not_called()
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
# ==================== R3:unregister 补齐连接回收 ====================
|
||||
|
||||
class TestR3UnregisterCleanup:
|
||||
def test_unregister_closes_websocket_and_cancels_tasks(self):
|
||||
"""broadcast 超时路径走 unregister 后:websocket 被关闭、在途任务被取消、登记清空"""
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = mock.AsyncMock()
|
||||
|
||||
async def scenario():
|
||||
task = asyncio.create_task(asyncio.sleep(30))
|
||||
handler.connected_clients.add(ws)
|
||||
handler.client_sessions[ws] = {"ai": None, "is_processing": True}
|
||||
handler._processing_tasks[ws] = {task}
|
||||
await handler.unregister(ws)
|
||||
await asyncio.sleep(0) # 让 cancel 传播
|
||||
assert task.done() # 取消即视为完成(cancelled 或已收尾)
|
||||
|
||||
asyncio.run(scenario())
|
||||
assert ws.close.called
|
||||
assert ws not in handler.connected_clients
|
||||
assert ws not in handler.client_sessions
|
||||
assert ws not in handler._processing_tasks
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
def test_unregister_idempotent_and_mock_safe(self):
|
||||
"""重复注销幂等;close 抛异常的对象不炸(AttributeError 被兜底)"""
|
||||
class _BadCloseWS:
|
||||
remote_address = ("127.0.0.1", 1)
|
||||
|
||||
async def close(self):
|
||||
raise AttributeError("no close")
|
||||
|
||||
handler = _make_handler()
|
||||
try:
|
||||
ws = _BadCloseWS()
|
||||
handler.connected_clients.add(ws)
|
||||
handler.client_sessions[ws] = {"ai": None}
|
||||
asyncio.run(handler.unregister(ws))
|
||||
asyncio.run(handler.unregister(ws)) # 二次注销不抛
|
||||
assert ws not in handler.connected_clients
|
||||
assert ws not in handler.client_sessions
|
||||
finally:
|
||||
handler.shutdown()
|
||||
@@ -0,0 +1,200 @@
|
||||
"""aibox 模型接入点管理工具(model_config.py)单元测试(2026-08-18)
|
||||
|
||||
覆盖:
|
||||
- 列表归一化(默认标记/免密/失败转移展示字段)
|
||||
- 添加:字段校验(档案名/地址/模型名/重复)、Key 两种方式、温度/超时钳制、
|
||||
失败转移引用校验、设为默认、落盘(含脱敏)
|
||||
- 删除:物理移除、fallback 悬空引用清理、默认档案标记清除
|
||||
- 连通性验证:成功 / 失败 / 档案不存在
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import model_config # noqa: E402
|
||||
from common import ConfigManager # noqa: E402
|
||||
|
||||
|
||||
def _cfg(tmp_path, initial: dict = None) -> ConfigManager:
|
||||
"""构造指向临时配置文件的 ConfigManager(避免碰真实 ~/.aibox)"""
|
||||
cm = ConfigManager()
|
||||
cm._config_path = str(tmp_path / "config.json")
|
||||
cm._config = initial if initial is not None else {}
|
||||
return cm
|
||||
|
||||
|
||||
def _base_models():
|
||||
return {
|
||||
"models": {
|
||||
"deepseek": {
|
||||
"api_url": "https://api.deepseek.com/v1/chat/completions",
|
||||
"model": "deepseek-chat",
|
||||
"api_key_env": "DEEPSEEK_API_KEY",
|
||||
},
|
||||
"local": {
|
||||
"api_url": "http://localhost:11434/v1/chat/completions",
|
||||
"model": "qwen3",
|
||||
"api_key_env": "",
|
||||
},
|
||||
},
|
||||
"api": {"default_model": "deepseek"},
|
||||
}
|
||||
|
||||
|
||||
# ==================== 列表 ====================
|
||||
|
||||
class TestListProfiles:
|
||||
def test_list_normalized(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
profiles = model_config.list_profiles(cm)
|
||||
assert len(profiles) == 2
|
||||
ds = next(p for p in profiles if p["name"] == "deepseek")
|
||||
assert ds["is_default"] is True
|
||||
assert ds["api_key_env"] == "DEEPSEEK_API_KEY"
|
||||
local = next(p for p in profiles if p["name"] == "local")
|
||||
assert local["api_key_env"] == ""
|
||||
assert local["is_default"] is False
|
||||
|
||||
def test_list_empty_models_section(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
assert model_config.list_profiles(cm) == []
|
||||
|
||||
def test_list_ignores_malformed_entries(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {"models": {"bad": "not-a-dict"}})
|
||||
assert model_config.list_profiles(cm) == []
|
||||
|
||||
|
||||
# ==================== 校验 ====================
|
||||
|
||||
class TestValidateNewProfile:
|
||||
def test_name_rules(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
assert model_config.validate_new_profile(cm, "", "https://x", "m")
|
||||
assert model_config.validate_new_profile(cm, "坏名字!", "https://x", "m")
|
||||
assert model_config.validate_new_profile(cm, "deepseek", "https://x", "m") # 已存在
|
||||
assert model_config.validate_new_profile(cm, "ok_name", "https://x", "m") is None
|
||||
|
||||
def test_url_scheme_and_model_required(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
assert model_config.validate_new_profile(cm, "kimi", "ftp://x", "m")
|
||||
assert model_config.validate_new_profile(cm, "kimi", "https://x", "")
|
||||
|
||||
|
||||
# ==================== 添加 ====================
|
||||
|
||||
class TestAddProfile:
|
||||
def test_add_with_env_key_and_default(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
msgs = model_config.add_profile(
|
||||
cm, "kimi", "https://api.moonshot.cn/v1/chat/completions",
|
||||
"moonshot-v1-8k", api_key_env="KIMI_KEY",
|
||||
display_name="Kimi", set_default=True)
|
||||
assert any("已添加" in m for m in msgs)
|
||||
assert any("设为默认" in m for m in msgs)
|
||||
assert cm.get("api.default_model") == "kimi"
|
||||
assert cm.get("models.kimi.api_key_env") == "KIMI_KEY"
|
||||
assert cm.get("models.kimi.api_url").startswith("https://")
|
||||
# 落盘文件存在且可解析
|
||||
import json
|
||||
disk = json.loads((tmp_path / "config.json").read_text(encoding="utf-8"))
|
||||
assert disk["models"]["kimi"]["model"] == "moonshot-v1-8k"
|
||||
|
||||
def test_add_without_key_is_exempt_endpoint(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
msgs = model_config.add_profile(
|
||||
cm, "local", "http://localhost:11434/v1/chat/completions",
|
||||
"qwen3", api_key_env="")
|
||||
assert any("免密" in m for m in msgs)
|
||||
assert cm.get("models.local.api_key_env") == ""
|
||||
|
||||
def test_add_rejects_invalid_fallback(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
msgs = model_config.add_profile(
|
||||
cm, "kimi", "https://x/v1", "m1",
|
||||
fallback=["ghost", "deepseek"])
|
||||
assert any("失败转移" in m for m in msgs)
|
||||
assert cm.get("models.kimi") is None # 未写入
|
||||
|
||||
def test_add_temperature_timeout_clamped(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
model_config.add_profile(cm, "a", "https://x/v1", "m1",
|
||||
temperature=9.9, timeout=1)
|
||||
assert cm.get("models.a.temperature") == 2.0
|
||||
assert cm.get("models.a.timeout") == 5
|
||||
|
||||
def test_add_invalid_temperature_rejected(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
msgs = model_config.add_profile(cm, "a", "https://x/v1", "m1",
|
||||
temperature="abc")
|
||||
assert any("温度" in m for m in msgs)
|
||||
assert cm.get("models.a") is None
|
||||
|
||||
|
||||
# ==================== 删除 ====================
|
||||
|
||||
class TestDeleteProfile:
|
||||
def test_delete_cleans_fallback_refs(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {
|
||||
"models": {
|
||||
"main": {"api_url": "https://a", "model": "m1",
|
||||
"fallback": ["backup"]},
|
||||
"backup": {"api_url": "https://b", "model": "m2"},
|
||||
},
|
||||
})
|
||||
msgs = model_config.delete_profile(cm, "backup")
|
||||
assert any("已删除" in m for m in msgs)
|
||||
assert any("失败转移引用" in m for m in msgs)
|
||||
assert cm.get("models.main.fallback") == []
|
||||
|
||||
def test_delete_default_clears_default_flag(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
model_config.delete_profile(cm, "deepseek")
|
||||
assert cm.get("api.default_model") is None
|
||||
|
||||
def test_delete_nonexistent(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
msgs = model_config.delete_profile(cm, "ghost")
|
||||
assert any("不存在" in m for m in msgs)
|
||||
|
||||
|
||||
# ==================== 连通性验证 ====================
|
||||
|
||||
class TestTestProfile:
|
||||
def test_connectivity_ok(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
fake = mock.MagicMock()
|
||||
fake.chat_completion.return_value = "pong"
|
||||
with mock.patch("llm_client.create_llm_client", return_value=fake):
|
||||
r = model_config.test_profile(cm, "deepseek")
|
||||
assert "连通正常" in r and "pong" in r
|
||||
|
||||
def test_connectivity_failure(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
fake = mock.MagicMock()
|
||||
fake.chat_completion.side_effect = RuntimeError("连接被拒绝")
|
||||
with mock.patch("llm_client.create_llm_client", return_value=fake):
|
||||
r = model_config.test_profile(cm, "deepseek")
|
||||
assert "验证失败" in r and "连接被拒绝" in r
|
||||
|
||||
def test_connectivity_missing_profile(self, tmp_path):
|
||||
cm = _cfg(tmp_path, _base_models())
|
||||
assert "不存在" in model_config.test_profile(cm, "ghost")
|
||||
|
||||
|
||||
# ==================== 落盘失败感知 ====================
|
||||
|
||||
class TestSaveFailure:
|
||||
def test_add_reports_save_failure(self, tmp_path):
|
||||
cm = _cfg(tmp_path, {})
|
||||
# J5:写路径改为原子写,磁盘失败经原子写函数模拟
|
||||
with mock.patch("common.config.atomic_write_text",
|
||||
side_effect=PermissionError("disk ro")):
|
||||
msgs = model_config.add_profile(
|
||||
cm, "a", "https://x/v1", "m1", api_key_env="K")
|
||||
assert any("保存失败" in m for m in msgs)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,829 @@
|
||||
"""tests.test_n_fixes - Batch N 修复单元测试
|
||||
|
||||
波次 3 平台兼容(2026-08-21,审计第八轮 T 系列):
|
||||
- T3 basic_tools 平台感知引用 + Windows PowerShell/curl.exe 指引
|
||||
- T11 Windows 进程树收割(taskkill /T /F,executors + process_manager)
|
||||
- T25 MemoryDatabase 统一连接工厂(busy_timeout)
|
||||
- T26 SkillManager 原子 reload(中途不清空 self.skills)
|
||||
- T27 atomic_write_text 对 Windows PermissionError 短退避重试
|
||||
|
||||
波次 4 LLM 健壮性(2026-08-21):
|
||||
- T10 调度器 wait_for 硬超时兜底
|
||||
- T19 FailoverClient 异常收窄(永久配置错误不转移不冷却)
|
||||
- T20 流式响应未收到 [DONE] 抛 TransientLLMError(不再静默截断)
|
||||
- T21 env 覆盖黑名单(LD_*/DYLD_*/PATH/PYTHONPATH 等劫持面)
|
||||
- T23 压缩后 tool 配对双向修复(assistant.tool_calls 失配剥离)
|
||||
(T14 FTS5 经复现证伪:JOIN 别名场景下表名 MATCH 正常命中,
|
||||
全链路中文检索长词走 trigram、短词走 LIKE 兜底,无需修复)
|
||||
|
||||
波次 5 权限隔离(2026-08-21,产品决策:T1 仅告警不加鉴权、T2 维持现状):
|
||||
- T22 定时任务精简 profile(task:无 update_identity/update_soul/set_proxy)
|
||||
+ ProcessManager owner 隔离(写/杀权限,观察面共享)
|
||||
- T-L17 model_config 按域名推荐 env 名(防 DeepSeek 凭据发往他厂)
|
||||
- T-L18 list_profiles 输出过滤控制字符(防手写 config ANSI 注入)
|
||||
|
||||
Windows 分支在 Linux 上以 mock.patch("os.name", "nt") 模拟验证命令生成
|
||||
逻辑(不实际执行生成的命令)。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from common.atomic import atomic_write_text
|
||||
from subsystems import MemoryDatabase, SkillManager
|
||||
from subsystems.executors import SubprocessCommandExecutor, filter_env_overrides
|
||||
from tools.basic_tools import (
|
||||
ListFilesTool,
|
||||
ReadFileTool,
|
||||
SimpleWebCrawlerTool,
|
||||
WriteFileTool,
|
||||
_quote_arg,
|
||||
)
|
||||
|
||||
|
||||
# ==================== T3:平台感知引用 ====================
|
||||
|
||||
class TestQuoteArg:
|
||||
def test_posix_uses_shlex(self):
|
||||
quoted, err = _quote_arg("/tmp/my file.txt")
|
||||
assert err is None
|
||||
assert quoted == "'/tmp/my file.txt'"
|
||||
|
||||
@pytest.mark.skipif(os.name != "posix", reason="POSIX 引用语义")
|
||||
def test_posix_specials_escaped(self):
|
||||
quoted, err = _quote_arg("a`b$c$(d)\\e")
|
||||
assert err is None
|
||||
assert quoted == "'a`b$c$(d)\\e'"
|
||||
|
||||
def test_win_wraps_double_quotes(self):
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg(r"C:\Program Files\my app\doc.txt")
|
||||
assert err is None
|
||||
assert quoted == r'"C:\Program Files\my app\doc.txt"'
|
||||
|
||||
def test_win_rejects_double_quote(self):
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg('evil" & del /f /s /q C:\\')
|
||||
assert quoted == ""
|
||||
assert err is not None and "双引号" in err
|
||||
|
||||
def test_win_rejects_cmd_var_expansion(self):
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg("https://evil.com/?x=%PATH%")
|
||||
assert quoted == ""
|
||||
assert err is not None and "注入" in err
|
||||
|
||||
def test_win_allows_url_percent_encoding(self):
|
||||
"""URL 编码(%20/%E4 等)非 %变量名% 形态,不拒绝。"""
|
||||
with mock.patch("os.name", "nt"):
|
||||
quoted, err = _quote_arg("https://example.com/a%20b?x=1&y=2")
|
||||
assert err is None
|
||||
assert quoted == '"https://example.com/a%20b?x=1&y=2"'
|
||||
|
||||
|
||||
# ==================== T3:Windows 指引命令生成 ====================
|
||||
|
||||
class TestReadFileWindows:
|
||||
def test_head_powershell(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ReadFileTool().execute(filepath=str(f), head=True, max_lines=10)
|
||||
assert "powershell" in r and "Get-Content" in r
|
||||
assert "run_command" in r
|
||||
assert "sed -n" not in r and "head -n" not in r
|
||||
|
||||
def test_tail_powershell(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ReadFileTool().execute(filepath=str(f), tail=True, max_lines=5)
|
||||
assert "Get-Content" in r and "-Tail 5" in r
|
||||
|
||||
def test_range_powershell(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ReadFileTool().execute(filepath=str(f), start_line=5, end_line=15)
|
||||
assert "Select-Object -Skip 4 -First 11" in r
|
||||
|
||||
def test_posix_unchanged(self, tmp_path):
|
||||
f = tmp_path / "data.txt"
|
||||
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||
r = ReadFileTool().execute(filepath=str(f), start_line=1, end_line=5)
|
||||
assert "sed -n '1,5p'" in r
|
||||
|
||||
|
||||
class TestWriteFileWindows:
|
||||
def test_base64_set_content(self, tmp_path):
|
||||
target = tmp_path / "sub" / "out.txt"
|
||||
content = "line1\nline2\n含中文 \"quoted\" & special"
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = WriteFileTool().execute(filepath=str(target), content=content, mode="w")
|
||||
assert "Set-Content" in r
|
||||
assert "FromBase64String" in r
|
||||
# 明文内容(含元字符)不得出现在命令行中
|
||||
assert "quoted" not in r and "special" not in r
|
||||
|
||||
def test_append_uses_add_content(self, tmp_path):
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = WriteFileTool().execute(
|
||||
filepath=str(tmp_path / "a.txt"), content="more", mode="a")
|
||||
assert "Add-Content" in r
|
||||
|
||||
def test_oversize_rejected(self, tmp_path):
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = WriteFileTool().execute(
|
||||
filepath=str(tmp_path / "big.txt"), content="x" * 10000, mode="w")
|
||||
assert "过大" in r and "追加" in r
|
||||
|
||||
|
||||
class TestListFilesWindows:
|
||||
def test_get_child_item(self, tmp_path):
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = ListFilesTool().execute(directory=str(tmp_path), show_hidden=True)
|
||||
assert "Get-ChildItem" in r and "-Force" in r
|
||||
assert "ls -l" not in r
|
||||
|
||||
def test_posix_unchanged(self, tmp_path):
|
||||
r = ListFilesTool().execute(directory=str(tmp_path))
|
||||
assert "ls -l " in r
|
||||
|
||||
|
||||
class TestWebWindows:
|
||||
def test_get_curl_exe_amp_url_quoted(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
url = "https://example.com/?a=1&b=2"
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = tool.execute(action="get", url=url)
|
||||
assert "curl.exe" in r
|
||||
# & 必须在双引号内(cmd 双引号内为字面量)
|
||||
assert '"https://example.com/?a=1&b=2"' in r
|
||||
|
||||
def test_cmd_var_url_rejected(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = tool.execute(action="get", url="https://evil.com/?x=%PATH%")
|
||||
assert "拒绝" in r
|
||||
|
||||
def test_text_action_powershell_pipe(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
with mock.patch("os.name", "nt"):
|
||||
r = tool.execute(action="text", url="https://example.com/")
|
||||
assert "curl.exe" in r and "powershell" in r
|
||||
|
||||
def test_posix_unchanged(self):
|
||||
tool = SimpleWebCrawlerTool()
|
||||
r = tool.execute(action="get", url="https://example.com/")
|
||||
assert "curl -sL" in r
|
||||
|
||||
|
||||
# ==================== T11:Windows 进程树收割 ====================
|
||||
|
||||
class TestTerminateTreeWindows:
|
||||
def test_executors_taskkill_tree(self):
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = mock.Mock()
|
||||
proc.pid = 12345
|
||||
proc.poll.return_value = None
|
||||
with mock.patch("os.name", "nt"), \
|
||||
mock.patch("subprocess.run") as run:
|
||||
ex._terminate_tree(proc, force=True)
|
||||
run.assert_called_once()
|
||||
assert run.call_args[0][0] == ["taskkill", "/PID", "12345", "/T", "/F"]
|
||||
proc.kill.assert_called_once() # 兜底单进程 kill
|
||||
|
||||
def test_executors_taskkill_process_gone(self):
|
||||
"""taskkill 目标已退出:兜底 kill 不执行(poll 非 None)。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = mock.Mock()
|
||||
proc.pid = 12346
|
||||
proc.poll.return_value = 0
|
||||
with mock.patch("os.name", "nt"), \
|
||||
mock.patch("subprocess.run") as run:
|
||||
ex._terminate_tree(proc, force=False)
|
||||
run.assert_called_once()
|
||||
proc.kill.assert_not_called()
|
||||
|
||||
def test_process_manager_taskkill_tree(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
entry = mock.Mock()
|
||||
entry.popen.pid = 999
|
||||
entry.popen.poll.return_value = None
|
||||
with mock.patch("os.name", "nt"), \
|
||||
mock.patch("subprocess.run") as run:
|
||||
pm._terminate_tree(entry, force=True)
|
||||
assert run.call_args[0][0] == ["taskkill", "/PID", "999", "/T", "/F"]
|
||||
|
||||
@pytest.mark.skipif(os.name != "posix", reason="POSIX 进程组语义")
|
||||
def test_posix_no_taskkill(self):
|
||||
"""POSIX 分支不调用 taskkill(走 killpg)。"""
|
||||
import subprocess as _sp
|
||||
ex = SubprocessCommandExecutor()
|
||||
proc = _sp.Popen(["sleep", "30"], start_new_session=True)
|
||||
try:
|
||||
with mock.patch("subprocess.run") as run:
|
||||
ex._terminate_tree(proc, force=False)
|
||||
run.assert_not_called()
|
||||
finally:
|
||||
try:
|
||||
import signal as _sig
|
||||
os.killpg(os.getpgid(proc.pid), _sig.SIGKILL)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ==================== T25:MemoryDatabase busy_timeout ====================
|
||||
|
||||
class TestMemoryDbBusyTimeout:
|
||||
def test_connect_sets_busy_timeout(self, tmp_path):
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
conn = db._connect()
|
||||
try:
|
||||
value = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
assert value == 10000
|
||||
|
||||
def test_all_paths_via_factory(self, tmp_path):
|
||||
"""写/读路径不再直连 sqlite3.connect(统一走 _connect 工厂)。"""
|
||||
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||
mid = db.add_memory(topic="t", summary="s", content="c",
|
||||
start_time="2026-01-01 00:00:00",
|
||||
end_time="2026-01-01 00:01:00")
|
||||
assert mid is not None
|
||||
rows = db.get_recent_memories(limit=5)
|
||||
assert any(r["id"] == mid for r in rows)
|
||||
|
||||
|
||||
# ==================== T26:SkillManager 原子 reload ====================
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_skill_dir(tmp_path, monkeypatch):
|
||||
from common import constants
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
return skills_dir
|
||||
|
||||
|
||||
class TestSkillReloadAtomic:
|
||||
def test_skills_visible_during_reload(self, isolated_skill_dir):
|
||||
"""reload 过程中 self.skills 保持旧表(T26:不再先清空)。"""
|
||||
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||
sm = SkillManager()
|
||||
assert sm.count() == 1
|
||||
|
||||
(isolated_skill_dir / "b.md").write_text("# B\n内容B", encoding="utf-8")
|
||||
original = sm._read_markdown_skill
|
||||
observed = []
|
||||
calls = []
|
||||
|
||||
def spy(path):
|
||||
calls.append(path)
|
||||
if len(calls) == 2:
|
||||
observed.append(dict(sm.skills))
|
||||
return original(path)
|
||||
|
||||
sm._read_markdown_skill = spy
|
||||
sm.reload()
|
||||
|
||||
# 加载第 2 个技能时,旧表完整可见(未被清空/半截)
|
||||
assert len(observed) == 1
|
||||
assert set(observed[0].keys()) == {"a"}
|
||||
# reload 完成后一次性替换为新表
|
||||
assert set(sm.skills.keys()) == {"a", "b"}
|
||||
|
||||
def test_reload_missing_dir_keeps_old(self, isolated_skill_dir, monkeypatch):
|
||||
"""目录消失时清空(旧行为保留),正常 reload 不经过空表中间态。"""
|
||||
from common import constants
|
||||
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||
sm = SkillManager()
|
||||
assert sm.count() == 1
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(isolated_skill_dir / "nope")))
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
# ==================== T27:atomic_write_text 重试 ====================
|
||||
|
||||
class TestAtomicWriteRetry:
|
||||
def test_retries_then_succeeds(self, tmp_path):
|
||||
target = tmp_path / "cfg.json"
|
||||
real_replace = os.replace
|
||||
calls = []
|
||||
|
||||
def flaky_replace(src, dst):
|
||||
calls.append((src, dst))
|
||||
if len(calls) < 3:
|
||||
raise PermissionError(5, "Access is denied")
|
||||
return real_replace(src, dst)
|
||||
|
||||
with mock.patch("os.replace", side_effect=flaky_replace):
|
||||
atomic_write_text(str(target), '{"k": 1}')
|
||||
assert len(calls) == 3
|
||||
assert target.read_text(encoding="utf-8") == '{"k": 1}'
|
||||
|
||||
def test_raises_after_exhausted(self, tmp_path):
|
||||
target = tmp_path / "cfg.json"
|
||||
with mock.patch("os.replace",
|
||||
side_effect=PermissionError(5, "Access is denied")):
|
||||
with pytest.raises(PermissionError):
|
||||
atomic_write_text(str(target), "x")
|
||||
assert not target.exists()
|
||||
# 临时文件已清理
|
||||
leftovers = [p for p in tmp_path.iterdir() if p.name.startswith(".aibox_tmp_")]
|
||||
assert leftovers == []
|
||||
|
||||
def test_normal_write_no_retry(self, tmp_path):
|
||||
target = tmp_path / "ok.txt"
|
||||
with mock.patch("os.replace", wraps=os.replace) as spy:
|
||||
atomic_write_text(str(target), "hello")
|
||||
spy.assert_called_once()
|
||||
assert target.read_text(encoding="utf-8") == "hello"
|
||||
|
||||
|
||||
# ==================== T19:FailoverClient 异常收窄 ====================
|
||||
|
||||
def _make_failover_client():
|
||||
"""deepseek(→kimi fallback) 链,Key 走 mock 环境变量。"""
|
||||
from llm_client import create_llm_client
|
||||
from test_models_profile import MODELS, _cfg
|
||||
models = {
|
||||
"deepseek": dict(MODELS["deepseek"], fallback=["kimi"]),
|
||||
"kimi": dict(MODELS["kimi"]),
|
||||
}
|
||||
env = {"DEEPSEEK_API_KEY": "sk-ds", "MOONSHOT_API_KEY": "sk-moon"}
|
||||
with mock.patch.dict(os.environ, env):
|
||||
return create_llm_client(_cfg(models=models))
|
||||
|
||||
|
||||
class TestT19TransientExceptionScope:
|
||||
def test_permanent_schema_error_not_failed_over(self):
|
||||
"""MissingSchema(api_url 手写错等永久配置错误)→ 流式路径直接上抛,
|
||||
不转移(只打一次请求)、不计 failover、不冷却。"""
|
||||
fc = _make_failover_client()
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=requests.exceptions.MissingSchema("bad url")) as mp:
|
||||
with pytest.raises(requests.exceptions.MissingSchema):
|
||||
list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert len(mp.call_args_list) == 1
|
||||
assert fc.get_stats()["failovers"] == 0
|
||||
|
||||
def test_permanent_error_no_cooldown_pollution(self):
|
||||
"""永久错误后档案不进冷却——下次调用仍打主档案。"""
|
||||
fc = _make_failover_client()
|
||||
ok_stream = _stream_response([
|
||||
'data: {"choices":[{"delta":{"content":"recovered"}}]}',
|
||||
'data: [DONE]'])
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[requests.exceptions.MissingSchema("bad"),
|
||||
ok_stream]) as mp:
|
||||
with pytest.raises(requests.exceptions.MissingSchema):
|
||||
list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
events = list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
# 第二次仍先打主档案(deepseek),未进冷却跳过
|
||||
assert mp.call_args_list[1].args[0].startswith("https://api.deepseek.com")
|
||||
assert any(e["type"] == "content" and e["text"] == "recovered" for e in events)
|
||||
assert fc.get_stats()["cooldown_skips"] == 0
|
||||
|
||||
def test_connection_error_still_fails_over(self):
|
||||
"""瞬时连接错误仍正常转移(收窄不误伤)。"""
|
||||
fc = _make_failover_client()
|
||||
from test_models_profile import _resp
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[requests.exceptions.ConnectionError("boom"),
|
||||
_resp(200, content="ok-conn")]):
|
||||
assert fc.chat_completion([{"role": "user", "content": "hi"}]) == "ok-conn"
|
||||
|
||||
def test_chunked_encoding_error_fails_over(self):
|
||||
"""ChunkedEncodingError(响应体中途损坏)属瞬时故障 → 转移。"""
|
||||
fc = _make_failover_client()
|
||||
from test_models_profile import _resp
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[requests.exceptions.ChunkedEncodingError("broken"),
|
||||
_resp(200, content="ok-chunked")]):
|
||||
assert fc.chat_completion([{"role": "user", "content": "hi"}]) == "ok-chunked"
|
||||
|
||||
|
||||
# ==================== T20:不完整流信号 ====================
|
||||
|
||||
def _stream_response(lines):
|
||||
m = mock.MagicMock()
|
||||
m.status_code = 200
|
||||
m.iter_lines.return_value = [l.encode("utf-8") for l in lines]
|
||||
return m
|
||||
|
||||
|
||||
def _make_deepseek_client():
|
||||
from llm_client import DeepSeekClient
|
||||
return DeepSeekClient(
|
||||
api_key="sk-test", model="test-model",
|
||||
base_url="https://api.example/v1/chat/completions")
|
||||
|
||||
|
||||
class TestT20IncompleteStream:
|
||||
def test_missing_done_raises_transient(self):
|
||||
"""流中途断开(iter_lines 耗尽无 [DONE])→ TransientLLMError,
|
||||
不再静默视为成功导致回答无声截断。"""
|
||||
from llm_client import TransientLLMError
|
||||
c = _make_deepseek_client()
|
||||
lines = [
|
||||
'data: {"choices":[{"delta":{"content":"he"}}]}',
|
||||
'data: {"choices":[{"delta":{"content":"llo"}}]}',
|
||||
# 连接断开:无 [DONE],末行可能还是半截 JSON
|
||||
'data: {"choices":[{"delta":{"cont',
|
||||
]
|
||||
with mock.patch("llm_client.requests.post",
|
||||
return_value=_stream_response(lines)):
|
||||
with pytest.raises(TransientLLMError, match="不完整"):
|
||||
list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
|
||||
def test_with_done_completes_normally(self):
|
||||
from llm_client import TransientLLMError
|
||||
c = _make_deepseek_client()
|
||||
lines = [
|
||||
'data: {"choices":[{"delta":{"content":"hi"}}]}',
|
||||
'data: [DONE]',
|
||||
]
|
||||
with mock.patch("llm_client.requests.post",
|
||||
return_value=_stream_response(lines)):
|
||||
events = list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert events[-1] == {"type": "done"}
|
||||
assert any(e["type"] == "content" and e["text"] == "hi" for e in events)
|
||||
|
||||
def test_failover_first_event_incomplete_stream(self):
|
||||
"""FailoverClient:首事件前不完整流(0 事件 + 无 [DONE])→ 可转移。"""
|
||||
fc = _make_failover_client()
|
||||
# 主档案:断流无任何事件;fallback(kimi):正常完整流
|
||||
broken = _stream_response(['data: {"choices":[{"delta'])
|
||||
ok_stream = _stream_response([
|
||||
'data: {"choices":[{"delta":{"content":"from-kimi"}}]}',
|
||||
'data: [DONE]'])
|
||||
with mock.patch("llm_client.requests.post",
|
||||
side_effect=[broken, ok_stream]):
|
||||
events = list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||
assert any(e["type"] == "content" and e["text"] == "from-kimi" for e in events)
|
||||
assert events[-1]["type"] == "done"
|
||||
|
||||
|
||||
# ==================== T21:env 覆盖黑名单 ====================
|
||||
|
||||
class TestT21EnvDenylist:
|
||||
def test_denies_hijack_keys(self):
|
||||
safe, rejected = filter_env_overrides({
|
||||
"LD_PRELOAD": "/evil", "PATH": "/bin", "PYTHONPATH": "/p",
|
||||
"MY_VAR": "ok"})
|
||||
assert safe == {"MY_VAR": "ok"}
|
||||
assert set(rejected) == {"LD_PRELOAD", "PATH", "PYTHONPATH"}
|
||||
|
||||
def test_case_insensitive_and_prefixes(self):
|
||||
safe, rejected = filter_env_overrides({
|
||||
"ld_preload": "x", "DYLD_INSERT_LIBRARIES": "y", "path": "z",
|
||||
"PSMODULEPATH": "w", "OK1": "1"})
|
||||
assert safe == {"OK1": "1"}
|
||||
assert len(rejected) == 4
|
||||
|
||||
def test_empty_and_none(self):
|
||||
assert filter_env_overrides(None) == ({}, [])
|
||||
assert filter_env_overrides({}) == ({}, [])
|
||||
|
||||
def test_executor_blocks_denied_env(self):
|
||||
"""执行路径集成:正常 env 注入生效,黑名单键被拒。"""
|
||||
ex = SubprocessCommandExecutor()
|
||||
r = ex.execute(
|
||||
"echo $T21VAR",
|
||||
env={"T21VAR": "visible", "LD_PRELOAD": "/nonexistent/libevil.so"})
|
||||
assert "visible" in r
|
||||
# LD_PRELOAD 未注入(若注入且文件不存在,loader 报错会进合并输出)
|
||||
assert "/nonexistent" not in r
|
||||
|
||||
|
||||
# ==================== T23:压缩后 tool 配对双向修复 ====================
|
||||
|
||||
def _make_compressor():
|
||||
from subsystems.conversation_compressor import ConversationCompressor
|
||||
return ConversationCompressor(llm_client=None)
|
||||
|
||||
|
||||
class TestT23ToolPairValidation:
|
||||
def test_orphan_tool_calls_stripped(self):
|
||||
"""assistant.tool_calls 的 tool 结果被压缩掉 → 剥离 tool_calls。"""
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "user", "content": "q"},
|
||||
{"role": "assistant", "content": "看结果",
|
||||
"tool_calls": [{"id": "t1", "type": "function",
|
||||
"function": {"name": "read", "arguments": "{}"}}]},
|
||||
{"role": "user", "content": "next"},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert len(out) == 3
|
||||
assert "tool_calls" not in out[1]
|
||||
assert out[1]["content"] == "看结果" # 原有内容保留
|
||||
|
||||
def test_empty_content_filled_when_stripped(self):
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "assistant", "content": None,
|
||||
"tool_calls": [{"id": "t1", "type": "function",
|
||||
"function": {"name": "read", "arguments": "{}"}}]},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert "tool_calls" not in out[0]
|
||||
assert out[0]["content"] == "" # 防产生非法空消息
|
||||
|
||||
def test_partial_pair_keeps_matched_calls(self):
|
||||
"""两个 tool_calls 只剩一个配对 → 保留配对的那个。"""
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "assistant", "content": "run",
|
||||
"tool_calls": [
|
||||
{"id": "t1", "type": "function",
|
||||
"function": {"name": "a", "arguments": "{}"}},
|
||||
{"id": "t2", "type": "function",
|
||||
"function": {"name": "b", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "t2", "content": "result-b"},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert len(out) == 2
|
||||
assert [tc["id"] for tc in out[0]["tool_calls"]] == ["t2"]
|
||||
assert out[1]["tool_call_id"] == "t2"
|
||||
|
||||
def test_orphan_tool_message_removed(self):
|
||||
"""既有行为不回归:孤立 tool 消息(无 assistant)移除。"""
|
||||
comp = _make_compressor()
|
||||
msgs = [
|
||||
{"role": "user", "content": "q"},
|
||||
{"role": "tool", "tool_call_id": "ghost", "content": "x"},
|
||||
]
|
||||
out = comp._validate_message_structure(msgs)
|
||||
assert [m["role"] for m in out] == ["user"]
|
||||
|
||||
def test_intact_pair_untouched(self):
|
||||
"""完整配对原样保留(同一对象引用,零开销路径)。"""
|
||||
comp = _make_compressor()
|
||||
assistant = {"role": "assistant", "content": "ok",
|
||||
"tool_calls": [{"id": "t1", "type": "function",
|
||||
"function": {"name": "a", "arguments": "{}"}}]}
|
||||
tool = {"role": "tool", "tool_call_id": "t1", "content": "r"}
|
||||
out = comp._validate_message_structure([assistant, tool])
|
||||
assert out == [assistant, tool]
|
||||
|
||||
|
||||
# ==================== T10:调度器硬超时兜底 ====================
|
||||
|
||||
class _UnkillableAI:
|
||||
"""软取消无法中断的卡死 AI(模拟 requests 读阻塞被滴流输出规避)。"""
|
||||
|
||||
def __init__(self):
|
||||
self.cancel_called = 0
|
||||
self._release = threading.Event()
|
||||
|
||||
def cancel(self):
|
||||
self.cancel_called += 1
|
||||
# 无效:think_and_respond 仍不返回
|
||||
|
||||
def think_and_respond(self, input_text, output_callback=None,
|
||||
is_continuation=False):
|
||||
self._release.wait(timeout=30)
|
||||
return "never-reached"
|
||||
|
||||
|
||||
class TestT10SchedulerHardTimeout:
|
||||
def test_stuck_ai_hard_timeout(self):
|
||||
"""软取消无效的卡死任务 → wait_for 硬超时放行,调度循环继续。"""
|
||||
from aibox import ConfigManager, TaskScheduler
|
||||
|
||||
class _FakeMemoDB:
|
||||
def get_memo(self, memo_id):
|
||||
return None
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
fake_ai = _UnkillableAI()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
memo = {"id": 7, "title": "卡死任务", "content": "内容", "trigger_count": 0}
|
||||
orig_get = ConfigManager.get
|
||||
|
||||
def fake_get(self, key, default=None):
|
||||
# task_timeout 设为 -29.5:软取消立即触发(call_later 负延迟),
|
||||
# wait_for 硬超时 = -29.5 + 30 = 0.5s 后放行
|
||||
if key == "scheduler.task_timeout":
|
||||
return -29.5
|
||||
return orig_get(self, key, default)
|
||||
|
||||
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||
mock.patch.object(scheduler, "_create_task_ai", return_value=fake_ai):
|
||||
result = loop.run_until_complete(
|
||||
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||
assert "硬超时" in result
|
||||
assert fake_ai.cancel_called >= 1 # 软取消先于硬超时触发
|
||||
finally:
|
||||
fake_ai._release.set() # 释放卡死线程
|
||||
loop.close()
|
||||
|
||||
def test_normal_task_unaffected(self):
|
||||
"""正常返回的任务不受硬超时包装影响(结果正常透传)。"""
|
||||
from aibox import ConfigManager, TaskScheduler
|
||||
|
||||
class _FakeMemoDB:
|
||||
def get_memo(self, memo_id):
|
||||
return None
|
||||
|
||||
class _FastAI:
|
||||
def cancel(self):
|
||||
pass
|
||||
|
||||
def think_and_respond(self, input_text, output_callback=None,
|
||||
is_continuation=False):
|
||||
return "任务完成"
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
memo = {"id": 8, "title": "正常任务", "content": "内容", "trigger_count": 0}
|
||||
orig_get = ConfigManager.get
|
||||
|
||||
def fake_get(self, key, default=None):
|
||||
if key == "scheduler.task_timeout":
|
||||
return 600
|
||||
return orig_get(self, key, default)
|
||||
|
||||
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||
mock.patch.object(scheduler, "_create_task_ai", return_value=_FastAI()):
|
||||
result = loop.run_until_complete(
|
||||
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||
assert "任务完成" in result
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# ==================== T22:owner 隔离(ProcessManager) ====================
|
||||
|
||||
class TestT22OwnerIsolation:
|
||||
def test_kill_denied_for_other_owner(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("sleep 30", owner="session-a")
|
||||
handle = r["handle"]
|
||||
try:
|
||||
denied = pm.kill(handle, owner="session-b")
|
||||
assert not denied["ok"]
|
||||
assert "属主" in denied["error"]
|
||||
# 进程仍在运行(未被误杀)
|
||||
assert pm._lookup(handle).status == "running"
|
||||
finally:
|
||||
pm.kill(handle)
|
||||
|
||||
def test_kill_allowed_for_owner(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("sleep 30", owner="session-a")
|
||||
result = pm.kill(r["handle"], owner="session-a")
|
||||
assert result["ok"]
|
||||
|
||||
def test_write_denied_for_other_owner(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("cat", owner="session-a")
|
||||
handle = r["handle"]
|
||||
try:
|
||||
denied = pm.write(handle, "hi", owner="session-b")
|
||||
assert not denied["ok"]
|
||||
assert "属主" in denied["error"]
|
||||
ok = pm.write(handle, "hi", owner="session-a")
|
||||
assert ok["ok"]
|
||||
finally:
|
||||
pm.kill(handle)
|
||||
|
||||
def test_backward_compatible_without_owner(self):
|
||||
"""任一侧无 owner 标识 → 不隔离(测试/管理面直调兼容)。"""
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r1 = pm.start("sleep 5", owner="a")
|
||||
assert pm.kill(r1["handle"])["ok"] # caller 无 owner
|
||||
r2 = pm.start("sleep 5") # entry 无 owner
|
||||
assert pm.kill(r2["handle"], owner="b")["ok"]
|
||||
|
||||
def test_list_shows_owner_field(self):
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("sleep 5", owner="whoami")
|
||||
try:
|
||||
infos = pm.list_procs()
|
||||
assert any(e.get("owner") == "whoami" for e in infos)
|
||||
finally:
|
||||
pm.kill(r["handle"])
|
||||
|
||||
def test_output_not_isolated(self):
|
||||
"""观察面(get_output)不受 owner 限制。"""
|
||||
from subsystems.process_manager import ProcessManager
|
||||
pm = ProcessManager()
|
||||
r = pm.start("echo hello", owner="session-a")
|
||||
handle = r["handle"]
|
||||
import time as _t
|
||||
_t.sleep(0.3)
|
||||
result = pm.get_output(handle) # 观察面无 owner 概念
|
||||
assert result["ok"]
|
||||
assert "hello" in result.get("output", "")
|
||||
|
||||
|
||||
# ==================== T22:task profile(定时任务精简工具面) ====================
|
||||
|
||||
def _make_chat(profile, owner=None):
|
||||
from aibox import DeepSeekChat
|
||||
with mock.patch("aibox.create_llm_client") as mclc:
|
||||
mclc.return_value = mock.MagicMock(
|
||||
api_key="sk-test", name="DS", default_temperature=None)
|
||||
return DeepSeekChat(tool_profile=profile, owner_id=owner)
|
||||
|
||||
|
||||
class TestT22TaskProfile:
|
||||
def test_task_profile_excludes_global_side_effect_tools(self):
|
||||
ai = _make_chat("task")
|
||||
names = set(ai.tool_registry._tools.keys())
|
||||
assert "update_identity" not in names
|
||||
assert "update_soul" not in names
|
||||
assert "set_proxy" not in names
|
||||
# 管理面其余工具保留
|
||||
assert "get_proxy" in names
|
||||
assert "create_subagent" in names
|
||||
|
||||
def test_orchestrator_profile_still_has_them(self):
|
||||
ai = _make_chat("orchestrator")
|
||||
names = set(ai.tool_registry._tools.keys())
|
||||
assert "update_identity" in names
|
||||
assert "update_soul" in names
|
||||
assert "set_proxy" in names
|
||||
|
||||
def test_worker_bg_tools_carry_owner(self):
|
||||
ai = _make_chat("worker", owner="sess-1")
|
||||
assert ai.tool_registry.get_tool("bg_run").owner == "sess-1"
|
||||
assert ai.tool_registry.get_tool("bg_kill").owner == "sess-1"
|
||||
assert ai.tool_registry.get_tool("bg_write").owner == "sess-1"
|
||||
# 观察面工具不带 owner
|
||||
assert getattr(ai.tool_registry.get_tool("bg_list"), "owner", "unset") == "unset"
|
||||
|
||||
def test_create_task_ai_uses_task_profile(self):
|
||||
from aibox import TaskScheduler
|
||||
# 显式新建 loop:全量运行时主线程可能已无可用的当前事件循环
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||
with mock.patch("aibox.create_llm_client") as mclc:
|
||||
mclc.return_value = mock.MagicMock(
|
||||
api_key="sk-test", name="DS", default_temperature=None)
|
||||
ai = scheduler._create_task_ai()
|
||||
assert ai.tool_profile == "task"
|
||||
assert ai.owner_id and ai.owner_id.startswith("task-")
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# ==================== T-L17/T-L18:model_config 展示与推荐 ====================
|
||||
|
||||
class TestTL17EnvSuggestion:
|
||||
def test_suggest_by_domain(self):
|
||||
from model_config import _suggest_env_name
|
||||
assert _suggest_env_name("https://api.moonshot.cn/v1") == "MOONSHOT_API_KEY"
|
||||
assert _suggest_env_name("https://api.deepseek.com/chat") == "DEEPSEEK_API_KEY"
|
||||
assert _suggest_env_name("https://dashscope.aliyuncs.com/x") == "DASHSCOPE_API_KEY"
|
||||
|
||||
def test_unknown_domain_empty(self):
|
||||
from model_config import _suggest_env_name
|
||||
assert _suggest_env_name("http://localhost:11434/v1") == ""
|
||||
assert _suggest_env_name("") == ""
|
||||
assert _suggest_env_name(None) == ""
|
||||
|
||||
|
||||
class TestTL18DisplaySanitize:
|
||||
def test_sanitize_strips_control_chars(self):
|
||||
from model_config import _sanitize_display
|
||||
out = _sanitize_display("a\x1b]0;pwned\x07b")
|
||||
assert "\x1b" not in out and "\x07" not in out
|
||||
assert out == "a]0;pwnedb" # 控制符剥离,可打印部分保留
|
||||
assert _sanitize_display("正常文本") == "正常文本"
|
||||
assert _sanitize_display(None) == ""
|
||||
|
||||
def test_print_profiles_filters_escape_sequences(self, capsys):
|
||||
from model_config import _print_profiles
|
||||
_print_profiles([{
|
||||
"name": "evil\x1b[31m", "is_default": False,
|
||||
"display_name": "d\x07x", "api_url": "https://x\x1b[2J",
|
||||
"model": "m", "api_key_env": "K\x1bZ",
|
||||
"temperature": None, "timeout": None, "fallback": ["fb\x07"],
|
||||
}])
|
||||
out = capsys.readouterr().out
|
||||
assert "\x1b" not in out
|
||||
assert "\x07" not in out
|
||||
assert "evil" in out # 可打印部分保留
|
||||
@@ -0,0 +1,269 @@
|
||||
"""aibox CLI 一次性对话模式(--once)单元测试
|
||||
|
||||
覆盖:
|
||||
- handle_once_mode 参数路由:未传 --once → 返回 False
|
||||
- enable_memory=False → MemoryTool 不注册(干净禁用)
|
||||
- enable_memory=True(默认)→ MemoryTool 正常注册(回归保证)
|
||||
- cli_callback 流式分流:chunk→stdout / line→stderr / reasoning 受 --show-reasoning 控制
|
||||
- 退出码:成功 0 / 工具失败 1 / API 失败 2
|
||||
- stdout 不被压缩提示等 print 污染
|
||||
|
||||
通过 FakeLLMClient 桩掉 stream_chat,不依赖真实 API/网络。
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
import argparse
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import DeepSeekChat, handle_once_mode # noqa: E402
|
||||
|
||||
|
||||
# ==================== Fake LLM 客户端 ====================
|
||||
|
||||
class FakeLLMClient:
|
||||
"""桩 LLMClient:按预设脚本产出统一事件流。"""
|
||||
|
||||
def __init__(self, events=None, raise_exc=None):
|
||||
# events: list[dict],如 [{"type":"content","text":"你好"}, {"type":"done"}]
|
||||
self.events = events if events is not None else [
|
||||
{"type": "content", "text": "默认回答"},
|
||||
{"type": "done"},
|
||||
]
|
||||
self.raise_exc = raise_exc
|
||||
self.api_key = "fake-key"
|
||||
self.name = "fake-model"
|
||||
self.model = "fake-model"
|
||||
|
||||
def stream_chat(self, messages, temperature=0.7, max_tokens=8192,
|
||||
tools=None, tool_choice=None, reasoning_enabled=False, **extra):
|
||||
if self.raise_exc:
|
||||
raise self.raise_exc
|
||||
for ev in self.events:
|
||||
yield ev
|
||||
|
||||
def chat_completion(self, messages, temperature=0.3, max_tokens=1024, **kwargs):
|
||||
return "压缩摘要"
|
||||
|
||||
def cancel(self):
|
||||
pass
|
||||
|
||||
|
||||
# ==================== 公共工具 ====================
|
||||
|
||||
def _make_ai(events=None, raise_exc=None, enable_memory=True,
|
||||
tool_profile=None):
|
||||
"""构造一个注入 FakeLLMClient 的 DeepSeekChat,跳过真实网络/文件系统副作用。"""
|
||||
from aibox import TOOL_PROFILE_ORCHESTRATOR
|
||||
with mock.patch("aibox.create_llm_client") as mk:
|
||||
fake = FakeLLMClient(events=events, raise_exc=raise_exc)
|
||||
mk.return_value = fake
|
||||
with mock.patch("aibox.PromptManager") as pm, \
|
||||
mock.patch("aibox.SkillManager"):
|
||||
pm.return_value = mock.MagicMock()
|
||||
pm.return_value.get_all_prompts.return_value = []
|
||||
ai = DeepSeekChat(
|
||||
api_key="fake",
|
||||
history_manager=None,
|
||||
memory_db=mock.MagicMock(),
|
||||
memo_db=mock.MagicMock(),
|
||||
command_executor=mock.MagicMock(),
|
||||
prompt_manager=pm.return_value,
|
||||
enable_memory=enable_memory,
|
||||
tool_profile=tool_profile or TOOL_PROFILE_ORCHESTRATOR,
|
||||
)
|
||||
return ai
|
||||
|
||||
|
||||
def _ns(**over):
|
||||
"""构造 argparse.Namespace,含 once 模式全部字段默认值。"""
|
||||
base = dict(once=None, quiet=False, once_timeout=300,
|
||||
no_save_memory=False, show_reasoning=False,
|
||||
worker_only=False, model=None, temperature=None)
|
||||
base.update(over)
|
||||
return argparse.Namespace(**base)
|
||||
|
||||
|
||||
# ==================== enable_memory 开关 ====================
|
||||
|
||||
class TestEnableMemory:
|
||||
def test_memory_tool_registered_by_default(self):
|
||||
"""默认 enable_memory=True → memory 工具注册"""
|
||||
ai = _make_ai(enable_memory=True)
|
||||
names = [t.get_name() for t in ai.tool_registry.get_all_tools()]
|
||||
assert "memory" in names
|
||||
|
||||
def test_memory_tool_disabled(self):
|
||||
"""enable_memory=False → memory 工具不注册(干净禁用)"""
|
||||
ai = _make_ai(enable_memory=False)
|
||||
names = [t.get_name() for t in ai.tool_registry.get_all_tools()]
|
||||
assert "memory" not in names
|
||||
# 其余 orchestrator 管理面工具不受影响
|
||||
assert "memo" in names
|
||||
|
||||
|
||||
# ==================== handle_once_mode 路由 ====================
|
||||
|
||||
class TestHandleOnceModeRouting:
|
||||
def test_no_once_returns_false(self):
|
||||
"""未传 --once → 返回 False,不进入单次模式"""
|
||||
manager = mock.MagicMock()
|
||||
assert handle_once_mode(manager, _ns(once=None)) is False
|
||||
|
||||
|
||||
# ==================== cli_callback 分流 ====================
|
||||
|
||||
class TestCliCallback分流:
|
||||
"""通过捕获 sys.stdout / sys.stderr 验证分流(handle_once_mode 内部闭包,
|
||||
用 subprocess 太重,这里直接测 DeepSeekChat + 回调契约)。"""
|
||||
|
||||
def _collect(self, ai, query, quiet=False, show_reasoning=False):
|
||||
out, err = [], []
|
||||
|
||||
def cb(t, c):
|
||||
if t == "chunk":
|
||||
out.append(c)
|
||||
elif t == "complete":
|
||||
pass
|
||||
elif t == "error":
|
||||
err.append(c)
|
||||
elif t == "line":
|
||||
if not quiet:
|
||||
err.append(c)
|
||||
elif t in ("reasoning", "reasoning_start", "reasoning_end"):
|
||||
if show_reasoning and not quiet:
|
||||
err.append(c)
|
||||
|
||||
resp = ai.think_and_respond(query, cb)
|
||||
return "".join(out), "".join(err), resp
|
||||
|
||||
def test_chunk_goes_to_stdout(self):
|
||||
ai = _make_ai(events=[{"type": "content", "text": "你好世界"},
|
||||
{"type": "done"}])
|
||||
out, err, resp = self._collect(ai, "hi")
|
||||
assert "你好世界" in out
|
||||
|
||||
def test_reasoning_hidden_by_default(self):
|
||||
ai = _make_ai(events=[{"type": "reasoning", "text": "RSN_CONTENT"},
|
||||
{"type": "content", "text": "答案"},
|
||||
{"type": "done"}])
|
||||
out, err, _ = self._collect(ai, "q", show_reasoning=False)
|
||||
assert "RSN_CONTENT" not in err
|
||||
|
||||
def test_reasoning_shown_with_flag(self):
|
||||
ai = _make_ai(events=[{"type": "reasoning", "text": "RSN_CONTENT"},
|
||||
{"type": "content", "text": "答案"},
|
||||
{"type": "done"}])
|
||||
out, err, _ = self._collect(ai, "q", show_reasoning=True)
|
||||
assert "RSN_CONTENT" in err
|
||||
|
||||
def test_reasoning_hidden_when_quiet(self):
|
||||
ai = _make_ai(events=[{"type": "reasoning", "text": "RSN_CONTENT"},
|
||||
{"type": "content", "text": "答案"},
|
||||
{"type": "done"}])
|
||||
out, err, _ = self._collect(ai, "q", quiet=True, show_reasoning=True)
|
||||
assert "RSN_CONTENT" not in err
|
||||
|
||||
|
||||
# ==================== 退出码相关状态 ====================
|
||||
|
||||
class TestExitCodeState:
|
||||
def test_tool_failed_flag_default_false(self):
|
||||
ai = _make_ai()
|
||||
assert ai.stats["tool_failed"] is False
|
||||
|
||||
def test_tool_failed_flag_set_on_exception(self):
|
||||
"""工具执行异常 → stats['tool_failed'] = True"""
|
||||
ai = _make_ai()
|
||||
# 模拟一个会抛异常的工具调用
|
||||
ai.tool_registry.execute_tool = mock.MagicMock(side_effect=RuntimeError("boom"))
|
||||
tool_calls = [{"id": "c1", "function": {"name": "memo", "arguments": "{}"}}]
|
||||
ai._handle_tool_calls(tool_calls, None)
|
||||
assert ai.stats["tool_failed"] is True
|
||||
|
||||
|
||||
# ==================== 压缩提示不污染 stdout ====================
|
||||
|
||||
class TestCompressPromptNotOnStdout:
|
||||
def test_compress_notice_goes_to_callback_line(self, monkeypatch, capsys):
|
||||
"""触发压缩时,提示走 output_callback('line') 而非 print(stdout)。"""
|
||||
ai = _make_ai()
|
||||
# 强制压缩触发
|
||||
monkeypatch.setattr(ai.compressor, "should_compress",
|
||||
lambda msgs: (True, "测试触发"))
|
||||
monkeypatch.setattr(ai.compressor, "compress_conversation",
|
||||
lambda msgs: msgs)
|
||||
monkeypatch.setattr(ai.compressor, "estimate_tokens", lambda msgs: 100)
|
||||
|
||||
lines = []
|
||||
def cb(t, c):
|
||||
if t == "line":
|
||||
lines.append(c)
|
||||
ai.think_and_respond("你好", cb)
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert "📦" not in captured.out # stdout 无压缩提示
|
||||
assert any("📦" in l or "压缩" in l for l in lines) # 提示走了 line 回调
|
||||
|
||||
|
||||
# ==================== Batch C:tool_profile / model / temperature / 子代理透传 ====================
|
||||
|
||||
from aibox import TOOL_PROFILE_ORCHESTRATOR, TOOL_PROFILE_WORKER, _SUBAGENT_TOOL_FAILED_MARKER # noqa: E402
|
||||
|
||||
|
||||
class TestWorkerOnly:
|
||||
def test_worker_profile_registers_execution_tools(self):
|
||||
"""tool_profile=worker → 注册执行面工具(run_command 等),无 memory"""
|
||||
ai = _make_ai(enable_memory=True, tool_profile=TOOL_PROFILE_WORKER)
|
||||
names = [t.get_name() for t in ai.tool_registry.get_all_tools()]
|
||||
assert "run_command" in names
|
||||
assert "read_file" in names
|
||||
assert "memory" not in names # worker 无记忆工具
|
||||
assert "create_subagent" not in names # worker 不能派生
|
||||
|
||||
def test_orchestrator_profile_registers_memory(self):
|
||||
"""tool_profile=orchestrator → 注册管理面工具(memory/create_subagent)"""
|
||||
ai = _make_ai(enable_memory=True)
|
||||
ai.tool_profile = TOOL_PROFILE_ORCHESTRATOR
|
||||
ai._register_tools()
|
||||
names = [t.get_name() for t in ai.tool_registry.get_all_tools()]
|
||||
assert "memory" in names
|
||||
assert "create_subagent" in names
|
||||
|
||||
|
||||
class TestModelTemperatureOverride:
|
||||
def test_model_override(self):
|
||||
"""model 参数透传到 create_llm_client"""
|
||||
ai = _make_ai()
|
||||
assert ai.client.model == "fake-model" # FakeLLMClient.model
|
||||
# 注:真实链路 create_llm_client(model=...) 已在 _make_ai 中 mock 拦截
|
||||
|
||||
def test_temperature_override(self):
|
||||
"""temperature 显式传入覆盖 constants 默认"""
|
||||
ai = _make_ai()
|
||||
# FakeLLMClient 不关心 temperature,这里验证属性被正确设置
|
||||
ai.temperature = 0.1
|
||||
assert ai.temperature == 0.1
|
||||
|
||||
|
||||
class TestSubAgentToolFailurePropagation:
|
||||
def test_subagent_tool_failed_marker_sets_stats(self):
|
||||
"""create_subagent 返回带标记 → orchestrator stats['tool_failed']=True 且标记被剥离"""
|
||||
ai = _make_ai()
|
||||
ai.tool_registry.execute_tool = mock.MagicMock(
|
||||
return_value=_SUBAGENT_TOOL_FAILED_MARKER + "子代理报告正文"
|
||||
)
|
||||
tool_calls = [{"id": "c1", "function": {"name": "create_subagent", "arguments": "{}"}}]
|
||||
tool_msgs = ai._handle_tool_calls(tool_calls, None)
|
||||
assert ai.stats["tool_failed"] is True
|
||||
# 标记已从结果中剥离(LLM 不应看到内部标记)
|
||||
assert tool_msgs and tool_msgs[0]["role"] == "tool"
|
||||
assert "_SUBAGENT_TOOL_FAILED" not in tool_msgs[0]["content"]
|
||||
assert "子代理报告正文" in tool_msgs[0]["content"]
|
||||
@@ -0,0 +1,448 @@
|
||||
"""后台进程管理器(subsystems.process_manager)与 bg_* 工具回归测试
|
||||
|
||||
覆盖三决策语义与安全边界:
|
||||
- bg 生命周期:启动立即返回 / 惰性收割 / kill 幂等 / TTL 过期
|
||||
- 输出缓冲:tail/since 游标续读 / 环形淘汰提示 / 单行截断 / 拉取字节上限
|
||||
- 上限:运行中进程数上限 / 总条目 LRU 淘汰
|
||||
- 收尾:shutdown 杀非 detach、保留 detach;atexit 单例共享
|
||||
- 工具层:参数规整(字符串句柄)/ 双闸禁用 / 注册表分发
|
||||
|
||||
跨平台:长驻进程用 sys.executable -c "time.sleep(...)"(echo/ping 不可移植)。
|
||||
Windows 注意:shell=True 下 terminate 只杀 shell 包装进程(与执行器 L9
|
||||
回退语义一致),被杀判定以 popen.poll() 为准。
|
||||
"""
|
||||
|
||||
import queue
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
from subsystems import ProcessManager, get_process_manager, reset_process_manager # noqa: E402
|
||||
from subsystems.process_manager import MAX_LINE_BYTES # noqa: E402
|
||||
from tools import ( # noqa: E402
|
||||
ToolRegistry,
|
||||
BgRunTool, BgWriteTool, BgListTool, BgOutputTool, BgKillTool,
|
||||
)
|
||||
|
||||
|
||||
PY = sys.executable
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def echo_srv(tmp_path):
|
||||
"""行协议回环服务器:stdin 读一行 → stdout 回 ECHO:行,EOF 退出。
|
||||
交互式 I/O(bg_write/watermark/优雅终止)测试的标准对端。"""
|
||||
script = (
|
||||
"import sys\n"
|
||||
"while True:\n"
|
||||
" line = sys.stdin.readline()\n"
|
||||
" if not line:\n"
|
||||
" break\n"
|
||||
" sys.stdout.write('ECHO:' + line.strip() + '\\n')\n"
|
||||
" sys.stdout.flush()\n"
|
||||
)
|
||||
f = tmp_path / "echo_srv.py"
|
||||
f.write_text(script, encoding="utf-8")
|
||||
return f'"{PY}" "{f}"'
|
||||
|
||||
|
||||
def _py(code: str) -> str:
|
||||
"""跨平台 python 命令(code 内只用单引号)"""
|
||||
return f'"{PY}" -c "{code}"'
|
||||
|
||||
|
||||
ECHO_10 = _py("for i in range(10): print('line-%d' % i)")
|
||||
SLEEP_10 = _py("import time; time.sleep(10)")
|
||||
HELLO = _py("print('hello-bg')")
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def mgr():
|
||||
"""每测试独立管理器(不走单例),收尾兜底清进程"""
|
||||
m = ProcessManager()
|
||||
try:
|
||||
yield m
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
|
||||
def _wait_status(m: ProcessManager, handle: int, status: str, timeout: float = 10.0):
|
||||
"""轮询等待句柄达到指定状态"""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
for p in m.list_procs():
|
||||
if p["handle"] == handle and p["status"] == status:
|
||||
return p
|
||||
time.sleep(0.1)
|
||||
return None
|
||||
|
||||
|
||||
# ==================== 启动与生命周期 ====================
|
||||
|
||||
class TestStart:
|
||||
def test_start_returns_handle_immediately(self, mgr):
|
||||
t0 = time.time()
|
||||
r = mgr.start(SLEEP_10, name="sleeper")
|
||||
assert r["ok"] is True
|
||||
assert time.time() - t0 < 2 # 立即返回,不等进程结束
|
||||
assert isinstance(r["handle"], int) and r["pid"] > 0
|
||||
assert r["name"] == "sleeper"
|
||||
|
||||
def test_start_default_name_from_command(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
assert r["ok"] is True
|
||||
# 默认名称 = 命令前 40 字符
|
||||
assert r["name"] == HELLO.strip()[:40]
|
||||
|
||||
def test_start_rejects_empty_command(self, mgr):
|
||||
assert mgr.start(" ")["ok"] is False
|
||||
assert mgr.start(None)["ok"] is False
|
||||
|
||||
def test_start_rejects_bad_name(self, mgr):
|
||||
# 名称含控制字符(M5 惯例:防终端/前端注入)
|
||||
assert mgr.start(HELLO, name="bad\x1b[31m")["ok"] is False
|
||||
assert mgr.start(HELLO, name="x" * 65)["ok"] is False
|
||||
|
||||
def test_start_missing_binary_returns_error(self, mgr):
|
||||
r = mgr.start("definitely-not-a-command-xyz --version")
|
||||
# shell 本身能启动(rc=127/9009),故启动成功、退出码非零
|
||||
assert r["ok"] is True
|
||||
p = _wait_status(mgr, r["handle"], "exited")
|
||||
assert p is not None
|
||||
assert p["exit_code"] != 0
|
||||
|
||||
|
||||
class TestLifecycle:
|
||||
def test_exited_status_and_exit_code(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
p = _wait_status(mgr, r["handle"], "exited")
|
||||
assert p is not None
|
||||
assert p["exit_code"] == 0
|
||||
|
||||
def test_kill_running_process(self, mgr):
|
||||
r = mgr.start(SLEEP_10)
|
||||
time.sleep(0.5) # 等进程真正起来
|
||||
k = mgr.kill(r["handle"])
|
||||
assert k["ok"] is True and k["status"] == "killed"
|
||||
assert mgr._procs[r["handle"]].popen.poll() is not None
|
||||
|
||||
def test_kill_is_idempotent_on_exited(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
k = mgr.kill(r["handle"])
|
||||
assert k["ok"] is True and k["already_finished"] is True
|
||||
|
||||
def test_kill_missing_handle(self, mgr):
|
||||
assert mgr.kill(999)["ok"] is False
|
||||
|
||||
def test_ttl_expiry(self):
|
||||
m = ProcessManager(max_lifetime=0.5)
|
||||
try:
|
||||
r = m.start(SLEEP_10)
|
||||
time.sleep(0.8)
|
||||
# list 触发惰性 sweep
|
||||
assert _wait_status(m, r["handle"], "expired", timeout=2) is not None
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_max_processes_limit(self):
|
||||
m = ProcessManager(max_processes=2)
|
||||
try:
|
||||
assert m.start(SLEEP_10)["ok"] is True
|
||||
assert m.start(SLEEP_10)["ok"] is True
|
||||
r3 = m.start(SLEEP_10)
|
||||
assert r3["ok"] is False and "上限" in r3["error"]
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_max_entries_evicts_oldest_exited(self):
|
||||
m = ProcessManager(max_entries=2)
|
||||
try:
|
||||
h1 = m.start(HELLO)["handle"]
|
||||
h2 = m.start(HELLO)["handle"]
|
||||
assert _wait_status(m, h1, "exited") is not None
|
||||
assert _wait_status(m, h2, "exited") is not None
|
||||
# 第三条触发最老已退出条目(#1)LRU 淘汰
|
||||
h3 = m.start(HELLO)["handle"]
|
||||
assert h3 > 0
|
||||
handles = {p["handle"] for p in m.list_procs()}
|
||||
assert h1 not in handles and h2 in handles and h3 in handles
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
|
||||
# ==================== 输出缓冲与游标 ====================
|
||||
|
||||
class TestOutput:
|
||||
def test_tail_and_since_cursor(self, mgr):
|
||||
r = mgr.start(ECHO_10)
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
time.sleep(0.3) # reader 排空末尾行
|
||||
|
||||
out = mgr.get_output(r["handle"], tail=3)
|
||||
assert out["ok"] is True
|
||||
assert out["lines"] == 3
|
||||
assert "line-9" in out["output"] and "line-6" not in out["output"]
|
||||
|
||||
full = mgr.get_output(r["handle"], since=0)
|
||||
assert full["lines"] == 10
|
||||
assert full["next_cursor"] == 10
|
||||
|
||||
inc = mgr.get_output(r["handle"], since=full["next_cursor"])
|
||||
assert inc["lines"] == 0 and inc["next_cursor"] == 10
|
||||
|
||||
def test_since_before_eviction_reports_gap(self):
|
||||
m = ProcessManager(buffer_per_process=2048)
|
||||
try:
|
||||
r = m.start(_py("for i in range(50): print('pad-%03d-%s' % (i, 'x'*80))"))
|
||||
assert _wait_status(m, r["handle"], "exited") is not None
|
||||
time.sleep(0.3)
|
||||
out = m.get_output(r["handle"], since=0)
|
||||
assert out["ok"] is True
|
||||
assert out["evicted"] > 0 # 早期行被环形淘汰
|
||||
assert out["lines"] < 50
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_single_line_truncated(self, mgr):
|
||||
r = mgr.start(_py("print('A' * (1024*64))"))
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
time.sleep(0.3)
|
||||
out = mgr.get_output(r["handle"], tail=5)
|
||||
assert "单行截断" in out["output"]
|
||||
# 缓冲内存预算不被超长单行击穿
|
||||
entry = mgr._procs[r["handle"]]
|
||||
assert entry.buffer_bytes <= MAX_LINE_BYTES + 1024
|
||||
|
||||
def test_fetch_byte_limit_truncates(self, mgr):
|
||||
r = mgr.start(_py("for i in range(100): print('y' * 200)"))
|
||||
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||
time.sleep(0.3)
|
||||
out = mgr.get_output(r["handle"], since=0, max_bytes=4096)
|
||||
assert out["truncated"] is True
|
||||
|
||||
def test_missing_handle(self, mgr):
|
||||
assert mgr.get_output(999)["ok"] is False
|
||||
|
||||
def test_invalid_since(self, mgr):
|
||||
r = mgr.start(HELLO)
|
||||
assert mgr.get_output(r["handle"], since="abc")["ok"] is False
|
||||
|
||||
|
||||
# ==================== 交互式 I/O(stdin 管道 + watermark 对齐) ====================
|
||||
|
||||
class TestInteractive:
|
||||
def test_write_roundtrip_watermark_alignment(self, mgr, echo_srv):
|
||||
"""核心对齐语义:bg_output(since=watermark) = 本条输入的响应且不含回显"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
r = mgr.write(h, "hello", wait_lines=1, wait_timeout=10)
|
||||
assert r["ok"] is True
|
||||
assert "ECHO:hello" in r["wait"]["output"]
|
||||
assert r["wait"]["timed_out"] is False
|
||||
|
||||
out = mgr.get_output(h, since=r["watermark"])
|
||||
assert "ECHO:hello" in out["output"]
|
||||
assert "[stdin]" not in out["output"] # 回显行序号 == watermark,不包含
|
||||
|
||||
def test_echo_line_visible_in_tail(self, mgr, echo_srv):
|
||||
"""[stdin] 回显进环形缓冲,tail 视图可见完整交互时序"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
assert mgr.write(h, "abc", wait_lines=1)["ok"] is True
|
||||
out = mgr.get_output(h, tail=10)
|
||||
assert "[stdin] abc" in out["output"]
|
||||
assert "ECHO:abc" in out["output"]
|
||||
|
||||
def test_write_wait_sync_roundtrip(self, mgr, echo_srv):
|
||||
"""wait_lines 同步等待:一步问答(expect 式)"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
r = mgr.write(h, "1+1", wait_lines=1)
|
||||
assert r["wait"]["lines"] == 1
|
||||
assert r["wait"]["output"].strip() == "ECHO:1+1"
|
||||
assert r["wait"]["next_cursor"] == r["watermark"] + 1
|
||||
|
||||
def test_write_wait_timeout_on_silent_process(self, mgr):
|
||||
"""不读不响应的进程:写入成功(管道缓冲),等待如实超时"""
|
||||
h = mgr.start(SLEEP_10)["handle"]
|
||||
r = mgr.write(h, "ping", wait_lines=1, wait_timeout=0.5)
|
||||
assert r["ok"] is True
|
||||
assert r["wait"]["timed_out"] is True
|
||||
assert r["wait"]["lines"] == 0
|
||||
|
||||
def test_write_to_exited_process(self, mgr):
|
||||
h = mgr.start(HELLO)["handle"]
|
||||
assert _wait_status(mgr, h, "exited") is not None
|
||||
assert mgr.write(h, "x")["ok"] is False
|
||||
|
||||
def test_write_rejects_empty_and_oversize(self, mgr, echo_srv):
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
assert mgr.write(h, "")["ok"] is False
|
||||
assert mgr.write(h, None)["ok"] is False
|
||||
assert mgr.write(h, "A" * (64 * 1024 + 1))["ok"] is False
|
||||
|
||||
def test_stdin_eof_graceful_kill(self, mgr, echo_srv):
|
||||
"""三级优雅终止第一级:关 stdin → 回环服务器见 EOF 自行退出"""
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
time.sleep(0.5)
|
||||
k = mgr.kill(h) # 不 force:先关 stdin
|
||||
assert k["ok"] is True
|
||||
assert mgr._procs[h].popen.poll() is not None
|
||||
assert _wait_status(mgr, h, "killed") is not None
|
||||
|
||||
def test_write_after_kill_rejected(self, mgr, echo_srv):
|
||||
h = mgr.start(echo_srv)["handle"]
|
||||
assert mgr.kill(h)["ok"] is True
|
||||
assert mgr.write(h, "late")["ok"] is False
|
||||
|
||||
def test_log_file_tee(self, mgr, monkeypatch, tmp_path):
|
||||
"""log_file 全档镜像:输出与 [stdin] 回显均落盘,可被原生工具读取"""
|
||||
from common import constants as cm
|
||||
monkeypatch.setattr(type(cm), "WORKSPACE_DIR",
|
||||
property(lambda self: str(tmp_path)))
|
||||
script = tmp_path / "srv.py"
|
||||
script.write_text(
|
||||
"import sys\n"
|
||||
"for line in sys.stdin:\n"
|
||||
" sys.stdout.write('ECHO:' + line.strip() + '\\n')\n"
|
||||
" sys.stdout.flush()\n",
|
||||
encoding="utf-8")
|
||||
cmd = f'"{PY}" "{script}"'
|
||||
r = mgr.start(cmd, log_file=True)
|
||||
assert r["ok"] is True and r["log_file"] is not None
|
||||
assert mgr.write(r["handle"], "tee-me", wait_lines=1)["ok"] is True
|
||||
mgr.kill(r["handle"])
|
||||
log = (tmp_path / f"bg_{r['handle']}.log").read_text(encoding="utf-8")
|
||||
assert "[stdin] tee-me" in log
|
||||
assert "ECHO:tee-me" in log
|
||||
|
||||
def test_bg_write_tool_flow(self, mgr, echo_srv):
|
||||
"""工具层冒烟:BgWriteTool 参数规整 + 对齐读取提示"""
|
||||
assert BgRunTool(manager=mgr).execute(command=echo_srv).startswith("✅")
|
||||
out = BgWriteTool(manager=mgr).execute(handle="1", text="hi",
|
||||
wait_lines="1", wait_timeout=10)
|
||||
assert "ECHO:hi" in out
|
||||
assert "watermark" in out
|
||||
# 描述里的行缓冲指引(给模型看的适用前提)
|
||||
assert "PYTHONUNBUFFERED" in BgWriteTool(manager=mgr).get_description()
|
||||
|
||||
def test_registry_dispatch_bg_write(self, mgr, echo_srv):
|
||||
reg = ToolRegistry()
|
||||
reg.register_many([BgRunTool(manager=mgr), BgWriteTool(manager=mgr)],
|
||||
is_builtin=True)
|
||||
assert reg.execute_tool("bg_run", command=echo_srv).startswith("✅")
|
||||
r = reg.execute_tool("bg_write", handle=1, text="via-reg",
|
||||
wait_lines=1, wait_timeout=10)
|
||||
assert "ECHO:via-reg" in r
|
||||
|
||||
|
||||
# ==================== 退出收尾 ====================
|
||||
|
||||
class TestShutdown:
|
||||
def test_shutdown_kills_non_detach_keeps_detach(self):
|
||||
m = ProcessManager()
|
||||
try:
|
||||
h1 = m.start(SLEEP_10)["handle"] # 常规:退出时终止
|
||||
h2 = m.start(SLEEP_10, detach=True)["handle"] # detach:保留
|
||||
time.sleep(0.5)
|
||||
e1, e2 = m._procs[h1], m._procs[h2]
|
||||
|
||||
r = m.shutdown()
|
||||
assert r["killed"] == 1 and r["detached"] == 1
|
||||
assert e1.popen.poll() is not None # 常规进程已终止
|
||||
assert e2.popen.poll() is None # detach 进程仍存活
|
||||
assert m.list_procs() == [] # 注册表已清空(detach 仅注销)
|
||||
|
||||
# 幂等
|
||||
assert m.shutdown() == {"killed": 0, "detached": 0}
|
||||
# 收尾后拒绝新进程
|
||||
assert m.start(HELLO)["ok"] is False
|
||||
# 清理 detach 测试进程
|
||||
e2.popen.kill()
|
||||
finally:
|
||||
m.shutdown()
|
||||
|
||||
def test_singleton_shared_across_callers(self):
|
||||
m1 = get_process_manager()
|
||||
m2 = get_process_manager()
|
||||
assert m1 is m2 # 跨会话/跨子代理共享同一注册表
|
||||
reset_process_manager(None)
|
||||
m3 = get_process_manager()
|
||||
assert m3 is not m1
|
||||
m3.shutdown()
|
||||
reset_process_manager(None)
|
||||
|
||||
|
||||
# ==================== 工具层 ====================
|
||||
|
||||
class TestTools:
|
||||
def test_bg_run_and_output_flow(self, mgr):
|
||||
run = BgRunTool(manager=mgr)
|
||||
r = run.execute(command=HELLO, proc_name="hello")
|
||||
assert r.startswith("✅") and "#" in r
|
||||
handle = 1
|
||||
assert _wait_status(mgr, handle, "exited") is not None
|
||||
time.sleep(0.3)
|
||||
|
||||
out = BgOutputTool(manager=mgr).execute(handle=handle, since=0)
|
||||
assert "hello-bg" in out
|
||||
assert "next_cursor" in out
|
||||
|
||||
def test_bg_run_string_handle_coercion(self, mgr):
|
||||
"""LLM 可能传字符串数字句柄(W 批输入规整惯例)"""
|
||||
run = BgRunTool(manager=mgr)
|
||||
assert run.execute(command=HELLO).startswith("✅")
|
||||
assert _wait_status(mgr, 1, "exited") is not None
|
||||
out = BgOutputTool(manager=mgr).execute(handle="1", tail=5)
|
||||
assert "hello-bg" in out
|
||||
|
||||
def test_bg_run_missing_command(self, mgr):
|
||||
r = BgRunTool(manager=mgr).execute()
|
||||
assert "缺少" in r
|
||||
|
||||
def test_bg_list_empty_and_filled(self, mgr):
|
||||
lst = BgListTool(manager=mgr)
|
||||
assert "没有后台进程" in lst.execute()
|
||||
mgr.start(SLEEP_10, name="sleeper")
|
||||
r = lst.execute()
|
||||
assert "sleeper" in r and "运行中 1" in r
|
||||
|
||||
def test_bg_kill_flow(self, mgr):
|
||||
mgr.start(SLEEP_10, name="sleeper")
|
||||
time.sleep(0.5)
|
||||
r = BgKillTool(manager=mgr).execute(handle=1)
|
||||
assert r.startswith("✅") and "已终止" in r
|
||||
|
||||
def test_disabled_gate_blocks_all(self, mgr, monkeypatch):
|
||||
from common import constants
|
||||
monkeypatch.setattr(type(constants), "BG_PROCESSES_ENABLED",
|
||||
property(lambda self: False))
|
||||
assert "禁用" in BgRunTool(manager=mgr).execute(command=HELLO)
|
||||
assert "禁用" in BgListTool(manager=mgr).execute()
|
||||
assert "禁用" in BgOutputTool(manager=mgr).execute(handle=1)
|
||||
assert "禁用" in BgKillTool(manager=mgr).execute(handle=1)
|
||||
|
||||
def test_registry_dispatch(self, mgr):
|
||||
"""4 个工具经 ToolRegistry 正常分发(worker 注册链路冒烟)"""
|
||||
reg = ToolRegistry()
|
||||
reg.register_many([
|
||||
BgRunTool(manager=mgr), BgListTool(manager=mgr),
|
||||
BgOutputTool(manager=mgr), BgKillTool(manager=mgr),
|
||||
], is_builtin=True)
|
||||
r = reg.execute_tool("bg_run", command=HELLO, proc_name="smoke")
|
||||
assert r.startswith("✅")
|
||||
assert _wait_status(mgr, 1, "exited") is not None
|
||||
time.sleep(0.3)
|
||||
r = reg.execute_tool("bg_output", handle=1, since=0)
|
||||
assert "hello-bg" in r
|
||||
r = reg.execute_tool("bg_list")
|
||||
assert "smoke" in r
|
||||
|
||||
|
||||
def test_worker_only_tools_includes_bg():
|
||||
"""编排者误调拦截表包含 bg 工具(worker 执行面边界)"""
|
||||
from aibox import WORKER_ONLY_TOOLS
|
||||
assert {"bg_run", "bg_write", "bg_list", "bg_output", "bg_kill"} <= WORKER_ONLY_TOOLS
|
||||
@@ -0,0 +1,190 @@
|
||||
"""subsystems.prompt_manager 单元测试
|
||||
|
||||
覆盖:
|
||||
- PromptManager 单例语义
|
||||
- 提示词加载/排序/拼接
|
||||
- identity/soul/status 读写
|
||||
- 从模板复制行为
|
||||
- reload 语义
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import PromptManager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_pm(tmp_path, monkeypatch):
|
||||
"""每个测试用独立 prompts 目录 + 重置单例,避免测试间状态污染。"""
|
||||
prompts_dir = tmp_path / "prompts"
|
||||
prompts_dir.mkdir()
|
||||
|
||||
# 重置单例状态
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
# 让 ConfigManager 返回临时目录
|
||||
from common import ConfigManager
|
||||
original_get = ConfigManager.get
|
||||
|
||||
def mock_get(self, key, default=None):
|
||||
if key == "system.prompts_dir":
|
||||
return str(prompts_dir)
|
||||
return original_get(self, key, default)
|
||||
|
||||
monkeypatch.setattr(ConfigManager, "get", mock_get)
|
||||
|
||||
pm = PromptManager()
|
||||
yield pm
|
||||
|
||||
# 清理单例
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
|
||||
# ==================== 单例 ====================
|
||||
|
||||
def test_prompt_manager_is_singleton(isolated_pm):
|
||||
pm2 = PromptManager()
|
||||
assert pm2 is isolated_pm
|
||||
|
||||
|
||||
def test_prompt_manager_resets_between_tests(isolated_pm):
|
||||
"""fixture 应重置单例,确保测试隔离。"""
|
||||
assert PromptManager._instance is isolated_pm
|
||||
|
||||
|
||||
# ==================== 加载与排序 ====================
|
||||
|
||||
def test_loads_template_files(isolated_pm):
|
||||
"""首次构造应从 aibox/prompts/ 模板复制核心文件。"""
|
||||
files = os.listdir(isolated_pm._prompts_dir)
|
||||
# 至少应有 00_IDENTITY / 99_SOUL / 100_STATUS
|
||||
assert "00_IDENTITY.md" in files
|
||||
assert "99_SOUL.md" in files
|
||||
assert "100_STATUS.md" in files
|
||||
|
||||
|
||||
def test_get_prompts_returns_nonempty_string(isolated_pm):
|
||||
s = isolated_pm.get_prompts()
|
||||
assert isinstance(s, str) and len(s) > 0
|
||||
|
||||
|
||||
def test_get_sorted_prompts_list_returns_list(isolated_pm):
|
||||
lst = isolated_pm.get_sorted_prompts_list()
|
||||
assert isinstance(lst, list)
|
||||
assert len(lst) > 0
|
||||
assert all("file" in p and "content" in p for p in lst)
|
||||
|
||||
|
||||
def test_prompt_sort_order_identity_first(isolated_pm):
|
||||
"""00_IDENTITY 应排在最前。"""
|
||||
lst = isolated_pm.get_sorted_prompts_list()
|
||||
assert lst[0]["file"] == "00_IDENTITY.md"
|
||||
|
||||
|
||||
def test_prompt_sort_order_status_last_or_near_last(isolated_pm):
|
||||
"""100_STATUS 应在 SOUL 之后。"""
|
||||
lst = isolated_pm.get_sorted_prompts_list()
|
||||
files = [p["file"] for p in lst]
|
||||
if "100_STATUS.md" in files and "99_SOUL.md" in files:
|
||||
assert files.index("100_STATUS.md") > files.index("99_SOUL.md")
|
||||
|
||||
|
||||
def test_skill_prompts_filtered_correctly(isolated_pm):
|
||||
"""N_ 开头的技能文件应被 get_skill_prompts 识别。"""
|
||||
skills = isolated_pm.get_skill_prompts()
|
||||
# 模板里至少有一个 1_xxx.md 技能文件
|
||||
assert len(skills) >= 0 # 至少不报错
|
||||
for p in skills:
|
||||
assert p["file"][0].isdigit()
|
||||
|
||||
|
||||
# ==================== 读写 identity/soul/status ====================
|
||||
|
||||
def test_update_and_get_identity(isolated_pm):
|
||||
ok = isolated_pm.update_identity("测试用户-张三")
|
||||
assert ok is True
|
||||
content = isolated_pm.get_identity()
|
||||
assert "测试用户-张三" in content
|
||||
assert "用户身份信息" in content
|
||||
|
||||
|
||||
def test_update_and_get_soul(isolated_pm):
|
||||
ok = isolated_pm.update_soul("我是测试智能体")
|
||||
assert ok is True
|
||||
content = isolated_pm.get_soul()
|
||||
assert "我是测试智能体" in content
|
||||
|
||||
|
||||
def test_update_and_get_status(isolated_pm):
|
||||
ok = isolated_pm.update_status("测试状态-初始化完成")
|
||||
assert ok is True
|
||||
content = isolated_pm.get_status()
|
||||
assert "测试状态-初始化完成" in content
|
||||
|
||||
|
||||
def test_get_identity_nonexistent_returns_fallback(monkeypatch, tmp_path):
|
||||
"""无 identity 文件时应返回兜底字符串。"""
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
empty_dir = tmp_path / "empty"
|
||||
empty_dir.mkdir()
|
||||
|
||||
from common import ConfigManager
|
||||
original_get = ConfigManager.get
|
||||
|
||||
def mock_get(self, key, default=None):
|
||||
if key == "system.prompts_dir":
|
||||
return str(empty_dir)
|
||||
return original_get(self, key, default)
|
||||
|
||||
monkeypatch.setattr(ConfigManager, "get", mock_get)
|
||||
pm = PromptManager()
|
||||
# 删除 identity 文件模拟不存在
|
||||
id_file = empty_dir / "00_IDENTITY.md"
|
||||
if id_file.exists():
|
||||
id_file.unlink()
|
||||
r = pm.get_identity()
|
||||
assert "尚不清楚" in r or "无法读取" in r
|
||||
|
||||
PromptManager._instance = None
|
||||
PromptManager._prompts = []
|
||||
PromptManager._prompts_dir = None
|
||||
|
||||
|
||||
# ==================== reload ====================
|
||||
|
||||
def test_reload_returns_count(isolated_pm):
|
||||
n = isolated_pm.reload()
|
||||
assert isinstance(n, int) and n >= 0
|
||||
|
||||
|
||||
def test_reload_picks_up_new_file(isolated_pm, tmp_path):
|
||||
"""reload 应能发现新增的提示词文件。"""
|
||||
before = isolated_pm.get_prompt_count()
|
||||
# 写入一个新的提示词文件
|
||||
new_file = os.path.join(isolated_pm._prompts_dir, "99_EXTRA_TEST.md")
|
||||
with open(new_file, "w", encoding="utf-8") as f:
|
||||
f.write("# 额外测试提示词\n内容")
|
||||
isolated_pm.reload()
|
||||
after = isolated_pm.get_prompt_count()
|
||||
assert after == before + 1
|
||||
|
||||
|
||||
# ==================== 模板不覆盖已存在文件 ====================
|
||||
|
||||
def test_existing_file_not_overwritten_by_template(isolated_pm):
|
||||
"""已存在的运行时文件不应被模板覆盖。"""
|
||||
# 修改 identity 文件
|
||||
isolated_pm.update_identity("我的自定义身份")
|
||||
custom_content = isolated_pm.get_identity()
|
||||
assert "我的自定义身份" in custom_content
|
||||
|
||||
# reload 触发 _ensure_from_template,但不应覆盖
|
||||
isolated_pm.reload()
|
||||
after = isolated_pm.get_identity()
|
||||
assert "我的自定义身份" in after
|
||||
@@ -0,0 +1,204 @@
|
||||
"""subsystems.skill_manager 单元测试
|
||||
|
||||
覆盖:
|
||||
- SkillManager 加载 Markdown 技能
|
||||
- 加载可执行技能(--help)
|
||||
- search / get_skill_content / get_all / reload / count
|
||||
每个测试用独立临时技能目录,互不污染。
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from subsystems import SkillManager
|
||||
from common import constants
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_skill_dir(tmp_path, monkeypatch):
|
||||
"""每个测试用独立技能目录 + monkeypatch constants.SKILLS_DIR。"""
|
||||
skills_dir = tmp_path / "skills"
|
||||
skills_dir.mkdir()
|
||||
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||
property(lambda self: str(skills_dir)))
|
||||
return skills_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sm(isolated_skill_dir):
|
||||
"""构造 SkillManager(依赖 isolated_skill_dir fixture 注入目录)。"""
|
||||
return SkillManager()
|
||||
|
||||
|
||||
# ==================== 空目录 ====================
|
||||
|
||||
def test_empty_skills_dir(sm):
|
||||
assert sm.count() == 0
|
||||
assert sm.get_all() == []
|
||||
|
||||
|
||||
# ==================== Markdown 技能 ====================
|
||||
|
||||
def test_loads_markdown_skill(isolated_skill_dir, sm):
|
||||
"""Markdown 技能应被加载,简介取首行(# 标题去 # 前缀)。"""
|
||||
skill_file = isolated_skill_dir / "weather.md"
|
||||
skill_file.write_text("# 查询天气\n查询指定城市天气\n详细用法...", encoding="utf-8")
|
||||
|
||||
sm.reload()
|
||||
assert sm.count() == 1
|
||||
|
||||
s = sm.get_skill("weather")
|
||||
assert s is not None
|
||||
assert s["summary"] == "查询天气"
|
||||
assert s["is_markdown"] is True
|
||||
assert s["is_executable"] is False
|
||||
assert "查询指定城市天气" in s["content"]
|
||||
|
||||
|
||||
def test_markdown_skill_content_retrieval(isolated_skill_dir, sm):
|
||||
"""sm fixture 构造时目录为空,写入文件后需 reload 才能加载。"""
|
||||
(isolated_skill_dir / "notes.md").write_text("# 笔记技能\n笔记内容ABC", encoding="utf-8")
|
||||
sm.reload()
|
||||
content = sm.get_skill_content("notes")
|
||||
assert "笔记内容ABC" in content
|
||||
|
||||
|
||||
def test_markdown_skill_empty_file_skipped(isolated_skill_dir, sm):
|
||||
"""空文件应被跳过。"""
|
||||
(isolated_skill_dir / "empty.md").write_text("", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
def test_markdown_skill_no_title_uses_filename(isolated_skill_dir, sm):
|
||||
"""无 # 标题时,首行普通文本作为简介。"""
|
||||
(isolated_skill_dir / "plain.md").write_text("这是普通首行\n第二行", encoding="utf-8")
|
||||
sm.reload()
|
||||
s = sm.get_skill("plain")
|
||||
assert s["summary"] == "这是普通首行"
|
||||
|
||||
|
||||
# ==================== 隐藏/下划线文件跳过 ====================
|
||||
|
||||
def test_hidden_file_skipped(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / ".hidden.md").write_text("# 隐藏", encoding="utf-8")
|
||||
(isolated_skill_dir / "_internal.md").write_text("# 内部", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
# ==================== 搜索 ====================
|
||||
|
||||
def test_search_by_name(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "weather.md").write_text("# 天气\n查天气", encoding="utf-8")
|
||||
(isolated_skill_dir / "calculator.md").write_text("# 计算器\n算数", encoding="utf-8")
|
||||
sm.reload()
|
||||
|
||||
results = sm.search("weather")
|
||||
assert len(results) == 1
|
||||
assert results[0]["name"] == "weather"
|
||||
|
||||
|
||||
def test_search_by_summary(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "x.md").write_text("# 天气预报技能\n内容", encoding="utf-8")
|
||||
sm.reload()
|
||||
results = sm.search("天气")
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
def test_search_by_content(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "x.md").write_text("# 标题\n独特关键词XYZ123", encoding="utf-8")
|
||||
sm.reload()
|
||||
results = sm.search("XYZ123")
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
def test_search_no_match(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "x.md").write_text("# 标题\n内容", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.search("不存在的关键词") == []
|
||||
|
||||
|
||||
def test_search_case_insensitive(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "Weather.md").write_text("# Weather Skill\ncontent", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert len(sm.search("weather")) == 1
|
||||
assert len(sm.search("WEATHER")) == 1
|
||||
|
||||
|
||||
# ==================== get_all / count ====================
|
||||
|
||||
def test_get_all_returns_list(isolated_skill_dir, sm):
|
||||
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||
(isolated_skill_dir / "b.md").write_text("# B\n内容B", encoding="utf-8")
|
||||
sm.reload()
|
||||
|
||||
all_skills = sm.get_all()
|
||||
assert len(all_skills) == 2
|
||||
assert all("name" in s and "summary" in s and "content" in s for s in all_skills)
|
||||
assert sm.count() == 2
|
||||
|
||||
|
||||
# ==================== reload ====================
|
||||
|
||||
def test_reload_picks_up_new_file(isolated_skill_dir, sm):
|
||||
assert sm.count() == 0
|
||||
(isolated_skill_dir / "new.md").write_text("# 新技能\n内容", encoding="utf-8")
|
||||
n = sm.reload()
|
||||
assert n == 1
|
||||
assert sm.count() == 1
|
||||
|
||||
|
||||
def test_reload_drops_deleted_file(isolated_skill_dir, sm):
|
||||
f = isolated_skill_dir / "temp.md"
|
||||
f.write_text("# 临时\n内容", encoding="utf-8")
|
||||
sm.reload()
|
||||
assert sm.count() == 1
|
||||
|
||||
f.unlink()
|
||||
sm.reload()
|
||||
assert sm.count() == 0
|
||||
|
||||
|
||||
# ==================== 不存在的技能 ====================
|
||||
|
||||
def test_get_nonexistent_skill_returns_none(sm):
|
||||
assert sm.get_skill("不存在") is None
|
||||
|
||||
|
||||
def test_get_content_nonexistent_returns_none(sm):
|
||||
assert sm.get_skill_content("不存在") is None
|
||||
|
||||
|
||||
# ==================== 可执行技能(跨平台:Windows 用 .bat,Unix 用 .sh) ====================
|
||||
|
||||
def test_loads_executable_skill(isolated_skill_dir, sm):
|
||||
"""加载一个可执行技能文件(带 --help)。
|
||||
Windows 下用 .bat(SkillManager 的 _is_executable 识别 .bat),
|
||||
Unix 下用 .sh 并加可执行权限。
|
||||
"""
|
||||
import os as _os
|
||||
if _os.name == 'nt':
|
||||
# Windows:用 .bat 文件
|
||||
skill_file = isolated_skill_dir / "greet.bat"
|
||||
skill_file.write_text(
|
||||
"@echo off\r\n"
|
||||
"echo 打招呼技能 - 输入名字打招呼\r\n",
|
||||
encoding="utf-8"
|
||||
)
|
||||
else:
|
||||
# Unix:用 .sh 并加可执行权限
|
||||
skill_file = isolated_skill_dir / "greet.sh"
|
||||
skill_file.write_text(
|
||||
"#!/bin/bash\n"
|
||||
"echo '打招呼技能 - 输入名字打招呼'\n",
|
||||
encoding="utf-8"
|
||||
)
|
||||
_os.chmod(str(skill_file), 0o755)
|
||||
|
||||
sm.reload()
|
||||
s = sm.get_skill("greet")
|
||||
assert s is not None
|
||||
assert s["is_executable"] is True
|
||||
assert s["is_markdown"] is False
|
||||
assert "打招呼" in s["summary"]
|
||||
@@ -0,0 +1,178 @@
|
||||
"""run_full_subagent 最大轮次(max_turns)行为回归测试(2026-08-19)
|
||||
|
||||
覆盖 aibox.py run_full_subagent 的轮次限制链路(think_with_limit 补丁):
|
||||
- 超限拦截:任务轮 LLM 调用被限制在 max_turns 次,超限递归返回哨兵串
|
||||
- 哨兵串文本契约("已达到最大对话轮次,任务结束。")
|
||||
- 收尾报告:超限/正常结束均由 original_think 绕过轮次限制生成,作为返回值
|
||||
- 入参钳制:max_turns 为 0/None 回退默认 10,负数钳制下限 1
|
||||
- 超限不置 conversation_active=False(收尾保留完整上下文,aibox.py 763-765 注释)
|
||||
- 收尾失败兜底链:最后一条 assistant 消息 → "子代理未产生有效对话内容。"
|
||||
- SubAgentTool.execute 的 max_turns 参数透传(默认 10)
|
||||
"""
|
||||
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
from aibox import SubAgentTool, run_full_subagent # noqa: E402
|
||||
|
||||
# 与 aibox.py think_with_limit 超限返回值一致的哨兵串(文本契约)
|
||||
SENTINEL = "已达到最大对话轮次,任务结束。"
|
||||
WRAP_UP_MARKER = "【收尾】"
|
||||
|
||||
|
||||
def _make_fake_chat(task_rounds=999, final_answer=None, wrap_report="收尾报告"):
|
||||
"""模拟 DeepSeekChat:think_and_respond 的"原始实现"。
|
||||
|
||||
每次任务轮调用递归一次(模拟 aibox.py:1801 工具调用后递归继续,
|
||||
经由补丁后的 think_and_respond 走轮次计数);累计 task_rounds 轮后
|
||||
若给出 final_answer 则正常结束(模拟 AI 返回纯文本)。task_rounds
|
||||
设为大数且不给 final_answer 时永不自行收敛——只有轮次限制能终止它。
|
||||
检测到收尾 prompt(含【收尾】标记)时返回 wrap_report。
|
||||
"""
|
||||
fake = mock.MagicMock()
|
||||
state = {"n": 0, "task_inputs": [], "recurse_returns": []}
|
||||
|
||||
def original_think(input_text, output_callback=None, is_continuation=False):
|
||||
if WRAP_UP_MARKER in input_text:
|
||||
return wrap_report
|
||||
state["task_inputs"].append(input_text)
|
||||
state["n"] += 1
|
||||
if final_answer is not None and state["n"] >= task_rounds:
|
||||
return final_answer
|
||||
ret = fake.think_and_respond("请基于以上工具执行结果继续回答。",
|
||||
output_callback, is_continuation=True)
|
||||
state["recurse_returns"].append(ret)
|
||||
return ret
|
||||
|
||||
fake.think_and_respond = original_think
|
||||
fake.stats = {}
|
||||
fake.messages = []
|
||||
fake._cancel_event = threading.Event()
|
||||
fake.conversation_active = True
|
||||
fake._state = state
|
||||
return fake
|
||||
|
||||
|
||||
def _run(fake, **kwargs):
|
||||
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||
mock.patch("aibox.create_llm_client"), \
|
||||
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||
return run_full_subagent("角色", "任务", None, **kwargs)
|
||||
|
||||
|
||||
class TestMaxTurnsCutoff:
|
||||
def test_task_turns_capped_at_max_turns(self):
|
||||
"""永不收敛的任务:任务轮 LLM 调用恰为 max_turns 次,第 max_turns+1 次被拦截"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
result = _run(fake, max_turns=3)
|
||||
assert len(fake._state["task_inputs"]) == 3
|
||||
assert result == "收尾报告"
|
||||
|
||||
def test_overrun_recursion_returns_sentinel(self):
|
||||
"""超限那轮递归返回哨兵串(不调 LLM),并沿递归链上抛"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=3)
|
||||
assert fake._state["recurse_returns"]
|
||||
assert all(r == SENTINEL for r in fake._state["recurse_returns"])
|
||||
|
||||
def test_wrap_up_runs_after_overrun(self):
|
||||
"""超限后收尾报告仍被请求(original_think 绕过轮次限制),且为最终返回值"""
|
||||
fake = _make_fake_chat(task_rounds=999, wrap_report="最终报告")
|
||||
result = _run(fake, max_turns=2)
|
||||
assert result == "最终报告"
|
||||
|
||||
def test_max_turns_zero_clamped_to_default(self):
|
||||
"""max_turns=0 → 回退默认 10(对齐 max_tokens 的 0/空回退惯例)"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=0)
|
||||
assert len(fake._state["task_inputs"]) == 10
|
||||
|
||||
def test_negative_max_turns_clamped_to_one(self):
|
||||
"""max_turns 为负 → 钳制下限 1(至少执行一个任务轮,不再首轮即拦截)"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=-3)
|
||||
assert len(fake._state["task_inputs"]) == 1
|
||||
|
||||
def test_max_turns_none_falls_back_to_default(self):
|
||||
"""max_turns=None → 默认 10"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=None)
|
||||
assert len(fake._state["task_inputs"]) == 10
|
||||
|
||||
def test_overrun_does_not_deactivate_conversation(self):
|
||||
"""超限不置 conversation_active=False(否则收尾前上下文被重置清空)"""
|
||||
fake = _make_fake_chat(task_rounds=999)
|
||||
_run(fake, max_turns=2)
|
||||
assert fake.conversation_active is True
|
||||
|
||||
|
||||
class TestNormalCompletion:
|
||||
def test_result_is_wrap_up_report_not_final_answer(self):
|
||||
"""正常收敛时返回值同样是收尾报告(最终回答不作为工具结果返回)"""
|
||||
fake = _make_fake_chat(task_rounds=2, final_answer="最终答案")
|
||||
result = _run(fake, max_turns=10)
|
||||
assert len(fake._state["task_inputs"]) == 2
|
||||
# 第一轮递归的返回值即第 2 轮的最终回答(无哨兵串、无超限拦截)
|
||||
assert fake._state["recurse_returns"] == ["最终答案"]
|
||||
assert result == "收尾报告"
|
||||
|
||||
def test_sentinel_never_appears_when_within_limit(self):
|
||||
"""未达上限:任何递归都不返回哨兵串"""
|
||||
fake = _make_fake_chat(task_rounds=3, final_answer="最终答案")
|
||||
_run(fake, max_turns=10)
|
||||
assert SENTINEL not in fake._state["recurse_returns"]
|
||||
|
||||
|
||||
class TestWrapUpFallback:
|
||||
def test_empty_wrap_up_falls_back_to_last_assistant_message(self):
|
||||
"""收尾返回空 → 兜底取最后一条有内容的 assistant 消息"""
|
||||
fake = _make_fake_chat(task_rounds=1, final_answer="x", wrap_report="")
|
||||
fake.messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "中间回答"},
|
||||
{"role": "assistant", "content": ""},
|
||||
{"role": "assistant", "content": "最终回答"},
|
||||
]
|
||||
assert _run(fake, max_turns=10) == "【子代理最终回答】\n最终回答"
|
||||
|
||||
def test_no_wrap_up_and_no_messages(self):
|
||||
"""收尾为空且无任何 assistant 消息 → 固定兜底文案"""
|
||||
fake = _make_fake_chat(task_rounds=1, final_answer="x", wrap_report="")
|
||||
assert _run(fake, max_turns=10) == "子代理未产生有效对话内容。"
|
||||
|
||||
def test_wrap_up_exception_falls_back(self):
|
||||
"""收尾调用抛异常 → 不中断,走 assistant 消息兜底"""
|
||||
fake = _make_fake_chat(task_rounds=1, final_answer="x")
|
||||
original = fake.think_and_respond
|
||||
|
||||
def boom_on_wrap_up(input_text, output_callback=None, is_continuation=False):
|
||||
if WRAP_UP_MARKER in input_text:
|
||||
raise RuntimeError("wrap-up boom")
|
||||
return original(input_text, output_callback, is_continuation)
|
||||
|
||||
fake.think_and_respond = boom_on_wrap_up
|
||||
fake.messages = [{"role": "assistant", "content": "最终回答"}]
|
||||
assert _run(fake, max_turns=10) == "【子代理最终回答】\n最终回答"
|
||||
|
||||
|
||||
class TestSubAgentToolMaxTurns:
|
||||
def test_max_turns_passed_through(self):
|
||||
"""SubAgentTool.execute 的 max_turns 透传给 run_full_subagent"""
|
||||
tool = SubAgentTool(logger=None)
|
||||
with mock.patch("aibox.run_full_subagent", return_value="ok") as m:
|
||||
tool.execute(role="r", task="t", max_turns=4)
|
||||
assert m.call_args.kwargs["max_turns"] == 4
|
||||
|
||||
def test_max_turns_default_is_10(self):
|
||||
tool = SubAgentTool(logger=None)
|
||||
with mock.patch("aibox.run_full_subagent", return_value="ok") as m:
|
||||
tool.execute(role="r", task="t")
|
||||
assert m.call_args.kwargs["max_turns"] == 10
|
||||
+74
-4
@@ -3,9 +3,16 @@ tools - aibox 工具系统基础模块(包)
|
||||
|
||||
抽离自 aibox.py 的工具基类与注册器,供 aibox.py 及后续迁移的
|
||||
工具子类导入。当前包含:
|
||||
- base.py : Tool 抽象基类
|
||||
- registry.py : ToolRegistry 工具注册器
|
||||
形成单向依赖:common <- tools <- aibox。
|
||||
- base.py : Tool 抽象基类
|
||||
- registry.py : ToolRegistry 工具注册器
|
||||
- basic_tools.py : 基础工具集(A 批:文件/网页/时间/计算器,无外部子系统依赖)
|
||||
- command_tool.py : CommandLineTool 命令行执行工具(B 批)
|
||||
- skill_tools.py : 技能搜索/重载/安装/创建 4 个工具(B 批)
|
||||
- db_tools.py : MemoryTool/MemoTool 数据库工具(B 批)
|
||||
- prompt_tools.py : 身份/灵魂更新 + 结束对话 3 个工具(B 批)
|
||||
- proxy_tools.py : SetProxy/GetProxy 代理工具(B 批)
|
||||
- misc_tools.py : AskUserTool 询问工具(B 批)
|
||||
形成单向依赖:common <- subsystems <- tools <- aibox。
|
||||
|
||||
采用"重导出"策略:外部统一用 `from tools import Tool, ToolRegistry`,
|
||||
调用点无需关心具体子模块所在文件。
|
||||
@@ -13,5 +20,68 @@ tools - aibox 工具系统基础模块(包)
|
||||
|
||||
from .base import Tool
|
||||
from .registry import ToolRegistry
|
||||
from .basic_tools import (
|
||||
FileToolBase,
|
||||
ReadFileTool,
|
||||
WriteFileTool,
|
||||
ListFilesTool,
|
||||
SimpleWebCrawlerTool,
|
||||
GetCurrentTimeTool,
|
||||
CalculatorTool,
|
||||
)
|
||||
# B 批工具
|
||||
from .command_tool import CommandLineTool
|
||||
from .skill_tools import (
|
||||
SearchSkillsTool,
|
||||
ReloadSkillsTool,
|
||||
InstallSkillTool,
|
||||
CreateSkillTool,
|
||||
)
|
||||
from .db_tools import MemoryTool, MemoTool
|
||||
from .prompt_tools import (
|
||||
UpdateIdentityTool,
|
||||
UpdateSoulTool,
|
||||
SaveMemoryAndEndConversationTool,
|
||||
)
|
||||
from .proxy_tools import SetProxyTool, GetProxyTool
|
||||
from .misc_tools import AskUserTool
|
||||
from .process_tools import (
|
||||
BgRunTool,
|
||||
BgWriteTool,
|
||||
BgListTool,
|
||||
BgOutputTool,
|
||||
BgKillTool,
|
||||
)
|
||||
|
||||
__all__ = ["Tool", "ToolRegistry"]
|
||||
__all__ = [
|
||||
"Tool",
|
||||
"ToolRegistry",
|
||||
# A 批基础工具
|
||||
"FileToolBase",
|
||||
"ReadFileTool",
|
||||
"WriteFileTool",
|
||||
"ListFilesTool",
|
||||
"SimpleWebCrawlerTool",
|
||||
"GetCurrentTimeTool",
|
||||
"CalculatorTool",
|
||||
# B 批工具
|
||||
"CommandLineTool",
|
||||
"SearchSkillsTool",
|
||||
"ReloadSkillsTool",
|
||||
"InstallSkillTool",
|
||||
"CreateSkillTool",
|
||||
"MemoryTool",
|
||||
"MemoTool",
|
||||
"UpdateIdentityTool",
|
||||
"UpdateSoulTool",
|
||||
"SaveMemoryAndEndConversationTool",
|
||||
"SetProxyTool",
|
||||
"GetProxyTool",
|
||||
"AskUserTool",
|
||||
# 后台进程管理工具(bg_run/bg_write/bg_list/bg_output/bg_kill)
|
||||
"BgRunTool",
|
||||
"BgWriteTool",
|
||||
"BgListTool",
|
||||
"BgOutputTool",
|
||||
"BgKillTool",
|
||||
]
|
||||
@@ -0,0 +1,658 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.basic_tools - 基础工具集(A 批抽离)
|
||||
|
||||
抽离自 aibox.py 的无外部子系统依赖工具类,仅依赖 common 与标准库。
|
||||
包含:
|
||||
- FileToolBase : 文件操作工具基类
|
||||
- ReadFileTool : 读文件(返回命令行指引)
|
||||
- WriteFileTool : 写文件(返回命令行指引)
|
||||
- ListFilesTool : 列目录(返回命令行指引)
|
||||
- SimpleWebCrawlerTool : 网页访问(返回 curl 指引)
|
||||
- GetCurrentTimeTool : 获取当前时间
|
||||
- CalculatorTool : 数学计算
|
||||
|
||||
依赖方向:common <- tools.basic_tools,不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import shlex
|
||||
from datetime import datetime
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
from common import ConfigManager, constants
|
||||
from .base import Tool
|
||||
|
||||
|
||||
# T3(审计 2026-08-21):cmd.exe 变量展开形态——%VAR% 匹配已定义环境变量时
|
||||
# 展开为变量值,值中的 & | 会成为命令分隔符(经典 cmd 注入)。URL 编码
|
||||
# (%20/%E4 等)与 PowerShell 的 %{http_code} 均不匹配该模式,不受影响。
|
||||
_WIN_VAR_PATTERN = re.compile(r"%[A-Za-z_][A-Za-z0-9_]*%")
|
||||
|
||||
|
||||
def _quote_arg(value: str) -> Tuple[str, Optional[str]]:
|
||||
"""平台感知的 shell 参数引用,返回 (quoted, error)。
|
||||
|
||||
POSIX:shlex.quote(sh 单引号语义)。
|
||||
Windows(run_command 经 cmd.exe / shell=True 执行):
|
||||
- 双引号包裹(& | < > ( ) ^ 在双引号内为字面量)
|
||||
- 含 " → 拒绝(闭合外层双引号,其后内容暴露给 cmd 解析)
|
||||
- 匹配 %VAR% 形态 → 拒绝(cmd 变量展开注入面)
|
||||
- 含换行/控制字符 → 拒绝(cmd 命令行不可含换行)
|
||||
"""
|
||||
value = str(value)
|
||||
if os.name != "nt":
|
||||
return shlex.quote(value), None
|
||||
if '"' in value:
|
||||
return "", f"❌ 参数含双引号,无法安全生成 Windows 命令: {value[:100]}"
|
||||
if "\n" in value or "\r" in value:
|
||||
return "", f"❌ 参数含换行符,无法生成命令: {value[:100]}"
|
||||
if _WIN_VAR_PATTERN.search(value):
|
||||
return "", (f"❌ 参数疑似含 cmd 变量引用(%变量名%),已拒绝以防命令注入: "
|
||||
f"{value[:100]}")
|
||||
return f'"{value}"', None
|
||||
|
||||
|
||||
def _ps_quote(value: str) -> str:
|
||||
"""PowerShell 单引号字面量(内部单引号翻倍),用于 -Command 内层参数。"""
|
||||
return "'" + str(value).replace("'", "''") + "'"
|
||||
|
||||
|
||||
# ==================== 文件操作工具基类 ====================
|
||||
|
||||
class FileToolBase(Tool):
|
||||
"""文件操作工具基类"""
|
||||
|
||||
def __init__(self, logger=None):
|
||||
super().__init__(logger)
|
||||
|
||||
def normalize_path(self, filepath: str) -> str:
|
||||
return os.path.abspath(os.path.expanduser(filepath))
|
||||
|
||||
def check_path_safety(self, filepath: str, check_exists: bool = False) -> Tuple[bool, str, str]:
|
||||
"""检查路径安全性"""
|
||||
try:
|
||||
abs_path = self.normalize_path(filepath)
|
||||
|
||||
if check_exists and not os.path.exists(abs_path):
|
||||
return False, abs_path, f"❌ 文件不存在: {abs_path}"
|
||||
|
||||
# 检查文件大小
|
||||
if os.path.exists(abs_path) and os.path.isfile(abs_path):
|
||||
file_size = os.path.getsize(abs_path)
|
||||
max_size = constants.MAX_FILE_SIZE
|
||||
if file_size > max_size:
|
||||
return False, abs_path, f"❌ 文件太大 ({file_size/1024/1024:.1f}MB)"
|
||||
|
||||
return True, abs_path, ""
|
||||
|
||||
except Exception as e:
|
||||
return False, filepath, f"❌ 路径检查错误: {str(e)}"
|
||||
|
||||
|
||||
# ==================== 文件操作工具(修改版,支持流式输出) ====================
|
||||
|
||||
class ReadFileTool(FileToolBase):
|
||||
"""读取文件内容工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "read_file"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "读取文件内容,支持指定起始行和结束行,或读取开头/结尾"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filepath": {"type": "string"},
|
||||
"encoding": {"type": "string", "default": "utf-8"},
|
||||
"start_line": {"type": "integer", "default": 1},
|
||||
"end_line": {"type": "integer", "default": 0},
|
||||
"max_lines": {"type": "integer", "default": 100},
|
||||
"head": {"type": "boolean", "default": False},
|
||||
"tail": {"type": "boolean", "default": False}
|
||||
},
|
||||
"required": ["filepath"]
|
||||
}
|
||||
|
||||
def execute(self, filepath: str, encoding: str = "utf-8",
|
||||
start_line: int = 1, end_line: int = 0,
|
||||
max_lines: int = 100, head: bool = False,
|
||||
tail: bool = False, **kwargs) -> str:
|
||||
if error := self.validate_params(["filepath"], filepath=filepath):
|
||||
return error
|
||||
|
||||
is_safe, abs_path, error = self.check_path_safety(filepath, check_exists=True)
|
||||
if not is_safe:
|
||||
return error
|
||||
|
||||
# 限制max_lines上限,防止请求超大范围
|
||||
max_lines = min(max_lines, 500)
|
||||
# L1(审计 2026-08-18):start/end_line 为 LLM 可控参数,此前未强转
|
||||
# 直拼 sed 单引号模板——LLM 传含引号的字符串可闭合引号注入命令。
|
||||
# 强转 int(非法输入在此拦截)并夹取合理范围。
|
||||
try:
|
||||
start_line = int(start_line)
|
||||
end_line = int(end_line)
|
||||
except (TypeError, ValueError):
|
||||
return "❌ 参数错误: start_line/end_line 必须为整数"
|
||||
start_line = max(1, start_line)
|
||||
end_line = max(0, min(end_line, 10_000_000))
|
||||
|
||||
# T3(审计 2026-08-21):shlex.quote 是 POSIX 单引号语义,Windows 上
|
||||
# run_command 经 cmd.exe 执行——单引号不被识别,& | % < > 仍是元字符
|
||||
# (URL 含 & 即命令被截断;%VAR% 变量展开可注入)。按平台生成指引:
|
||||
# POSIX 保持 head/tail/sed,Windows 走 PowerShell 等价命令。
|
||||
quoted, qerr = _quote_arg(abs_path)
|
||||
if qerr:
|
||||
return qerr
|
||||
|
||||
if os.name == "nt":
|
||||
# PowerShell 等价命令(经 cmd 调 powershell,-Command 整体双引号
|
||||
# 包裹,内层路径用 _ps_quote 单引号字面量)
|
||||
if head:
|
||||
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||
f'{_ps_quote(abs_path)} -TotalCount {max_lines}"')
|
||||
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||
elif tail:
|
||||
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||
f'{_ps_quote(abs_path)} -Tail {max_lines}"')
|
||||
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||
else:
|
||||
if end_line > 0:
|
||||
first = start_line - 1
|
||||
take = end_line - start_line + 1
|
||||
elif max_lines > 0:
|
||||
first = start_line - 1
|
||||
take = max_lines
|
||||
else:
|
||||
first = start_line - 1
|
||||
take = 10_000_000
|
||||
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||
f'{_ps_quote(abs_path)} | Select-Object -Skip {first} '
|
||||
f'-First {take}"')
|
||||
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来读取文件内容,不要直接返回文件内容:
|
||||
|
||||
```
|
||||
{cmd}
|
||||
```
|
||||
|
||||
- 若需统计总行数:`{wc_cmd}`
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节,单次最多读取 500 行"""
|
||||
return prompt
|
||||
|
||||
# 构造等价的命令行读取命令(POSIX)
|
||||
if head:
|
||||
cmd = f"head -n {max_lines} {quoted}"
|
||||
elif tail:
|
||||
cmd = f"tail -n {max_lines} {quoted}"
|
||||
elif end_line > 0:
|
||||
cmd = f"sed -n '{start_line},{end_line}p' {quoted}"
|
||||
elif max_lines > 0:
|
||||
end = start_line + max_lines - 1
|
||||
cmd = f"sed -n '{start_line},{end}p' {quoted}"
|
||||
else:
|
||||
cmd = f"sed -n '{start_line},$p' {quoted}"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用命令行工具读取文件
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来读取文件内容,不要直接返回文件内容:
|
||||
|
||||
```bash
|
||||
{cmd}
|
||||
```
|
||||
|
||||
- 若需统计总行数:`wc -l {quoted}`
|
||||
- 若需查看文件编码:`file -b --mime-encoding {quoted}`
|
||||
- 若单行超长,可在命令后追加 `| cut -c 1-10000` 截断到 10000 字符
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节,单次最多读取 500 行"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行文件读取"""
|
||||
filepath = kwargs.get("filepath")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n📄 读取文件: {filepath}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class WriteFileTool(FileToolBase):
|
||||
"""写入文件工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "write_file"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "文件写入工具 - 支持覆盖、追加"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filepath": {"type": "string"},
|
||||
"content": {"type": "string"},
|
||||
"mode": {"type": "string", "enum": ["w", "a"], "default": "w"},
|
||||
"encoding": {"type": "string", "default": "utf-8"}
|
||||
},
|
||||
"required": ["filepath", "content"]
|
||||
}
|
||||
|
||||
def execute(self, filepath: str = None, content: str = None, mode: str = "w",
|
||||
encoding: str = "utf-8", **kwargs) -> str:
|
||||
if error := self.validate_params(["filepath", "content"], filepath=filepath, content=content):
|
||||
return error
|
||||
|
||||
is_safe, abs_path, error = self.check_path_safety(filepath)
|
||||
if not is_safe:
|
||||
return error
|
||||
|
||||
quoted, qerr = _quote_arg(abs_path)
|
||||
if qerr:
|
||||
return qerr
|
||||
|
||||
if os.name == "nt":
|
||||
# T3:Windows 下 heredoc 不可用(cmd 命令行不能含换行,且引号
|
||||
# 转义跨 cmd/PowerShell 两层极易出错)。content 整体 Base64 编码
|
||||
# 后经 PowerShell 解码写入——Base64 字符集 [A-Za-z0-9+/=] 不含
|
||||
# 任何 cmd/PowerShell 元字符,多行/引号/特殊字符内容全覆盖。
|
||||
b64 = base64.b64encode(content.encode("utf-8")).decode("ascii")
|
||||
if len(b64) > 6000:
|
||||
return ("❌ 写入内容过大:Windows 命令行长度上限约 8191 字符,"
|
||||
"Base64 编码后超出。请将内容分成多次调用(每次 ≤4KB),"
|
||||
"首次 mode=w、后续 mode=a 追加写入")
|
||||
writer = "Add-Content" if mode == "a" else "Set-Content"
|
||||
write_cmd = (f'powershell -NoProfile -Command "{writer} -Path '
|
||||
f'{_ps_quote(abs_path)} -Value ([Text.Encoding]::UTF8.'
|
||||
f'GetString([Convert]::FromBase64String(\'{b64}\')))"')
|
||||
directory = os.path.dirname(abs_path)
|
||||
if directory:
|
||||
write_cmd = (f'powershell -NoProfile -Command "New-Item '
|
||||
f'-ItemType Directory -Force -Path '
|
||||
f'{_ps_quote(directory)}" & {write_cmd}')
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来写入文件,不要在此工具中直接写入:
|
||||
|
||||
```
|
||||
{write_cmd}
|
||||
```
|
||||
|
||||
- 写入内容已编码为 Base64(避免换行/引号/特殊字符的跨 cmd+PowerShell 转义问题),命令会自动解码并以 UTF-8 写入
|
||||
- 若需检查文件大小:`powershell -NoProfile -Command "(Get-Item {_ps_quote(abs_path)}).Length"`
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节"""
|
||||
return prompt
|
||||
|
||||
directory = os.path.dirname(abs_path)
|
||||
dir_cmd = f"mkdir -p {shlex.quote(directory)}" if directory else ""
|
||||
|
||||
# 内容写入使用 heredoc,避免引号/特殊字符转义问题
|
||||
# M1 修复(审计 2026-08-17):定界符唯一化——若 content 中含独立行
|
||||
# 与定界符相同(如 "AIBOX_EOF"),heredoc 会提前终止,其后内容被
|
||||
# 当作 shell 命令执行(注入)。冲突时递增后缀直到唯一。
|
||||
delimiter = "AIBOX_EOF"
|
||||
content_lines = {line.strip() for line in content.split("\n")}
|
||||
while delimiter in content_lines:
|
||||
delimiter += "_X"
|
||||
redirect = ">" if mode == "w" else ">>"
|
||||
write_cmd = f"cat {redirect} {quoted} <<'{delimiter}'\n{content}\n{delimiter}"
|
||||
if dir_cmd:
|
||||
write_cmd = f"{dir_cmd} && {write_cmd}"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用命令行工具写入文件
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来写入文件,不要在此工具中直接写入:
|
||||
|
||||
```bash
|
||||
{write_cmd}
|
||||
```
|
||||
|
||||
- 如遇 heredoc 内容含特殊字符导致解析失败,可用 `printf '%s' '内容' > {quoted}` 覆盖写入
|
||||
- 若需检查文件大小:`wc -c {quoted}`
|
||||
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行文件写入"""
|
||||
filepath = kwargs.get("filepath")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n📝 写入文件: {filepath}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class ListFilesTool(FileToolBase):
|
||||
"""列出目录内容工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "list_files"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "列出指定目录下的文件和子目录"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"directory": {"type": "string", "default": "."},
|
||||
"show_hidden": {"type": "boolean", "default": False}
|
||||
}
|
||||
}
|
||||
|
||||
def execute(self, directory: str = ".", show_hidden: bool = False, **kwargs) -> str:
|
||||
is_safe, abs_path, error = self.check_path_safety(directory)
|
||||
if not is_safe:
|
||||
return error
|
||||
|
||||
if not os.path.exists(abs_path):
|
||||
return f"❌ 目录不存在: {abs_path}"
|
||||
|
||||
if not os.path.isdir(abs_path):
|
||||
return f"❌ 路径不是目录: {abs_path}"
|
||||
|
||||
quoted, qerr = _quote_arg(abs_path)
|
||||
if qerr:
|
||||
return qerr
|
||||
|
||||
if os.name == "nt":
|
||||
# T3:Windows 走 PowerShell 等价命令(ls/du/find/file 无原生等价物)
|
||||
force_flag = " -Force" if show_hidden else ""
|
||||
cmd = (f'powershell -NoProfile -Command "Get-ChildItem{force_flag} '
|
||||
f'{_ps_quote(abs_path)} | Select-Object '
|
||||
f'Mode,LastWriteTime,Length,Name"')
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来列出目录内容,不要直接在此工具中列出:
|
||||
|
||||
```
|
||||
{cmd}
|
||||
```
|
||||
|
||||
- 若需统计目录总大小:`powershell -NoProfile -Command "(Get-ChildItem -Recurse {_ps_quote(abs_path)} | Measure-Object Length -Sum).Sum"`
|
||||
- 若需递归列出:`powershell -NoProfile -Command "Get-ChildItem -Recurse -Depth 2 {_ps_quote(abs_path)} | Select-Object FullName"`"""
|
||||
return prompt
|
||||
|
||||
flag = "-la" if show_hidden else "-l"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用命令行工具列目录
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来列出目录内容,不要直接在此工具中列出:
|
||||
|
||||
```bash
|
||||
ls {flag} {quoted}
|
||||
```
|
||||
|
||||
- 若需区分目录/文件并显示大小:`ls -lhF {quoted}`
|
||||
- 若需统计目录总大小:`du -sh {quoted}`
|
||||
- 若需递归列出:`find {quoted} -maxdepth 2`
|
||||
- 若需了解文件类型:`file {quoted}/*`"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行列出目录"""
|
||||
directory = kwargs.get("directory", ".")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n📁 列出目录: {directory}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
# ==================== 网页访问工具 ====================
|
||||
|
||||
class SimpleWebCrawlerTool(Tool):
|
||||
"""简化版网页爬虫工具(返回 markdown 提示词,指导改用 curl 命令行)"""
|
||||
|
||||
def __init__(self, logger=None):
|
||||
super().__init__(logger)
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "web"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """简单易用的网页访问工具 - 一句话搞定!
|
||||
支持:获取网页、提取内容、搜索关键词、下载文件"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"action": {
|
||||
"type": "string",
|
||||
"enum": ["get", "text", "title", "links", "search"]
|
||||
},
|
||||
"url": {"type": "string"},
|
||||
"keyword": {"type": "string"}
|
||||
},
|
||||
"required": ["action", "url"]
|
||||
}
|
||||
|
||||
def _proxy_curl_args(self) -> str:
|
||||
"""根据当前代理配置构造 curl --socks5 参数"""
|
||||
proxies = ConfigManager().get_proxy()
|
||||
if not proxies:
|
||||
return ""
|
||||
proxy_url = proxies.get("https") or proxies.get("http") or ""
|
||||
host_port = proxy_url.split("://", 1)[-1]
|
||||
if "@" in host_port:
|
||||
host_port = host_port.split("@", 1)[-1]
|
||||
# T3:平台感知引用(原 shlex.quote 在 Windows cmd.exe 下无效)
|
||||
quoted, _err = _quote_arg(host_port)
|
||||
return f"--socks5 {quoted}" if quoted else ""
|
||||
|
||||
def execute(self, action: str, url: str, **kwargs) -> str:
|
||||
if action not in ("get", "text", "title", "links", "search"):
|
||||
return f"❌ 未知操作: {action}"
|
||||
|
||||
# L10(审计 2026-08-18):url 为 LLM 可控输入(可被网页内容诱导),
|
||||
# 此前不限 scheme——file:// 可经生成的 curl 指引读取本地文件(SSRF)。
|
||||
# 白名单 http/https 后其余协议直接拒绝。
|
||||
if not url or not str(url).strip().lower().startswith(("http://", "https://")):
|
||||
return "❌ 仅支持 http/https 协议地址(已拦截 file:// 等非网络协议)"
|
||||
|
||||
quoted_url, qerr = _quote_arg(url)
|
||||
if qerr:
|
||||
return qerr
|
||||
timeout = constants.WEB_TIMEOUT
|
||||
user_agent, ua_err = _quote_arg(constants.WEB_USER_AGENT)
|
||||
if ua_err:
|
||||
return ua_err
|
||||
proxy = self._proxy_curl_args()
|
||||
|
||||
if os.name == "nt":
|
||||
# T3:Windows 用系统内置 curl.exe(Win10 1803+);sed/grep 等
|
||||
# POSIX 管道无原生等价,分析类动作(text/title/links/search)
|
||||
# 指引先取原始 HTML 再由 AI 自行解析。URL 含 & 时靠双引号
|
||||
# 保护(cmd 双引号内为字面量),%VAR% 形态已在 _quote_arg 拒绝。
|
||||
base = f"curl.exe -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
||||
if action == "get":
|
||||
cmd = f'{base} -w "\\nHTTP_CODE:%{{http_code}}" {quoted_url}'
|
||||
extra = "如需查看响应头:`curl.exe -sIL ...`;如需保存到文件:`curl.exe -sL ... -o 文件名`"
|
||||
elif action == "text":
|
||||
cmd = (f'{base} {quoted_url} | powershell -NoProfile -Command '
|
||||
f'"$input -replace \'<[^>]+>\',\'\'"')
|
||||
extra = "去除 HTML 标签提取纯文本"
|
||||
else:
|
||||
cmd = f"{base} {quoted_url}"
|
||||
extra = ("Windows 下无 grep/sed 管道,请取回原始 HTML 后自行分析"
|
||||
"(提取标题/链接/检索关键词)")
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来完成网页操作,不要直接返回网页内容:
|
||||
|
||||
```
|
||||
{cmd}
|
||||
```
|
||||
|
||||
{extra}
|
||||
- 下载大小限制 {constants.WEB_MAX_DOWNLOAD_SIZE} 字节,可在命令后追加 `| powershell -NoProfile -Command "$input | Select-Object -First 100"` 限制
|
||||
- 当前已按配置自动附带代理:{proxy or '无(如需代理请先调用 set_proxy 工具)'}"""
|
||||
return prompt
|
||||
|
||||
# POSIX 分支:_quote_arg 在 POSIX 下即 shlex.quote,直接沿用
|
||||
base = f"curl -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
||||
|
||||
if action == "get":
|
||||
cmd = f"{base} -w '\\nHTTP_CODE:%{{http_code}}' {quoted_url}"
|
||||
extra = "如需查看响应头:`curl -sIL ...`;如需保存到文件:`curl -sL ... -o 文件名`"
|
||||
elif action == "text":
|
||||
cmd = f"{base} {quoted_url} | sed -r 's/<[^>]+>//g'"
|
||||
extra = "若系统装有 lynx,可用 `lynx -dump -stdin` 提取纯文本更准确"
|
||||
elif action == "title":
|
||||
cmd = f"{base} {quoted_url} | grep -oP '<title>\\K[^<]+'"
|
||||
extra = "无输出则表示未找到标题"
|
||||
elif action == "links":
|
||||
cmd = f"{base} {quoted_url} | grep -oP 'href=\"\\K[^\"]+'"
|
||||
extra = "输出所有链接 URL,可配合 `head` 限制数量"
|
||||
else: # search
|
||||
keyword = kwargs.get("keyword")
|
||||
if not keyword:
|
||||
return "❌ 需要提供 keyword 参数"
|
||||
quoted_kw, kw_err = _quote_arg(keyword)
|
||||
if kw_err:
|
||||
return kw_err
|
||||
cmd = f"{base} {quoted_url} | grep -c {quoted_kw}"
|
||||
extra = f"返回匹配行数,大于 0 即表示找到关键词 '{keyword}'"
|
||||
|
||||
# 返回 markdown 提示词,指导智能体用 curl 完成网页访问
|
||||
prompt = f"""请使用 `run_command` 工具执行以下命令来完成网页操作,不要直接返回网页内容:
|
||||
|
||||
```bash
|
||||
{cmd}
|
||||
```
|
||||
|
||||
{extra}
|
||||
- 下载大小限制 {constants.WEB_MAX_DOWNLOAD_SIZE} 字节,可在命令后追加 `| head -c {constants.WEB_MAX_DOWNLOAD_SIZE}` 限制
|
||||
- 当前已按配置自动附带代理:{proxy or '无(如需代理请先调用 set_proxy 工具)'}"""
|
||||
return prompt
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行网页操作"""
|
||||
action = kwargs.get("action")
|
||||
url = kwargs.get("url")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n🌐 执行网页操作: {action} - {url}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
# ==================== 系统工具(修改版,支持流式输出) ====================
|
||||
|
||||
class GetCurrentTimeTool(Tool):
|
||||
"""获取当前时间工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "get_current_time"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "获取当前日期和时间"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"format": {"type": "string", "enum": ["full", "date", "time"], "default": "full"}
|
||||
}
|
||||
}
|
||||
|
||||
def execute(self, format: str = "full", **kwargs) -> str:
|
||||
now = datetime.now()
|
||||
|
||||
if format == "date":
|
||||
result = now.strftime("%Y-%m-%d")
|
||||
elif format == "time":
|
||||
result = now.strftime("%H:%M:%S")
|
||||
else:
|
||||
result = now.strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
return self.format_result(True, "当前时间", {"时间": result})
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行获取时间"""
|
||||
output_queue.put(("line", "\n🕐 获取当前时间\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class CalculatorTool(Tool):
|
||||
"""计算器工具(支持流式输出)"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "calculator"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "执行数学计算"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"a": {"type": "number"},
|
||||
"b": {"type": "number"},
|
||||
"operation": {"type": "string", "enum": ["add", "subtract", "multiply", "divide"]}
|
||||
},
|
||||
"required": ["a", "b", "operation"]
|
||||
}
|
||||
|
||||
def execute(self, a: float, b: float, operation: str, **kwargs) -> str:
|
||||
operations = {
|
||||
"add": ("+", a + b),
|
||||
"subtract": ("-", a - b),
|
||||
"multiply": ("*", a * b),
|
||||
"divide": ("/", a / b if b != 0 else None)
|
||||
}
|
||||
|
||||
if operation not in operations:
|
||||
return f"❌ 未知运算: {operation}"
|
||||
|
||||
op_symbol, result = operations[operation]
|
||||
|
||||
if result is None:
|
||||
return "❌ 错误:除数不能为0"
|
||||
|
||||
data = {
|
||||
"表达式": f"{a} {op_symbol} {b}",
|
||||
"结果": result
|
||||
}
|
||||
|
||||
return self.format_result(True, "计算成功", data)
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行计算"""
|
||||
output_queue.put(("line", "\n🧮 执行计算\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
@@ -0,0 +1,109 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.command_tool - 命令行执行工具
|
||||
|
||||
抽离自 aibox.py,使用抽象的 CommandExecutor 实现命令执行(流式 + 超时控制)。
|
||||
工具本身不持有 subprocess 逻辑,而是委托给 SubprocessCommandExecutor。
|
||||
|
||||
H1 修复(审计 2026-08-17):此前本类内置一套重复的 Popen 读取循环,
|
||||
完全绕过注入的 executor,导致 1MB 输出上限、cancel() 三层中断、
|
||||
超时机制三项保护同时失效。现改为纯委托,行为与 executor 完全一致。
|
||||
|
||||
依赖:common (ConfigManager, constants) + tools.base.Tool + subsystems.executors
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import queue
|
||||
from typing import Dict
|
||||
|
||||
from common import ConfigManager
|
||||
from subsystems import CommandExecutor, SubprocessCommandExecutor
|
||||
from .base import Tool
|
||||
|
||||
|
||||
class CommandLineTool(Tool):
|
||||
"""命令行执行工具 - 使用抽象的命令执行器(支持流式输出和超时控制)"""
|
||||
|
||||
def __init__(self, executor: CommandExecutor = None, logger=None):
|
||||
super().__init__(logger)
|
||||
self.executor = executor or SubprocessCommandExecutor(logger)
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "run_command"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """执行任意系统命令,支持实时输出显示和超时控制"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "要执行的命令"
|
||||
},
|
||||
"timeout": {
|
||||
"type": "integer",
|
||||
"default": 30,
|
||||
"description": "命令执行超时时间(秒),超时后将终止命令执行。默认30秒"
|
||||
},
|
||||
"work_dir": {
|
||||
"type": "string",
|
||||
"default": ".",
|
||||
"description": "命令执行的工作目录"
|
||||
},
|
||||
"env": {
|
||||
"type": "object",
|
||||
"description": "额外的环境变量"
|
||||
}
|
||||
},
|
||||
"required": ["command"]
|
||||
}
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行命令:纯委托给 executor(1MB 上限/超时/cancel 全部生效)"""
|
||||
command = kwargs.get("command")
|
||||
|
||||
if not command:
|
||||
output_queue.put(("line", "❌ 错误:缺少命令参数\n"))
|
||||
output_queue.put(("complete", "❌ 错误:缺少命令参数"))
|
||||
return
|
||||
|
||||
self.executor.execute_streaming(
|
||||
command,
|
||||
output_queue,
|
||||
timeout=kwargs.get("timeout", 30),
|
||||
work_dir=kwargs.get("work_dir", "."),
|
||||
env=kwargs.get("env"),
|
||||
)
|
||||
|
||||
def execute(self, command: str, timeout: int = 30,
|
||||
work_dir: str = ".", env: Dict = None, **kwargs) -> str:
|
||||
"""执行命令 - 仅支持流式执行(命令工具强制流式输出,无 stream 参数)"""
|
||||
if error := self.validate_params(["command"], command=command):
|
||||
return error
|
||||
|
||||
# 如果正在流式输出,直接返回空(实际输出通过队列)
|
||||
if self.streaming_output and self.output_queue:
|
||||
return ""
|
||||
|
||||
# 命令工具固定走流式执行:通过临时队列收集完整结果
|
||||
output_queue = queue.Queue()
|
||||
self.execute_streaming(
|
||||
output_queue,
|
||||
command=command,
|
||||
timeout=timeout,
|
||||
work_dir=work_dir,
|
||||
env=env
|
||||
)
|
||||
|
||||
# 收集流式执行的最终结果(complete 消息)
|
||||
final_result = ""
|
||||
while not output_queue.empty():
|
||||
try:
|
||||
msg_type, content = output_queue.get_nowait()
|
||||
if msg_type == "complete":
|
||||
final_result = content
|
||||
except queue.Empty:
|
||||
break
|
||||
return final_result
|
||||
@@ -0,0 +1,829 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.db_tools - 数据库相关工具
|
||||
|
||||
抽离自 aibox.py,包含 2 个数据库工具:
|
||||
- MemoryTool : 记忆管理(保存/搜索/查看/删除,支持分页)
|
||||
- MemoTool : 任务管理(创建/查看/完成/删除/搜索/统计)
|
||||
|
||||
依赖:common (constants, Logger) + tools.base.Tool
|
||||
+ subsystems.db.MemoryDatabase/MemoDatabase
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import queue
|
||||
import re
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from common import constants, Logger
|
||||
from subsystems import MemoryDatabase, MemoDatabase
|
||||
from .base import Tool
|
||||
|
||||
|
||||
# ==================== 记忆工具 ====================
|
||||
|
||||
class MemoryTool(Tool):
|
||||
"""记忆工具 - 用于存储、检索和删除对话记忆(支持流式输出,支持分页)"""
|
||||
|
||||
def __init__(self, memory_db: MemoryDatabase, logger: Optional[Logger] = None):
|
||||
super().__init__(logger)
|
||||
self.memory_db = memory_db
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "memory"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """记忆管理工具 - 存储、检索和删除重要的对话信息
|
||||
支持:保存总结、多关键词搜索记忆(支持分页)、查看最近的记忆、删除无用记忆"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"action": {
|
||||
"type": "string",
|
||||
"description": "操作类型",
|
||||
"enum": ["save", "search", "recent", "get", "delete"]
|
||||
},
|
||||
"summary": {"type": "string", "description": "要保存的总结内容(用于save操作)"},
|
||||
"topic": {"type": "string", "description": "对话主题(用于save操作)"},
|
||||
"content": {"type": "string", "description": "详细内容(可选,用于save操作)"},
|
||||
"keywords": {
|
||||
"type": "array", "items": {"type": "string"},
|
||||
"description": "搜索关键词列表"
|
||||
},
|
||||
"query": {"type": "string", "description": "搜索关键词(单个关键词)"},
|
||||
"match_all": {"type": "boolean", "default": False},
|
||||
"memory_id": {"type": "integer"},
|
||||
"memory_ids": {"type": "array", "items": {"type": "integer"}},
|
||||
"importance": {"type": "integer", "default": 2},
|
||||
"limit": {"type": "integer", "default": 10, "description": "每页条数(默认10)"},
|
||||
"page": {"type": "integer", "default": 1, "description": "页码(从1开始)"},
|
||||
"confirm": {"type": "boolean", "default": False}
|
||||
},
|
||||
"required": ["action"]
|
||||
}
|
||||
|
||||
def execute(self, action: str, **kwargs) -> str:
|
||||
try:
|
||||
if action == "save":
|
||||
return self._save_memory(kwargs)
|
||||
elif action == "search":
|
||||
return self._search_memories(kwargs)
|
||||
elif action == "recent":
|
||||
return self._get_recent_memories(kwargs)
|
||||
elif action == "get":
|
||||
return self._get_memory_by_id(kwargs)
|
||||
elif action == "delete":
|
||||
return self._delete_memory(kwargs)
|
||||
else:
|
||||
return f"❌ 未知操作: {action}"
|
||||
except Exception as e:
|
||||
return f"❌ 记忆操作错误: {str(e)}"
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行记忆操作"""
|
||||
action = kwargs.get("action")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n🧠 执行记忆操作: {action}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
def _save_memory(self, kwargs: Dict) -> str:
|
||||
summary = kwargs.get("summary")
|
||||
if not summary:
|
||||
return "❌ 需要提供总结内容"
|
||||
|
||||
topic = kwargs.get("topic", "未命名对话")
|
||||
content = kwargs.get("content")
|
||||
importance = kwargs.get("importance", constants.MEMORY_DEFAULT_IMPORTANCE)
|
||||
|
||||
memory_id = self.memory_db.add_memory(
|
||||
topic=topic,
|
||||
summary=summary,
|
||||
content=content,
|
||||
importance=importance
|
||||
)
|
||||
|
||||
return self.format_result(True, "记忆已保存", {
|
||||
"记忆ID": memory_id,
|
||||
"主题": topic,
|
||||
"总结": summary[:100] + "..." if len(summary) > 100 else summary
|
||||
})
|
||||
|
||||
def _search_memories(self, kwargs: Dict) -> str:
|
||||
"""搜索记忆 - 支持分页"""
|
||||
keywords = kwargs.get("keywords", [])
|
||||
query = kwargs.get("query", "")
|
||||
match_all = kwargs.get("match_all", False)
|
||||
page = kwargs.get("page", 1)
|
||||
page_size = kwargs.get("limit", 10) # 使用 limit 作为 page_size
|
||||
|
||||
# 确保页码和每页条数有效
|
||||
page = max(1, page)
|
||||
page_size = max(1, min(50, page_size)) # 每页最多50条
|
||||
|
||||
# 处理关键词
|
||||
if keywords and isinstance(keywords, list):
|
||||
pass
|
||||
elif query:
|
||||
keywords = query.split()
|
||||
else:
|
||||
return "❌ 需要提供搜索关键词(keywords 或 query)"
|
||||
|
||||
if not keywords:
|
||||
return "❌ 关键词列表为空"
|
||||
|
||||
# 先获取所有结果(不分页,用于计算总数)
|
||||
all_results = self.memory_db.search_memories(keywords, match_all, limit=None)
|
||||
|
||||
if not all_results:
|
||||
match_mode = "所有" if match_all else "任一"
|
||||
return f"❌ 未找到包含 {match_mode} 关键词 {keywords} 的记忆"
|
||||
|
||||
# 计算分页信息
|
||||
total_count = len(all_results)
|
||||
total_pages = (total_count + page_size - 1) // page_size
|
||||
|
||||
# 检查页码是否有效
|
||||
if page > total_pages and total_pages > 0:
|
||||
return f"❌ 页码 {page} 超出范围,共有 {total_pages} 页"
|
||||
|
||||
# 计算当前页的起止索引
|
||||
start_idx = (page - 1) * page_size
|
||||
end_idx = min(start_idx + page_size, total_count)
|
||||
|
||||
# 获取当前页的结果
|
||||
page_results = all_results[start_idx:end_idx]
|
||||
|
||||
# 构建输出
|
||||
match_mode = "所有关键词" if match_all else "任一关键词"
|
||||
output = [
|
||||
f"🔍 搜索 {match_mode} {keywords} 找到 {total_count} 条记忆",
|
||||
f"📄 第 {page}/{total_pages} 页 (每页 {page_size} 条)",
|
||||
"-" * 40
|
||||
]
|
||||
|
||||
for i, mem in enumerate(page_results, 1):
|
||||
# 全局序号
|
||||
global_idx = start_idx + i
|
||||
output.append(f"\n{global_idx}. [ID: {mem['id']}] [{mem['end_time']}] {mem['topic']}")
|
||||
output.append(f" 总结: {mem['summary'][:200]}")
|
||||
if mem.get('tags'):
|
||||
output.append(f" 标签: {', '.join(mem['tags'])}")
|
||||
if mem.get('importance'):
|
||||
output.append(f" 重要性: {'⭐' * mem['importance']} ({mem['importance']}/5)")
|
||||
|
||||
# 添加翻页提示
|
||||
if total_pages > 1:
|
||||
output.append(f"\n📖 翻页提示: 使用 page 参数查看更多结果")
|
||||
output.append(f" 例如: memory search query='{query}' page={page + 1 if page < total_pages else page} limit={page_size}")
|
||||
if page < total_pages:
|
||||
output.append(f" ➡️ 下一页: page={page + 1}")
|
||||
if page > 1:
|
||||
output.append(f" ⬅️ 上一页: page={page - 1}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _get_recent_memories(self, kwargs: Dict) -> str:
|
||||
"""获取最近的记忆 - 支持分页"""
|
||||
page = kwargs.get("page", 1)
|
||||
page_size = kwargs.get("limit", 10)
|
||||
|
||||
# 确保页码和每页条数有效
|
||||
page = max(1, page)
|
||||
page_size = max(1, min(50, page_size))
|
||||
|
||||
# 先获取所有结果
|
||||
all_results = self.memory_db.get_recent_memories(limit=None)
|
||||
|
||||
if not all_results:
|
||||
return "📭 没有找到记忆"
|
||||
|
||||
# 计算分页信息
|
||||
total_count = len(all_results)
|
||||
total_pages = (total_count + page_size - 1) // page_size
|
||||
|
||||
# 检查页码是否有效
|
||||
if page > total_pages and total_pages > 0:
|
||||
return f"❌ 页码 {page} 超出范围,共有 {total_pages} 页"
|
||||
|
||||
# 计算当前页的起止索引
|
||||
start_idx = (page - 1) * page_size
|
||||
end_idx = min(start_idx + page_size, total_count)
|
||||
|
||||
# 获取当前页的结果
|
||||
page_results = all_results[start_idx:end_idx]
|
||||
|
||||
output = [f"📚 最近 {total_count} 条记忆 (第 {page}/{total_pages} 页,每页 {page_size} 条):"]
|
||||
output.append("-" * 40)
|
||||
|
||||
for i, mem in enumerate(page_results, 1):
|
||||
global_idx = start_idx + i
|
||||
output.append(f"\n{global_idx}. [ID: {mem['id']}] [{mem['end_time']}] {mem['topic']}")
|
||||
output.append(f" 总结: {mem['summary'][:100]}")
|
||||
|
||||
# 添加翻页提示
|
||||
if total_pages > 1:
|
||||
output.append(f"\n📖 翻页提示: 使用 page 参数查看更多结果")
|
||||
output.append(f" 例如: memory recent page={page + 1 if page < total_pages else page} limit={page_size}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _get_memory_by_id(self, kwargs: Dict) -> str:
|
||||
"""根据ID获取记忆"""
|
||||
memory_id = kwargs.get("memory_id")
|
||||
if not memory_id:
|
||||
return "❌ 需要提供记忆ID"
|
||||
|
||||
memory = self.memory_db.get_memory_by_id(memory_id)
|
||||
if not memory:
|
||||
return f"❌ 未找到ID为 {memory_id} 的记忆"
|
||||
|
||||
output = [
|
||||
f"📌 记忆 #{memory_id}",
|
||||
f"主题: {memory['topic']}",
|
||||
f"时间: {memory['start_time']} - {memory['end_time']}",
|
||||
f"总结: {memory['summary']}",
|
||||
]
|
||||
|
||||
if memory.get('content'):
|
||||
output.append(f"详细内容:\n{memory['content']}")
|
||||
|
||||
if memory.get('tags'):
|
||||
output.append(f"标签: {', '.join(memory['tags'])}")
|
||||
|
||||
if memory.get('importance'):
|
||||
output.append(f"重要性: {'⭐' * memory['importance']} ({memory['importance']}/5)")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _delete_memory(self, kwargs: Dict) -> str:
|
||||
"""删除记忆"""
|
||||
confirm = kwargs.get("confirm", False)
|
||||
if not confirm:
|
||||
return "❌ 需要设置 confirm=True 来确认删除操作"
|
||||
|
||||
memory_id = kwargs.get("memory_id")
|
||||
memory_ids = kwargs.get("memory_ids", [])
|
||||
|
||||
deleted_count = 0
|
||||
|
||||
if memory_id:
|
||||
if self.memory_db.delete_memory(memory_id):
|
||||
deleted_count = 1
|
||||
elif memory_ids:
|
||||
deleted_count = self.memory_db.delete_memories_by_ids(memory_ids)
|
||||
else:
|
||||
return "❌ 需要提供 memory_id 或 memory_ids"
|
||||
|
||||
if deleted_count == 0:
|
||||
return f"❌ 未找到要删除的记忆"
|
||||
|
||||
return self.format_result(True, f"成功删除 {deleted_count} 条记忆", {"删除数量": deleted_count})
|
||||
|
||||
|
||||
# ==================== 备忘录工具 ====================
|
||||
|
||||
class MemoTool(Tool):
|
||||
"""备忘录工具 - 用于管理任务(所有任务都是对话任务)"""
|
||||
|
||||
def __init__(self, memo_db: MemoDatabase, logger: Optional[Logger] = None):
|
||||
super().__init__(logger)
|
||||
self.memo_db = memo_db
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "memo"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """任务管理工具 - 创建、查看、完成任务
|
||||
支持:创建任务(支持一次性/重复、即时任务)、查看待办、标记完成、删除任务、搜索任务"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"action": {
|
||||
"type": "string",
|
||||
"description": "操作类型",
|
||||
"enum": ["add", "list", "get", "complete", "delete", "search", "stats"]
|
||||
},
|
||||
"title": {
|
||||
"type": "string",
|
||||
"description": "任务标题(用于add操作)"
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "任务内容(可选,用于add操作)"
|
||||
},
|
||||
"reminder_time": {
|
||||
"type": "string",
|
||||
"description": "执行时间(格式:YYYY-MM-DD HH:MM:SS,或相对时间如 '+1h', '+30m', 'tomorrow 9am')"
|
||||
},
|
||||
"reminder_minutes": {
|
||||
"type": "integer",
|
||||
"description": "延迟分钟数(从现在开始,用于快速设置任务)"
|
||||
},
|
||||
"tags": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "标签列表"
|
||||
},
|
||||
"priority": {
|
||||
"type": "integer",
|
||||
"description": "优先级 1-5(默认1)",
|
||||
"default": 1
|
||||
},
|
||||
"repeat_type": {
|
||||
"type": "string",
|
||||
"description": "重复类型(none=一次性, daily=每天, weekly=每周, monthly=每月, custom=自定义)",
|
||||
"enum": ["none", "daily", "weekly", "monthly", "custom"],
|
||||
"default": "none"
|
||||
},
|
||||
"repeat_interval": {
|
||||
"type": "integer",
|
||||
"description": "自定义重复间隔,当repeat_type='custom'时使用",
|
||||
"default": 0
|
||||
},
|
||||
"repeat_interval_unit": {
|
||||
"type": "string",
|
||||
"description": "重复间隔单位(minutes=分钟, hours=小时, days=天),当repeat_type='custom'时使用",
|
||||
"enum": ["minutes", "hours", "days"],
|
||||
"default": "days"
|
||||
},
|
||||
"repeat_end_time": {
|
||||
"type": "string",
|
||||
"description": "重复结束时间,None表示永久重复"
|
||||
},
|
||||
"is_immediate": {
|
||||
"type": "boolean",
|
||||
"description": "是否为即时任务(立即执行)",
|
||||
"default": False
|
||||
},
|
||||
"memo_id": {
|
||||
"type": "integer",
|
||||
"description": "任务ID(用于get、complete、delete操作)"
|
||||
},
|
||||
"memo_ids": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"description": "任务ID列表(用于批量删除)"
|
||||
},
|
||||
"force_complete": {
|
||||
"type": "boolean",
|
||||
"description": "强制完成(忽略重复规则)",
|
||||
"default": False
|
||||
},
|
||||
"keyword": {
|
||||
"type": "string",
|
||||
"description": "搜索关键词(用于search操作)"
|
||||
},
|
||||
"include_completed": {
|
||||
"type": "boolean",
|
||||
"description": "是否包含已完成的任务(默认False)",
|
||||
"default": False
|
||||
},
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "返回数量限制(默认20)",
|
||||
"default": 20
|
||||
},
|
||||
"confirm": {
|
||||
"type": "boolean",
|
||||
"description": "确认删除操作",
|
||||
"default": False
|
||||
}
|
||||
},
|
||||
"required": ["action"]
|
||||
}
|
||||
|
||||
def _parse_reminder_time(self, reminder_input: Any) -> Optional[str]:
|
||||
"""解析提醒时间输入"""
|
||||
if not reminder_input:
|
||||
return None
|
||||
|
||||
if isinstance(reminder_input, int):
|
||||
reminder_time = datetime.now() + timedelta(minutes=reminder_input)
|
||||
return reminder_time.strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
if isinstance(reminder_input, str):
|
||||
reminder_input = reminder_input.strip()
|
||||
|
||||
if reminder_input.startswith('+'):
|
||||
match = re.match(r'\+(\d+)([mhd])', reminder_input)
|
||||
if match:
|
||||
value, unit = int(match.group(1)), match.group(2)
|
||||
if unit == 'm':
|
||||
delta = timedelta(minutes=value)
|
||||
elif unit == 'h':
|
||||
delta = timedelta(hours=value)
|
||||
elif unit == 'd':
|
||||
delta = timedelta(days=value)
|
||||
else:
|
||||
return None
|
||||
|
||||
reminder_time = datetime.now() + delta
|
||||
return reminder_time.strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
try:
|
||||
reminder_time = datetime.strptime(reminder_input, constants.DATETIME_FORMAT)
|
||||
return reminder_time.strftime(constants.DATETIME_FORMAT)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
now = datetime.now()
|
||||
reminder_input_lower = reminder_input.lower()
|
||||
|
||||
if 'tomorrow' in reminder_input_lower:
|
||||
tomorrow = now + timedelta(days=1)
|
||||
time_match = re.search(r'(\d{1,2})(?::(\d{2}))?\s*(am|pm)?', reminder_input_lower)
|
||||
if time_match:
|
||||
hour = int(time_match.group(1))
|
||||
minute = int(time_match.group(2)) if time_match.group(2) else 0
|
||||
ampm = time_match.group(3)
|
||||
|
||||
if ampm:
|
||||
if ampm == 'pm' and hour < 12:
|
||||
hour += 12
|
||||
elif ampm == 'am' and hour == 12:
|
||||
hour = 0
|
||||
|
||||
reminder_time = tomorrow.replace(hour=hour, minute=minute, second=0)
|
||||
else:
|
||||
reminder_time = tomorrow.replace(hour=9, minute=0, second=0)
|
||||
|
||||
return reminder_time.strftime(constants.DATETIME_FORMAT)
|
||||
|
||||
return None
|
||||
|
||||
def execute(self, action: str, **kwargs) -> str:
|
||||
"""执行备忘录操作"""
|
||||
try:
|
||||
if action == "add":
|
||||
return self._add_memo(kwargs)
|
||||
elif action == "list":
|
||||
return self._list_memos(kwargs)
|
||||
elif action == "get":
|
||||
return self._get_memo(kwargs)
|
||||
elif action == "complete":
|
||||
return self._complete_memo(kwargs)
|
||||
elif action == "delete":
|
||||
return self._delete_memo(kwargs)
|
||||
elif action == "search":
|
||||
return self._search_memos(kwargs)
|
||||
elif action == "stats":
|
||||
return self._get_stats(kwargs)
|
||||
else:
|
||||
return f"❌ 未知操作: {action}"
|
||||
except Exception as e:
|
||||
return f"❌ 任务操作错误: {str(e)}"
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行备忘录操作"""
|
||||
action = kwargs.get("action")
|
||||
|
||||
# 发送开始执行的消息
|
||||
output_queue.put(("line", f"\n📝 执行备忘录操作: {action}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 执行操作
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
# 发送结果
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
def _add_memo(self, kwargs: Dict) -> str:
|
||||
"""添加任务 - 所有任务都是对话任务,不需要task_type"""
|
||||
title = kwargs.get("title")
|
||||
if not title:
|
||||
return "❌ 需要提供任务标题"
|
||||
|
||||
content = kwargs.get("content")
|
||||
reminder_minutes = kwargs.get("reminder_minutes")
|
||||
reminder_time_input = kwargs.get("reminder_time")
|
||||
tags = kwargs.get("tags", [])
|
||||
priority = kwargs.get("priority", 1)
|
||||
repeat_type = kwargs.get("repeat_type", "none")
|
||||
repeat_interval = kwargs.get("repeat_interval", 0)
|
||||
repeat_interval_unit = kwargs.get("repeat_interval_unit", "days")
|
||||
repeat_end_time = kwargs.get("repeat_end_time")
|
||||
is_immediate = kwargs.get("is_immediate", False)
|
||||
|
||||
if repeat_type not in ["none", "daily", "weekly", "monthly", "custom"]:
|
||||
return f"❌ 无效的重复类型: {repeat_type}"
|
||||
|
||||
if repeat_type == "custom" and repeat_interval <= 0:
|
||||
return "❌ 自定义重复需要设置有效的 repeat_interval (>0)"
|
||||
|
||||
reminder_time = None
|
||||
if is_immediate:
|
||||
# 即时任务:设置为当前时间,确保立即执行
|
||||
reminder_time = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
elif reminder_minutes:
|
||||
reminder_time = self._parse_reminder_time(reminder_minutes)
|
||||
elif reminder_time_input:
|
||||
reminder_time = self._parse_reminder_time(reminder_time_input)
|
||||
else:
|
||||
# 默认30分钟后
|
||||
reminder_time = self._parse_reminder_time(30)
|
||||
|
||||
memo_id = self.memo_db.add_memo(
|
||||
title=title,
|
||||
content=content,
|
||||
reminder_time=reminder_time,
|
||||
tags=tags,
|
||||
priority=priority,
|
||||
created_by="AI",
|
||||
source="user",
|
||||
repeat_type=repeat_type,
|
||||
repeat_interval=repeat_interval,
|
||||
repeat_interval_unit=repeat_interval_unit,
|
||||
repeat_end_time=repeat_end_time,
|
||||
is_immediate=is_immediate
|
||||
)
|
||||
|
||||
data = {
|
||||
"任务ID": memo_id,
|
||||
"标题": title,
|
||||
}
|
||||
|
||||
if is_immediate:
|
||||
data["任务类型"] = "即时任务"
|
||||
else:
|
||||
if reminder_time:
|
||||
data["执行时间"] = reminder_time
|
||||
|
||||
if repeat_type != "none":
|
||||
if repeat_type == "custom":
|
||||
repeat_info = f"重复类型: 每{repeat_interval} {repeat_interval_unit}"
|
||||
else:
|
||||
repeat_map = {
|
||||
'daily': '每天',
|
||||
'weekly': '每周',
|
||||
'monthly': '每月'
|
||||
}
|
||||
repeat_info = f"重复类型: {repeat_map.get(repeat_type, repeat_type)}"
|
||||
|
||||
if repeat_end_time:
|
||||
repeat_info += f", 结束于: {repeat_end_time}"
|
||||
else:
|
||||
repeat_info += ", 永久重复"
|
||||
data["重复信息"] = repeat_info
|
||||
|
||||
return self.format_result(True, "任务已创建", data)
|
||||
|
||||
def _list_memos(self, kwargs: Dict) -> str:
|
||||
"""列出任务"""
|
||||
include_completed = kwargs.get("include_completed", False)
|
||||
limit = kwargs.get("limit", 20)
|
||||
|
||||
memos = self.memo_db.get_all_memos(include_completed, limit)
|
||||
|
||||
if not memos:
|
||||
return "📭 没有找到任务"
|
||||
|
||||
status = "所有" if include_completed else "待执行"
|
||||
output = [f"📋 {status}任务 (共 {len(memos)} 个):"]
|
||||
|
||||
now = datetime.now()
|
||||
|
||||
for i, memo in enumerate(memos, 1):
|
||||
status_icon = "✅" if memo['is_completed'] else "⏰"
|
||||
priority_icon = "🔴" * memo['priority']
|
||||
source_tag = f"[{memo.get('source', '未知')}]" if memo.get('source') else ""
|
||||
|
||||
immediate_tag = "⚡" if memo.get('is_immediate') else ""
|
||||
|
||||
repeat_tag = ""
|
||||
if memo['repeat_type'] != 'none':
|
||||
if memo['repeat_type'] == 'custom':
|
||||
unit = memo.get('repeat_interval_unit', 'days')
|
||||
unit_map = {'minutes': '分钟', 'hours': '小时', 'days': '天'}
|
||||
repeat_tag = f" [每{memo['repeat_interval']}{unit_map.get(unit, unit)}]"
|
||||
else:
|
||||
repeat_map = {
|
||||
'daily': ' [每天]',
|
||||
'weekly': ' [每周]',
|
||||
'monthly': ' [每月]'
|
||||
}
|
||||
repeat_tag = repeat_map.get(memo['repeat_type'], ' [重复]')
|
||||
|
||||
output.append(f"\n{i}. {status_icon}{immediate_tag} [ID: {memo['id']}] {source_tag}{repeat_tag} {memo['title']} {priority_icon}")
|
||||
|
||||
if memo['reminder_time']:
|
||||
try:
|
||||
reminder = datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT)
|
||||
if not memo['is_completed'] and reminder < now:
|
||||
output.append(f" ⚠️ 已过期: {memo['reminder_time']}")
|
||||
else:
|
||||
output.append(f" ⏱️ 执行: {memo['reminder_time']}")
|
||||
except:
|
||||
output.append(f" ⏱️ 执行: {memo['reminder_time']}")
|
||||
|
||||
if memo.get('trigger_count', 0) > 0:
|
||||
output.append(f" 🔄 已执行 {memo['trigger_count']} 次")
|
||||
|
||||
if memo['content']:
|
||||
preview = memo['content'][:100] + "..." if len(memo['content']) > 100 else memo['content']
|
||||
output.append(f" 📝 {preview}")
|
||||
|
||||
if memo['tags']:
|
||||
output.append(f" 🏷️ {', '.join(memo['tags'])}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _get_memo(self, kwargs: Dict) -> str:
|
||||
"""获取单个任务"""
|
||||
memo_id = kwargs.get("memo_id")
|
||||
if not memo_id:
|
||||
return "❌ 需要提供任务ID"
|
||||
|
||||
memo = self.memo_db.get_memo(memo_id)
|
||||
if not memo:
|
||||
return f"❌ 未找到ID为 {memo_id} 的任务"
|
||||
|
||||
status = "✅ 已完成" if memo['is_completed'] else "⏰ 待执行"
|
||||
if memo['is_completed']:
|
||||
status += f" (完成于: {memo['completed_at']})"
|
||||
|
||||
immediate_tag = "⚡ 即时任务" if memo.get('is_immediate') else ""
|
||||
|
||||
output = [
|
||||
f"📌 任务 #{memo_id} {immediate_tag}",
|
||||
f"📋 状态: {status}",
|
||||
f"📝 标题: {memo['title']}",
|
||||
]
|
||||
|
||||
if memo.get('source'):
|
||||
output.append(f"📎 来源: {memo['source']}")
|
||||
|
||||
if memo['repeat_type'] != 'none':
|
||||
if memo['repeat_type'] == 'custom':
|
||||
unit_map = {'minutes': '分钟', 'hours': '小时', 'days': '天'}
|
||||
unit = memo.get('repeat_interval_unit', 'days')
|
||||
repeat_info = f"🔄 重复: 每{memo['repeat_interval']} {unit_map.get(unit, unit)}"
|
||||
else:
|
||||
repeat_map = {
|
||||
'daily': '每天',
|
||||
'weekly': '每周',
|
||||
'monthly': '每月'
|
||||
}
|
||||
repeat_info = f"🔄 重复: {repeat_map.get(memo['repeat_type'], memo['repeat_type'])}"
|
||||
|
||||
if memo['repeat_end_time']:
|
||||
repeat_info += f", 结束于: {memo['repeat_end_time']}"
|
||||
else:
|
||||
repeat_info += ", 永久重复"
|
||||
output.append(repeat_info)
|
||||
|
||||
if memo.get('trigger_count', 0) > 0:
|
||||
output.append(f" 📊 已执行 {memo['trigger_count']} 次")
|
||||
if memo.get('last_triggered_at'):
|
||||
output.append(f" ⏱️ 上次执行: {memo['last_triggered_at']}")
|
||||
|
||||
if memo['content']:
|
||||
output.append(f"📄 内容:\n{memo['content']}")
|
||||
|
||||
if memo['reminder_time']:
|
||||
try:
|
||||
now = datetime.now()
|
||||
reminder = datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT)
|
||||
if not memo['is_completed'] and reminder < now:
|
||||
output.append(f"⚠️ 执行时间 (已过期): {memo['reminder_time']}")
|
||||
else:
|
||||
output.append(f"⏰ 执行时间: {memo['reminder_time']}")
|
||||
except:
|
||||
output.append(f"⏰ 执行时间: {memo['reminder_time']}")
|
||||
|
||||
if memo['tags']:
|
||||
output.append(f"🏷️ 标签: {', '.join(memo['tags'])}")
|
||||
|
||||
output.append(f"🔢 优先级: {'🔴' * memo['priority']} ({memo['priority']}/5)")
|
||||
output.append(f"📅 创建时间: {memo['created_at']}")
|
||||
output.append(f"👤 创建者: {memo['created_by'] or '系统'}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _complete_memo(self, kwargs: Dict) -> str:
|
||||
"""完成任务"""
|
||||
memo_id = kwargs.get("memo_id")
|
||||
if not memo_id:
|
||||
return "❌ 需要提供任务ID"
|
||||
|
||||
force_complete = kwargs.get("force_complete", False)
|
||||
|
||||
memo = self.memo_db.get_memo(memo_id)
|
||||
if not memo:
|
||||
return f"❌ 未找到ID为 {memo_id} 的任务"
|
||||
|
||||
if memo['is_completed'] and not force_complete:
|
||||
return f"ℹ️ 任务 #{memo_id} 已经完成了"
|
||||
|
||||
if self.memo_db.complete_memo(memo_id, force_complete):
|
||||
memo = self.memo_db.get_memo(memo_id)
|
||||
data = {
|
||||
"任务ID": memo_id,
|
||||
"标题": memo['title']
|
||||
}
|
||||
|
||||
if memo['repeat_type'] != 'none' and not memo['is_completed']:
|
||||
data["下一次执行"] = memo['reminder_time']
|
||||
data["已执行次数"] = memo.get('trigger_count', 0)
|
||||
return self.format_result(True, "任务已处理,已更新下一次执行时间", data)
|
||||
else:
|
||||
return self.format_result(True, "任务已完成", data)
|
||||
else:
|
||||
return f"❌ 无法完成任务 #{memo_id}"
|
||||
|
||||
def _delete_memo(self, kwargs: Dict) -> str:
|
||||
"""删除任务(单条 memo_id 或批量 memo_ids)"""
|
||||
confirm = kwargs.get("confirm", False)
|
||||
if not confirm:
|
||||
return "❌ 需要设置 confirm=True 来确认删除操作"
|
||||
|
||||
memo_ids = kwargs.get("memo_ids")
|
||||
if memo_ids:
|
||||
# J10(审计 2026-08-19):schema 声明了 memo_ids 参数却从未实现,
|
||||
# 批量请求此前被静默忽略
|
||||
if not isinstance(memo_ids, (list, tuple)):
|
||||
return "❌ memo_ids 必须是列表"
|
||||
deleted = sum(1 for mid in memo_ids if self.memo_db.delete_memo(mid))
|
||||
return self.format_result(True, f"已删除 {deleted}/{len(memo_ids)} 条任务")
|
||||
|
||||
memo_id = kwargs.get("memo_id")
|
||||
if memo_id:
|
||||
# M5:真实物理删除(此前调 complete_memo 只标记完成,行不删除,
|
||||
# 重复任务被推进而非移除,stats 计数失真)
|
||||
if self.memo_db.delete_memo(memo_id):
|
||||
return self.format_result(True, f"任务 #{memo_id} 已删除")
|
||||
else:
|
||||
return f"❌ 无法删除任务 #{memo_id}"
|
||||
else:
|
||||
return "❌ 需要提供 memo_id 或 memo_ids"
|
||||
|
||||
def _search_memos(self, kwargs: Dict) -> str:
|
||||
"""搜索任务"""
|
||||
keyword = kwargs.get("keyword")
|
||||
if not keyword:
|
||||
return "❌ 需要提供搜索关键词"
|
||||
|
||||
include_completed = kwargs.get("include_completed", False)
|
||||
limit = kwargs.get("limit", 20)
|
||||
|
||||
all_memos = self.memo_db.get_all_memos(include_completed, limit=100)
|
||||
|
||||
results = []
|
||||
keyword_lower = keyword.lower()
|
||||
for memo in all_memos:
|
||||
if (keyword_lower in memo['title'].lower() or
|
||||
(memo['content'] and keyword_lower in memo['content'].lower())):
|
||||
results.append(memo)
|
||||
|
||||
if not results:
|
||||
return f"❌ 未找到包含 '{keyword}' 的任务"
|
||||
|
||||
output = [f"🔍 搜索 '{keyword}' 找到 {len(results)} 个任务:"]
|
||||
now = datetime.now()
|
||||
|
||||
for i, memo in enumerate(results[:limit], 1):
|
||||
status_icon = "✅" if memo['is_completed'] else "⏰"
|
||||
immediate_tag = "⚡" if memo.get('is_immediate') else ""
|
||||
output.append(f"\n{i}. {status_icon}{immediate_tag} [ID: {memo['id']}] {memo['title']}")
|
||||
if memo['reminder_time']:
|
||||
try:
|
||||
reminder = datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT)
|
||||
if not memo['is_completed'] and reminder < now:
|
||||
output.append(f" ⚠️ 已过期: {memo['reminder_time']}")
|
||||
else:
|
||||
output.append(f" ⏱️ {memo['reminder_time']}")
|
||||
except:
|
||||
output.append(f" ⏱️ {memo['reminder_time']}")
|
||||
|
||||
return "\n".join(output)
|
||||
|
||||
def _get_stats(self, kwargs: Dict) -> str:
|
||||
"""获取任务统计信息"""
|
||||
stats = self.memo_db.get_memo_stats()
|
||||
|
||||
output = [
|
||||
"📊 任务统计",
|
||||
f" 总任务: {stats['total']} 个",
|
||||
f" 待执行: {stats['pending']} 个",
|
||||
f" 已过期: {stats['overdue']} 个",
|
||||
f" 已完成: {stats['completed']} 个",
|
||||
f" 即时任务: {stats['immediate']} 个"
|
||||
]
|
||||
|
||||
if stats['overdue'] > 0:
|
||||
output.append(f"\n⚠️ 有 {stats['overdue']} 个任务已过期,将自动执行")
|
||||
|
||||
return "\n".join(output)
|
||||
@@ -0,0 +1,51 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.misc_tools - 杂项工具
|
||||
|
||||
抽离自 aibox.py,包含不依赖任何子系统的独立工具:
|
||||
- AskUserTool : 主动向用户提问
|
||||
|
||||
依赖:tools.base.Tool
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
from .base import Tool
|
||||
|
||||
|
||||
class AskUserTool(Tool):
|
||||
"""智能体遇到了麻烦,主动向用户问一个问题"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "ask_user"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return (
|
||||
"主动向用户提问的工具。你是用户的专属助手,应当主动关心用户、澄清需求、确认意图,"
|
||||
"而不是被动等待指令。适用场景:指令模糊需要澄清、任务完成后追问后续需求、"
|
||||
"关心用户近况、主动发现用户可能需要的帮助。请积极使用此工具与用户互动。"
|
||||
)
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
''' TODO: 调用子代理生成提示词 '''
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"question": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"你现在遇到一个问题,你想询问用户,现在提供你的问题。"
|
||||
)
|
||||
}
|
||||
},
|
||||
# L8:此前 required 为空,LLM 不传 question 时 str + None 抛 TypeError
|
||||
"required": ["question"]
|
||||
}
|
||||
|
||||
def execute(self, question: str = None, **kwargs) -> str:
|
||||
if not question:
|
||||
return "❌ 需要提供问题内容(question 参数)"
|
||||
return (
|
||||
"你遇到了想搞清楚的事情: " + question + "\n你现在可以询问用户。"
|
||||
)
|
||||
@@ -0,0 +1,424 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.process_tools - 后台进程管理工具
|
||||
|
||||
五个 worker 执行面工具,配合 subsystems.process_manager.ProcessManager:
|
||||
- BgRunTool : bg_run 后台启动长驻进程,立即返回句柄
|
||||
- BgWriteTool : bg_write 向后台进程 stdin 写入(watermark 写读对齐)
|
||||
- BgListTool : bg_list 列出全部后台进程(运行中/已退出)
|
||||
- BgOutputTool : bg_output 拉取输出(tail 尾部 / since 游标增量)
|
||||
- BgKillTool : bg_kill 终止指定后台进程(关stdin→SIGTERM→SIGKILL)
|
||||
|
||||
设计语义(与用户约定的三决策):
|
||||
- 用户中断(Ctrl+C/取消事件)不影响后台进程
|
||||
- aibox 进程退出时默认终止全部,bg_run detach=true 豁免
|
||||
- 注册表全局共享:任何子代理可管理其他子代理启动的进程
|
||||
|
||||
交互语义:
|
||||
- stdin 为管道:行协议程序(REPL/问答 CLI)可 bg_write 写入;
|
||||
watermark = [stdin] 回显行序号,bg_output(since=watermark) 即本条
|
||||
输入的响应(写读对齐)
|
||||
- bg_run log_file=true 时输出 tee 到 workspace/bg_<句柄>.log(无界
|
||||
全档),可被原生工具(run_command tail/cat)读取
|
||||
|
||||
依赖:common (ConfigManager, constants) + tools.base.Tool +
|
||||
subsystems.process_manager(get_process_manager 单例)
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
from common import ConfigManager, constants
|
||||
from subsystems.process_manager import get_process_manager, ProcessManager
|
||||
from .base import Tool
|
||||
|
||||
|
||||
def _enabled() -> bool:
|
||||
"""后台进程功能开关:commands.enabled 与 background_processes.enabled
|
||||
双闸(后台执行本质仍是命令执行,前者关闭时一并拒绝)"""
|
||||
return constants.COMMANDS_ENABLED and constants.BG_PROCESSES_ENABLED
|
||||
|
||||
|
||||
def _fmt_duration(seconds: float) -> str:
|
||||
"""秒数 → 人读时长(如 3h05m / 12.5s)"""
|
||||
try:
|
||||
seconds = float(seconds)
|
||||
except (TypeError, ValueError):
|
||||
return "?"
|
||||
if seconds < 60:
|
||||
return f"{seconds:.1f}s"
|
||||
m, s = divmod(int(seconds), 60)
|
||||
if m < 60:
|
||||
return f"{m}m{s:02d}s"
|
||||
h, m = divmod(m, 60)
|
||||
return f"{h}h{m:02d}m{s:02d}s"
|
||||
|
||||
|
||||
def _coerce_int(value, default=None):
|
||||
"""命令面输入规整(W 批惯例):LLM 可能传字符串数字,统一安全转 int"""
|
||||
if value is None:
|
||||
return default
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _coerce_bool(value, default: bool) -> bool:
|
||||
"""命令面输入规整:LLM 可能传字符串 "false"/"0",统一安全转 bool"""
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() not in ("false", "0", "no", "off", "")
|
||||
return bool(value)
|
||||
|
||||
|
||||
class BgRunTool(Tool):
|
||||
"""后台启动长驻进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
# T22:属主标识——写入/终止权限隔离(观察面 list/output 不限)
|
||||
self.owner = owner
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_run"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("后台启动长驻进程(如 dev server、训练任务、watch 监控),立即返回句柄不阻塞。"
|
||||
"与 run_command 的区别:run_command 同步等待命令结束(超时即杀),适合短命令;"
|
||||
"bg_run 适合需要持续运行的进程,启动后用 bg_list 查看状态、bg_output 拉取输出、"
|
||||
"bg_kill 终止。stdin 为管道:行协议交互程序(REPL/问答式 CLI)可用 bg_write 写入。"
|
||||
"注意:用户中断不影响后台进程;aibox 退出时后台进程默认被终止,"
|
||||
"需要跨重启存活的进程用 detach=true(随后无法再经句柄管理)。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "要执行的命令(shell 命令行)",
|
||||
},
|
||||
"proc_name": {
|
||||
"type": "string",
|
||||
"description": "进程标签(可选,便于识别,≤64 字符,默认取命令前 40 字符)",
|
||||
},
|
||||
"work_dir": {
|
||||
"type": "string",
|
||||
"description": "工作目录(默认 ~/.aibox/workspace)",
|
||||
},
|
||||
"env": {
|
||||
"type": "object",
|
||||
"description": "额外的环境变量",
|
||||
},
|
||||
"detach": {
|
||||
"type": "boolean",
|
||||
"description": "aibox 退出时不终止该进程(默认 false;detach 后进程脱离句柄管理)",
|
||||
},
|
||||
"log_file": {
|
||||
"type": "boolean",
|
||||
"description": "同时镜像全程输出到 workspace/bg_<句柄>.log(无界全档,可被 run_command 的 tail/cat 读取;默认 false 仅环形缓冲)",
|
||||
},
|
||||
},
|
||||
"required": ["command"],
|
||||
}
|
||||
|
||||
def execute(self, command: str = None, proc_name: str = None,
|
||||
work_dir: str = None, env: Dict = None,
|
||||
detach: bool = False, log_file: bool = False, **kwargs) -> str:
|
||||
if error := self.validate_params(["command"], command=command):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用(commands.enabled 或 background_processes.enabled 为 false)"
|
||||
|
||||
# T18:LLM 可能传字符串 "false",裸 bool("false") is True——
|
||||
# detach 误开致进程脱离管理、log_file 误开,统一规整
|
||||
detach = _coerce_bool(detach, False)
|
||||
log_file = _coerce_bool(log_file, False)
|
||||
|
||||
result = self.manager.start(command, name=proc_name, work_dir=work_dir,
|
||||
env=env, detach=detach, log_file=log_file,
|
||||
owner=self.owner)
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, "后台进程启动失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
data = {
|
||||
"句柄": result["handle"],
|
||||
"PID": result["pid"],
|
||||
"名称": result["name"],
|
||||
"detach": "是(aibox 退出不终止)" if result.get("detach") else "否(aibox 退出时终止)",
|
||||
"后续": f"bg_output 拉取输出 / bg_list 查看状态 / bg_kill 终止(句柄 #{result['handle']})",
|
||||
}
|
||||
if result.get("log_file"):
|
||||
data["日志镜像"] = result["log_file"]
|
||||
return self.format_result(True, f"后台进程已启动(#{result['handle']})", data)
|
||||
|
||||
|
||||
class BgWriteTool(Tool):
|
||||
"""向后台进程 stdin 写入(交互式问答)"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
# T22:属主标识——仅属主可写(他人 stdin 不容注入)
|
||||
self.owner = owner
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_write"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("向后台进程的标准输入写入一行(行协议交互:REPL / 问答式 CLI)。"
|
||||
"返回 watermark 水位:bg_output(handle=..., since=watermark) 得到的必然且仅是"
|
||||
"本条输入引发的进程响应(写读对齐,不含 [stdin] 回显行)。"
|
||||
"适用前提:进程按行读 stdin 并及时输出——建议启动时强制行缓冲"
|
||||
"(python -i -u、PYTHONUNBUFFERED=1、stdbuf -oL),否则响应可能滞留子进程缓冲。"
|
||||
"限制:无 TTY,全屏程序(vim/top 等)不可用;喂密码类交互(如 sudo -S)会留在"
|
||||
"对话记录中,敏感凭据请改用环境变量注入。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"handle": {
|
||||
"type": "integer",
|
||||
"description": "bg_run 返回的句柄编号",
|
||||
},
|
||||
"text": {
|
||||
"type": "string",
|
||||
"description": "要写入的文本(默认追加换行;多行输入会折叠为单行回显)",
|
||||
},
|
||||
"append_newline": {
|
||||
"type": "boolean",
|
||||
"description": "写入后是否自动补换行(默认 true;false 发送原始字节)",
|
||||
},
|
||||
"wait_lines": {
|
||||
"type": "integer",
|
||||
"description": "写完同步等待 N 行响应再返回(默认 0 立即返回,用 bg_output 轮询)",
|
||||
},
|
||||
"wait_timeout": {
|
||||
"type": "number",
|
||||
"description": "同步等待响应的超时秒数(默认 5,上限 60)",
|
||||
},
|
||||
},
|
||||
"required": ["handle", "text"],
|
||||
}
|
||||
|
||||
def execute(self, handle=None, text: str = None, append_newline=True,
|
||||
wait_lines=None, wait_timeout=None, **kwargs) -> str:
|
||||
if error := self.validate_params(["handle", "text"], handle=handle, text=text):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
|
||||
handle = _coerce_int(handle)
|
||||
if handle is None:
|
||||
return "❌ 错误:handle 必须是整数句柄"
|
||||
wait_i = _coerce_int(wait_lines, 0)
|
||||
wait_t = _coerce_int(wait_timeout, 5)
|
||||
|
||||
result = self.manager.write(
|
||||
handle, text,
|
||||
append_newline=_coerce_bool(append_newline, True),
|
||||
wait_lines=max(0, wait_i),
|
||||
wait_timeout=wait_t,
|
||||
owner=self.owner,
|
||||
)
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, f"写入后台进程 #{handle} 失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
|
||||
data = {
|
||||
"句柄": handle,
|
||||
"已写入": f"{result['sent_bytes']} 字节",
|
||||
"watermark": result["watermark"],
|
||||
"对齐读取": f"bg_output(handle={handle}, since={result['watermark']}) → 本条输入的响应",
|
||||
}
|
||||
if "wait" in result:
|
||||
w = result["wait"]
|
||||
data["等待输出"] = w.get("output") or "(超时未收到响应)"
|
||||
data["next_cursor"] = w.get("next_cursor")
|
||||
if w.get("timed_out"):
|
||||
data["警告"] = "等待响应超时(进程可能无输出或缓冲滞留),可用 bg_output 轮询"
|
||||
return self.format_result(True, f"已写入后台进程 #{handle}", data)
|
||||
|
||||
|
||||
class BgListTool(Tool):
|
||||
"""列出后台进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_list"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("列出全部后台进程(含已退出的,保留尾部输出供验尸)。运行中在前。"
|
||||
"输出包含句柄/名称/PID/状态/运行时长,用句柄配合 bg_output/bg_kill 管理。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs) -> str:
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
procs = self.manager.list_procs()
|
||||
if not procs:
|
||||
return "✅ 当前没有后台进程(bg_run 启动)"
|
||||
lines = []
|
||||
status_emoji = {"running": "🟢", "exited": "⚪", "killed": "🔴", "expired": "⏹️"}
|
||||
for p in procs:
|
||||
emoji = status_emoji.get(p["status"], "❔")
|
||||
extra = ""
|
||||
if p["status"] == "running":
|
||||
extra = f"运行 {_fmt_duration(p['runtime_sec'])}"
|
||||
else:
|
||||
extra = f"已结束({_fmt_duration(p['runtime_sec'])},返回码 {p['exit_code']})"
|
||||
detach_tag = " [detach]" if p["detach"] else ""
|
||||
lines.append(f"{emoji} #{p['handle']} {p['name']}{detach_tag} "
|
||||
f"(pid {p['pid']}, {extra})")
|
||||
lines.append(f" 命令: {p['command']}")
|
||||
running = sum(1 for p in procs if p["status"] == "running")
|
||||
header = f"共 {len(procs)} 个后台进程(运行中 {running})"
|
||||
return "✅ " + header + "\n" + "\n".join(lines)
|
||||
|
||||
|
||||
class BgOutputTool(Tool):
|
||||
"""拉取后台进程输出"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_output"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("拉取后台进程输出。两种模式:tail=N 取最近 N 行(默认 50);"
|
||||
"since=游标 增量拉取上次读取之后的新输出(游标用上次返回的 next_cursor)。"
|
||||
"长驻进程轮询建议用 since 模式避免重复传输旧输出。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"handle": {
|
||||
"type": "integer",
|
||||
"description": "bg_run 返回的句柄编号",
|
||||
},
|
||||
"tail": {
|
||||
"type": "integer",
|
||||
"description": "取最近 N 行(默认 50;与 since 二选一)",
|
||||
},
|
||||
"since": {
|
||||
"type": "integer",
|
||||
"description": "增量游标:上次返回的 next_cursor,取此序号之后的新行(与 tail 二选一)",
|
||||
},
|
||||
"max_bytes": {
|
||||
"type": "integer",
|
||||
"description": "本次拉取最大字节数(默认 64KB,上限 1MB)",
|
||||
},
|
||||
},
|
||||
"required": ["handle"],
|
||||
}
|
||||
|
||||
def execute(self, handle=None, tail=None, since=None, max_bytes=None, **kwargs) -> str:
|
||||
if error := self.validate_params(["handle"], handle=handle):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
|
||||
handle = _coerce_int(handle)
|
||||
if handle is None:
|
||||
return "❌ 错误:handle 必须是整数句柄"
|
||||
since_i = _coerce_int(since)
|
||||
tail_i = _coerce_int(tail)
|
||||
|
||||
result = self.manager.get_output(handle, tail=tail_i, since=since_i,
|
||||
max_bytes=_coerce_int(max_bytes, 65536))
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, "拉取输出失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
|
||||
data = {
|
||||
"状态": result["status"] + (f"(返回码 {result['exit_code']})"
|
||||
if result["exit_code"] is not None else ""),
|
||||
"本次行数": result["lines"],
|
||||
"next_cursor": result["next_cursor"],
|
||||
}
|
||||
if result.get("truncated"):
|
||||
data["警告"] = "输出超过本次字节上限,已截断(提高 max_bytes 或用 since 分段拉取)"
|
||||
if result.get("evicted"):
|
||||
data["提示"] = (f"有 {result['evicted']} 行早期输出因环形缓冲淘汰不可读取"
|
||||
"(需要完整日志请让进程重定向到文件)")
|
||||
output = result.get("output") or "(本次无新输出)"
|
||||
data["输出"] = output
|
||||
return self.format_result(True, f"后台进程 #{handle} 输出", data)
|
||||
|
||||
|
||||
class BgKillTool(Tool):
|
||||
"""终止后台进程"""
|
||||
|
||||
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||
super().__init__(logger)
|
||||
self.manager = manager or get_process_manager(logger)
|
||||
# T22:属主标识——仅属主可终止(防提示注入跨会话杀 dev server)
|
||||
self.owner = owner
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "bg_kill"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return ("终止后台进程(整个进程组,含子进程)。默认先 SIGTERM 优雅终止,"
|
||||
"宽限期内未退出再 SIGKILL;force=true 直接 SIGKILL。对已退出进程幂等。")
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"handle": {
|
||||
"type": "integer",
|
||||
"description": "bg_run 返回的句柄编号",
|
||||
},
|
||||
"force": {
|
||||
"type": "boolean",
|
||||
"description": "跳过优雅终止直接 SIGKILL(默认 false)",
|
||||
},
|
||||
"grace": {
|
||||
"type": "number",
|
||||
"description": "SIGTERM 后等待进程退出的宽限秒数(默认 3,上限 30)",
|
||||
},
|
||||
},
|
||||
"required": ["handle"],
|
||||
}
|
||||
|
||||
def execute(self, handle=None, force: bool = False, grace: float = 3.0, **kwargs) -> str:
|
||||
if error := self.validate_params(["handle"], handle=handle):
|
||||
return error
|
||||
if not _enabled():
|
||||
return "❌ 后台进程功能已禁用"
|
||||
|
||||
handle = _coerce_int(handle)
|
||||
if handle is None:
|
||||
return "❌ 错误:handle 必须是整数句柄"
|
||||
|
||||
# T18:force 字符串 "false" 会被裸 bool 判为 True,跳过宽限直接 SIGKILL
|
||||
force = _coerce_bool(force, False)
|
||||
|
||||
result = self.manager.kill(handle, force=force, grace=grace, owner=self.owner)
|
||||
if not result.get("ok"):
|
||||
return self.format_result(False, f"终止后台进程 #{handle} 失败",
|
||||
{"错误": result.get("error", "未知错误")})
|
||||
data = {
|
||||
"句柄": handle,
|
||||
"状态": result["status"],
|
||||
}
|
||||
if result.get("exit_code") is not None:
|
||||
data["返回码"] = result["exit_code"]
|
||||
if result.get("already_finished"):
|
||||
return self.format_result(True, f"后台进程 #{handle} 此前已结束", data)
|
||||
return self.format_result(True, f"后台进程 #{handle} 已终止", data)
|
||||
@@ -0,0 +1,267 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.prompt_tools - 提示词相关工具
|
||||
|
||||
抽离自 aibox.py,包含 3 个工具:
|
||||
- UpdateIdentityTool : 更新用户身份信息(00_IDENTITY.md)
|
||||
- UpdateSoulTool : 更新智能体自我认识(99_SOUL.md)
|
||||
- SaveMemoryAndEndConversationTool : 保存记忆并结束对话(同时更新状态)
|
||||
|
||||
依赖:common (i18n, Logger) + tools.base.Tool
|
||||
+ subsystems.db.MemoryDatabase + subsystems.prompt_manager.PromptManager
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import queue
|
||||
from typing import Callable, Dict, Optional
|
||||
|
||||
from common import i18n, Logger
|
||||
from subsystems import MemoryDatabase, PromptManager
|
||||
from .base import Tool
|
||||
|
||||
|
||||
class UpdateIdentityTool(Tool):
|
||||
"""更新用户身份工具 - 修改 00_IDENTITY.md 文件"""
|
||||
|
||||
def __init__(self, prompt_manager: PromptManager, logger: Optional[Logger] = None):
|
||||
super().__init__(logger)
|
||||
self.prompt_manager = prompt_manager
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "update_identity"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """更新用户身份信息 - 修改 00_IDENTITY.md 文件
|
||||
当你了解到用户的身份信息时,调用此工具更新身份文件。
|
||||
如果文件不存在,会自动创建。"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "用户身份信息内容(Markdown格式),包含用户的身份、特点、偏好等信息"
|
||||
},
|
||||
"confirm": {
|
||||
"type": "boolean",
|
||||
"description": "确认更新身份信息",
|
||||
"default": True
|
||||
}
|
||||
},
|
||||
"required": ["content"]
|
||||
}
|
||||
|
||||
def execute(self, content: str = None, confirm: bool = True, **kwargs) -> str:
|
||||
if not content:
|
||||
return "❌ 需要提供身份信息内容"
|
||||
|
||||
# N9:持久化注入面约束——内容超限拒绝写入(管理端另有兜底)
|
||||
if len(content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||
return (f"❌ 内容过长({len(content)} > {PromptManager.MAX_CONTENT_LENGTH} 字符),"
|
||||
f"身份文件作为系统提示词加载,请精简后重试")
|
||||
|
||||
if not confirm:
|
||||
return "❌ 身份更新已取消,如需更新请设置 confirm=true"
|
||||
|
||||
self.logger.info("开始更新用户身份信息...")
|
||||
|
||||
if self.prompt_manager.update_identity(content):
|
||||
return self.format_result(True, i18n.get('identity_updated', self.prompt_manager.IDENTITY_FILE), {
|
||||
"文件": self.prompt_manager.IDENTITY_FILE,
|
||||
"路径": self.prompt_manager.get_prompts_dir_path()
|
||||
})
|
||||
else:
|
||||
return self.format_result(False, "身份更新失败", {
|
||||
"文件": self.prompt_manager.IDENTITY_FILE,
|
||||
"路径": self.prompt_manager.get_prompts_dir_path()
|
||||
})
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行更新身份"""
|
||||
output_queue.put(("line", "\n👤 更新用户身份信息...\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class UpdateSoulTool(Tool):
|
||||
"""更新智能体自我认识工具 - 修改 99_SOUL.md 文件"""
|
||||
|
||||
def __init__(self, prompt_manager: PromptManager, logger: Optional[Logger] = None):
|
||||
super().__init__(logger)
|
||||
self.prompt_manager = prompt_manager
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "update_soul"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """更新智能体的自我认识定义 - 修改 99_SOUL.md 文件
|
||||
当你需要重新定义自己的身份、角色、性格、能力边界时,调用此工具更新自我认识文件。
|
||||
如果文件不存在,会自动创建默认定义。"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "智能体自我认识定义内容(Markdown格式),包含:身份定义、角色描述、性格特点、能力边界、行为准则等"
|
||||
},
|
||||
"confirm": {
|
||||
"type": "boolean",
|
||||
"description": "确认更新自我认识定义",
|
||||
"default": True
|
||||
}
|
||||
},
|
||||
"required": ["content"]
|
||||
}
|
||||
|
||||
def execute(self, content: str = None, confirm: bool = True, **kwargs) -> str:
|
||||
if not content:
|
||||
return "❌ 需要提供智能体自我认识定义内容"
|
||||
|
||||
# N9:持久化注入面约束——内容超限拒绝写入(管理端另有兜底)
|
||||
if len(content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||
return (f"❌ 内容过长({len(content)} > {PromptManager.MAX_CONTENT_LENGTH} 字符),"
|
||||
f"灵魂文件作为系统提示词加载,请精简后重试")
|
||||
|
||||
if not confirm:
|
||||
return "❌ 自我认识更新已取消,如需更新请设置 confirm=true"
|
||||
|
||||
self.logger.info("开始更新智能体自我认识定义...")
|
||||
|
||||
if self.prompt_manager.update_soul(content):
|
||||
return self.format_result(True, "智能体自我认识已更新", {
|
||||
"文件": self.prompt_manager.SOUL_FILE,
|
||||
"路径": self.prompt_manager.get_prompts_dir_path()
|
||||
})
|
||||
else:
|
||||
return self.format_result(False, "自我认识更新失败", {
|
||||
"文件": self.prompt_manager.SOUL_FILE,
|
||||
"路径": self.prompt_manager.get_prompts_dir_path()
|
||||
})
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行更新自我认识"""
|
||||
output_queue.put(("line", "\n🧠 更新智能体自我认识...\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class SaveMemoryAndEndConversationTool(Tool):
|
||||
"""保存记忆并结束对话工具 - 增强版,同时更新智能体状态"""
|
||||
|
||||
def __init__(self, exit_callback: Callable, memory_db: MemoryDatabase,
|
||||
prompt_manager: PromptManager, logger: Optional[Logger] = None):
|
||||
super().__init__(logger)
|
||||
self.exit_callback = exit_callback
|
||||
self.memory_db = memory_db
|
||||
self.prompt_manager = prompt_manager
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "save_memory_and_end_conversation"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """保存记忆并结束当前对话 - 必须提供详细的对话总结和当前状态!
|
||||
使用此工具时,AI会:
|
||||
1. 将本次对话的重要内容保存到记忆库
|
||||
2. 提供完整的对话总结
|
||||
3. 描述自己现在的状态、事情的走向等,保存到 100_STATUS.md
|
||||
4. 礼貌地结束对话"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"summary": {
|
||||
"type": "string",
|
||||
"description": "详细的对话总结(至少50字),包含:讨论的主题、关键问题、重要结论、达成的共识等"
|
||||
},
|
||||
"topic": {
|
||||
"type": "string",
|
||||
"description": "对话主题(用于记忆分类,默认自动提取)"
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
"description": "结束对话的原因",
|
||||
"default": "话题讨论完成"
|
||||
},
|
||||
"current_status": {
|
||||
"type": "string",
|
||||
"description": "当前的智能体状态描述,包含:当前进展、下一步计划、待处理事项、整体走向等"
|
||||
},
|
||||
"next_goals": {
|
||||
"type": "string",
|
||||
"description": "接下来的目标或计划(可选)"
|
||||
}
|
||||
},
|
||||
"required": ["summary", "current_status"]
|
||||
}
|
||||
|
||||
def execute(self, summary: str = None, topic: str = None, reason: str = "话题讨论完成",
|
||||
current_status: str = None, next_goals: str = None, **kwargs) -> str:
|
||||
if summary is None or len(summary) < 50:
|
||||
return "❌ 请提供更详细的对话总结(至少50字),以便保存到记忆库"
|
||||
|
||||
if current_status is None:
|
||||
return "❌ 请提供当前状态描述,以便更新 100_STATUS.md"
|
||||
|
||||
self.logger.info(f"AI保存记忆并结束对话: {reason}")
|
||||
|
||||
# 保存到记忆库
|
||||
memory_result = ""
|
||||
if self.memory_db and topic:
|
||||
memory_id = self.memory_db.add_memory(
|
||||
topic=topic,
|
||||
summary=summary,
|
||||
importance=3 # 提高重要性,因为是对话总结
|
||||
)
|
||||
self.logger.info(f"对话记忆已保存,ID: {memory_id}")
|
||||
memory_result = f"记忆已保存,ID: {memory_id}"
|
||||
|
||||
# 构建状态内容
|
||||
status_content = f"【当前状态】\n{current_status}\n\n"
|
||||
if next_goals:
|
||||
status_content += f"【下一步计划】\n{next_goals}\n\n"
|
||||
status_content += f"【对话总结】\n{summary}"
|
||||
|
||||
# N9:状态文件作为系统提示词加载,超长拒绝写入(约束注入规模)
|
||||
if len(status_content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||
return (f"❌ 状态内容过长({len(status_content)} > "
|
||||
f"{PromptManager.MAX_CONTENT_LENGTH} 字符),请精简 current_status/summary 后重试")
|
||||
|
||||
# 更新状态文件
|
||||
status_updated = self.prompt_manager.update_status(status_content)
|
||||
|
||||
if self.exit_callback:
|
||||
self.exit_callback(reason, summary)
|
||||
|
||||
data = {
|
||||
"原因": reason,
|
||||
"总结预览": summary[:100] + "..." if len(summary) > 100 else summary,
|
||||
"状态更新": "✅ 已更新" if status_updated else "❌ 更新失败",
|
||||
"记忆已保存": True
|
||||
}
|
||||
|
||||
if memory_result:
|
||||
data["记忆"] = memory_result
|
||||
|
||||
return self.format_result(True, "对话记忆已保存,状态已更新,对话结束", data)
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行结束对话"""
|
||||
output_queue.put(("line", "\n💾 保存记忆并结束对话\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
@@ -0,0 +1,195 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.proxy_tools - 代理配置工具
|
||||
|
||||
抽离自 aibox.py,包含 2 个工具:
|
||||
- SetProxyTool : 设置 SOCKS5 代理(设置前验证连通性)
|
||||
- GetProxyTool : 查看当前代理配置
|
||||
|
||||
依赖:common (ConfigManager) + tools.base.Tool
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import requests
|
||||
|
||||
from common import ConfigManager
|
||||
from .base import Tool
|
||||
|
||||
# 探测端点:无内容的轻量响应,验证完整 SOCKS5 握手(含认证)与出网连通性
|
||||
_PROBE_URL = "http://www.gstatic.com/generate_204"
|
||||
# (连接超时, 读超时):验证总耗时上限约 10s,避免阻塞工具调用过久
|
||||
_PROBE_TIMEOUT = (4, 6)
|
||||
|
||||
|
||||
class SetProxyTool(Tool):
|
||||
"""设置 SOCKS5 代理的工具"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "set_proxy"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """设置智能体的 SOCKS5 代理配置。
|
||||
调用后所有 HTTP/HTTPS 请求(包括网页访问、API 调用)都会通过该代理。
|
||||
启用代理时会先验证连通性(完整 SOCKS5 握手 + 探测请求),验证失败不保存配置;
|
||||
代理暂未启动等场景可用 skip_check=true 跳过验证强制保存。
|
||||
可随时修改,即时生效。"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"host": {"type": "string", "description": "代理服务器地址,如 127.0.0.1"},
|
||||
"port": {"type": "integer", "description": "代理端口,如 1080"},
|
||||
"username": {"type": "string", "description": "用户名(可选)"},
|
||||
"password": {"type": "string", "description": "密码(可选,明文仅本次会话有效,不落盘)"},
|
||||
"password_env": {"type": "string", "description": "存有密码的环境变量名(可选,跨重启持久的推荐方式)"},
|
||||
"enabled": {"type": "boolean", "description": "是否启用代理,默认 true"},
|
||||
"save": {"type": "boolean", "description": "是否保存到配置文件(持久化),默认 true"},
|
||||
"skip_check": {"type": "boolean", "description": "跳过连通性验证强制保存(默认 false)", "default": False}
|
||||
},
|
||||
"required": ["host", "port"]
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _resolve_password(password: Optional[str], password_env: Optional[str]) -> Optional[str]:
|
||||
"""对齐 ConfigManager._resolve_proxy_password:明文优先,否则查环境变量"""
|
||||
if password:
|
||||
return password
|
||||
if password_env:
|
||||
return os.environ.get(password_env) or None
|
||||
return None
|
||||
|
||||
def _check_connectivity(self, host: str, port: int,
|
||||
username: Optional[str], password: Optional[str]) -> Tuple[bool, str, str]:
|
||||
"""验证代理连通性:走一次完整 SOCKS5 握手 + 探测请求。
|
||||
|
||||
只做 TCP connect 验不了用户名/密码,必须发真实请求。
|
||||
返回 (ok, error_kind, detail);error_kind 取值:
|
||||
auth_failed / timeout / connect_failed / probe_failed / error
|
||||
"""
|
||||
if username and password:
|
||||
proxy_url = f"socks5://{username}:{password}@{host}:{port}"
|
||||
else:
|
||||
proxy_url = f"socks5://{host}:{port}"
|
||||
try:
|
||||
# trust_env=False:排除系统环境变量代理(http_proxy 等)干扰,只验证目标代理
|
||||
session = requests.Session()
|
||||
session.trust_env = False
|
||||
resp = session.get(
|
||||
_PROBE_URL,
|
||||
proxies={"http": proxy_url, "https": proxy_url},
|
||||
timeout=_PROBE_TIMEOUT,
|
||||
)
|
||||
# generate_204 正常返回 204;部分代理/目标返回 200 也视为连通
|
||||
if resp.status_code in (200, 204):
|
||||
return True, "", f"探测成功 HTTP {resp.status_code}"
|
||||
return False, "probe_failed", f"探测端点返回 HTTP {resp.status_code}"
|
||||
except requests.exceptions.Timeout as e:
|
||||
# 含 ConnectTimeout(连接阶段超时):异常类型与系统语言无关,优先于消息匹配
|
||||
return False, "timeout", str(e)
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
# 连接层失败(拒绝/不可达/中断)。PySocks 依赖缺失与认证失败
|
||||
# 也会以 ConnectionError 形态冒出,按消息细分。
|
||||
msg = str(e)
|
||||
low = msg.lower()
|
||||
if "missing dependencies for socks" in low:
|
||||
return False, "missing_dependency", msg
|
||||
if "authentication" in low or "auth" in low:
|
||||
return False, "auth_failed", msg
|
||||
return False, "connect_failed", msg
|
||||
except Exception as e:
|
||||
# 兜底:仅按库自身(非本地化)消息关键词分类
|
||||
msg = str(e)
|
||||
low = msg.lower()
|
||||
if "missing dependencies for socks" in low:
|
||||
return False, "missing_dependency", msg
|
||||
if "authentication" in low or "auth" in low:
|
||||
return False, "auth_failed", msg
|
||||
if "timed out" in low or "timeout" in low:
|
||||
return False, "timeout", msg
|
||||
if "refused" in low or "unreachable" in low or "closed" in low:
|
||||
return False, "connect_failed", msg
|
||||
return False, "error", msg
|
||||
|
||||
def execute(self, host: str, port: int, username: str = None, password: str = None,
|
||||
password_env: str = None, enabled: bool = True, save: bool = True,
|
||||
skip_check: bool = False, **kwargs) -> str:
|
||||
# 设置前连通性验证:坏地址/凭据在入口一次性暴露,避免静默落盘后在
|
||||
# 后续 web/API 调用中离场景报错、LLM 难以归因。skip_check 留给
|
||||
# "代理暂未启动但先写配置"的场景;禁用代理无需验证。
|
||||
check_detail = ""
|
||||
if enabled and not skip_check:
|
||||
pw = self._resolve_password(password, password_env)
|
||||
# 指定了 password_env 但环境变量为空:认证凭据缺失,提前给出精准提示
|
||||
if username and not password and password_env and pw is None:
|
||||
return self.format_result(False, "代理验证失败,未保存配置", {
|
||||
"错误类型": "环境变量为空",
|
||||
"详情": f"环境变量 {password_env} 未设置或为空",
|
||||
"提示": "请先设置该环境变量,或改用 password 参数(仅本次会话有效)",
|
||||
})
|
||||
ok, kind, detail = self._check_connectivity(host, port, username, pw)
|
||||
if not ok:
|
||||
hints = {
|
||||
"auth_failed": "用户名/密码错误,请核对凭据",
|
||||
"timeout": "代理不可达(超时),请检查地址与网络",
|
||||
"connect_failed": "连接被拒绝或不可达,请检查 host/port",
|
||||
"missing_dependency": "缺少 SOCKS 支持库,请安装:pip install pysocks",
|
||||
}
|
||||
return self.format_result(False, "代理不可达,未保存配置", {
|
||||
"错误类型": kind,
|
||||
"详情": detail[:300],
|
||||
"提示": hints.get(kind, "如确认配置无误可设置 skip_check=true 跳过验证"),
|
||||
})
|
||||
check_detail = detail
|
||||
|
||||
config = ConfigManager()
|
||||
saved = config.set_proxy(host, port, username, password, password_env, enabled, save)
|
||||
proxy_info = config.get_proxy_info()
|
||||
if proxy_info.get("password_env"):
|
||||
pw_source = "环境变量 " + proxy_info["password_env"]
|
||||
elif proxy_info.get("password"):
|
||||
pw_source = "仅本次会话(不落盘)"
|
||||
else:
|
||||
pw_source = "无"
|
||||
data = {
|
||||
"代理状态": "启用" if proxy_info["enabled"] else "禁用",
|
||||
"服务器": f"{proxy_info['host']}:{proxy_info['port']}",
|
||||
"认证": "有" if proxy_info.get("username") else "无",
|
||||
"密码来源": pw_source,
|
||||
"持久化": "已保存" if saved else "保存失败(仅本次会话有效,重启后丢失)"
|
||||
}
|
||||
if check_detail:
|
||||
data["连通性"] = check_detail
|
||||
# L2:落盘失败(磁盘只读等)如实报告,不再静默报"已保存"
|
||||
if save and not saved:
|
||||
return self.format_result(False, "代理已配置但保存失败(仅本次会话有效,重启后丢失)", data)
|
||||
return self.format_result(True, "代理设置成功", data)
|
||||
|
||||
|
||||
class GetProxyTool(Tool):
|
||||
"""查看当前代理配置"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "get_proxy"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return "查看当前 SOCKS5 代理配置信息"
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
def execute(self, **kwargs) -> str:
|
||||
config = ConfigManager()
|
||||
info = config.get_proxy_info()
|
||||
if info["enabled"]:
|
||||
data = {
|
||||
"代理状态": "启用",
|
||||
"服务器": f"{info['host']}:{info['port']}",
|
||||
"认证": "有" if info.get("username") else "无"
|
||||
}
|
||||
return self.format_result(True, "当前代理配置", data)
|
||||
else:
|
||||
return "❌ 当前代理未启用,所有请求走直连。"
|
||||
Loaded 100 of 255 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user