Compare commits
103
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
121e155248 | ||
|
|
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 | ||
|
|
d3a672a2a8 | ||
|
|
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 | ||
|
|
96d9f96e6d | ||
|
|
f606f7ec56 | ||
|
|
02d95a2bd9 | ||
|
|
82875b3021 | ||
|
|
4d6e1e1fd4 | ||
|
|
8fe2afdfde | ||
|
|
4b7773daec | ||
|
|
6309358262 |
No files matched your search
@@ -25,3 +25,6 @@ waf/certs/
|
||||
|
||||
# 运行时日志(HailoRT 等)
|
||||
hailort.log
|
||||
|
||||
# 录制产物(可复现:mtc DSL → 编译 → pybot --record)
|
||||
actpred/data_v3
|
||||
@@ -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(核心模块 · 进程托管服务)
|
||||
@@ -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` 失败转移),`--once` 支持无 WebSocket 的 CLI 一次性对话;详见 [aibox/README.md](aibox/README.md)。
|
||||
- **四轮安全审计修复(2026-08,Batch E/F/G/H/I 共 60+ 项)**:命令执行保护链(1MB/超时/中断)、凭据脱敏(代理密码不落盘、日志打码、Bearer 脱敏)、取消传播(子代理/情绪子代理/调度器超时)、DB 事务与 FTS5 转义、POSIX 进程组终止等;审计报告见 [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) |
|
||||
|
||||
|
||||
+133
-38
@@ -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` | 4 | 批次大小(默认分辨率 1920×1080 下 CPU/Pi 内存受限;小分辨率可调大) |
|
||||
| `--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,13 +226,16 @@ 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 实测(224×224,2026-08-13 复测)11.4ms/次含预处理
|
||||
`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。
|
||||
@@ -189,7 +266,8 @@ reshape)**无法被 DFC 解析**,必须用 `export_deploy.py` 的部署版
|
||||
|
||||
当前 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),形状与量化参数由 `hailo_infer.py` 从 vstream
|
||||
pred_token (1,1,1,128)(`--imu-dim`>0 时追加 imu (1,1,1,imu_dim),共 6 输入),
|
||||
形状与量化参数由 `hailo_infer.py` 从 vstream
|
||||
info 自动推导,不硬编码。
|
||||
|
||||
```bash
|
||||
@@ -208,25 +286,42 @@ python3 dfc_compile.py # 产物 checkpoints/actpred_pretrain.hef
|
||||
量化、编译全部通过;hailo10h 目标。
|
||||
编译要点与踩坑记录见 `dfc_compile.py` 文档字符串。
|
||||
|
||||
> **警告(2026-08-13 树莓派实测)**:首版 HEF 在 NPU 上**输出对全部
|
||||
> 5 个输入均不敏感**(全 0 / 全 255 / 真实数据输出完全一致,对照
|
||||
> resnet HEF 同 API 验证输入通道有效)。驱动升级 5.3.0 后复测依旧,
|
||||
> 判定为 DFC 编译侧问题(输入在编译图中被断开/常量化),并非推理封装
|
||||
> 问题。**该 HEF 不可用于推理**。已做如下修复,待 x86 重编译后复验:
|
||||
> - `_onnx_validate.py`:编译前数值验证 `actpred_deploy.onnx` 输入
|
||||
> 敏感性与 PyTorch 一致性(本次已 PASS,差异 1.5e-5)
|
||||
> - `dfc_compile.py`:parser/optimize 后自动检查输入层连通性,断开即报错
|
||||
> - `gen_calibration.py`:pred_token 常量校准退化(scale=18571.7,任何
|
||||
> 输入都被量化成 0)已改抖动校准,推理端喂学习值(`--pred-token`)
|
||||
> - `hailo_infer.py`:帧输入按 input_layerN 排序绑定(vstream 报告顺序
|
||||
> 实测为 1,5,4,3,2,按报告顺序绑定会 frame1/frame2 错位)
|
||||
> **验证记录(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 注释
|
||||
>
|
||||
> 重编译后用 `_npu_test.py` 做 4 角输入敏感性复验(6.3ms/次的封装
|
||||
> 性能基线已在新驱动 5.3.0 下测通)。
|
||||
> 验证链:`export_deploy.py`(token 量级)→ `_onnx_validate.py`(ONNX
|
||||
> 数值)→ `dfc_compile.py`(输入连通性)→ `_npu_test.py`(4 角敏感性)。
|
||||
> 遗留:全零历史角(z≈-3.5)超出校准输入范围(-2.7)会被截断,
|
||||
> 该角仅作连通性验证,非真实部署场景。
|
||||
|
||||
`hailo_infer.py` 的 `HailoPredictor` 适配 5 输入 HEF,输入布局与量化
|
||||
参数(含负零偏→ImageNet 归一化识别)从 vstream info 自动推导,单元测试
|
||||
含 fake hailo_platform 的 5 输入与乱序报告顺序用例。CLI 支持
|
||||
`--pred-token` 喂入学习值、`--stats` 喂入归一化统计。
|
||||
`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)
|
||||
+89
-23
@@ -4,7 +4,7 @@
|
||||
|
||||
4 角测试(真实/全零帧 × 真实/全零历史)+ 与 PyTorch 参考对比。
|
||||
HEF 尺寸/布局从 vstream info 自动推导,checkpoint 用
|
||||
checkpoints/pretrain_best_inference.pt。
|
||||
checkpoints/best_inference.pt(v4 onehot,fixed_range)。
|
||||
"""
|
||||
import sys
|
||||
import time
|
||||
@@ -20,8 +20,8 @@ from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
from hailo_infer import HailoPredictor
|
||||
|
||||
HEF_PATH = "checkpoints/actpred_pretrain.hef"
|
||||
CKPT = "checkpoints/pretrain_best_inference.pt"
|
||||
HEF_PATH = "checkpoints/actpred_v4.hef"
|
||||
CKPT = "checkpoints/best_inference.pt"
|
||||
IMG_DIR = "data/episode_0000/images"
|
||||
|
||||
hef = HEF(HEF_PATH)
|
||||
@@ -40,11 +40,11 @@ imgs_real = [np.array(Image.open(f"{IMG_DIR}/000{i}.jpg").convert("RGB"))
|
||||
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],
|
||||
[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)
|
||||
|
||||
@@ -55,6 +55,25 @@ 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([
|
||||
@@ -64,27 +83,56 @@ tf = transforms.Compose([
|
||||
])
|
||||
|
||||
|
||||
def torch_pred(imgs, hist, use_learned_pred=True):
|
||||
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():
|
||||
return model(x, h).squeeze(0).numpy()
|
||||
saved = model.pred_token.data.clone()
|
||||
model.pred_token.data.zero_()
|
||||
with torch.no_grad():
|
||||
out = model(x, h).squeeze(0).numpy()
|
||||
model.pred_token.data.copy_(saved)
|
||||
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
|
||||
|
||||
|
||||
with HailoPredictor(HEF_PATH) as p:
|
||||
# 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)
|
||||
b = p.predict(imgs_zero, hist_real, pred_token=pred_learned)
|
||||
c = p.predict(imgs_real, hist_zero, pred_token=pred_learned)
|
||||
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/次")
|
||||
@@ -96,13 +144,31 @@ with HailoPredictor(HEF_PATH) as p:
|
||||
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)
|
||||
ref_zero_pred = torch_pred(imgs_real, hist_real, False)
|
||||
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}")
|
||||
print(f"HEF vs torch(learned) max_diff={np.abs(a - ref).max():.4f}"
|
||||
f"(原始动作空间)")
|
||||
|
||||
ok = img_sens > 0.1 and hist_sens > 0.1
|
||||
print("\nPASS: HEF 输入连通,可继续与 torch 对比量化误差"
|
||||
if ok else "\nFAIL: HEF 输出对输入不敏感,编译仍有问题")
|
||||
+56
-10
@@ -24,6 +24,23 @@ 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)
|
||||
@@ -48,9 +65,28 @@ hist_real = np.array([
|
||||
], 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"""
|
||||
@@ -58,35 +94,45 @@ def frames_nhwc(imgs):
|
||||
return [chw[i].transpose(1, 2, 0)[None] for i in range(len(imgs))]
|
||||
|
||||
|
||||
def run_onnx(imgs, hist, pred):
|
||||
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):
|
||||
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).squeeze(0).numpy()
|
||||
return model(x, hh, ii).squeeze(0).numpy()
|
||||
|
||||
|
||||
a = run_onnx(imgs_real, hist_real, pred_learned)
|
||||
b = run_onnx(imgs_zero, hist_real, pred_learned)
|
||||
c = run_onnx(imgs_real, hist_zero, pred_learned)
|
||||
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)
|
||||
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}")
|
||||
|
||||
ref = run_torch(imgs_real, hist_real)
|
||||
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 = (np.abs(a - b).max() > 1e-4 and np.abs(a - c).max() > 1e-4
|
||||
and np.abs(a - ref).max() < 1e-4)
|
||||
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 分类目标互转
|
||||
|
||||
分箱设计(actpred_onehot_multihead_dev_plan_20260814.md):
|
||||
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
|
||||
+180
-13
@@ -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,18 +58,32 @@ class ModelConfig:
|
||||
|
||||
@dataclass
|
||||
class TrainConfig:
|
||||
# 默认输入分辨率 1920×1080,batch=32 在 CPU/Pi 上内存不可行;
|
||||
# 高分辨率训练取小 batch,小分辨率可用 --batch-size 调大
|
||||
# 默认输入分辨率 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
|
||||
@@ -70,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)
|
||||
@@ -82,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,
|
||||
@@ -108,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),
|
||||
@@ -129,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),
|
||||
@@ -139,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:
|
||||
@@ -150,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)}")
|
||||
@@ -186,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"),
|
||||
@@ -198,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:
|
||||
|
||||
+67
-42
@@ -25,6 +25,7 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
@@ -33,8 +34,16 @@ import tempfile
|
||||
|
||||
|
||||
NET_NAME = "actpred_deploy"
|
||||
START_NODES = ["frame0", "frame1", "frame2", "action_history", "pred_token"]
|
||||
INPUT_FORMATS = ["NHWC", "NHWC", "NHWC", "NHWC", "NHWC"]
|
||||
|
||||
|
||||
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):
|
||||
@@ -43,56 +52,63 @@ def hailo(*args, venv_bin: str):
|
||||
return subprocess.run(cmd, env=env, check=True)
|
||||
|
||||
|
||||
def har_input_layer_names(har_path: str) -> list:
|
||||
"""从 HAR(tar 包)中读 hn 文件,按 input_layerN 排序返回层名"""
|
||||
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")]
|
||||
with open(os.path.join(tmp, hn_files[0])) as f:
|
||||
hn = json.load(f)
|
||||
layers = [(name, l) for name, l in hn["layers"].items()
|
||||
if l["type"] == "input_layer"]
|
||||
layers.sort(key=lambda x: x[0])
|
||||
return [name for name, _ in layers]
|
||||
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 中输入层是否都有下游边
|
||||
"""检查 HAR 图中每个输入层是否可达任一输出层
|
||||
|
||||
输入层被优化断开时编译照常成功,但 HEF 输出恒定值且无任何报错
|
||||
(2026-08-13 实测踩坑),必须在编译前拦截。.edges 文件缺失或格式
|
||||
无法识别时仅告警不阻断。
|
||||
(2026-08-13 实测踩坑),必须在编译前拦截。DFC 5.3 的 HAR 不含
|
||||
.edges 文件,连通性改从主 hn 图的 layer output 边 BFS 判定。
|
||||
"""
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
with tarfile.open(har_path) as tf:
|
||||
tf.extractall(tmp)
|
||||
edge_files = [f for f in os.listdir(tmp) if f.endswith(".edges")]
|
||||
if not edge_files:
|
||||
print(f"WARNING: {har_path} 无 .edges 文件,跳过输入连通性检查")
|
||||
return
|
||||
with open(os.path.join(tmp, edge_files[0])) as f:
|
||||
edges = json.load(f)
|
||||
|
||||
if not isinstance(edges, list):
|
||||
print(f"WARNING: .edges 不是列表({type(edges).__name__}),"
|
||||
f"跳过连通性检查")
|
||||
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
|
||||
adjacency = set()
|
||||
for e in edges:
|
||||
if isinstance(e, dict) and "from" in e:
|
||||
adjacency.add(e["from"])
|
||||
elif isinstance(e, (list, tuple)) and len(e) == 2:
|
||||
adjacency.add(e[0])
|
||||
else:
|
||||
print(f"WARNING: 无法识别 .edges 条目格式: {e!r},跳过连通性检查")
|
||||
return
|
||||
disconnected = [n for n in input_layers if n not in adjacency]
|
||||
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"错误: {har_path} 中输入层不可达输出层(输入被优化断开,"
|
||||
f"编译出的 HEF 将输出恒定值): {disconnected}")
|
||||
print(f"输入连通性检查通过: {len(input_layers)} 个输入层均有下游边")
|
||||
print(f"输入连通性检查通过: {len(input_layers)} 个输入层均可达输出层")
|
||||
|
||||
|
||||
def main():
|
||||
@@ -103,6 +119,12 @@ def main():
|
||||
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()
|
||||
@@ -111,6 +133,9 @@ def main():
|
||||
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")
|
||||
@@ -121,16 +146,16 @@ def main():
|
||||
print("[1/3] parser: ONNX → HAR")
|
||||
hailo("parser", "onnx", args.onnx, "--hw-arch", args.hw_arch, "-y",
|
||||
"--har-path", har_raw,
|
||||
"--start-node-names", *START_NODES,
|
||||
"--start-node-names", *nodes,
|
||||
"--end-node-names", "predicted_action",
|
||||
"--input-format",
|
||||
*[f"{n}={f}" for n, f in zip(START_NODES, INPUT_FORMATS)],
|
||||
*[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(START_NODES):
|
||||
sys.exit(f"错误: HAR 输入层数 {len(names)} != {len(START_NODES)},解析可能被截断")
|
||||
if len(names) != len(nodes):
|
||||
sys.exit(f"错误: HAR 输入层数 {len(names)} != {len(nodes)},解析可能被截断")
|
||||
check_input_connectivity(har_raw, names)
|
||||
|
||||
print("[2/3] optimize: 量化校准")
|
||||
|
||||
@@ -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()
|
||||
+143
-12
@@ -12,13 +12,26 @@ Reshape 序列。此脚本逐项改造(权重数值完全不变):
|
||||
"""
|
||||
|
||||
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)"""
|
||||
@@ -78,6 +91,9 @@ class DeployWrapper(nn.Module):
|
||||
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)])
|
||||
@@ -91,11 +107,23 @@ class DeployWrapper(nn.Module):
|
||||
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())
|
||||
self.pool = nn.AvgPool2d((m._image_height // 16, m._image_width // 16))
|
||||
# 固定 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)
|
||||
@@ -109,21 +137,86 @@ class DeployWrapper(nn.Module):
|
||||
x = self.proj_convs[i](x) # 1x1 conv ≡ 池化后的 img_proj
|
||||
return self.pool(x) # (1, d, 1, 1) NCHW
|
||||
|
||||
def forward(self, frame0: torch.Tensor, frame1: torch.Tensor, frame2: torch.Tensor,
|
||||
action_history: torch.Tensor, pred_token: torch.Tensor) -> torch.Tensor:
|
||||
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 = [frame0, frame1, frame2]
|
||||
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)
|
||||
tokens = torch.cat([img_tokens, act_tokens, pred_token], dim=2) # (1, 1, S, 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")
|
||||
@@ -143,21 +236,48 @@ def main():
|
||||
|
||||
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)
|
||||
ref = wrapper(*dummy_frames, dummy_history, dummy_pred, *extra_inputs)
|
||||
|
||||
torch.onnx.export(
|
||||
wrapper, (*dummy_frames, dummy_history, dummy_pred),
|
||||
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"],
|
||||
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)")
|
||||
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
|
||||
@@ -166,8 +286,13 @@ def main():
|
||||
|
||||
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))
|
||||
print(f"wrapper 与原模型最大差异: {(ref.squeeze(1) - orig_out).abs().max():.6e}")
|
||||
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
|
||||
@@ -176,8 +301,14 @@ def main():
|
||||
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]
|
||||
print(f"ONNX 一致性验证通过(最大差异: {np.abs(ort_out - ref.numpy()).max():.6e})")
|
||||
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
|
||||
|
||||
|
||||
+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)")
|
||||
|
||||
@@ -22,6 +22,8 @@
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -53,15 +55,38 @@ def build_action_history_samples(mt_dir, n_samples, history_len, action_dim, rng
|
||||
return np.stack([windows[i] for i in idx]).astype(np.float32)
|
||||
|
||||
|
||||
def build_image_samples(n_samples, frame_len, height, width, rng):
|
||||
# 存 uint8 原始像素值(0~255),加载时再做 ToTensor+Normalize,节省磁盘
|
||||
return rng.integers(0, 256,
|
||||
size=(n_samples, frame_len, 3, height, width),
|
||||
dtype=np.uint8)
|
||||
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)
|
||||
@@ -77,6 +102,34 @@ def build_deploy_set(args):
|
||||
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)
|
||||
@@ -92,14 +145,24 @@ def build_deploy_set(args):
|
||||
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([*frames, history, pred])})
|
||||
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():
|
||||
@@ -133,6 +196,9 @@ def main():
|
||||
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}")
|
||||
|
||||
+335
-55
@@ -6,9 +6,10 @@ HEF 由 actpred ONNX 经 Hailo DFC 编译生成(量化校准后)。本封装
|
||||
vstream 名称/形状/量化参数,全部从 HEF 的 input/output vstream info 动态推导,
|
||||
支持两种布局:
|
||||
|
||||
新版 5 输入(export_deploy.py + dfc_compile.py 产物):
|
||||
新版 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)
|
||||
@@ -30,27 +31,51 @@ import contextlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
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)
|
||||
|
||||
# 与 predict.parse_action_history 相同逻辑;此处内联以保持
|
||||
# 本模块零 torch 依赖(HAT 部署环境只装 hailo_platform + numpy + PIL)
|
||||
# 异步推理缓冲数:≥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:
|
||||
@@ -69,9 +94,39 @@ class HailoPredictor:
|
||||
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())
|
||||
@@ -87,13 +142,13 @@ class HailoPredictor:
|
||||
raise ValueError(
|
||||
f"HEF 应只有一个输出 vstream,实际 {len(self._out_infos)}")
|
||||
self._out_info = self._out_infos[0]
|
||||
self._frame_infos, self._history_info, self._pred_info = \
|
||||
self._classify_inputs()
|
||||
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)
|
||||
@@ -110,14 +165,14 @@ class HailoPredictor:
|
||||
# ---------------- 输入分类与形状推导 ----------------
|
||||
|
||||
def _classify_inputs(self):
|
||||
"""将 HEF 输入分为 帧列表 / 动作历史 / pred_token(可能为 None)
|
||||
"""将 HEF 输入分为 帧列表 / 动作历史 / pred_token / imu(后两者可为 None)
|
||||
|
||||
多帧输入的绑定顺序按 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)
|
||||
out_dim = tuple(self._out_info.shape)[-1]
|
||||
# 新版:每帧独立 (H,W,3) 输入
|
||||
frames = [i for i in ins if len(i.shape) == 3
|
||||
and i.shape[2] == 3 and i.shape[0] > 3]
|
||||
@@ -127,24 +182,27 @@ class HailoPredictor:
|
||||
# 旧版:单一 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 = None, None
|
||||
history, pred, imu = None, None, None
|
||||
for i in rest:
|
||||
s = tuple(i.shape)
|
||||
if len(s) == 3 and s[-1] == out_dim:
|
||||
history = i
|
||||
elif len(s) == 3 and s[1] == 1:
|
||||
pred = i
|
||||
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))")
|
||||
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(
|
||||
f"HEF 输入中未找到动作历史输入(输出维度 {out_dim}),"
|
||||
f"HEF 输入中未找到动作历史输入(输出维度 "
|
||||
f"{tuple(self._out_info.shape)}),"
|
||||
f"实际: {[(i.name, tuple(i.shape)) for i in ins]}")
|
||||
return frames, history, pred
|
||||
return frames, history, pred, imu
|
||||
|
||||
def _derive_dims(self):
|
||||
# 帧输入布局
|
||||
@@ -174,19 +232,48 @@ class HailoPredictor:
|
||||
self.frame_len = len(self._frame_infos)
|
||||
self._image_layout = "nhwc_frames"
|
||||
self._frame_quant = self._frame_infos[0].quant_info
|
||||
# 帧量化零偏为负 ⇒ 校准数据是 ImageNet 归一化后的浮点,
|
||||
# 推理端需先归一化再量化;零偏为 0 则为原始像素直通
|
||||
# 帧输入是否需 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 = \
|
||||
self._frame_quant.qp_zp * self._frame_quant.qp_scale < 0
|
||||
(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} 不一致")
|
||||
|
||||
# ---------------- 生命周期 ----------------
|
||||
|
||||
@@ -202,13 +289,8 @@ class HailoPredictor:
|
||||
|
||||
# ---------------- 推理 ----------------
|
||||
|
||||
def predict(self, images, action_history: np.ndarray,
|
||||
pred_token: np.ndarray = None) -> np.ndarray:
|
||||
"""images: 长度为 frame_len 的 RGB 图像列表(旧→新,末帧为当前帧)
|
||||
|
||||
返回 (action_dim,) float32 动作向量。pred_token 仅新版 5 输入 HEF
|
||||
需要;默认喂零,抖动校准的新 HEF 建议喂 checkpoint 学习值。
|
||||
"""
|
||||
def _prepare_inputs(self, images, action_history):
|
||||
"""输入校验 + 动作历史归一化(predict/predict_async 共用)"""
|
||||
if not isinstance(images, (list, tuple)):
|
||||
images = [images]
|
||||
if len(images) != self.frame_len:
|
||||
@@ -221,14 +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()
|
||||
|
||||
frame_float = np.stack([self._frame_float(img) for img in images])
|
||||
self._bind_frames(bindings, frame_float)
|
||||
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:
|
||||
@@ -241,6 +335,8 @@ class HailoPredictor:
|
||||
|
||||
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:
|
||||
@@ -251,6 +347,16 @@ class HailoPredictor:
|
||||
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)
|
||||
@@ -258,19 +364,120 @@ 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 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
|
||||
@@ -281,6 +488,43 @@ class HailoPredictor:
|
||||
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
|
||||
@@ -318,23 +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()
|
||||
|
||||
pred_token = np.load(args.pred_token) if args.pred_token else None
|
||||
with HailoPredictor(args.hef, stats_path=args.stats) as predictor:
|
||||
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, pred_token=pred_token)
|
||||
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])
|
||||
@@ -30,15 +30,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
|
||||
+1927
-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)
|
||||
+273
-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,74 @@ 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 解耦)
|
||||
├── tools/ # 工具系统包(A 批 + B 批共 17 个工具类)
|
||||
│ ├── 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
|
||||
├── tests/ # 单元测试套件(441 项全过)
|
||||
│ ├── 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_model_config.py # 模型接入点管理工具(17 项)
|
||||
├── skills/ # 内置扩展技能(*.py 可执行脚本 / *.md Markdown 技能)
|
||||
│ ├── hello-world.py
|
||||
│ ├── aes-crypto.py
|
||||
@@ -79,6 +134,7 @@ aibox/
|
||||
│ ├── ssh-exec.py
|
||||
│ ├── sftp-transfer.py
|
||||
│ └── ...
|
||||
├── prompts/ # 系统提示词源文件(运行期复制到 ~/.aibox/prompts/)
|
||||
├── SKILLS_FORMAT.md # 技能文件格式说明
|
||||
├── _run.sh # 启动脚本(被 core 守护时调用)
|
||||
├── _runcli.sh # 命令行启动脚本
|
||||
@@ -87,6 +143,25 @@ 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
|
||||
# 当前:264 项全过
|
||||
```
|
||||
|
||||
运行期数据存放于 `~/.aibox/`:
|
||||
|
||||
```
|
||||
@@ -136,7 +211,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 +221,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)工具集
|
||||
@@ -452,15 +527,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,
|
||||
@@ -512,11 +641,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 +656,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 +679,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 +708,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 +798,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 +831,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 +911,85 @@ elif command == "/mycommand":
|
||||
|
||||
## 📄 更新日志
|
||||
|
||||
### 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`)**(分三批落地,开发文档见 [docs/reports/aibox_cli_once_mode_dev_doc.md](../docs/reports/aibox_cli_once_mode_dev_doc.md)):
|
||||
- **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 分组装配工具集
|
||||
|
||||
+632
-4662
File diff suppressed because it is too large.
Load diff
+83
-18
@@ -8,6 +8,7 @@ common.config - aibox 配置管理器
|
||||
|
||||
import os
|
||||
import json
|
||||
import copy
|
||||
from typing import Dict, Optional
|
||||
|
||||
from .i18n import i18n
|
||||
@@ -64,14 +65,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, # 最大历史消息数量
|
||||
@@ -165,6 +188,7 @@ class ConfigManager:
|
||||
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 _merge_configs(self, default: Dict, user: Dict) -> Dict:
|
||||
result = default.copy()
|
||||
@@ -198,6 +222,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 +252,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 +274,45 @@ 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
|
||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
||||
with open(config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(safe_config, f, 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 ""
|
||||
}
|
||||
+490
-35
@@ -22,27 +22,48 @@ 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._~+/=-]+")
|
||||
|
||||
|
||||
def _sanitize_log_msg(msg: str) -> str:
|
||||
"""日志脱敏:替换错误消息中的 Bearer 令牌。"""
|
||||
if not msg:
|
||||
return msg
|
||||
return _AUTH_HEADER_RE.sub("Bearer ***", 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 +76,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 +103,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 +126,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 +145,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 +162,35 @@ 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()
|
||||
if response.status_code == 429 or response.status_code >= 500:
|
||||
# L4a:瞬时故障显式抛出——FailoverClient 据此转移下一档案;
|
||||
# 无包装器时由调用方 catch-all 接住,下游仍走失败分支(语义等价)
|
||||
raise TransientLLMError(
|
||||
f"LLM 请求瞬时失败: HTTP {response.status_code}: {response.text[:200]}",
|
||||
retry_after=_retry_after_of(response))
|
||||
self._log("error", f"LLM 压缩请求失败: HTTP {response.status_code}")
|
||||
return None
|
||||
except TransientLLMError:
|
||||
raise # L4a:瞬时故障上抛供 FailoverClient 转移,不吞
|
||||
except (requests.exceptions.ConnectionError, requests.exceptions.Timeout):
|
||||
raise # L4a:连接类瞬时故障同样上抛(原被吞成 None,无法转移)
|
||||
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,14 +219,19 @@ class DeepSeekClient(LLMClient):
|
||||
self._current_response = response
|
||||
|
||||
if response.status_code != 200:
|
||||
if response.status_code == 429 or response.status_code >= 500:
|
||||
# L4b:瞬时故障显式分类(此时尚未产出任何事件,可安全转移重放)
|
||||
raise TransientLLMError(
|
||||
f"LLM 流式请求瞬时失败: HTTP {response.status_code}: {response.text[:200]}",
|
||||
retry_after=_retry_after_of(response))
|
||||
raise RuntimeError(
|
||||
f"HTTP {response.status_code}: {response.text}")
|
||||
f"HTTP {response.status_code}: {response.text[:200]}")
|
||||
|
||||
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:]
|
||||
@@ -198,11 +252,19 @@ class DeepSeekClient(LLMClient):
|
||||
finally:
|
||||
self._current_response = None
|
||||
|
||||
# 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,17 +281,251 @@ 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
|
||||
|
||||
def _cooling(self, client) -> bool:
|
||||
entry = self._fail_at.get(client.name)
|
||||
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[client.name] = (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, requests.exceptions.RequestException) 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:流式转移——仅首事件前故障可安全重放(对用户不可见)。"""
|
||||
self._cancelled = False # 新一轮请求,清除上一轮中断标志
|
||||
self._call_started = time.monotonic()
|
||||
self._stats["requests"] += 1
|
||||
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, requests.exceptions.RequestException) 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"])
|
||||
key_path = os.path.join(home, ".cert", "deepseek.key")
|
||||
@@ -240,26 +536,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 +721,7 @@ def create_llm_client(config, api_key: Optional[str] = None,
|
||||
config=config,
|
||||
logger=logger,
|
||||
)
|
||||
return client
|
||||
|
||||
# 其他 OpenAI 兼容厂商:可直接复用 DeepSeekClient(继承其协议解析),
|
||||
# 或在此接入专有实现。新增厂商时在此扩展即可。
|
||||
|
||||
@@ -0,0 +1,415 @@
|
||||
#!/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
|
||||
|
||||
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 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
|
||||
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": profile.get("fallback") or [],
|
||||
})
|
||||
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}"]
|
||||
|
||||
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:
|
||||
config._config.setdefault("api", {})["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") or []
|
||||
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:
|
||||
config._config.setdefault("api", {}).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:
|
||||
return f"❌ 档案 '{name}' 验证失败:{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)})")
|
||||
|
||||
|
||||
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"环境变量 {p['api_key_env']}" if p["api_key_env"] else "免密"
|
||||
fb = f",失败转移 → {', '.join(p['fallback'])}" if p["fallback"] else ""
|
||||
print(f" {p['name']}{mark} [{p['display_name']}]")
|
||||
print(f" 地址: {p['api_url']} 模型: {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":
|
||||
api_key_env = _ask("环境变量名", "DEEPSEEK_API_KEY")
|
||||
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
|
||||
|
||||
msgs = add_profile(config, name, api_url, model,
|
||||
api_key_env=api_key_env, display_name=display_name,
|
||||
temperature=float(temperature_raw) if temperature_raw else None,
|
||||
timeout=int(timeout_raw) if timeout_raw else 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:
|
||||
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")
|
||||
|
||||
|
||||
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,37 @@
|
||||
"""
|
||||
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
|
||||
|
||||
__all__ = [
|
||||
"CommandExecutor",
|
||||
"SubprocessCommandExecutor",
|
||||
"PromptManager",
|
||||
"MemoryDatabase",
|
||||
"MemoDatabase",
|
||||
"SkillManager",
|
||||
"ConversationHistory",
|
||||
"ConversationCompressor",
|
||||
]
|
||||
@@ -0,0 +1,524 @@
|
||||
#!/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 _is_range_safe(self, messages: List[Dict], start: int, end: int, protected: set) -> bool:
|
||||
"""检查消息范围是否安全 - 简化版,不检查tool完整性"""
|
||||
return True # 不保护tool消息,允许压缩
|
||||
|
||||
def _get_protected_indices(self, messages: List[Dict]) -> set:
|
||||
"""
|
||||
找出所有受保护的消息索引 - 简化版,只保护最近的对话轮次
|
||||
"""
|
||||
protected_indices = set()
|
||||
|
||||
# ========== 1. 保护最近的消息(由配置决定) ==========
|
||||
keep_recent = constants.COMPRESS_KEEP_RECENT
|
||||
if keep_recent > 0:
|
||||
recent_start = max(0, len(messages) - keep_recent)
|
||||
for i in range(recent_start, len(messages)):
|
||||
protected_indices.add(i)
|
||||
|
||||
# ========== 2. 保护前2个对话轮次 ==========
|
||||
conversation_turns = self._build_conversation_turns(messages)
|
||||
|
||||
if len(conversation_turns) >= 2:
|
||||
# 保护第1轮
|
||||
first_turn = conversation_turns[0]
|
||||
protected_indices.add(first_turn["user_index"])
|
||||
if first_turn["assistant_index"] is not None:
|
||||
protected_indices.add(first_turn["assistant_index"])
|
||||
for tool_idx in first_turn["tool_indices"]:
|
||||
protected_indices.add(tool_idx)
|
||||
|
||||
# 保护第2轮
|
||||
second_turn = conversation_turns[1]
|
||||
protected_indices.add(second_turn["user_index"])
|
||||
if second_turn["assistant_index"] is not None:
|
||||
protected_indices.add(second_turn["assistant_index"])
|
||||
for tool_idx in second_turn["tool_indices"]:
|
||||
protected_indices.add(tool_idx)
|
||||
|
||||
# 调试日志
|
||||
if self.logger.level >= 3:
|
||||
self.logger.debug(f"受保护的消息索引: {sorted(protected_indices)}")
|
||||
|
||||
return protected_indices
|
||||
|
||||
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 消息"""
|
||||
if not messages:
|
||||
return messages
|
||||
|
||||
# 找出所有 assistant 消息的 tool_call_ids
|
||||
assistant_tool_ids = set()
|
||||
for msg in messages:
|
||||
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
||||
for tc in msg.get("tool_calls", []):
|
||||
tc_id = tc.get("id")
|
||||
if tc_id:
|
||||
assistant_tool_ids.add(tc_id)
|
||||
|
||||
# 过滤掉没有对应 assistant 的 tool 消息
|
||||
filtered = []
|
||||
for msg in messages:
|
||||
if msg.get("role") == "tool":
|
||||
tool_call_id = msg.get("tool_call_id")
|
||||
if tool_call_id and tool_call_id not in assistant_tool_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,371 @@
|
||||
#!/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
|
||||
|
||||
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,历史内容丢失)"""
|
||||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||||
with open(tmp, 'w', encoding='utf-8') as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
os.replace(tmp, path)
|
||||
|
||||
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,775 @@
|
||||
#!/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 _init_database(self):
|
||||
"""初始化数据库表"""
|
||||
with sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('DELETE FROM memories WHERE id = ?', (memory_id,))
|
||||
deleted = cursor.rowcount > 0
|
||||
|
||||
if deleted:
|
||||
cursor.execute('DELETE FROM memories_fts WHERE rowid = ?', (memory_id,))
|
||||
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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
placeholders = ','.join(['?'] * len(memory_ids))
|
||||
cursor.execute(f'DELETE FROM memories WHERE id IN ({placeholders})', memory_ids)
|
||||
deleted_count = cursor.rowcount
|
||||
|
||||
if deleted_count > 0:
|
||||
cursor.execute(f'DELETE FROM memories_fts WHERE rowid IN ({placeholders})', memory_ids)
|
||||
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 _init_database(self):
|
||||
"""初始化备忘录表"""
|
||||
with sqlite3.connect(self.db_path) as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
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和任务类型"""
|
||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||
tags_str = ",".join(tags) if tags else ""
|
||||
|
||||
with sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.cursor()
|
||||
# I 批(审计 2026-08-18):多进程同库(CLI --memo-add + 服务调度)
|
||||
# 写竞争时默认 5s 后直接报 database is locked——busy_timeout 让
|
||||
# SQLite 等待锁释放(最长 10s)而不是立即失败
|
||||
cursor.execute("PRAGMA busy_timeout = 10000")
|
||||
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']:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
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' and memo['repeat_interval'] > 0:
|
||||
unit = memo.get('repeat_interval_unit', 'days')
|
||||
if unit == 'minutes':
|
||||
# 每N分钟重复:当前时间 + N分钟
|
||||
next_time = current_time + timedelta(minutes=memo['repeat_interval'])
|
||||
elif unit == 'hours':
|
||||
# 每N小时重复:当前时间 + N小时
|
||||
next_time = current_time + timedelta(hours=memo['repeat_interval'])
|
||||
else: # days
|
||||
# 每N天重复:当前时间 + N天
|
||||
next_time = current_time + timedelta(days=memo['repeat_interval'])
|
||||
else:
|
||||
return None
|
||||
|
||||
# 检查是否超过重复结束时间
|
||||
if memo['repeat_end_time']:
|
||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||
if 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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 sqlite3.connect(self.db_path) 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', 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,
|
||||
created_by="user",
|
||||
source="websocket",
|
||||
is_immediate=is_immediate
|
||||
)
|
||||
|
||||
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": datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT) < now if memo['reminder_time'] else False,
|
||||
"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,413 @@
|
||||
#!/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
|
||||
|
||||
|
||||
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()
|
||||
|
||||
def _terminate_tree(self, proc, force: bool = False):
|
||||
"""终止整个进程组(L9:POSIX 下 shell=True 只杀 shell 时,孙进程
|
||||
(`(x &)` 等)仍持管道写端——reader 阻塞无 EOF、fd/线程悬挂至孙
|
||||
进程退出,写满 64KB 管道缓冲后孙进程也卡死。start_new_session 使
|
||||
shell 成为会话首进程,按组信号可一次覆盖整棵进程树)。
|
||||
组 ID 在 Popen 后立即保存:shell 死亡后 os.getpgid(pid) 会抛
|
||||
ProcessLookupError,届时无法再从进程反查组。
|
||||
Windows 无进程组概念,退回单进程 terminate/kill。"""
|
||||
try:
|
||||
if os.name == "posix":
|
||||
sig = signal.SIGKILL if force else signal.SIGTERM
|
||||
with self._process_lock:
|
||||
pgid = self._current_pgid
|
||||
if pgid is None:
|
||||
pgid = os.getpgid(proc.pid)
|
||||
os.killpg(pgid, sig)
|
||||
elif force:
|
||||
proc.kill()
|
||||
else:
|
||||
proc.terminate()
|
||||
except Exception:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def cancel(self):
|
||||
"""取消当前正在运行的命令(kill子进程组)。线程安全。"""
|
||||
with self._process_lock:
|
||||
proc = self._current_process
|
||||
if proc is not None:
|
||||
try:
|
||||
if proc.poll() is None:
|
||||
self._terminate_tree(proc, force=True)
|
||||
self.logger.info("已终止正在运行的子进程组")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=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 阻塞不再影响本判断)
|
||||
if timeout > 0 and time.time() - start_time > timeout:
|
||||
self._terminate_tree(process)
|
||||
timeout_occurred = True
|
||||
break
|
||||
|
||||
# 检查输出大小
|
||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
||||
truncated = True
|
||||
self._terminate_tree(process)
|
||||
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)
|
||||
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)
|
||||
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 "❌ 命令执行功能已禁用"
|
||||
|
||||
# 默认工作目录为workspace,避免在启动目录散落文件
|
||||
if work_dir is None or work_dir == ".":
|
||||
work_dir = constants.WORKSPACE_DIR
|
||||
|
||||
self.logger.info(f"执行命令: {command} (work_dir={work_dir})")
|
||||
|
||||
# 准备环境变量
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
exec_env.update(env)
|
||||
|
||||
process = 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,
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
if os.name == "posix":
|
||||
try:
|
||||
self._current_pgid = os.getpgid(process.pid)
|
||||
except Exception:
|
||||
self._current_pgid = None
|
||||
|
||||
start_time = time.time()
|
||||
output_str, timeout_occurred, truncated = self._read_with_limit(process, timeout, start_time)
|
||||
|
||||
# 确保进程终止
|
||||
if process.poll() is None:
|
||||
self._terminate_tree(process, force=timeout_occurred or truncated)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
self._terminate_tree(process, force=True)
|
||||
|
||||
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)
|
||||
# 清除当前进程引用
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
self._current_pgid = None
|
||||
|
||||
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"))
|
||||
|
||||
# 准备环境变量
|
||||
exec_env = os.environ.copy()
|
||||
if env:
|
||||
exec_env.update(env)
|
||||
|
||||
process = 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,
|
||||
env=exec_env,
|
||||
cwd=os.path.expanduser(work_dir),
|
||||
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||
start_new_session=(os.name == "posix")
|
||||
)
|
||||
self.process = process
|
||||
# 保存进程引用以便cancel()可以kill它
|
||||
with self._process_lock:
|
||||
self._current_process = process
|
||||
if os.name == "posix":
|
||||
try:
|
||||
self._current_pgid = os.getpgid(process.pid)
|
||||
except Exception:
|
||||
self._current_pgid = None
|
||||
|
||||
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)
|
||||
|
||||
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)
|
||||
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)
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except:
|
||||
self._terminate_tree(process, force=True)
|
||||
self.process = None
|
||||
with self._process_lock:
|
||||
self._current_process = None
|
||||
self._current_pgid = None
|
||||
@@ -0,0 +1,325 @@
|
||||
#!/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
|
||||
|
||||
|
||||
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):
|
||||
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}
|
||||
"""
|
||||
with open(identity_path, 'w', encoding='utf-8') as f:
|
||||
f.write(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}
|
||||
"""
|
||||
with open(status_path, 'w', encoding='utf-8') as f:
|
||||
f.write(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}
|
||||
"""
|
||||
with open(soul_path, 'w', encoding='utf-8') as f:
|
||||
f.write(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,230 @@
|
||||
#!/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):第一行为简介(# 标题或普通文本),其余为技能内容
|
||||
"""
|
||||
|
||||
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)。"""
|
||||
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()
|
||||
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 文件)。"""
|
||||
skills_dir = constants.SKILLS_DIR
|
||||
self.skills = {}
|
||||
|
||||
if not os.path.exists(skills_dir):
|
||||
self.logger.info(f"技能目录不存在: {skills_dir}")
|
||||
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('.')
|
||||
]
|
||||
|
||||
for skill_file in skill_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:
|
||||
# Markdown 格式技能:直接读取内容
|
||||
summary, content = self._read_markdown_skill(skill_path)
|
||||
if content:
|
||||
self.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}"
|
||||
|
||||
self.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.logger.info(f"共加载 {len(self.skills)} 个技能({len([s for s in self.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,408 @@
|
||||
"""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")
|
||||
with mock.patch("builtins.open", 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)
|
||||
with mock.patch("builtins.open", 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,227 @@
|
||||
"""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"))
|
||||
mid = db.add_memo(title="格式异常", content="c", reminder_time="不是时间")
|
||||
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,198 @@
|
||||
"""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, {})
|
||||
with mock.patch("builtins.open", 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,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,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"]
|
||||
+61
-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,55 @@ 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
|
||||
|
||||
__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",
|
||||
]
|
||||
@@ -0,0 +1,486 @@
|
||||
#!/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 os
|
||||
import queue
|
||||
import shlex
|
||||
from datetime import datetime
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
from common import ConfigManager, constants
|
||||
from .base import Tool
|
||||
|
||||
|
||||
# ==================== 文件操作工具基类 ====================
|
||||
|
||||
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))
|
||||
quoted = shlex.quote(abs_path)
|
||||
|
||||
# 构造等价的命令行读取命令
|
||||
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 = shlex.quote(abs_path)
|
||||
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 = shlex.quote(abs_path)
|
||||
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]
|
||||
return f"--socks5 {shlex.quote(host_port)}"
|
||||
|
||||
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 = shlex.quote(url)
|
||||
timeout = constants.WEB_TIMEOUT
|
||||
user_agent = shlex.quote(constants.WEB_USER_AGENT)
|
||||
proxy = self._proxy_curl_args()
|
||||
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 参数"
|
||||
cmd = f"{base} {quoted_url} | grep -c {shlex.quote(keyword)}"
|
||||
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,820 @@
|
||||
#!/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:
|
||||
"""删除任务"""
|
||||
confirm = kwargs.get("confirm", False)
|
||||
if not confirm:
|
||||
return "❌ 需要设置 confirm=True 来确认删除操作"
|
||||
|
||||
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"
|
||||
|
||||
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,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 "❌ 当前代理未启用,所有请求走直连。"
|
||||
+72
-11
@@ -11,10 +11,26 @@ import threading
|
||||
import time
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import Logger, i18n
|
||||
from common import Logger, constants, i18n
|
||||
|
||||
from .base import Tool
|
||||
|
||||
# 工具参数名命中这些关键词时,debug 日志里打码(防密码/密钥进日志文件)
|
||||
# L11(审计 2026-08-18):精确匹配漏掉 apiKey/access_token/passwd 等变体,
|
||||
# 改包含匹配——日志场景下过度打码可接受,漏打码才是问题。
|
||||
_SENSITIVE_PARAM_KEYS = (
|
||||
"password", "passwd", "api_key", "apikey", "token", "secret",
|
||||
"authorization", "credential",
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_kwargs(kwargs: Dict) -> Dict:
|
||||
"""复制一份参数并把敏感键的值替换为 ***(仅用于日志输出)"""
|
||||
return {
|
||||
k: ("***" if any(s in str(k).lower() for s in _SENSITIVE_PARAM_KEYS) else v)
|
||||
for k, v in kwargs.items()
|
||||
}
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
"""工具注册器 - 支持技能热重载"""
|
||||
@@ -25,6 +41,8 @@ class ToolRegistry:
|
||||
self.logger = logger or Logger("ToolRegistry")
|
||||
self.current_output_queue = None # 用于流式输出的队列
|
||||
self.skill_manager: Optional['SkillManager'] = None # 技能管理器引用
|
||||
# 单工具流式执行总上限(工具自身秒级超时之外的兜底;测试可调小)
|
||||
self.stream_tool_timeout = 3600
|
||||
|
||||
def register(self, tool: Tool, is_builtin: bool = False):
|
||||
"""注册工具,可以标记是否为内置工具"""
|
||||
@@ -63,7 +81,7 @@ class ToolRegistry:
|
||||
if not tool:
|
||||
return f"❌ 错误:找不到工具 '{name}'"
|
||||
|
||||
self.logger.debug(f"工具调用: {name}({kwargs})")
|
||||
self.logger.debug(f"工具调用: {name}({_sanitize_kwargs(kwargs)})")
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
@@ -84,16 +102,52 @@ class ToolRegistry:
|
||||
)
|
||||
thread.daemon = True
|
||||
thread.start()
|
||||
|
||||
# 从队列收集输出并发送到当前队列
|
||||
|
||||
# M2 修复(审计 2026-08-17):转发循环加总超时。
|
||||
# 此前 while thread.is_alive() 在工具内部卡死(如子进程
|
||||
# 等待 stdin)时永不退出;超时后放弃等待,避免本调用
|
||||
# 线程被连带挂起。
|
||||
# N6 修复(审计 2026-08-18):超时路径此前直接 return,
|
||||
# 跳过下方清理——daemon 线程继续持有 tool.output_queue
|
||||
# 实例属性,属性被后续调用覆盖到新会话队列后,僵尸输出
|
||||
# 会混入后续对话流。现在统一在出口清理。
|
||||
tool_timeout = self.stream_tool_timeout
|
||||
deadline = time.time() + tool_timeout
|
||||
final_result = None
|
||||
timed_out = False
|
||||
# N8:流式转发同样受 TOOL_MAX_RESULT_LENGTH 约束——此前仅非流式
|
||||
# 返回值受限,超长内容(如记忆全文)经流式通道可完整推到前端
|
||||
line_budget = constants.TOOL_MAX_RESULT_LENGTH
|
||||
forwarded = 0
|
||||
truncation_noted = False
|
||||
while thread.is_alive() or not output_queue.empty():
|
||||
if time.time() > deadline:
|
||||
self.logger.warning(
|
||||
f"工具 {name} 流式执行超过 {tool_timeout}s 总上限,放弃等待")
|
||||
timed_out = True
|
||||
break
|
||||
try:
|
||||
msg_type, content = output_queue.get(timeout=0.1)
|
||||
if msg_type == "line":
|
||||
piece = str(content) if content is not None else ""
|
||||
if forwarded >= line_budget:
|
||||
# 预算耗尽:停止转发,仅提示一次
|
||||
if not truncation_noted:
|
||||
truncation_noted = True
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(
|
||||
("line", f"\n⚠️ 流式输出超过 {line_budget} 字符上限,已停止转发剩余内容\n"))
|
||||
continue
|
||||
remaining = line_budget - forwarded
|
||||
if len(piece) > remaining:
|
||||
piece = piece[:remaining] + \
|
||||
f"\n⚠️ 流式输出超过 {line_budget} 字符上限,已截断"
|
||||
forwarded = line_budget
|
||||
else:
|
||||
forwarded += len(piece)
|
||||
# 发送到当前队列(会通过WebSocket发送)
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(("line", content))
|
||||
self.current_output_queue.put(("line", piece))
|
||||
elif msg_type == "complete":
|
||||
final_result = content
|
||||
# 也发送完整结果
|
||||
@@ -101,14 +155,21 @@ class ToolRegistry:
|
||||
self.current_output_queue.put(("complete", content))
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
thread.join(timeout=5)
|
||||
if thread.is_alive():
|
||||
return "❌ 工具执行超时"
|
||||
|
||||
result = final_result or ""
|
||||
|
||||
if not timed_out:
|
||||
thread.join(timeout=5)
|
||||
if thread.is_alive():
|
||||
timed_out = True
|
||||
|
||||
# 统一清理(正常与超时路径都必须执行):断开实例属性,
|
||||
# 僵尸 daemon 线程的 send_line 落到 None 上不再外写
|
||||
tool.streaming_output = False
|
||||
tool.output_queue = None
|
||||
|
||||
if timed_out:
|
||||
result = f"❌ 工具执行超时(>{tool_timeout}s,已放弃等待)"
|
||||
else:
|
||||
result = final_result or ""
|
||||
else:
|
||||
# 普通执行
|
||||
result = tool.execute(**kwargs)
|
||||
|
||||
@@ -0,0 +1,405 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.skill_tools - 技能相关工具
|
||||
|
||||
抽离自 aibox.py,包含 4 个技能管理工具:
|
||||
- SearchSkillsTool : 按关键字搜索扩展技能
|
||||
- ReloadSkillsTool : 热重载技能
|
||||
- InstallSkillTool : 返回技能安装指引
|
||||
- CreateSkillTool : 返回创建技能的指令提示词
|
||||
|
||||
依赖:common (constants, i18n, Logger) + tools.base.Tool + tools.base.ToolRegistry
|
||||
+ subsystems.skill_manager.SkillManager
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import shlex
|
||||
from typing import Dict, Optional
|
||||
|
||||
from common import constants, i18n, Logger
|
||||
from subsystems import SkillManager
|
||||
from .base import Tool
|
||||
from .registry import ToolRegistry
|
||||
|
||||
|
||||
class SearchSkillsTool(Tool):
|
||||
"""扩展技能搜索工具 - 按关键字搜索 ~/.aibox/skills/ 下的技能文件。
|
||||
|
||||
支持两种技能格式:
|
||||
1. 可执行脚本/程序:通过 --help 获取简介和完整使用说明
|
||||
2. Markdown 文件 (.md):第一行为简介,其余为技能内容
|
||||
|
||||
工作流程:
|
||||
1. 用关键字调用本工具,获取匹配技能的名称和简介列表
|
||||
2. 对于 Markdown 技能,可以直接读取 content 字段获取完整技能内容
|
||||
3. 对于可执行技能,通过 create_subagent 派生子代理执行 "技能名 --help" 查看完整命令手册
|
||||
4. 根据内容,使用相应方式调用技能
|
||||
"""
|
||||
|
||||
def __init__(self, skill_manager: SkillManager, logger: Optional[Logger] = None):
|
||||
super().__init__(logger)
|
||||
self.skill_manager = skill_manager
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "search_skills"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return (
|
||||
"扩展技能搜索工具。\n\n"
|
||||
"用关键字搜索 ~/.aibox/skills/ 目录下的技能文件。\n"
|
||||
"支持两种技能格式:\n"
|
||||
" 1. 可执行脚本/程序:通过 --help 获取简介和完整使用说明\n"
|
||||
" 2. Markdown 文件 (.md):第一行为简介,其余为技能内容\n\n"
|
||||
"使用方法:\n"
|
||||
"1. 调用本工具搜索关键字,获取匹配技能列表(空字符串列出全部)。\n"
|
||||
"2. 对于 Markdown 技能:设置 show_content=true 读取返回的内容获取完整技能说明。\n"
|
||||
"3. 对于可执行技能:用 create_subagent 派生子代理执行 \"技能名 --help\" 获取完整命令手册,再按手册调用该技能。\n"
|
||||
"4. 按照内容说明,使用相应方式调用该技能(执行类操作一律经 create_subagent 派生子代理完成)。"
|
||||
)
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"keyword": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"搜索关键字,在技能名称、简介和内容中进行匹配。"
|
||||
"传入空字符串可列出全部可用技能。"
|
||||
)
|
||||
},
|
||||
"show_content": {
|
||||
"type": "boolean",
|
||||
"description": (
|
||||
"是否在结果中显示技能内容。"
|
||||
"对于 Markdown 技能,这会影响返回结果的详细程度。"
|
||||
"默认 False(只显示名称和简介)。"
|
||||
),
|
||||
"default": False
|
||||
}
|
||||
},
|
||||
"required": ["keyword"]
|
||||
}
|
||||
|
||||
def execute(self, keyword: str = "", show_content: bool = False, **kwargs) -> str:
|
||||
kw = keyword.strip()
|
||||
if kw:
|
||||
results = self.skill_manager.search(kw)
|
||||
header = f"🔍 搜索 '{kw}' 的结果(共 {len(results)} 个匹配技能):"
|
||||
empty_msg = (
|
||||
f"❌ 未找到与 '{kw}' 相关的技能。\n"
|
||||
"提示:传入空字符串可列出所有可用技能。"
|
||||
)
|
||||
else:
|
||||
results = self.skill_manager.get_all()
|
||||
header = f"🔧 所有可用技能(共 {len(results)} 个):"
|
||||
empty_msg = "❌ ~/.aibox/skills/ 目录下暂无可用技能。"
|
||||
|
||||
if not results:
|
||||
return empty_msg
|
||||
|
||||
lines = [header]
|
||||
for item in results:
|
||||
skill_type = "📄 Markdown" if item.get('is_markdown') else "⚙️ 可执行"
|
||||
lines.append(f"\n {skill_type} • {item['name']}")
|
||||
lines.append(f" 简介: {item['summary']}")
|
||||
|
||||
if show_content and item.get('content'):
|
||||
content = item['content']
|
||||
if item.get('is_markdown'):
|
||||
preview = content[:500]
|
||||
if len(content) > 500:
|
||||
preview += "...\n (内容较长,已截断,请直接读取技能文件获取完整内容)"
|
||||
lines.append(f" 内容预览:\n {preview}")
|
||||
else:
|
||||
preview = content[:200] if content else "(无帮助输出)"
|
||||
if len(content) > 200:
|
||||
preview += "..."
|
||||
lines.append(f" 帮助预览: {preview}")
|
||||
else:
|
||||
if item.get('is_markdown'):
|
||||
lines.append(f" 💡 提示: 设置 show_content=true 可查看技能内容")
|
||||
else:
|
||||
lines.append(f" 💡 提示: 用 create_subagent 派生子代理执行 \"{item['name']} --help\" 查看完整手册")
|
||||
|
||||
lines.append("")
|
||||
lines.append("📖 技能调用方式:")
|
||||
lines.append(" • Markdown 技能: 设置 show_content=true 读取完整技能说明")
|
||||
lines.append(" • 可执行技能: 用 create_subagent 派生子代理执行(先 '技能名 --help' 查手册,再按手册调用)")
|
||||
return "\n".join(lines)
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行搜索技能"""
|
||||
keyword = kwargs.get("keyword", "")
|
||||
show_content = kwargs.get("show_content", False)
|
||||
|
||||
output_queue.put(("line", f"\n🔍 搜索技能: {keyword if keyword else '(全部)'}\n"))
|
||||
output_queue.put(("line", "-" * 50 + "\n"))
|
||||
|
||||
# 直接传递 kwargs,避免重复参数
|
||||
result = self.execute(**kwargs)
|
||||
|
||||
output_queue.put(("line", result + "\n"))
|
||||
output_queue.put(("complete", result))
|
||||
|
||||
|
||||
class ReloadSkillsTool(Tool):
|
||||
"""技能热重载工具 - 重新加载所有扩展技能"""
|
||||
|
||||
def __init__(self, tool_registry: ToolRegistry, logger: Optional[Logger] = None):
|
||||
super().__init__(logger)
|
||||
self.tool_registry = tool_registry
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "reload_skills"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return """技能热重载工具 - 重新加载所有扩展技能
|
||||
当你添加、修改或删除了 ~/.aibox/skills/ 目录下的技能文件后,
|
||||
调用此工具可以立即生效,无需重启系统。"""
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"confirm": {
|
||||
"type": "boolean",
|
||||
"description": "确认重新加载技能",
|
||||
"default": True
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
def execute(self, confirm: bool = True, **kwargs) -> str:
|
||||
"""执行技能热重载"""
|
||||
if not confirm:
|
||||
return "❌ 技能重载已取消,如需重载请设置 confirm=true"
|
||||
|
||||
self.logger.info("开始执行技能热重载...")
|
||||
|
||||
try:
|
||||
count = self.tool_registry.reload_skills()
|
||||
|
||||
# 从 SkillManager 获取更新后的技能列表(含简介)
|
||||
skill_manager = self.tool_registry.skill_manager
|
||||
if skill_manager:
|
||||
all_skills = skill_manager.get_all()
|
||||
skill_lines = "\n".join(
|
||||
f" • {s['name']} — {s['summary']}"
|
||||
for s in all_skills[:10]
|
||||
)
|
||||
data = {
|
||||
"重新加载的技能数量": count,
|
||||
"技能目录": constants.SKILLS_DIR,
|
||||
"技能列表(含简介)": skill_lines or "(无技能)",
|
||||
}
|
||||
if count > 10:
|
||||
data["提示"] = f"还有 {count - 10} 个技能未显示"
|
||||
else:
|
||||
data = {
|
||||
"重新加载的技能数量": count,
|
||||
"技能目录": constants.SKILLS_DIR,
|
||||
}
|
||||
|
||||
return self.format_result(True, i18n.get('skills_reloaded', count), data)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"技能热重载失败: {e}")
|
||||
return i18n.get('skills_reload_failed', str(e))
|
||||
|
||||
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 InstallSkillTool(Tool):
|
||||
"""技能安装指引工具 - 返回如何将可执行文件/脚本安装为扩展技能的说明"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "install_skill"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return (
|
||||
"技能安装指引工具。\n"
|
||||
"返回将 Markdown 技能文件安装为扩展技能的完整步骤说明,"
|
||||
"包括文件检查、格式校验和部署方法(目标文件名统一带 .md 扩展名)。"
|
||||
)
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"skill_path": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"待安装的技能文件路径(可选)。"
|
||||
"若提供,说明中会带入该路径作为示例;"
|
||||
"不提供则返回通用说明。"
|
||||
)
|
||||
}
|
||||
},
|
||||
"required": []
|
||||
}
|
||||
|
||||
def execute(self, skill_path: str = "", **kwargs) -> str:
|
||||
skills_dir = constants.SKILLS_DIR
|
||||
|
||||
# 如果给了具体路径,在说明里带入,否则用占位符
|
||||
raw_src = skill_path.strip() if skill_path else ""
|
||||
# 去掉扩展名作为技能名称,再统一补 .md 作为目标文件名
|
||||
base = os.path.splitext(os.path.basename(raw_src))[0] if raw_src else "your_skill"
|
||||
dst = os.path.join(skills_dir, base + ".md")
|
||||
# L7(审计 2026-08-18):src 为 LLM 可控输入,含空格/元字符时
|
||||
# 未经引用的命令会被 shell 拆词/注入(指引由子代理经 run_command 执行)
|
||||
src = shlex.quote(raw_src) if raw_src else "/path/to/your_skill.md"
|
||||
dst_q = shlex.quote(dst)
|
||||
|
||||
guide = f"""📦 扩展技能安装指引
|
||||
{"=" * 55}
|
||||
|
||||
要将一个 Markdown 技能文件安装为扩展技能,请按以下步骤操作:
|
||||
|
||||
─────────────────────────────────────────────────────
|
||||
步骤 1:确认文件存在且格式正确
|
||||
─────────────────────────────────────────────────────
|
||||
调用 create_subagent 派生子代理执行以下命令,确认文件存在:
|
||||
|
||||
ls -l {src}
|
||||
|
||||
技能文件格式要求:
|
||||
• 文件名为纯文本 Markdown,第一行是一级标题(# 技能名),作为技能摘要
|
||||
• 其余内容为技能的执行说明、使用方法和示例
|
||||
• 文件应带 .md 扩展名(系统以 .md 文件作为技能加载)
|
||||
|
||||
─────────────────────────────────────────────────────
|
||||
步骤 2:复制到技能目录并统一为 .md 扩展名
|
||||
─────────────────────────────────────────────────────
|
||||
技能目录:{skills_dir}
|
||||
|
||||
调用 create_subagent 派生子代理执行复制命令(目标文件名统一补 .md 扩展名):
|
||||
|
||||
cp {src} {dst_q}
|
||||
|
||||
⚠️ 注意:目标文件名必须带 .md 扩展名,
|
||||
系统以 .md 文件作为技能加载,且技能名称取文件名(不含 .md)。
|
||||
|
||||
─────────────────────────────────────────────────────
|
||||
步骤 3:让系统识别新技能
|
||||
─────────────────────────────────────────────────────
|
||||
复制完成后,调用热重载工具使其立即生效:
|
||||
|
||||
reload_skills confirm=true
|
||||
|
||||
之后可用 search_skills 验证技能已加载:
|
||||
|
||||
search_skills keyword="{base}"
|
||||
|
||||
{"=" * 55}
|
||||
完成以上步骤后,即可通过该技能执行相应操作。"""
|
||||
|
||||
return guide
|
||||
|
||||
|
||||
class CreateSkillTool(Tool):
|
||||
"""创建扩展技能工具 - 返回提示词,驱动智能体自行编写并部署 Markdown 技能文件。
|
||||
|
||||
工具本身不调用 LLM,只输出一段结构化的指令提示词。
|
||||
智能体读到提示词后,会自行:
|
||||
1. 编写满足格式要求的 Markdown 技能文件
|
||||
2. 经 create_subagent 派生子代理写入 ~/.aibox/skills/<skill_name>.md
|
||||
3. 调用 reload_skills 热重载
|
||||
"""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "create_skill"
|
||||
|
||||
def get_description(self) -> str:
|
||||
return (
|
||||
"创建扩展技能工具。\n"
|
||||
"返回一段指令提示词,驱动你自行编写并部署一个 Markdown 技能文件(带 .md 扩展名)。\n\n"
|
||||
"技能文件格式要求:\n"
|
||||
" • 文件名必须带 .md 扩展名(如 fetch-weather.md、calc-bmi.md)\n"
|
||||
" • 第一行为一级标题(# 技能名),作为该技能的一句话简介\n"
|
||||
" • 其余内容为技能的执行说明、使用方法和示例,供智能体参考执行"
|
||||
)
|
||||
|
||||
def get_parameters(self) -> Dict:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"skill_name": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"技能名称(部署后的文件名,不含扩展名;系统会自动补齐 .md)。"
|
||||
"只能包含字母、数字、连字符和下划线,例如:fetch-weather、calc-bmi。"
|
||||
)
|
||||
},
|
||||
"description": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"对该技能的详细需求描述:它应该做什么、接受哪些输入参数、"
|
||||
"返回什么输出、有哪些典型使用场景。描述越详细,生成质量越高。"
|
||||
)
|
||||
}
|
||||
},
|
||||
"required": ["skill_name", "description"]
|
||||
}
|
||||
|
||||
def execute(self, skill_name: str = "", description: str = "", **kwargs) -> str:
|
||||
# ---- 参数校验 ----
|
||||
if not skill_name or not skill_name.strip():
|
||||
return "❌ 错误:skill_name 不能为空"
|
||||
if not description or not description.strip():
|
||||
return "❌ 错误:description 不能为空"
|
||||
if not re.match(r'^[A-Za-z0-9_\-]+$', skill_name.strip()):
|
||||
return "❌ 错误:skill_name 只能包含字母、数字、连字符(-)和下划线(_)"
|
||||
|
||||
skill_name = skill_name.strip()
|
||||
skills_dir = constants.SKILLS_DIR
|
||||
# 技能文件名统一带 .md 扩展名
|
||||
skill_path = os.path.join(skills_dir, skill_name + ".md")
|
||||
|
||||
# ---- 返回给智能体的指令提示词 ----
|
||||
prompt = f"""【任务】创建扩展技能:{skill_name}
|
||||
|
||||
请按照以下步骤,自行完成该扩展技能的编写和部署:
|
||||
|
||||
─────────────────────────────────────────────────
|
||||
步骤 1:编写 Markdown 技能文件
|
||||
─────────────────────────────────────────────────
|
||||
根据以下需求,编写一个 Markdown 技能文件:
|
||||
|
||||
需求描述:
|
||||
{description.strip()}
|
||||
|
||||
技能文件必须满足的格式约定:
|
||||
1. 第一行为一级标题 `# {skill_name}`,用于对该技能功能做一句话简介(系统用第一行作为技能摘要)
|
||||
2. 其余内容清晰说明该技能的用途、使用方式、注意事项和典型示例
|
||||
3. 内容完整可执行(即按内容操作即可完成该技能要完成的事)
|
||||
|
||||
─────────────────────────────────────────────────
|
||||
步骤 2:写入技能目录
|
||||
─────────────────────────────────────────────────
|
||||
调用 create_subagent 派生子代理,将你编写的完整内容写入(文件名必须带 .md 扩展名):
|
||||
{skill_path}
|
||||
(在 task 中附上完整文件内容,并要求子代理写入后用 head -5 验证首行格式)
|
||||
|
||||
─────────────────────────────────────────────────
|
||||
步骤 3:热重载
|
||||
─────────────────────────────────────────────────
|
||||
调用 reload_skills 工具使新技能立即生效:
|
||||
reload_skills confirm=true
|
||||
|
||||
─────────────────────────────────────────────────
|
||||
完成后请告知用户技能已创建成功,并展示技能的第一行简介。"""
|
||||
|
||||
return prompt
|
||||
@@ -0,0 +1,53 @@
|
||||
# jump.py —— jump.mt 的 mtc DSL 源码(新格式)
|
||||
# 跳跃:四腿 1 号关节(s2/s4/s6/s8)快速下蹲 10°→20°→30°(各 20ms)后 200ms 弹回站立。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/jump.py -o ../bot4cj/scripts/jump.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 3×20 + 200,T=460ms) ----
|
||||
stand()
|
||||
Beep(10, 460) # spk.F=10 贯穿全程
|
||||
Delay(200)
|
||||
for _ang in (10, 20, 30):
|
||||
for _dev in ("s2", "s4", "s6", "s8"):
|
||||
SetTo(_dev, _ang, 20)
|
||||
Delay(20)
|
||||
stand(200) # 弹回站立
|
||||
Delay(200)
|
||||
@@ -0,0 +1,57 @@
|
||||
# left-hand.py —— left-hand.mt 的 mtc DSL 源码(新格式)
|
||||
# 左手:s1/s8 展开维持 1s 后收回;蜂鸣 10→400→1000→10 Hz 阶跃。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/left-hand.py -o ../bot4cj/scripts/left-hand.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 200 + 1000 + 20,T=1420ms) ----
|
||||
stand()
|
||||
Beep(10, 200) # r1: F=10
|
||||
Delay(200)
|
||||
Beep(400, 200) # r2: 仅变频率(原行无舵机变化)
|
||||
Delay(200)
|
||||
SetTo("s1", 0, 1000) # r3: s1 45→0,s8 0→150
|
||||
SetTo("s8", 150, 1000)
|
||||
Beep(1000, 1000) # r3: F=1000
|
||||
Delay(1000)
|
||||
SetTo("s1", 45, 20) # r4: 收回
|
||||
SetTo("s8", 0, 20)
|
||||
Beep(10, 20) # r4: F=10
|
||||
Delay(20)
|
||||
@@ -0,0 +1,71 @@
|
||||
# left-turn.py —— left-turn.mt 的 mtc DSL 源码(新格式)
|
||||
# 左转步态:对角腿轮换抬摆(s2/s6 ↔ s4/s8 为一组,s1/s5、s3/s7 配合)。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/left-turn.py -o ../bot4cj/scripts/left-turn.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
def legs(dev_a, dev_b, ang, dur):
|
||||
SetTo(dev_a, ang, dur)
|
||||
SetTo(dev_b, ang, dur)
|
||||
|
||||
# ---- 动作序列(200+100+50+50+100+10+50+50+100+10+50,T=770ms) ----
|
||||
stand()
|
||||
Beep(10, 770) # spk.F=10 贯穿全程
|
||||
Delay(200)
|
||||
legs("s2", "s6", 45, 100) # r2
|
||||
Delay(100)
|
||||
legs("s1", "s5", 90, 50) # r3
|
||||
Delay(50)
|
||||
legs("s2", "s6", 0, 50) # r4
|
||||
Delay(50)
|
||||
legs("s4", "s8", 45, 100) # r5
|
||||
Delay(100)
|
||||
legs("s1", "s5", 45, 10) # r6
|
||||
Delay(10)
|
||||
legs("s3", "s7", 0, 50) # r7
|
||||
Delay(50)
|
||||
legs("s4", "s8", 0, 50) # r8
|
||||
Delay(50)
|
||||
legs("s2", "s6", 45, 100) # r9
|
||||
Delay(100)
|
||||
legs("s3", "s7", 45, 10) # r10
|
||||
Delay(10)
|
||||
legs("s2", "s6", 0, 50) # r11
|
||||
Delay(50)
|
||||
@@ -0,0 +1,51 @@
|
||||
# lie-down.py —— lie-down.mt 的 mtc DSL 源码(新格式)
|
||||
# 趴下:四腿 1 号关节(s2/s4/s6/s8)10°→15°→20°→25°→30° 渐增(各 20ms)。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/lie-down.py -o ../bot4cj/scripts/lie-down.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 5×20,T=300ms) ----
|
||||
stand()
|
||||
Beep(10, 300)
|
||||
Delay(200)
|
||||
for _ang in (10, 15, 20, 25, 30):
|
||||
for _dev in ("s2", "s4", "s6", "s8"):
|
||||
SetTo(_dev, _ang, 20)
|
||||
Delay(20)
|
||||
@@ -0,0 +1,56 @@
|
||||
# 录制用脚本:stand → 多姿态序列(微动作分解 + 合成帧 + IMU 占位)
|
||||
# 首行不分解(1 步),后续 4 行 delay=200 → 各 20 微步,共 81 步
|
||||
@Alias
|
||||
name alias;
|
||||
A angle;
|
||||
F frequency;
|
||||
@
|
||||
|
||||
@I2C
|
||||
name path;
|
||||
i2c-0 /dev/i2c-0;
|
||||
i2c-1 /dev/i2c-1;
|
||||
@
|
||||
|
||||
@PCA9685
|
||||
name I2C address;
|
||||
pca9685-0 i2c-1 0x40;
|
||||
@
|
||||
|
||||
@MPU6050_BIND
|
||||
name I2C addr;
|
||||
mpu6050-0 i2c-1 0x68;
|
||||
@
|
||||
|
||||
@Speaker
|
||||
name channels sample_rate volumn max_f min_f;
|
||||
spk 2 48000 0.6 5000 -;
|
||||
@
|
||||
|
||||
@Servo
|
||||
name max min offset controller channel;
|
||||
s1 90 - - pca9685-0 0; # rb0
|
||||
s2 180 - - - 1; # rb1
|
||||
s3 -90 - - - 2; # lb0
|
||||
s4 -180 - - - 3; # lb1
|
||||
s5 90 - - - 4; # lf0
|
||||
s6 180 - - - 5; # lf1
|
||||
s7 -90 - - - 6; # rf0
|
||||
s8 -180 - - - 7; # rf1
|
||||
s9 180 - - - 8; # cam0
|
||||
s10 180 - - - 9; # cam1
|
||||
@
|
||||
|
||||
@Camera
|
||||
name width height framerate capacity;
|
||||
cam0 640 480 30 8;
|
||||
@
|
||||
|
||||
@Action
|
||||
s1.A s2.A s3.A s4.A s5.A s6.A s7.A s8.A s9.A s10.A spk.F spk.delay;
|
||||
45 - 45 - 45 - 45 - 90 90 10 200;
|
||||
60 - 60 - 60 - 60 - 90 90 - 200;
|
||||
45 - 45 - 45 - 45 - 90 90 - 200;
|
||||
30 - 30 - 30 - 30 - 90 90 - 200;
|
||||
45 - 45 - 45 - 45 - 90 90 - 200;
|
||||
@
|
||||
@@ -0,0 +1,49 @@
|
||||
# record_stand.py —— record_stand.mt 的 mtc DSL 源码(新格式)
|
||||
# 录制用:stand → s1/s3/s5/s7 依次到 60°/45°/30°/45° 姿态序列(各 200ms),
|
||||
# 含 @Camera(合成帧)与激活的 @MPU6050_BIND(imu.npy 录制)。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/record_stand.py -o ../bot4cj/scripts/record_stand.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
MPU6050_BIND("mpu6050-0", "i2c-1", "0x68")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0) # min_f 原为占位符(首行解析为 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(5×200,T=1000ms → 100 栅格行) ----
|
||||
# 注:原文件注释「首行不分解,共 81 步」已过时——99c1bc1 起首行也分解为
|
||||
# 20 个恒等微步,原版/编译版在现行 pybot 语义下均为 100 步且逐元素相等。
|
||||
stand()
|
||||
Beep(10, 1000) # spk.F=10 贯穿全程(r2-r5 占位符继承首行)
|
||||
Delay(200)
|
||||
for _ang in (60, 45, 30, 45):
|
||||
for _dev in ("s1", "s3", "s5", "s7"):
|
||||
SetTo(_dev, _ang, 200)
|
||||
Delay(200)
|
||||
@@ -0,0 +1,60 @@
|
||||
# rght-hand.py —— rght-hand.mt 的 mtc DSL 源码(新格式)
|
||||
# 右手(右前腿 s7/s8 为 rf,本脚本动 s3/s6 = 左后/左前?保持原脚本语义):
|
||||
# s6 收回 → s3/s6 展开维持 1s → 收回;蜂鸣 10→400→1000→10 Hz 阶跃。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/rght-hand.py -o ../bot4cj/scripts/rght-hand.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 200 + 1000 + 20,T=1420ms) ----
|
||||
# spk.F 行级阶跃(pybot 微步对频率列不插值,直接取目标值),与 Beep 恒值段一致
|
||||
stand()
|
||||
Beep(10, 200) # r1: F=10
|
||||
Delay(200)
|
||||
SetTo("s6", 0, 200) # r2: s6 45→0
|
||||
Beep(400, 200) # r2: F=400
|
||||
Delay(200)
|
||||
SetTo("s3", 0, 1000) # r3: s3 45→0,s6 0→150
|
||||
SetTo("s6", 150, 1000)
|
||||
Beep(1000, 1000) # r3: F=1000
|
||||
Delay(1000)
|
||||
SetTo("s3", 45, 20) # r4: 收回
|
||||
SetTo("s6", 0, 20)
|
||||
Beep(10, 20) # r4: F=10
|
||||
Delay(20)
|
||||
@@ -0,0 +1,56 @@
|
||||
# std.py —— std.mt 的 mtc DSL 源码(新格式)
|
||||
# Dashun 蜘蛛机器人站立姿态:10 舵机到位 + spk 10Hz 节拍 200ms。
|
||||
#
|
||||
# 编译(产物覆盖同目录 std.mt):
|
||||
# cd mtc && python -m mtc ../bot4cj/scripts/std.py -o ../bot4cj/scripts/std.mt
|
||||
#
|
||||
# 与原手写 std.mt 的执行等价性:
|
||||
# 原 @Action 单行 delay=200,pybot 分解为 20 个恒等微步(首行不插值);
|
||||
# 本脚本编译为 20 行 delay=10 的常量栅格(servo 段 10ms 到位后保持,
|
||||
# spk.frequency 段覆盖全程),两者微步序列完全一致
|
||||
# (tests 验证:两版本 recorder 录制的 actions.npy 逐元素相等)。
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用(对应原 std.mt 中的注释行):
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0) # min_f 原为占位符 '-'(首行解析为 0)
|
||||
|
||||
Servo("s1", 90, 0, 0, "pca9685-0", 0) # rb0
|
||||
Servo("s2", 180, 0, 0, "pca9685-0", 1) # rb1
|
||||
Servo("s3", -90, 0, 0, "pca9685-0", 2) # lb0
|
||||
Servo("s4", -180, 0, 0, "pca9685-0", 3) # lb1
|
||||
Servo("s5", 90, 0, 0, "pca9685-0", 4) # lf0
|
||||
Servo("s6", 180, 0, 0, "pca9685-0", 5) # lf1
|
||||
Servo("s7", -90, 0, 0, "pca9685-0", 6) # rf0
|
||||
Servo("s8", -180, 0, 0, "pca9685-0", 7) # rf1
|
||||
Servo("s9", 180, 0, 0, "pca9685-0", 8) # cam0
|
||||
Servo("s10", 180, 0, 0, "pca9685-0", 9) # cam1
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
# stand:10ms(一个栅格步)内到位并保持;首行即目标值,无爬升
|
||||
SetTo("s1", 45, 10)
|
||||
SetTo("s2", 0, 10)
|
||||
SetTo("s3", 45, 10)
|
||||
SetTo("s4", 0, 10)
|
||||
SetTo("s5", 45, 10)
|
||||
SetTo("s6", 0, 10)
|
||||
SetTo("s7", 45, 10)
|
||||
SetTo("s8", 0, 10)
|
||||
SetTo("s9", 90, 10)
|
||||
SetTo("s10", 90, 10)
|
||||
|
||||
# spk 10Hz 持续 200ms —— 对齐原行 spk.F=10 + delay=200;T=200(20 栅格行)
|
||||
Beep(10, 200)
|
||||
@@ -0,0 +1,55 @@
|
||||
# stdd.py —— stdd.mt 的 mtc DSL 源码(新格式)
|
||||
# 站立 → 四腿 1 号关节(s2/s4/s6/s8)10°→20°→30°→20°→10° 渐摆 → 收回站立。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/stdd.py -o ../bot4cj/scripts/stdd.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0) # min_f 原为占位符(首行解析为 0)
|
||||
|
||||
# 名 max ch rb=右后 lb=左后 lf=左前 rf=右前 cam=云台
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
"""全体回站立姿态(dur ms 线性到位);首行用默认 10ms = 瞬时跳变"""
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(对应原 @Action 7 行:200 + 5×50 + 200,T=650ms) ----
|
||||
stand() # 首行瞬时到位(原首行分解为恒等微步)
|
||||
Beep(10, 650) # spk.F=10 贯穿全程(原各行占位符继承首行)
|
||||
Delay(200)
|
||||
for _ang in (10, 20, 30, 20, 10):
|
||||
for _dev in ("s2", "s4", "s6", "s8"):
|
||||
SetTo(_dev, _ang, 50)
|
||||
Delay(50)
|
||||
for _dev in ("s2", "s4", "s6", "s8"):
|
||||
SetTo(_dev, 0, 200) # 收回站立
|
||||
Delay(200)
|
||||
@@ -0,0 +1,63 @@
|
||||
# t0.py —— t0.mt 的 mtc DSL 源码(新格式)
|
||||
# 对角步态(trot)单周期:左后/右前/左前/右后依次抬摆,spk.F 10→0→200→0 Hz。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/t0.py -o ../bot4cj/scripts/t0.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
# I2C("i2c-0", "/dev/i2c-0") # 原 t0.mt 中此行为注释
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0) # min_f 原为占位符(首行解析为 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200+20 + 500+20 + 500+20 + 500+20 + 500+20,T=2300ms) ----
|
||||
# 频率通道:Beep 段外 rest=0 = 静音(对齐原脚本 F=0 / 占位符继承 0 的行)
|
||||
stand()
|
||||
Beep(10, 220) # r1-r2: F=10(r2 占位符继承首行)
|
||||
Delay(200)
|
||||
SetTo("s1", 90, 20) # r2: left-behind
|
||||
Delay(20)
|
||||
SetTo("s1", 45, 500) # r3: stand(F=0 → 段间空隙回落静音)
|
||||
Delay(500)
|
||||
SetTo("s3", 90, 20) # r4: right-behind
|
||||
SetTo("s7", 45, 20)
|
||||
Delay(20)
|
||||
SetTo("s3", 45, 500) # r5: stand + F=200
|
||||
SetTo("s10", 45, 500)
|
||||
Beep(200, 500)
|
||||
Delay(500)
|
||||
SetTo("s5", 90, 20) # r6: right-front(F=0)
|
||||
Delay(20)
|
||||
SetTo("s5", 45, 500) # r7: stand
|
||||
Delay(500)
|
||||
SetTo("s7", 90, 20) # r8: left-front
|
||||
Delay(20)
|
||||
SetTo("s7", 45, 500) # r9: stand
|
||||
Delay(500)
|
||||
SetTo("s10", 90, 20) # r10: 云台回 90
|
||||
Delay(20)
|
||||
@@ -1,21 +0,0 @@
|
||||
# 无硬件情况下的测试脚本
|
||||
|
||||
|
||||
# 在 @Action 矩阵中,硬件变量的参数可以使用别名
|
||||
@Alias
|
||||
name alias;
|
||||
A angle;
|
||||
F frequency;
|
||||
@
|
||||
|
||||
@TEST
|
||||
name parameter0 patameter1;
|
||||
t0 p0 p1;
|
||||
t1 p2 p3;
|
||||
@
|
||||
|
||||
|
||||
@Action
|
||||
t0.A t1.F;
|
||||
10 20;
|
||||
@
|
||||
@@ -0,0 +1,54 @@
|
||||
# tilt-b.py —— tilt-b.mt 的 mtc DSL 源码(新格式)
|
||||
# 后倾:s2/s4(后腿 1 号关节)10°→30° 渐摆后收回。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/tilt-b.py -o ../bot4cj/scripts/tilt-b.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 5×50 + 200,T=650ms) ----
|
||||
stand()
|
||||
Beep(10, 650)
|
||||
Delay(200)
|
||||
for _ang in (10, 20, 30, 20, 10):
|
||||
for _dev in ("s2", "s4"):
|
||||
SetTo(_dev, _ang, 50)
|
||||
Delay(50)
|
||||
for _dev in ("s2", "s4"):
|
||||
SetTo(_dev, 0, 200)
|
||||
Delay(200)
|
||||
@@ -0,0 +1,54 @@
|
||||
# tilt-f.py —— tilt-f.mt 的 mtc DSL 源码(新格式)
|
||||
# 前倾:s6/s8(前腿 1 号关节)10°→30° 渐摆后收回。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/tilt-f.py -o ../bot4cj/scripts/tilt-f.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 5×50 + 200,T=650ms) ----
|
||||
stand()
|
||||
Beep(10, 650)
|
||||
Delay(200)
|
||||
for _ang in (10, 20, 30, 20, 10):
|
||||
for _dev in ("s6", "s8"):
|
||||
SetTo(_dev, _ang, 50)
|
||||
Delay(50)
|
||||
for _dev in ("s6", "s8"):
|
||||
SetTo(_dev, 0, 200)
|
||||
Delay(200)
|
||||
@@ -0,0 +1,54 @@
|
||||
# tilt-l.py —— tilt-l.mt 的 mtc DSL 源码(新格式)
|
||||
# 左倾:s2/s8(右侧腿 1 号关节)10°→30° 渐摆后收回。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/tilt-l.py -o ../bot4cj/scripts/tilt-l.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 5×50 + 200,T=650ms) ----
|
||||
stand()
|
||||
Beep(10, 650)
|
||||
Delay(200)
|
||||
for _ang in (10, 20, 30, 20, 10):
|
||||
for _dev in ("s2", "s8"):
|
||||
SetTo(_dev, _ang, 50)
|
||||
Delay(50)
|
||||
for _dev in ("s2", "s8"):
|
||||
SetTo(_dev, 0, 200)
|
||||
Delay(200)
|
||||
@@ -0,0 +1,54 @@
|
||||
# tilt-r.py —— tilt-r.mt 的 mtc DSL 源码(新格式)
|
||||
# 右倾:s4/s6(左侧腿 1 号关节)10°→30° 渐摆后收回。
|
||||
# 编译:cd mtc && python -m mtc ../bot4cj/scripts/tilt-r.py -o ../bot4cj/scripts/tilt-r.mt
|
||||
|
||||
Alias("A", "angle")
|
||||
Alias("F", "frequency")
|
||||
|
||||
I2C("i2c-0", "/dev/i2c-0")
|
||||
I2C("i2c-1", "/dev/i2c-1")
|
||||
|
||||
PCA9685("pca9685-0", "i2c-1", "0x40")
|
||||
|
||||
# 姿态传感器按需启用:
|
||||
# MPU6050_BIND("mpu6050-0", "i2c-1", "0x69")
|
||||
# MPU6050_BIND("mpu6050-1", "i2c-1", "0x68")
|
||||
# MPU6050_BIND("mpu6050-2", "i2c-0", "0x68")
|
||||
# MPU6050_BIND("mpu6050-3", "i2c-0", "0x69")
|
||||
|
||||
Speaker("spk", 2, 48000, 0.6, 5000, 0)
|
||||
|
||||
# 名 max ch
|
||||
_SERVOS = [("s1", 90, 0), # rb0
|
||||
("s2", 180, 1), # rb1
|
||||
("s3", -90, 2), # lb0
|
||||
("s4", -180, 3), # lb1
|
||||
("s5", 90, 4), # lf0
|
||||
("s6", 180, 5), # lf1
|
||||
("s7", -90, 6), # rf0
|
||||
("s8", -180, 7), # rf1
|
||||
("s9", 180, 8), # cam0
|
||||
("s10", 180, 9)] # cam1
|
||||
for _name, _max, _ch in _SERVOS:
|
||||
Servo(_name, _max, 0, 0, "pca9685-0", _ch)
|
||||
|
||||
# 录制源(pybot --record 需要 @Camera 块)
|
||||
Camera("cam0", 640, 480, 30, 8)
|
||||
|
||||
_STAND = [45, 0, 45, 0, 45, 0, 45, 0, 90, 90]
|
||||
|
||||
def stand(dur=10):
|
||||
for _i, _ang in enumerate(_STAND, 1):
|
||||
SetTo(f"s{_i}", _ang, dur)
|
||||
|
||||
# ---- 动作序列(200 + 5×50 + 200,T=650ms) ----
|
||||
stand()
|
||||
Beep(10, 650)
|
||||
Delay(200)
|
||||
for _ang in (10, 20, 30, 20, 10):
|
||||
for _dev in ("s4", "s6"):
|
||||
SetTo(_dev, _ang, 50)
|
||||
Delay(50)
|
||||
for _dev in ("s4", "s6"):
|
||||
SetTo(_dev, 0, 200)
|
||||
Delay(200)
|
||||
@@ -36,6 +36,25 @@ services:
|
||||
depends_on: []
|
||||
sleep_before: 0
|
||||
|
||||
# --- 内存缓存 ---
|
||||
- name: memdb
|
||||
cmd: bash /opt/micro_servers/memdb/_run.sh
|
||||
args:
|
||||
- "--instance"
|
||||
- "memdb"
|
||||
- "--log-level"
|
||||
- "info"
|
||||
- "--host"
|
||||
- "127.0.0.1"
|
||||
- "--port"
|
||||
- "8011"
|
||||
- "--workers"
|
||||
- "4"
|
||||
- "--max-time"
|
||||
- "600"
|
||||
depends_on: []
|
||||
sleep_before: 0
|
||||
|
||||
# --- AI 审核 ---
|
||||
- name: ai_review
|
||||
cmd: bash /opt/micro_servers/ai_review/_run.sh
|
||||
|
||||
+16
-8
@@ -2,7 +2,7 @@
|
||||
|
||||
> 权威架构参考。本文以 `core/services.yaml`(实际运行清单)与 `git HEAD` 为事实来源。
|
||||
> 与根目录 `README.md` 的偏差见文末「§10 与 README 的偏差」。
|
||||
> 最后核对于:2026-08-10。
|
||||
> 最后核对于:2026-08-17。
|
||||
|
||||
---
|
||||
|
||||
@@ -31,8 +31,8 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Blackbird 系统边界 │
|
||||
│ waf → ca → gateway → {fsys, dnfs2, core} │
|
||||
│ 内部:db2 / ai_review / smtp / captcha2 / core(confs) │
|
||||
│ 独立:aibox(+BotOps) / bot4cj / lci │
|
||||
│ 内部:db2 / memdb / ai_review / smtp / captcha2 / core(confs) │
|
||||
│ 独立:aibox(+BotOps) / bot4cj+pybot / lci │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│ pymsg/TCP (内部 RPC) │ SMTP (外部邮件)
|
||||
▼ ▼
|
||||
@@ -85,14 +85,17 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
└─────────────────────────────────────┘
|
||||
|
||||
独立 / 外设(不在主链路):
|
||||
aibox :8765 (Python WebSocket, 独立产品)
|
||||
aibox :8765 (Python WebSocket, 独立产品; --once CLI 单次对话 + models 多模型档案含 fallback 失败转移)
|
||||
BotOps (aibox 的 TUI 前端)
|
||||
bot4cj (Cangjie + 树莓派 I2C, 硬件)
|
||||
pybot (Python, bot4cj 的迁移版: .mt 解析/录制/实时闭环)
|
||||
mtc (Python, 命令式动作脚本 → .mt 编译器, 开发期工具)
|
||||
lci (C++ /dev/fb0, 嵌入式 GUI)
|
||||
actpred / memdb / libtry (实验 / 支撑库)
|
||||
actpred / libtry (实验 / 支撑库)
|
||||
|
||||
日志链路: 各服务 log_client ──pymsg──► db2:8005 Logs 表 (fail-safe, db 不可用时静默丢弃)
|
||||
配置链路: gateway ──pymsg get──► core 内嵌 confs :8010 (confs/configs/gateway.ini, mtime 自动重载;其余服务配置纯 CLI 传递)
|
||||
内存缓存: memdb :8011 (libmemdb + libmsg RPC, TTL 惰性驱逐;暂无业务调用方,预留热点数据/会话缓存)
|
||||
```
|
||||
|
||||
> 注:`core` 启动参数为 `./nidus.out -i . --serve 5200 --confs 8010`(`-i .` 指向 services.yaml 所在目录),`--confs` 启用内嵌配置服务(由 confs2 合并而来,与 Python confs 协议兼容)。
|
||||
@@ -138,15 +141,17 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
| **fsys** | C++20 | 9200 | 127.0.0.1 | 前端托管 + dnfs 代理(static 的 C++ 重写) | 生产 |
|
||||
| **dnfs2** | C++20 | 8006 | 127.0.0.1 | 树状笔记 CRUD | 生产(dnfs 的 C++ 重写) |
|
||||
| **db2** | C++20 | 8005 | 127.0.0.1 | 通用 DB(pymsg execute) + 日志汇聚 | 生产(db 的 C++ 重写) |
|
||||
| **memdb** | C++20 | 8011 | 127.0.0.1 | 内存缓存数据库(libmemdb + libmsg RPC,TTL 惰性驱逐,9 路由) | 生产(2026-08 合入 services.yaml) |
|
||||
| **ai_review** | Python | 8007 | 127.0.0.1 | AI 内容审核(DeepSeek + ChromaDB RAG) | 生产 |
|
||||
| **smtp** | Python | 8008 | 127.0.0.1 | 邮件验证码后端(pymsg) | 生产(smtp3 C++ 重写待上线) |
|
||||
| **captcha2** | C++20 | 8009 | 127.0.0.1 | 图形验证码(pymsg) | 生产(captcha 的 C++ 重写) |
|
||||
| **aibox** | Python | 8765 | 127.0.0.1 | 独立 AI 智能体(WebSocket) | 独立产品 |
|
||||
| **BotOps** | Python | — | — | aibox 的 TUI 前端 | 独立 |
|
||||
| **bot4cj** | Cangjie | — | 树莓派 I2C | 机器人脚本解析器 | 独立硬件 |
|
||||
| **pybot** | Python | — | 树莓派 | bot4cj 的 Python 迁移版(.mt 解析、actpred episode 录制、--loop 实时预测闭环;HEF 后端自动异步,推理与执行上一步动作并行) | 独立硬件 |
|
||||
| **mtc** | Python | — | — | 命令式动作脚本(Python DSL)→ .mt 编译器(非阻塞时间轴 IR + 10ms 栅格化),pybot/bot4cj 的脚本 authoring 工具 | 开发期工具 |
|
||||
| **lci** | C++ | — | /dev/fb0 | 嵌入式帧缓冲 GUI | 独立硬件 |
|
||||
| **actpred** | Python | — | — | ML 预训练/推理 | 实验 |
|
||||
| **memdb** | C++20 | — | — | 内存数据库 | 实验/支撑 |
|
||||
| **actpred** | Python | — | — | 动作预测(图像+动作历史→下一动作),bot4cj/pybot 的「大脑」配套 | 实验 |
|
||||
| **libtry** | C | — | — | 嵌入式脚本解释器 | 支撑 |
|
||||
|
||||
### 已被替代、不应再部署的遗留模块(不在 services.yaml)
|
||||
@@ -170,6 +175,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
| libcrypt | C++20 | AES_CBC / RSA / SHA256 / Base64 |
|
||||
| libsqlite3 | C++20 | SQLite3 header-only 封装 |
|
||||
| libmsg | C++20 | pymsg 协议 C++ 实现(4B 大端长度前缀 + UTF-8 JSON),与 Python pymsg 字节级兼容,提供 RpcClient/RpcServer/AsyncQueue |
|
||||
| libmemdb | C++20 | header-only 内存数据库模板库:三索引存储(list 主存 + 键索引 + TTL 索引),TTL 惰性驱逐;memdb 服务基于此库 |
|
||||
| liblogc | C++20 | 统一集中式日志客户端(`us/LogC.hpp`,对应 Python `log_client.py`):pymsg/TCP 异步批量写 db2 Logs 表,队列上限 1 万条 + 级别过滤 + fail-safe,C++ 服务共用 |
|
||||
| libtry | C | 嵌入式脚本解释器(递归下降,中文标识符) |
|
||||
| pymsg | Python | 自研 RPC 基础库(`[4B 大端长度][UTF-8 JSON]`),Python 侧 db/log_client 依赖 |
|
||||
@@ -195,7 +201,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
- 三层覆盖:confs 默认值 → 环境变量 → 命令行参数。
|
||||
|
||||
### 7.4 弹性 / 容错
|
||||
- 启动顺序敏感:`core` 拉起顺序为 db2→ai_review→dnfs2→smtp→fsys→waf→ca→gateway→captcha2→aibox→redirect80;pymsg 客户端带 `retry_on_refused` 等依赖就绪。
|
||||
- 启动顺序敏感:`core` 拉起顺序为 db2→memdb→ai_review→dnfs2→smtp→fsys→waf→ca→gateway→captcha2→aibox→redirect80(共 12 个);pymsg 客户端带 `retry_on_refused` 等依赖就绪。
|
||||
- 日志 fail-safe:db2 不可用时静默丢弃,不影响业务。
|
||||
- 有状态单点:`smtp` / `captcha2` 验证码存进程内存,**不支持多实例水平扩展**(需换 Redis 才能扩)。
|
||||
|
||||
@@ -230,6 +236,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
- **smtp3 待上线**:C++ 重写已完成并通过安全审计,但 services.yaml 仍指向 Python `smtp`,建议择机平替以消除最后一块 Python 有状态服务。
|
||||
- **可观测性缺口**:缺 tracing/metrics/告警,建议后续引入轻量方案(或复用 db2 Logs + 定时聚合)。
|
||||
- **水平扩展限制**:smtp/captcha 内存态为单点,若需多实例需引入 Redis 等共享存储。
|
||||
- **panel 面板微服务(规划中)**:设计稿见 [docs/panel_service.md](panel_service.md)——背景视图 + 各服务自持控件片段(iframe 沙箱 + BFF 转发),解决整页应用互相改前端代码的耦合问题;落地后需三处同步(本文件 / MODULES.md / 根 README)。
|
||||
|
||||
---
|
||||
|
||||
@@ -244,6 +251,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
||||
5. **fsys README 图示滞后**:✅ 已修正(2026-08-10)。`fsys/README.md` 架构图与配置章节已更新为 `ca:8443 → gateway:8445`、纯 CLI 配置。
|
||||
6. **工作记忆偏差**:✅ 已同步。`ca:8443 → gateway:8445`、`dnfs2`、confs 内嵌 core 均已写入各文档。
|
||||
7. **dnfs2 README 图示滞后**:✅ 已修正(2026-08-10)。`dnfs2/README.md` 架构图与验签描述已更新为 ca:8443(CLI 参数 `--gateway-url` 指向 ca)。
|
||||
8. **memdb / pybot / mtc 未进根 README**:✅ 已回填(2026-08-15)。`memdb:8011`、`pybot`(bot4cj Python 迁移版)、`mtc`(Python DSL → .mt 编译器,2026-08-15 落地)连同 `actpred` 已入根 `README.md` 模块表与端口一览;本文件与 MODULES.md 已同步。
|
||||
|
||||
---
|
||||
|
||||
|
||||
+10
-6
@@ -1,7 +1,7 @@
|
||||
# 模块清单(blackbird / micro_servers)
|
||||
|
||||
> 一套 DIY 的轻量级「进程托管式微服务」(UniServe Framework)。
|
||||
> 混合语言:C++20(core/gateway/ca/dnfs2/db2/fsys/captcha2/lib*)、Python(static/dnfs/db/ai_review/smtp/captcha/confs/redirect80/BotOps/aibox/actpred/pymsg,多为遗留)、Go(waf)、仓颉/Cangjie(bot4cj)。
|
||||
> 混合语言:C++20(core/gateway/ca/dnfs2/db2/memdb/fsys/captcha2/lib*)、Python(static/dnfs/db/ai_review/smtp/captcha/confs/redirect80/BotOps/aibox/actpred/pybot/mtc/pymsg,多为遗留)、Go(waf)、仓颉/Cangjie(bot4cj)。
|
||||
> 由 `core` 读 `services.yaml` 拉起并守护各子进程;`waf:443 → ca:8443 → gateway:8445 → 业务服务` 为对外主链路。
|
||||
|
||||
## 一、核心与编排
|
||||
@@ -38,12 +38,13 @@
|
||||
| smtp | Python/pymsg | 8008 | 邮件服务,注册邮箱验证码后端(路由 send_email/verify_email/health),6 位码内存存储、冷却/限流/熔断(smtp3 C++ 重写待上线) |
|
||||
| captcha | Python/pymsg | 8009 | 验证码服务,图形验证码(一次性、5 分钟过期),经 ca `GET /captcha/image` 代理给浏览器(遗留,已由 captcha2 平替) |
|
||||
| captcha2 | C++20 | 8009 | captcha 的 C++20 重实现(pymsg),路由 image/verify/health 对齐 |
|
||||
| memdb | C++20 | 8011 | 内存缓存数据库(libmemdb 模板库 + libmsg RPC),TTL 惰性驱逐,9 路由(ping/set/get/del/update/exists/size/show/expire),线程安全 + SIGINT/SIGTERM 优雅退出 |
|
||||
|
||||
## 五、AI 智能体与前端
|
||||
|
||||
| 模块 | 语言 | 端口 | 职责 |
|
||||
|------|------|------|------|
|
||||
| aibox | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度/前端与语音客户端 |
|
||||
| aibox | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度/前端与语音客户端;`--once` CLI 单次对话(stdout/stderr 分流),`models` 配置节支持多 OpenAI 兼容端点按档案名切换 |
|
||||
| BotOps | Python/Textual | — | aibox 的终端 TUI 前端,WebSocket 连接 aibox,Markdown 增量流式渲染、思考模式、任务侧栏、自动重连 |
|
||||
| actpred | Python/PyTorch | — | 动作预测(图像 + 动作历史 → 预测下一个动作),bot4cj 的「大脑」配套;支持从 `.mt` 脚本预训练 |
|
||||
|
||||
@@ -52,6 +53,8 @@
|
||||
| 模块 | 语言 | 端口 | 职责 |
|
||||
|------|------|------|------|
|
||||
| bot4cj | 仓颉/Cangjie | — (树莓派 I2C) | 机器人硬件脚本解析器,解析 `.mt` 驱动 PCA9685 舵机 / MPU6050 陀螺仪 / 扬声器,面向 Dashun 蜘蛛机器人 |
|
||||
| pybot | Python | — (树莓派) | bot4cj 的 Python 迁移版:`.mt` 动作脚本解析器(零第三方依赖,与仓颉版逐语义对齐);支持执行 / `--dry-run` / `--record` 录制 actpred episode(含 imu.npy)/ `--loop` 实时预测闭环(HEF 后端自动异步:推理与执行上一步动作并行) |
|
||||
| mtc | Python | — | 命令式动作脚本(Python DSL)→ `.mt` 编译器:`exec` 注入命名空间 → 时间轴 IR(非阻塞语义、重叠截断)→ 10ms 栅格化;产物可由 pybot/bot4cj 直接执行,零第三方依赖 |
|
||||
| lci | C++ | — (/dev/fb0) | 嵌入式帧缓冲图形界面(无需 X11/Wayland),矢量 GB2312 中文字库、图片渲染、实时数据可视化 |
|
||||
| 3d | — (资产) | — | 机器人 / 机械臂 CAD 模型(STEP/STL/SLDPRT),bot4cj 配套硬件资产 |
|
||||
|
||||
@@ -64,17 +67,17 @@
|
||||
| libcrypt | C++ | 加密库:`AES_CBC`、`MyRSA`、`SHA256`、`Base64`、`Crypto_Basic` |
|
||||
| libsqlite3 | C++ | SQLite3 header-only 封装 |
|
||||
| libmsg | C++ | pymsg 协议 C++ 实现(4B 大端长度前缀 + UTF-8 载荷),`MsgClient/Server`、`RpcClient/Server`、`AsyncQueue`,与 Python pymsg 字节级兼容 |
|
||||
| libmemdb | C++20 | header-only 内存数据库模板库:三索引存储(`std::list` 主存 + `unordered_map` 键索引 + `multiset` TTL 索引),insert/remove/update/search 等操作,TTL 惰性驱逐;memdb 服务即基于此库 |
|
||||
| liblogc | C++ | 统一集中式日志客户端(`us/LogC.hpp`,对应 Python `log_client.py`):pymsg/TCP 异步批量写 db2 Logs 表,队列上限 1 万条 + 级别过滤 + fail-safe,ca/core/gateway/captcha2/smtp3 等 C++ 服务共用 |
|
||||
| libtry | C | 嵌入式脚本解释器(~741 行),手写递归下降,支持中文标识符,用于嵌入式操作界面 shell |
|
||||
| pymsg | Python | 自研 RPC 基础库(`[4B 长度][UTF-8 JSON]`),db 服务端与各 `log_client.py` 的基础依赖(setuptools 包) |
|
||||
| memdb | C++ | 内存数据库(include/src),内部组件,未在主 README 模块表列出 |
|
||||
|
||||
## 八、测试与文档
|
||||
|
||||
| 模块 | 说明 |
|
||||
|------|------|
|
||||
| tests | 黑盒测试套件(`test_db.py`/`test_gateway.py`/`test_waf.py`/`test_dnfs.py`/`test_captcha.py`/`test_smtp.py`/`test_ai_review.py`/`test_static.py`/`test_core_confs.py`),纯标准库 HTTP/TCP 用户旅程模拟,经 `test_runner.py` 统一调度(详见 [tests/blackbox/README.md](../tests/blackbox/README.md)) |
|
||||
| docs | 文档聚合(security/ 审计、reviews/ 审查、reports/ 分析),57 篇 md |
|
||||
| docs | 文档聚合(security/ 审计、reviews/ 审查、reports/ 分析),76 篇 md |
|
||||
|
||||
## 九、根级构建与部署脚本
|
||||
|
||||
@@ -88,5 +91,6 @@
|
||||
- **请求主链路**:client → `redirect80`(:80→443) → `waf`(:443 TLS + WAF) → `ca`(:8443 认证+前置转发) → `gateway`(:8445 纯转发) → `fsys`(:9200) / `dnfs2`(:8006)
|
||||
- **日志链路**:各服务 `log_client.py`(异步批量)→ `db2:8005` Logs 表(fail-safe,db2 不可用时静默丢弃)
|
||||
- **配置链路**:`gateway` → `core` 内嵌 confs:8010(内存缓存 + mtime 重载,仅 gateway.ini);其余服务配置纯 CLI 传递
|
||||
- **数据/审核/邮件**:`dnfs2` → `db2:8005`(pymsg 持久化);`dnfs2` → `ai_review:8007`(审核);注册邮件 `ca` → `smtp:8008`;登录验证码 `ca` → `captcha2:8009`
|
||||
- **端口一览**:80 redirect80 · 443 waf · 8443 ca · 8445 gateway · 5200/8010 core · 9200 fsys · 8006 dnfs2 · 8005 db2 · 8007 ai_review · 8009 captcha2 · 8008 smtp · 8765 aibox
|
||||
- **数据/审核/邮件/缓存**:`dnfs2` → `db2:8005`(pymsg 持久化);`dnfs2` → `ai_review:8007`(审核);注册邮件 `ca` → `smtp:8008`;登录验证码 `ca` → `captcha2:8009`;`memdb:8011`(内存缓存,libmsg RPC,暂无业务调用方,预留热点数据/会话缓存)
|
||||
- **机器人数据链**:`mtc` 编译 Python DSL → `.mt`(或手写 `.mt`)→ `pybot` 执行 / `--record` 录制 episode(images + actions.npy + imu.npy)→ `actpred` 两阶段训练 → ONNX/Hailo HEF 部署(`predict_async` 异步流水线)→ `pybot --loop` 实时预测闭环(bot4cj 语义对齐,HEF 后端推理与执行并行)
|
||||
- **端口一览**:80 redirect80 · 443 waf · 8443 ca · 8445 gateway · 5200/8010 core · 9200 fsys · 8006 dnfs2 · 8005 db2 · 8007 ai_review · 8009 captcha2 · 8008 smtp · 8011 memdb · 8765 aibox
|
||||
@@ -0,0 +1,254 @@
|
||||
# 面板微服务(panel)开发文档
|
||||
|
||||
> 状态:设计稿 v0.1(开发前评审)
|
||||
> 关联:[gateway/README.md](../gateway/README.md)、[docs/dnfs_custom_themes.md](dnfs_custom_themes.md)、[core/services.yaml](../core/services.yaml)、[libhttp/README.md](../libhttp/README.md)
|
||||
|
||||
## 1. 背景与目标
|
||||
|
||||
### 1.1 问题
|
||||
|
||||
现有前端页面(fsys 托管的 `index.html` / `dnfs.html` 等)都是"整页应用":每个页面由单个服务完整提供。当多个服务都想在首页/工作台展示自己的能力时,只能互相改动前端代码或各自开新页面,服务间耦合、发布相互牵连。
|
||||
|
||||
### 1.2 目标
|
||||
|
||||
新增**面板微服务**(panel),提供一张**背景视图**(工作台/仪表盘)。视图内的元素与控件**不由面板自己实现**,而是由各业务微服务以**前端代码片段**的形式提供,面板负责聚合、渲染与托管:
|
||||
|
||||
- 各服务**自持控件代码**(代码维护在各自模块),通过约定协议暴露片段
|
||||
- 面板**拉取并聚合**片段,以 iframe 沙箱插入背景视图
|
||||
- 面板只提供统一外壳(布局 / 主题 / 通信 / 身份透传),**不感知控件业务**
|
||||
|
||||
### 1.3 非目标(v1)
|
||||
|
||||
- 不做控件可视化拖拽编排(v2 候选)
|
||||
- 不做用户自定义控件 / 第三方插件市场(控件仅来自内部服务,白名单)
|
||||
- 不做控件代码的 AI 内容审核——控件是内部服务代码而非用户输入,信任边界与 dnfs2 用户样式节点不同(见 5.3)
|
||||
|
||||
## 2. 架构设计
|
||||
|
||||
### 2.1 拓扑
|
||||
|
||||
```
|
||||
浏览器
|
||||
│ https + q4_token Cookie / Authorization
|
||||
▼
|
||||
waf:443 ──► ca:8443(验签,注入 X-Auth-User / X-Auth-Admin)
|
||||
│ Target-Url: http://127.0.0.1:9201
|
||||
▼
|
||||
gateway:8445 ──► panel:9201(Proxy.opened 白名单)
|
||||
│
|
||||
┌────────────────┼──────────────────────┐
|
||||
│ 内部 HTTP(--providers 白名单) │ BFF 转发
|
||||
▼ ▼ ▼
|
||||
dnfs2:8006 fsys:9200 其他 provider 服务
|
||||
/panel-widget/x /panel-widget/y
|
||||
```
|
||||
|
||||
- 面板对浏览器只暴露**一个 origin**(面板自己);控件 iframe 由面板端点在**同一 origin** 下渲染,不存在跨域、自定义头问题(浏览器无法设置 `Target-Url`,因此控件回源统一走面板 BFF,对齐 fsys→dnfs2 的服务端代理先例)。
|
||||
- 面板加入 `Proxy.opened` 白名单后走 ca 受保护链路:只有登录用户能打开面板;BFF 转发时透传身份头。
|
||||
|
||||
### 2.2 核心概念
|
||||
|
||||
| 概念 | 说明 |
|
||||
|---|---|
|
||||
| Panel | 面板服务本身:背景视图 + 聚合 / 渲染 / 托管 |
|
||||
| Widget(控件) | 一个可独立渲染的前端单元,由某微服务提供 |
|
||||
| Provider(提供方) | 拥有控件代码的微服务,实现 `/panel-widget/<id>` 协议 |
|
||||
| Manifest | 控件描述:元数据 + html/css/js 片段 + 布局 / 主题声明 |
|
||||
|
||||
### 2.3 控件生命周期
|
||||
|
||||
1. **声明**:服务在自身实现 `GET /panel-widget/<id>`(内部 HTTP)
|
||||
2. **拉取**:面板按 `--widgets` 清单拉取各 manifest(缓存 + TTL)
|
||||
3. **聚合**:面板把 manifest 渲染为独立子页面 `/widget/<id>`(内联片段 + 白名单 CSP)
|
||||
4. **插入**:面板页用 `<iframe sandbox="allow-scripts">` 加载子页面,完成"背景视图 + 控件"组合
|
||||
5. **通信**:控件与面板用 postMessage 协议交互(就绪 / 调尺寸 / 主题 / 导航)
|
||||
6. **回源**:控件业务数据经面板 BFF 转发到所属服务(透传身份头)
|
||||
|
||||
### 2.4 端口与技术栈
|
||||
|
||||
| 项 | 值 | 理由 |
|
||||
|---|---|---|
|
||||
| 端口 | 9201 | 与 fsys:9200 相邻的 HTTP 对外段;8xxx 保留给 pymsg/TCP 内部服务 |
|
||||
| 语言 | C++20 + libhttp | 与 db2 / dnfs2 / fsys / captcha2 生态一致 |
|
||||
| 日志 | liblogc(`us/LogC.hpp`) | 统一异步推送 db2 Logs 表 |
|
||||
|
||||
## 3. 服务接口
|
||||
|
||||
### 3.1 对外(浏览器,经 gateway)
|
||||
|
||||
| 路由 | 方法 | 说明 |
|
||||
|---|---|---|
|
||||
| `/` | GET | 面板页(背景视图 + 控件网格) |
|
||||
| `/widgets` | GET | 控件清单 JSON(面板缓存内的 manifest 元数据 + 状态),供页面 JS 渲染网格 |
|
||||
| `/widget/<id>` | GET | 控件子页面(内联 html/css/js),供 iframe 加载;响应带独立 CSP(见 5.1) |
|
||||
| `/api/w/<id>/<path...>` | 任意 | BFF:转发到 `<id>` 所属 provider,透传身份头 |
|
||||
| `/health` | GET | 健康检查 |
|
||||
|
||||
### 3.2 对 Provider(内部 HTTP 协议)
|
||||
|
||||
```
|
||||
GET /panel-widget/<widget_id>
|
||||
|
||||
200 OK, application/json
|
||||
{
|
||||
"id": "dnfs2-minimap", // 必须与路径一致
|
||||
"title": "树状菜单迷你图", // 面板标题(≤ 64 字符)
|
||||
"version": 3, // 递增版本,作为面板缓存键
|
||||
"size": {"w": 2, "h": 1}, // 网格单位(1 ≤ w,h ≤ 6)
|
||||
"themeable": true, // 是否响应面板主题事件
|
||||
"html": "<div class=\"w-root\">…</div>",
|
||||
"css": ".w-root { … }",
|
||||
"js": "window.wx.ready(); …"
|
||||
}
|
||||
```
|
||||
|
||||
面板强制校验(失败视为该控件不可用,返回占位):
|
||||
- `html` / `css` / `js` 各 ≤ 256KB;`id` 必须与 URL 一致;所有字段必须存在
|
||||
- 拉取超时 3s;TTL 默认 60s(`--cache-ttl`)
|
||||
- provider 不可达 / 校验失败 → 占位卡片「控件暂不可用」,**fail-soft**:面板是展示性服务,一个控件失败不拖垮整页(区别于 ai_review 审核类的 fail-closed 语义)
|
||||
|
||||
### 3.3 面板配置(core/services.yaml)
|
||||
|
||||
```yaml
|
||||
- name: panel
|
||||
cmd: bash /opt/micro_servers/panel/_run.sh
|
||||
args:
|
||||
- "--instance"
|
||||
- "panel"
|
||||
- "--log-level"
|
||||
- "info"
|
||||
- "--host"
|
||||
- "127.0.0.1"
|
||||
- "--port"
|
||||
- "9201"
|
||||
# 控件可见清单(由面板控制;未列出的控件不渲染)
|
||||
- "--widgets"
|
||||
- "dnfs2-minimap;dnfs2-recent;fsys-storage"
|
||||
# 提供方映射(白名单:面板只向这些地址拉取 / 转发)
|
||||
- "--providers"
|
||||
- "dnfs2=http://127.0.0.1:8006;fsys=http://127.0.0.1:9200"
|
||||
- "--cache-ttl"
|
||||
- "60"
|
||||
depends_on: [dnfs2, fsys]
|
||||
sleep_before: 0
|
||||
```
|
||||
|
||||
同时修改 `core/configs/gateway.ini` 白名单:
|
||||
|
||||
```ini
|
||||
[Proxy]
|
||||
opened = http://127.0.0.1:9200;http://127.0.0.1:8006;http://127.0.0.1:5200;http://127.0.0.1:9201;
|
||||
```
|
||||
|
||||
## 4. 控件协议(前端侧)
|
||||
|
||||
### 4.1 运行环境约定
|
||||
|
||||
控件运行在 `<iframe sandbox="allow-scripts">` 中:
|
||||
|
||||
- **无 `allow-same-origin`**:iframe 内为 opaque origin,读不到主页面 DOM / Cookie / localStorage
|
||||
- 控件只能:渲染自己的片段、发 postMessage、`fetch` 面板同源路径(经 BFF)
|
||||
- 控件不得假设自己就是 provider 域:所有业务调用走 `fetch('/api/w/<id>/…')`
|
||||
|
||||
### 4.2 postMessage 事件
|
||||
|
||||
面板 → 控件:
|
||||
|
||||
| 事件 | payload | 说明 |
|
||||
|---|---|---|
|
||||
| `panel.theme` | `{theme:"dark"\|"light"}` | 主题切换(仅 `themeable: true` 的控件) |
|
||||
| `panel.user` | `{user:"user@example.com"}` | 当前登录用户(控件需要个性化时) |
|
||||
|
||||
控件 → 面板:
|
||||
|
||||
| 事件 | payload | 说明 |
|
||||
|---|---|---|
|
||||
| `widget.ready` | `{id}` | 控件渲染完成,面板收起 loading |
|
||||
| `widget.resize` | `{id, w, h}` | 请求调整网格占位(面板限幅在声明 size 的 ±2 内) |
|
||||
| `widget.navigate` | `{id, url}` | 请求面板同源跳转(url 必须同 origin,面板校验) |
|
||||
| `widget.toast` | `{id, message}` | 轻提示 |
|
||||
| `widget.log` | `{id, level, message}` | 前端日志(面板转发 db2 Logs) |
|
||||
|
||||
面板侧校验:优先校验 `e.origin === location.origin`;`payload.id` 必须与绑定的 iframe 一致;必要时以 `e.source === iframe.contentWindow` 兜底校验(opaque origin 场景)。
|
||||
|
||||
### 4.3 布局与主题
|
||||
|
||||
- v1 用 CSS Grid(12 列),按 `manifest.size` 自动流式排布(`grid-auto-flow: dense`),控件间留隙
|
||||
- 面板提供统一控件外壳(标题栏 = title + 版本 + 刷新按钮),内容区为 iframe
|
||||
- 主题变量沿用 dnfs2 风格(`--bg` / `--fg` / `--accent` / …),切换时经 `panel.theme` 同步给 themeable 控件
|
||||
|
||||
## 5. 安全设计
|
||||
|
||||
### 5.1 控件隔离(核心决策)
|
||||
|
||||
**用 iframe 子页面端点 `/widget/<id>`,而非 `srcdoc`。** 原因:
|
||||
|
||||
- `srcdoc` 文档**继承父页面 CSP**;面板页若保持严格 CSP(`script-src 'self'`),会连带封死控件的内联脚本
|
||||
- 独立端点可携带**控件文档自己的 CSP**,实现"父页面严格、控件文档自管"的分层
|
||||
|
||||
```http
|
||||
HTTP/1.1 200 OK
|
||||
Content-Type: text/html
|
||||
Content-Security-Policy: default-src 'none'; style-src 'unsafe-inline'; script-src 'unsafe-inline'; img-src 'self' data:; connect-src 'self'; sandbox allow-scripts
|
||||
```
|
||||
|
||||
配合 iframe 属性双保险:`<iframe sandbox="allow-scripts" src="/widget/<id>">`。即使 CSP 头缺失,`sandbox` 属性仍保证 opaque origin + 禁止顶层导航。
|
||||
|
||||
### 5.2 信任边界
|
||||
|
||||
| 边界 | 策略 |
|
||||
|---|---|
|
||||
| provider 拉取 | 仅 `--providers` 白名单地址,内部 HTTP;manifest 全字段校验 + 大小上限 |
|
||||
| 浏览器直连 | 面板在 `Proxy.opened`,必须经 ca 验签;身份头存在性由 gateway 不变量保证 |
|
||||
| BFF 转发 | 目标仅限该 widget 所属 provider;透传 `X-Auth-User` / `X-Auth-Admin` / `Authorization`,且 strip-then-set 剥离客户端伪造值(对齐 ca 契约) |
|
||||
| `widget.navigate` | url 必须同 origin(`new URL(u).origin === location.origin`),否则拒绝 |
|
||||
|
||||
### 5.3 日志与审计
|
||||
|
||||
- 拉取失败 / 超时 / 校验失败 → db2 Logs(warn 级别)
|
||||
- manifest `version` 变化 → info 日志(审计"谁升级了控件")
|
||||
- 明确不做内容审核:控件来自内部服务而非用户输入,信任边界在 provider 白名单处收口(与 dnfs2 用户样式节点区分,如需对用户内容把关仍走 dnfs2 的 ai_review 链路)
|
||||
|
||||
## 6. 目录结构(建议)
|
||||
|
||||
```
|
||||
panel/
|
||||
├── Makefile
|
||||
├── _set.sh / _run.sh
|
||||
├── src/
|
||||
│ ├── main.cpp # libhttp HttpServer + 路由注册 + 日志
|
||||
│ ├── widget_store.cpp # provider 拉取 / manifest 校验 / TTL 缓存
|
||||
│ └── bff.cpp # /api/w/<id> 转发(白名单 + 身份头透传)
|
||||
└── public/
|
||||
├── panel.html # 背景视图 + 网格外壳
|
||||
└── panel.js # 拉 /widgets → 建 iframe → postMessage 协议
|
||||
```
|
||||
|
||||
## 7. 实现阶段
|
||||
|
||||
| 阶段 | 内容 | 验收 |
|
||||
|---|---|---|
|
||||
| M1 | 骨架:`_set.sh` / `_run.sh`、libhttp 服务、`/health`、静态面板页、widget_store 拉取 + 缓存、`/widgets`、`/widget/<id>` | curl 可见清单与子页面 |
|
||||
| M2 | 页面 JS:网格渲染、iframe sandbox 加载、postMessage 协议、loading / 占位态 | 浏览器两个示例控件可显示 |
|
||||
| M3 | BFF 转发 + 主题同步 + `widget.log` → db2 | 控件调通自己后端接口 |
|
||||
| M4 | 安全加固审计(CSP / 沙箱 / 白名单复核)+ `tests/blackbox/test_panel.py` + 本文档定稿 | 黑盒测试通过 |
|
||||
|
||||
## 8. 测试(tests/blackbox/test_panel.py 思路)
|
||||
|
||||
- 起一个 mock provider(Python 标准库 HTTP server)提供两个 manifest
|
||||
- 断言项:`/widgets` 返回清单;`/widget/<id>` 响应含 CSP `sandbox` 头与内联片段;provider 下线后 `/widgets` 状态为不可用占位;BFF 转发透传 `X-Auth-User`;超大 manifest 被拒;未列入 `--widgets` 的控件不渲染
|
||||
|
||||
## 9. 风险与取舍
|
||||
|
||||
| 取舍 | 说明 |
|
||||
|---|---|
|
||||
| iframe 开销 | 每个控件一个 iframe(含独立文档 + CSP),内存与加载略重;换来的是控件间样式/脚本彻底隔离,控件可独立演进 |
|
||||
| 聚合单点 | 面板是控件编排的单点;控件业务能力仍分布在各自服务,面板故障只影响聚合视图,不影响各服务自身 |
|
||||
| 无在线编排 | v1 布局静态声明(manifest.size + 面板清单顺序);拖拽/自定义布局留给 v2 |
|
||||
|
||||
## 10. 参考
|
||||
|
||||
- gateway 转发白名单与身份契约:[gateway/README.md](../gateway/README.md)
|
||||
- 注入机制先例(js 特效 iframe 沙箱 / style 主题):[docs/dnfs_custom_themes.md](dnfs_custom_themes.md)
|
||||
- 服务注册与 CLI 配置约定:[core/services.yaml](../core/services.yaml)、[ca/README.md](../ca/README.md)(strip-then-set 契约)
|
||||
- HTTP 服务端库:[libhttp/README.md](../libhttp/README.md)
|
||||
@@ -0,0 +1,101 @@
|
||||
# actpred 动作归一化与逐列权重开发方案
|
||||
|
||||
> 文档类型:开发方案
|
||||
> 模块:actpred(联动 pybot 数据契约)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-14
|
||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`
|
||||
|
||||
---
|
||||
|
||||
## 1. 问题陈述
|
||||
|
||||
闭环评估(R4)暴露:z-score 归一化把 18 列方差拉平,隐式实现"全列等权",导致**舵机精度 ↔ spk.F 频率跟踪的权衡**——未归一化模型舵机准但 spk.F 崩(21k MSE),归一化模型 spk.F 改善 3× 但简单动作恶化 1.3-25×。
|
||||
|
||||
本质:**缩放方式(数值稳定性)与列重要性(优化方向)是两个正交维度**,当前被隐式耦合。方案目标:固定物理范围归一化(数值稳定、免 stats 文件)+ 显式逐列权重(舵机优先),两者解耦。
|
||||
|
||||
## 2. 方案设计
|
||||
|
||||
### 2.1 固定物理范围 min-max 归一化(actpred 侧)
|
||||
|
||||
位置决策:**归一化保留在 actpred 侧,pybot 继续记录原始量纲**。理由:
|
||||
|
||||
- **raw at rest**:episode 是数据资产,原始角度/频率可复用于逆运动学、碰撞检测等分析,归一化不可逆
|
||||
- **反归一化链条**:若 pybot 侧归一化,`--loop` 执行前需反归一化,min/max 表要同时存在于 recorder 与 loop 两端;现状 `predict()` 返回原始量纲,pybot 无感知,边界更干净
|
||||
- **代码库先例**:IMU 即"pybot 记录原始、actpred `imu_norm.py` 按物理量程归一化",动作走同一条路保持一致
|
||||
- **pybot 零依赖承诺**:归一化逻辑入 recorder 会加深 pybot 对 numpy 的依赖面
|
||||
|
||||
固定范围常量表(新增 `actpred/action_norm.py`,类比 `imu_norm.py`):
|
||||
|
||||
| 列 | 物理范围 | 依据 |
|
||||
|---|---|---|
|
||||
| 0-9 舵机角 | **[0, 180]** | `pybot/hd/servo.py` `MAX_ANGLE=180.0`(校准后全行程 0-180°) |
|
||||
| 10 spk.F | [0, 5000] | `@Speaker max_f=5000` |
|
||||
| 11 spk.delay | [0, 2000] | 实用上限(脚本用 10-500ms) |
|
||||
| 12-17 保留位 | 恒 0 | 不参与归一化,损失权重 0 |
|
||||
|
||||
min-max 映射:`x_norm = (x - lo) / (hi - lo)` → [0,1](或映射到 [-1,1],与 IMU 归一化输出区间一致,待实现时定)。
|
||||
|
||||
优点:免 action_stats.npy(消除 stats 路径类 bug 的整类风险)、零状态部署、数值区间有界。风险:min-max 对离群值敏感——若真实数据某列越界,需 clamp 到 [0,1] 并告警(防御口径与 `_action_bounds` 一致)。
|
||||
|
||||
可选细化(后续):按各舵机 `|max|`(s1=90°、s2=180°…)分列归一化,90° 舵机能用满 [0,1],分辨率更高;需要读取 @Servo 块的 max 列构建表。
|
||||
|
||||
### 2.2 逐列加权损失
|
||||
|
||||
```python
|
||||
COLUMN_WEIGHTS = np.array([
|
||||
1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, # 0-9 舵机(安全关键)
|
||||
0.2, 0.2, # 10-11 spk.F / spk.delay(声音与节奏,不影响运动)
|
||||
0.0, 0.0, 0.0, 0.0, 0.0, 0.0, # 12-17 保留位(恒 0,纯噪声)
|
||||
], dtype=np.float32)
|
||||
# loss = mean(w * (pred - target)^2) / mean(w)
|
||||
```
|
||||
|
||||
实施点:
|
||||
|
||||
- `train.py` / `pretrain_mt.py` 的 `train_one_epoch` / `validate`:MSE 处乘权重向量(权重存 DataConfig,随 checkpoint 序列化,推理端无需感知)
|
||||
- 权重可经 CLI `--column-weights` 覆盖(默认用上表)
|
||||
- 训练集/验证集用同一权重;`evaluate.py` 的逐列误差报告不受影响(报告仍按无权重口径,便于对照)
|
||||
|
||||
### 2.3 归一化与权重的组合矩阵
|
||||
|
||||
| 配置 | 数值稳定性 | 舵机精度 | spk.F 跟踪 | 预期适用 |
|
||||
|---|---|---|---|---|
|
||||
| raw + 等权(模型 A) | 差(loss 数百量级) | 好 | 崩 | 仅演示/对照 |
|
||||
| z-score + 等权(模型 B) | 好 | 差 | 好 | 现状 |
|
||||
| **固定范围 + 加权(本方案)** | 好 | 好 | 中(w=0.2 保留梯度) | **目标** |
|
||||
|
||||
## 3. 配套改进(同批实施)
|
||||
|
||||
### 3.1 R4 评估器修正(指标失效修复)
|
||||
|
||||
1. **oob_rate 恒值列退化**:span<阈值(如 1e-3)的列不参与越界判定;或改用固定物理界(舵机 [0,180]、spk.F [0,5000])。推荐后者——与 §2.1 常量表共用,且"越物理界"比"越经验界"更符合安全语义
|
||||
2. **error_growth 长度混叠**:按 episode 归一化后聚合(每 episode 自身首末段比值取中位数),避免短 episode 尖峰主导曲线
|
||||
|
||||
### 3.2 R1-B rollout + scheduled sampling
|
||||
|
||||
闭环误差累积源于 teacher forcing 训练与自预测推理的不匹配。已实现的 `--rollout-len K --ss-prob p` 直接对标:建议 `K=4, p=0.5`,与加权损失同批训练对照。
|
||||
|
||||
### 3.3 数据扩展
|
||||
|
||||
- **实机录制**(根本瓶颈):dry-run 帧为噪声、IMU 全零,视觉/姿态通路无真实信号。实机 `--record` 后复跑本实验
|
||||
- 验证集 1 个 episode 选模噪声大:val_split 0.15→0.3
|
||||
- 补充站立/保持类长序列(当前 std 类仅 1-81 步)
|
||||
|
||||
## 4. 实施步骤
|
||||
|
||||
| 步骤 | 内容 | 产出 |
|
||||
|---|---|---|
|
||||
| 1 | 新增 `actpred/action_norm.py`:固定范围常量表 + normalize/denormalize + clamp 防御 | 独立纯 numpy 模块 + 单测 |
|
||||
| 2 | DataConfig 增 `normalize_mode`(`zscore`/`fixed_range`)与 `column_weights`;dataset/train 接入 | 训练链路支持 + 单测 |
|
||||
| 3 | `evaluate.py` oob 改物理界、error_growth 按 episode 归一化 | 评估器修正 + 单测 |
|
||||
| 4 | 训练对照:固定范围+加权 vs z-score+等权(同种子、同数据) | 模型 C/D checkpoint |
|
||||
| 5 | 复跑 R4 闭环评估,验证权衡解耦 | 对比报告 |
|
||||
|
||||
步骤 1-3 可并行;步骤 4-5 依赖前置完成。预计步骤 1-3 改动集中在 `actpred/`(`action_norm.py` 新增、`config.py`/`dataset.py`/`train.py`/`evaluate.py` 小改),不影响 pybot 与部署链(归一化在推理端内部完成,`predict()` 输出仍为原始量纲)。
|
||||
|
||||
## 5. 决策记录
|
||||
|
||||
- **归一化位置**:讨论过"移入 pybot 按列 min/max 归一化"。结论:不采纳(raw at rest、反归一化链条变长、IMU 先例不一致),但采纳其"固定物理范围"内核,落在 actpred 侧常量表实现
|
||||
- **舵机范围**:±180° 为误述,实际全行程 0-180°(`servo.py MAX_ANGLE=180.0`),归一化表按 [0,180]
|
||||
- **spk.F 对数尺度**:备选未定。若加权后 spk.F 跟踪仍不足,再考虑 log 变换(频率感知是对数的,10→1000 跨两个数量级)
|
||||
@@ -0,0 +1,95 @@
|
||||
# actpred Hailo 推理速度测试报告:cv2.LUT 预处理优化验收
|
||||
|
||||
> 文档类型:测试报告
|
||||
> 模块:actpred(hailo_infer.py cv2.LUT 快路径)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-17
|
||||
> 前置:`actpred_hailo_preprocess_optimization_plan_20260817.md`(优化方案,已实施)、
|
||||
> `actpred_inference_speed_benchmark_20260813.md`(旧基准)
|
||||
|
||||
---
|
||||
|
||||
## 1. 测试环境
|
||||
|
||||
| 项 | 值 |
|
||||
|---|---|
|
||||
| 设备 | Raspberry Pi 5(4×A76)+ AI HAT 2(Hailo-10H,/dev/h1x-0) |
|
||||
| 软件 | hailo_platform 5.3.0、opencv-python-headless 5.0.0、torch CPU 2.x |
|
||||
| 模型 | actpred_v4.hef(onehot 1864 logits,fixed_range,640×480×3 帧输入) |
|
||||
| 数据 | 新录 640×480 数据集 episode_0000(图像与模型输入同尺寸,免 resize) |
|
||||
| 输入 | 3 帧 640×480 RGB + 5×18 动作历史(stand 姿态,spk 列 0/10 契约) |
|
||||
|
||||
## 2. 测试方法
|
||||
|
||||
- 预热 10 次后连续 200 次取均值(端到端 `predict` 调用,含预处理/绑定/NPU/解码)
|
||||
- 分段计时:预处理(LUT 量化 + 绑定)/ NPU run / 后处理(反量化 + 1864 logits 解码)
|
||||
- 对照:PyTorch CPU 同模型同输入(200 次均值)
|
||||
|
||||
## 3. 结果
|
||||
|
||||
### 3.1 端到端延迟
|
||||
|
||||
| 后端 | 延迟 | 相对优化前 |
|
||||
|---|---|---|
|
||||
| HEF(cv2.LUT 快路径) | **16.1 ms** | 52.4 → 16.1 ms(**3.3×**) |
|
||||
| PyTorch CPU | 107.9 ms | — |
|
||||
| 加速比(HEF vs CPU) | **6.7×** | 优化前 2.1× |
|
||||
|
||||
### 3.2 流水线分段(HEF)
|
||||
|
||||
| 阶段 | 耗时 | 占比 |
|
||||
|---|---|---|
|
||||
| 预处理(cv2.LUT + 输入绑定) | 3.0 ms | 19% |
|
||||
| NPU run(HEF 纯推理) | 12.5 ms | 78% |
|
||||
| 后处理(反量化 + decode) | 0.6 ms | 4% |
|
||||
|
||||
**瓶颈已从 Python 预处理转移到 NPU 本身**——预处理由优化前 26.5ms 降至 3.0ms(8.8×)。
|
||||
|
||||
### 3.3 闭环控制预算
|
||||
|
||||
| 控制率 | 周期预算 | HEF 占用 |
|
||||
|---|---|---|
|
||||
| 10 Hz | 100 ms | 16% |
|
||||
| 20 Hz | 50 ms | 32% |
|
||||
| 30 Hz | 33 ms | 48% |
|
||||
|
||||
10Hz 闭环下 NPU 仅占 16% 预算;上探 30Hz 也只占一半,余量充足。
|
||||
|
||||
## 4. 一致性验收
|
||||
|
||||
### 4.1 cv2.LUT vs float 路径输出(NPU 实测,max_diff)
|
||||
|
||||
| 输入类型 | max_diff |
|
||||
|---|---|
|
||||
| 640×480 RGB(主路径直通) | 0.0000 |
|
||||
| 224×224 RGB(PIL resize) | 0.0000 |
|
||||
| 全零帧 / 全 255 帧 | 0.0000 |
|
||||
| 灰度 (H,W)(PIL 转 RGB) | 0.0000 |
|
||||
| RGBA (H,W,4)(PIL 合成) | 0.0000 |
|
||||
| 奇数尺寸 333×777(PIL resize) | 0.0000 |
|
||||
| 随机帧 + 随机历史 × 8 seed | 0.0000(worst) |
|
||||
|
||||
### 4.2 无 cv2 回退
|
||||
|
||||
`HAS_CV2=False` 时 `_frame_luts=None` → float 路径,输出与 cv2 路径 max_diff = 0.0。
|
||||
|
||||
### 4.3 回归验收
|
||||
|
||||
| 项 | 结果 |
|
||||
|---|---|
|
||||
| `_npu_test.py` 4 角(真实/全零帧 × 真实/全零历史) | PASS(img/hist 敏感度 > 0.1) |
|
||||
| HEF vs torch(原始动作空间) | 0.30°(预算 ≤1°) |
|
||||
| 全套单测(test_actpred.py) | 188 tests OK |
|
||||
|
||||
## 5. 结论
|
||||
|
||||
1. cv2.LUT 快路径端到端 **16.1ms**,相比优化前 52.4ms 提升 3.3×;预处理不再是瓶颈
|
||||
2. 与 float 路径**逐点一致**(7 类输入 + 8 seed 随机全部 0 差异),无隐藏分支
|
||||
3. 无 cv2 环境自动回退,行为不变;量化误差(vs torch 0.30°)未回归
|
||||
4. 10Hz 闭环预算占用 16%,支持控制率上探至 30Hz
|
||||
|
||||
## 6. 附注
|
||||
|
||||
- 优化前数据(52.4ms/26.5ms)来自 2026-08-17 上午 640×480 数据集基准
|
||||
- cv2.LUT 的"归一化+量化"合并表数学上与 float 路径等价(0 LSB 取整差)
|
||||
- 依赖:opencv-python-headless>=4.9(已入 requirements.txt;缺失仅损失加速,功能正常)
|
||||
@@ -0,0 +1,96 @@
|
||||
# actpred 模型 C 实验报告:固定范围归一化 + 逐列加权
|
||||
|
||||
> 文档类型:实验报告
|
||||
> 模块:actpred
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-14
|
||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`(模型 A/B)
|
||||
> 开发方案:`actpred_action_norm_weighting_plan_20260814.md`(本实验为其步骤 4)
|
||||
> 后续实验:模型 D(fixed_range + 加权 + R1-B rollout),训练中
|
||||
|
||||
---
|
||||
|
||||
## 1. 实验设置
|
||||
|
||||
按开发方案步骤 4 训练对照模型 C:固定物理范围归一化 + 默认逐列权重。
|
||||
|
||||
| 项 | 值 |
|
||||
|---|---|
|
||||
| normalize_mode | fixed_range([0,1] min-max,`action_norm.py` 常量表) |
|
||||
| column_weights | 默认:舵机 1.0 / spk 0.2 / 保留位 0.0 |
|
||||
| 数据 / 种子 / 增强 | 与模型 A/B 相同(763/41 样本,seed 42,ColorJitter) |
|
||||
| epochs / patience | 100 / 15 |
|
||||
| checkpoint | `actpred/checkpoints/recorded_run_fixed/` |
|
||||
|
||||
## 2. 训练结果(teacher forcing 口径)
|
||||
|
||||
- 34 epochs 触发 early stopping
|
||||
- **best val_loss = 0.000394**(epoch 21)——三方最优(A: 19.02 原始量纲、B: 0.0288)
|
||||
- train loss 收敛至 ~0.0004-0.0008
|
||||
|
||||
**teacher forcing 口径下,方案目标完全达成**:loss 数值稳定(<0.001)、保留位列剔除、spk 列仍保留梯度(w=0.2)。
|
||||
|
||||
## 3. 闭环评估(R4)——目标失败
|
||||
|
||||
| 指标 | 模型 A(raw) | 模型 B(zscore) | 模型 C(fixed+加权) |
|
||||
|---|---|---|---|
|
||||
| overall_mse | 6446.7 | **2445.4** | 19843.5 |
|
||||
| error_growth | 0.005 | 0.082 | **13.43(发散)** |
|
||||
| smoothness_ratio | 0.147 | 1.047 | 3.19 |
|
||||
| oob_rate(物理界) | — | 28.1% | 100% |
|
||||
|
||||
**teacher forcing 指标与闭环行为完全背离**:C 的训练/验证 loss 最好,闭环却首次出现发散(error_growth 13.4 > 1)。
|
||||
|
||||
### 3.1 逐 episode 闭环 mean_mse
|
||||
|
||||
| Episode | A | B | C | C 相对 B |
|
||||
|---|---|---|---|---|
|
||||
| 0000 stand | 25.5 | 134.0 | 52882.2 | 395× 恶化 |
|
||||
| 0001 jump | 20.5 | 221.1 | 35228.6 | 159× 恶化 |
|
||||
| 0002 left-hand | 21252.2 | 6849.7 | **2420.3** | 2.8× 改善 |
|
||||
| 0003 left-turn | 276.0 | 641.0 | 3630.5 | 5.7× 恶化 |
|
||||
| 0004 lie-down | 35.2 | 46.9 | 470.8 | 10× 恶化 |
|
||||
| 0005 rght-hand | 21340.0 | 8022.9 | **3691.2** | 2.2× 改善 |
|
||||
| 0007 stdd | 49.8 | 115.9 | 47180.4 | 407× 恶化 |
|
||||
| 0008 t0 | 644.5 | 553.3 | 3927.6 | 7.1× 恶化 |
|
||||
| 0009-0012 tilt-* | 29.8-34.2 | 59.6-751.7 | 45111-46268 | 61-775× 恶化 |
|
||||
|
||||
C 在 hand 类 episode(spk.F 大跨度)上**三方最优**(2.2-2.8× 优于 B),证明加权确实教好了频率跟踪;但其余 episode 全部严重发散。
|
||||
|
||||
## 4. 根因分析
|
||||
|
||||
### 4.1 逐列预测定位(episode_0000,76 步闭环滚动)
|
||||
|
||||
| 列 | 预测范围 | 真值 | 判定 |
|
||||
|---|---|---|---|
|
||||
| col 7 s8 舵机 | 1.6 → **159.0** | 0 | 舵机漂移 159° |
|
||||
| col 10 spk.F | 27.6 → **1080.2** | 10 | 频率漂移 1080Hz |
|
||||
| col 11 spk.delay | **-9.8** → 12.6 | 10 | 负时延越界 |
|
||||
| col 12-17 保留位 | ±0.18 | 0 | 正常(排除权重 0 列失控假设) |
|
||||
|
||||
### 4.2 发散机制(尺度放大 + exposure bias)
|
||||
|
||||
1. **尺度放大**:fixed_range 下 1 个归一化单位 = spk.F 5000Hz 原始量(z-score 下各列单位方差,单位误差对应原始量级相当)。spk.F 归一化空间漂移 ±0.2 → 原始 ±1000Hz
|
||||
2. **权重放松**:w=0.2 使训练对 spk.F 施加 1/5 压力,配合 28 倍尺度放大,spk.F 成为闭环最脆弱列
|
||||
3. **反馈滚雪球**:漂移后的 spk.F 回写历史 → 注意力被污染 → 舵机列(col 7)跟漂 → 误差逐轮放大(episode 内 growth 3.6-430×)
|
||||
4. **teacher forcing 掩盖一切**:训练/验证喂真值历史,上述回路从未出现,故 val loss 0.000394 最优而闭环最差
|
||||
|
||||
### 4.3 直接诱因确认
|
||||
|
||||
- 保留位列(权重 0)预测值正常(±0.18)——**排除**"零权重列无梯度失控"假设
|
||||
- 推理端归一/反归一逻辑正确(`predict.py` fixed_range 分支逐行核对)——**排除**实现 bug
|
||||
- `denormalize_action` 不 clamp("保留模型外推能力"的设计决策)放大了越界值的回传,为次要因素
|
||||
|
||||
## 5. 结论
|
||||
|
||||
1. **逐列加权方案部分有效**:spk.F 跟踪达到三方最优(hand 类 2.2-2.8× 优于 B),方案目标"频率也要跟踪"达成
|
||||
2. **方案不完整**:缩放方式(fixed_range)与列权重正交解耦后,暴露了**尺度放大 × exposure bias** 的叠加风险——单步 teacher forcing 训练无法抵御闭环反馈
|
||||
3. **教训**:teacher forcing val loss 不能作为方案验收指标,必须过 R4 闭环评估;本实验正是 R4 的价值证明
|
||||
|
||||
## 6. 后续
|
||||
|
||||
- **模型 D**(训练中):fixed_range + 加权 + R1-B rollout(`rollout_len=4, ss_prob=0.5`),让模型在训练期面对自预测反馈,直击 exposure bias
|
||||
- 若 D 闭环仍不稳,备选路线(已在讨论中评估):
|
||||
- `denormalize_action` 输出 clamp 到物理界(快速缓解)
|
||||
- one-hot 离散化 + CE(对数分箱 + 软标签):消除尺度放大、内置物理界、NPU 友好;代价是序数信息与数据需求
|
||||
- spk.F 对数尺度(z-score 前 log 变换)
|
||||
@@ -146,6 +146,11 @@ resnet HEF 同 API 验证输入通道有效,排除封装问题)。
|
||||
3. 若仍不敏感:检查 `_optimized.har` 输入层下游连通性,定位
|
||||
optimize 断连的具体算子。
|
||||
|
||||
> **更新(同日)**:上述重编译复验已完成,输出仍不敏感;根因已定位
|
||||
> 为预训练 checkpoint 的 BN 统计退化导致图像特征 ±1e6、uint8 量化
|
||||
> 压死全部信号(非编译侧断连)。详见
|
||||
> [根因分析报告](actpred_hef_insensitivity_rootcause_20260813.md)。
|
||||
|
||||
## 6. 结论与产出
|
||||
|
||||
1. AI HAT 2 软件栈完整可用,新 API 推理链路实测打通(resnet50 307 FPS)。
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
# actpred Hailo 推理预处理优化:cv2.LUT 查表合并归一化与量化(已实施)
|
||||
|
||||
> 文档类型:开发方案(已实施)
|
||||
> 模块:actpred(部署链 hailo_infer.py)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-17(原型实测 + 集成实施)
|
||||
> 前置:`actpred_inference_speed_benchmark_20260813.md`(速度基准)、
|
||||
> 2026-08-17 新 640×480 数据集重测(本方案 §2 数据)
|
||||
|
||||
> 📌 **版本历程**
|
||||
> 1. 初版:numpy 查表法为核心,预期端到端 18-20ms
|
||||
> 2. 原型实测(§3.4):**numpy 高级索引查表在 Pi5 上无加速**(fancy
|
||||
> indexing 的 gather 开销与逐元素 float 运算相当),预期不成立;
|
||||
> 合并系数乘加仅 -5ms;**cv2.LUT 实测 15.8× 加速**(§3.4)
|
||||
> 3. **已实施(§5)**:`hailo_infer.py` 集成 cv2.LUT 快路径(可选依赖,
|
||||
> 无 cv2 自动回退 float 路径),端到端 52.4 → **16.0 ms**,输出与
|
||||
> float 路径逐点一致(max_diff = 0)
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景与动机
|
||||
|
||||
2026-08-17 用新录制的 640×480 数据集(免 resize 训练样本)重测 HEF 推理速度:
|
||||
|
||||
| 阶段 | 640×480 输入 | 224×224 输入(旧数据) |
|
||||
|---|---|---|
|
||||
| 端到端 predict | **52.4 ms** | 60.0 ms |
|
||||
| 预处理 | **35.9 ms** | 43.1 ms |
|
||||
| NPU run | 12.5 ms | 12.9 ms |
|
||||
| 后处理 | 0.6 ms | 0.6 ms |
|
||||
|
||||
输入尺寸已与模型一致(resize 变 no-op),**预处理仍是端到端延迟的主导项**(68%)。
|
||||
旧基准报告(2026-08-13)已记录同类问题:1080p 采集的预处理缩放约 45ms 主导端到端延迟,
|
||||
当时通过把默认分辨率改回 640×480 缓解;本次分析证明 640×480 下 Python 端预处理
|
||||
仍有约 36ms 的优化空间——但空间比初版估计小(见 §3.4 实测)。
|
||||
|
||||
## 2. 瓶颈分析(100 次均值分段计时)
|
||||
|
||||
| 子步骤 | 耗时 | 说明 |
|
||||
|---|---|---|
|
||||
| `_frame_float`(3 帧) | **27.2 ms** | PIL fromarray/convert + resize(同尺寸 no-op 仍有开销)+ `(x/255-mean)/std` float 运算(92 万像素×3 通道) |
|
||||
| `create_bindings` + 帧量化绑定 | **9.1 ms** | `_bind_frames`:float 归一化帧 → `round(x/q.scale+q.zp)` 再量化回 uint8 |
|
||||
| history/pred 绑定 | 0.07 ms | 可忽略 |
|
||||
|
||||
**关键冗余**:当前流程先做 float 归一化(`_frame_float`),随后又量化回 uint8(`_bind_frames`)——
|
||||
中间 float 帧是纯中间态,归一化与量化可以**合并为一步**。
|
||||
|
||||
## 3. 优化方案
|
||||
|
||||
### 3.1 主方案:合并系数乘加(纯 numpy,无新依赖)
|
||||
|
||||
HEF 帧输入是 uint8 量化值:`q = round(((x/255 - mean)/std) / q_scale + q_zp)`。
|
||||
把 4 步运算合并为**一次乘加**(系数预计算,按通道):
|
||||
|
||||
```
|
||||
q = round(x * A[c] + B[c])
|
||||
A[c] = 1 / (255 * std[c] * q_scale)
|
||||
B[c] = -mean[c] / (std[c] * q_scale) + q_zp
|
||||
```
|
||||
|
||||
- 现流程 4 次全数组运算(`/255`、`-mean`、`/std`、`/q_scale+q_zp`)→ 1 次乘加
|
||||
- 数值完全等价(实测 max_diff = 0 LSB,见 §3.4)
|
||||
- 直通模式(`_normalize_frames=False`)同理:`A=1/q_scale, B=q_zp`
|
||||
|
||||
### 3.2 附带优化(与主方案正交)
|
||||
|
||||
1. **免 PIL 路径**:输入已是 np.ndarray 时直接运算,跳过 `Image.fromarray/convert`
|
||||
2. **resize no-op 跳过**:`image.size == (W,H)` 时不调用 `resize()`
|
||||
3. **缓冲复用**:`np.multiply/np.add` 带 `out=` 参数,减少临时数组分配
|
||||
|
||||
### 3.3 已实施快路径:cv2.LUT(可选依赖,无 cv2 自动回退)
|
||||
|
||||
`cv2.LUT(img, lut)` 是 C 实现的 uint8→uint8 查表,单次遍历,无 float 中间态——
|
||||
把"归一化+量化"合并为一步。集成要点(`hailo_infer.py`):
|
||||
|
||||
- 惰性导入 cv2(`HAS_CV2`);**无 cv2 或非逐帧输入布局(ncfhw/nhwc_fused)自动回退 float 路径**
|
||||
- 加载时按每帧 `quant_info` 构建 256 项表(归一化/直通两种,`_normalize_frames` 判定)
|
||||
- **OpenCV 5.0 的 cv2.LUT 只接受单通道表**((256,3) 会断言失败),逐通道应用后 `cv2.merge`
|
||||
- 输入尺寸与模型不一致时先 PIL resize(仅此情况走 PIL)
|
||||
- 依赖安装:`opencv-python-headless`(5.0.0,清华镜像;官方源过慢)
|
||||
- 注:本方案实测后发现**合并系数乘加(§3.1)未实施**——cv2.LUT 收益已远超,
|
||||
避免冗余代码;§3.1 保留作为无 cv2 环境下的潜在备选
|
||||
|
||||
### 3.4 原型实测评估(2026-08-17,100-200 次均值,640×480×3 帧)
|
||||
|
||||
| 方案 | 3 帧耗时 | 对比现 float 28.5ms | 数值差异 |
|
||||
|---|---|---|---|
|
||||
| 现 float 路径 | 28.5 ms | 基线 | — |
|
||||
| numpy 高级索引查表 | 26.6 ms | **无加速(1.0x)** | 0 LSB |
|
||||
| 合并系数乘加 | 23.2 ms | 快 19%(-5ms) | 0 LSB |
|
||||
| **cv2.LUT(逐通道)** | **1.68 ms** | **快 15.8×** | **0 LSB** |
|
||||
|
||||
**结论**:
|
||||
- numpy fancy indexing 的 gather 开销与逐元素 float 运算相当,查表法**无收益**
|
||||
- 纯 numpy 可达的最佳是合并乘加(-5ms),受限于 92 万像素 float32 数组的内存吞吐
|
||||
- cv2.LUT 的 C 查表把帧预处理降一个数量级,且数值完全等价(0 LSB)
|
||||
|
||||
### 3.5 实测收益(已实施,集成后 NPU 实测)
|
||||
|
||||
| 指标 | 优化前 | 优化后(cv2.LUT) |
|
||||
|---|---|---|
|
||||
| 帧预处理(3 帧) | 26.5 ms | **3.0 ms(分段实测含绑定)** |
|
||||
| 端到端 predict | 52.4 ms | **16.1 ms(3.3×)**(200 次均值,见测试报告) |
|
||||
| 输出一致性 | — | 与 float 路径 max_diff = **0.0000**(7 类输入 + 8 seed 随机) |
|
||||
| HEF vs torch | 0.50° | 0.30°(4 角复验 PASS) |
|
||||
| vs PyTorch CPU | 2.1× | **6.7×**(16.1 vs 107.9 ms) |
|
||||
|
||||
闭环 10Hz(100ms 预算)占用率 52% → **16%**,20Hz 32%,30Hz(33ms)48%。
|
||||
|
||||
> 完整验收数据见 `actpred_cv2_lut_benchmark_report_20260817.md`(测试报告)。
|
||||
|
||||
## 4. 一致性约束(关键)
|
||||
|
||||
- **训练侧不动**:`dataset.py` 的 `Resize+ToTensor+Normalize` 保持 float 路径——
|
||||
合并乘加只是推理侧对"归一化+量化"的合并实现,数值等价(实测 0 LSB)
|
||||
- **三后端口径**:`predict.py`(PyTorch,float 输入不量化)、`onnx_infer.py`(ORT,float)
|
||||
不受影响;优化仅作用于 `hailo_infer.py` 的 uint8 量化输入路径
|
||||
- **帧输入判定不变**:`_normalize_frames`(量化网格负零偏判定)决定走
|
||||
归一化还是直通,逻辑与现状一致
|
||||
- 旧 5 输入 / 旧版单图像 HEF 布局同样适用(按各自 quant_info 算系数)
|
||||
|
||||
## 5. 实施步骤(已完成 ✅)
|
||||
|
||||
| 步骤 | 内容 | 状态 |
|
||||
|---|---|---|
|
||||
| 1 | `hailo_infer.py`:加载时按每帧 quant_info 构建 256 项 LUT(归一化/直通两种) | ✅ `_build_frame_luts` |
|
||||
| 2 | `predict` 的 nhwc_frames 分支走 cv2.LUT 快路径(尺寸不匹配时 PIL resize);无 cv2 回退 float | ✅ `_frame_lut_quant` + `HAS_CV2` |
|
||||
| 3 | NPU 实测:640×480 数据集 100 次均值 | ✅ 16.0 ms(52.4 → 16.0) |
|
||||
| 4 | `_npu_test.py` 4 角 + HEF vs torch 对比复验 | ✅ PASS,0.30°(≤1° 预算) |
|
||||
| 5 | 全套单测(TestHailoPredictor 等) | ✅ 188 tests OK |
|
||||
| 6 | 依赖安装:opencv-python-headless 5.0.0(清华镜像) | ✅ venv 已装 |
|
||||
|
||||
## 6. 风险与对策
|
||||
|
||||
| 风险 | 对策 |
|
||||
|---|---|
|
||||
| 查表取整与 float 路径存在差异 | 实测 0 LSB(表项 float32 与现路径同精度);NPU 输出 vs float 路径 max_diff=0,vs torch 0.30° |
|
||||
| 无 cv2 环境 | 惰性导入 + 自动回退 float 路径,行为与优化前完全一致 |
|
||||
| OpenCV 5.0 多通道 LUT 断言失败 | 逐通道单表 + `cv2.merge`(实测无性能损失) |
|
||||
| 非 nhwc_frames 布局(旧 HEF) | 仅该布局启用 cv2 路径,其余走原 float 路径 |
|
||||
| 未来输入尺寸变更 | LUT 与尺寸无关(逐像素映射);尺寸不一致时才走 PIL resize |
|
||||
|
||||
## 7. 决策记录
|
||||
|
||||
- **查表法(初版核心)实测证伪**:numpy 高级索引 gather 在 Pi5 上与逐元素 float
|
||||
运算同速(26.6 vs 28.5ms),"查表降为内存索引即快"的假设不成立——教训:
|
||||
性能方案须原型实测,不能凭直觉估计
|
||||
- **最终方案选 cv2.LUT**:C 实现查表把帧预处理从 26.5ms 降到 1.7ms(15.8×),
|
||||
端到端 52.4 → 16.0ms;数值 0 LSB 等价。合并系数乘加(-5ms)收益被碾压,
|
||||
不再实施(§3.1 保留为无 cv2 环境的备选)
|
||||
- **依赖取舍**:opencv-python-headless 作为可选依赖(惰性导入 + 自动回退),
|
||||
有 cv2 环境收益 3.3×,无 cv2 环境行为不变——收益远大于依赖成本
|
||||
- **训练侧不改**:训练用 float 归一化(增强/数值精度),推理侧查表实现——两侧
|
||||
口径在"归一化值"层面一致,量化差异由部署链固有预算吸收(HEF vs torch 0.30° 余量充足)
|
||||
@@ -0,0 +1,138 @@
|
||||
# actpred HEF 输出不敏感问题根因分析报告
|
||||
|
||||
日期:2026-08-13
|
||||
设备:Raspberry Pi 5 + AI HAT 2(Hailo-10H);编译端 x86 + Hailo DFC 5.3
|
||||
前情:[actpred_hailo_hat2_inference_research_20260813.md](actpred_hailo_hat2_inference_research_20260813.md) §5
|
||||
|
||||
## 1. 背景
|
||||
|
||||
首版 HEF 在 NPU 上输出对全部 5 个输入不敏感(全 0 / 全 255 / 真实数据
|
||||
输出完全一致)。修复连通性检查后 x86 重编译,新 HEF(224×224,含
|
||||
`_optimized.har` 连通性 PASS)拷回树莓派复验,**输出仍对输入不敏感**。
|
||||
本报告记录复验实验与最终根因定位。
|
||||
|
||||
## 2. 复验实验
|
||||
|
||||
### 2.1 新 HEF 扩展敏感性测试(`_npu_test.py` + 变体)
|
||||
|
||||
对 6 种输入变体逐组测试,与基线输出对比(`max_diff`):
|
||||
|
||||
| 变体 | max_diff | 结论 |
|
||||
|---|---|---|
|
||||
| 全零帧 | 0.0000 | 帧路径无效 |
|
||||
| 全 255 帧 | 0.0000 | 帧路径无效(极端值同样无效) |
|
||||
| 全零历史 | 0.0000 | 历史路径无效 |
|
||||
| 全 500 历史 | 0.0000 | 历史路径无效 |
|
||||
| pred_token=0 | 0.0000 | pred_token 无效 |
|
||||
| pred_token=±0.09 | 0.0000 | pred_token 无效 |
|
||||
|
||||
输出为恒定向量 `[11.814, -1.575, 9.451, 4.725, 11.814, 3.15, ...]`
|
||||
(uint8 q=`[17,0,14,8,17,6]`,zp=2,scale=0.7876)。
|
||||
|
||||
### 2.2 新 HEF vstream 量化参数
|
||||
|
||||
| 输入 | 形状 | zp | scale | 判定 |
|
||||
|---|---|---|---|---|
|
||||
| input_layer1~3(帧) | (224,224,3) | 114 | 0.0187 | 正常(≈归一化像素范围/255) |
|
||||
| input_layer4(历史) | (1,5,18) | 0 | 1.961 | 正常(=500/255,与数据范围 0~500 一致) |
|
||||
| input_layer5(pred_token) | (1,1,128) | 129 | **1965.92** | 退化(数据 ±0.09,合理 scale 应 ~0.0007) |
|
||||
|
||||
与 x86 备份 tgz 中旧 HEF(64×64)对比:pred_token scale
|
||||
`18571.7 → 1965.92`。**scale 随校准数据变化,说明抖动校准确实生效**——
|
||||
推翻了"抖动校准未生效"的初始假设。
|
||||
|
||||
### 2.3 控制组(排除环境因素)
|
||||
|
||||
resnet_v1_18.hef 同 API、同驱动实测:真实图 top1=984 vs 全零图
|
||||
top1=111,输出随输入变化。**硬件、驱动 5.3.0、hailo_platform 封装
|
||||
全部正常,问题在 HEF 本身。**
|
||||
|
||||
## 3. 根因:BN 统计退化 → 特征量级 ±1e6 → 量化压死全部信号
|
||||
|
||||
### 3.1 证据链
|
||||
|
||||
1. `pretrain_best_inference.pt` 中全部 4 个 BatchNorm 的 `running_var`
|
||||
均为 `3.7e-4`(每通道逐位相同)。预训练用全零图像占位,BN 从未
|
||||
见到真实数据,running_var 经 ~75 次零方差更新指数衰减至此。
|
||||
2. 推理(eval 模式)下每级 BN 放大 `1/sqrt(3.7e-4 + 1e-5) ≈ 51` 倍,
|
||||
4 级串联放大 ~7e6 倍。实测(torch 复算部署版图):
|
||||
- CNN 末级输出:**±1.4e6 ~ 2.5e6**(真实帧 / 全零帧均是)
|
||||
- `img_tokens`:±2.2e5 ~ 9.6e5
|
||||
- `act_tokens`:±537 ~ 694(历史 0~500 × act_proj 权重,量级正常)
|
||||
3. torch 浮点推理为何不受影响:pre-LN 架构的 LayerNorm 在注意力前
|
||||
归一化 token 幅度,±1e6 的残差流在 float32 下不损失精度。但 uint8
|
||||
量化无法表达这种动态范围。
|
||||
|
||||
### 3.2 失效机理(量化视角)
|
||||
|
||||
DFC 在 token merge 点(`concat(img, act, pred) + pos_embed` 的
|
||||
ew_add)按**最大幅度分配共享量化 scale**:
|
||||
|
||||
```
|
||||
merge scale ≈ 250k / 127.5 ≈ 1965.92 ← 与 HEF 边缘实测一致
|
||||
```
|
||||
|
||||
该 scale 下各支路命运:
|
||||
|
||||
| 支路 | 幅度 | 量化结果 |
|
||||
|---|---|---|
|
||||
| img_tokens | ±235k | 勉强存活(~120 级),但中间激活 ±1e6 超出 uint8 校准范围 → 饱和,任何输入差异被抹平 |
|
||||
| act_tokens | ±537 | 537/1965.92 ≈ 0.27 → 全部量化到 zp(**信号归零**) |
|
||||
| pred_token | ±0.09 | 全部量化到 zp(**信号归零**) |
|
||||
| pos_embed | ±0.06 | 全部量化到 zp(**信号归零**) |
|
||||
|
||||
图像路径经 CNN 内各级退化激活的饱和 + 其余支路归零,最终网络退化为
|
||||
对任何输入输出同一常量——与 2.1 的实测完全吻合。
|
||||
|
||||
### 3.3 为什么连通性检查通过仍然失败
|
||||
|
||||
`dfc_compile.py` 的检查验证的是**图拓扑**(输入层 BFS 可达输出层),
|
||||
拓扑确实连通。问题是**数值量级**:量化把信号压成 0,等价于数值上的
|
||||
断连。两者不矛盾,但都必须在编译前拦截。
|
||||
|
||||
## 4. 结论
|
||||
|
||||
- 当前 pretrain checkpoint 的图像支路数学上**无法 uint8 量化**
|
||||
(特征 ±1e6 远超表达范围),无论重编译多少次都不可用。
|
||||
- 必须先用真实/随机图像数据微调,恢复 BN 统计(特征回到 O(1)),
|
||||
再重走导出→编译链路。
|
||||
- pred_token 抖动校准本身没有 bug:merge 点 scale 由 ±235k 的图像
|
||||
token 主导,pred_token 无论怎么抖都会被压死。
|
||||
|
||||
## 5. 修复与后续
|
||||
|
||||
### 5.1 已落地:导出前量级前置检查
|
||||
|
||||
`actpred/export_deploy.py` 新增 `check_token_ranges()`:导出前用随机
|
||||
噪声帧 + 0~500 均匀历史实测 token 量级,`|img| > 100` 直接中止导出
|
||||
并提示微调;`act/img` 比例 > 32 时警告 merge 会粗量化图像路径。
|
||||
|
||||
实测:退化 checkpoint `|img|=2.46e5` → 正确拦截;模拟健康 BN
|
||||
(var=1)`|img|=0.117` → 放行。全量单测 76/76 通过。
|
||||
|
||||
### 5.2 建议路径
|
||||
|
||||
1. x86 微调:随机图像 + 混入真实帧(本机 6 episode × 30 帧可拷贝
|
||||
过去),动作目标继续用 .mt 动力学;微调后 BN 统计恢复。
|
||||
2. 重走链路:`export_deploy.py`(自带量级检查)→ `gen_calibration.py
|
||||
--deploy` → `dfc_compile.py`(连通性检查)→ 拷回 HEF。
|
||||
3. `_npu_test.py` 4 角敏感性复验作为验收门槛。
|
||||
|
||||
### 5.3 遗留风险
|
||||
|
||||
- 健康场景下 `act_tokens`(±700)仍远大于 `img_tokens`(O(1)):
|
||||
merge 共享 scale 会把图像路径粗量化到 ~2 级。后续可在部署图把
|
||||
`action_history` 输入除以 255 并将 `act_w` 放大 255 倍(线性等价,
|
||||
数学不变),把 act_tokens 压到 O(1) 量级。
|
||||
- 随机图像微调后 BN 统计匹配噪声分布,部署端真实摄像头帧的统计
|
||||
略有差异,特征量级可能偏大——导出时的新检查会兜底拦截,混入
|
||||
真实帧可进一步降低风险。
|
||||
|
||||
### 5.4 验证链(防再犯)
|
||||
|
||||
| 阶段 | 检查 | 脚本 |
|
||||
|---|---|---|
|
||||
| 导出时 | token 量级 | `export_deploy.py` `check_token_ranges` |
|
||||
| 编译前 | ONNX 数值正确性 | `_onnx_validate.py` |
|
||||
| 编译时 | 输入连通性 | `dfc_compile.py` `check_input_connectivity` |
|
||||
| 部署后 | 4 角输入敏感性 | `_npu_test.py` |
|
||||
@@ -4,15 +4,15 @@
|
||||
设备:Raspberry Pi 5(4×A76)+ AI HAT 2(Hailo-10H)
|
||||
模型:actpred(~450K 参数,frame_len=3,history_len=5,action_dim=18)
|
||||
|
||||
> ⚠️ **基线分辨率说明(2026-08-12 补注,N-3 修复)**
|
||||
> ⚠️ **基线分辨率说明(2026-08-13 更新)**
|
||||
>
|
||||
> 本报告 §1~§4 的全部数据均在 **640×480** 输入分辨率下测得。
|
||||
> 提交 `91b3004`(2026-08-12)已将模块默认输入分辨率改为
|
||||
> **1920×1080**(`config.py` 默认 `image_width/image_height`,与 pybot
|
||||
> 相机采集分辨率对齐)。**推理延迟随输入像素量近似线性增长
|
||||
> (1920×1080 ≈ 640×480 的 6.75 倍像素),本报告数字不能直接外推到
|
||||
> 新默认分辨率。** x86 参考机上的 1920×1080 复测数据见 §5;
|
||||
> Pi 5 / Hailo 板载在新分辨率下的数据待补测。
|
||||
> 提交 `91b3004`(2026-08-12)曾将模块默认输入分辨率改为
|
||||
> **1920×1080**(与 pybot 相机采集分辨率对齐);2026-08-13 实机测速
|
||||
> 发现 1080p 采集的预处理缩放(约 45ms)主导端到端延迟(§6.1),
|
||||
> 默认分辨率**改回 640×480**(`config.py` 默认
|
||||
> `image_width/image_height`),§1~§4 数据重新适用。
|
||||
> 1920×1080 的 x86 复测数据见 §5,仅作分辨率影响趋势参考。
|
||||
|
||||
## 1. 测试方法
|
||||
|
||||
@@ -93,11 +93,33 @@ NPU 部署轮次以 224×224 的 `pretrain_best_inference.pt` 为准(HEF 编
|
||||
| ONNX Runtime CPU | 224×224 | 10.7 ms | 纯 `session.run`(预缩放输入) |
|
||||
| Hailo-10 NPU(actpred HEF,首版) | 64×64 | 6.3 ms | 含预处理 + 量化/反量化,驱动 5.3.0 |
|
||||
|
||||
注:本组 ONNX 纯推理(10.7ms)与 §2 的 5.9ms(同为 224 模型)存在
|
||||
差异,未追根(torch≥2.12 dynamo 导出路径、系统负载、ORT 线程调度均
|
||||
可能影响),以本组实测为准;ONNX 仍比 PyTorch 快约 2 倍。NPU 首版
|
||||
HEF 因编译侧输入断开输出恒定,仅作封装性能基线,数值有效性待重编译
|
||||
后复验(见 hailo 调研报告)。
|
||||
### 6.1 归一化轮次 HEF 复测(2026-08-13 同日,输入连通性验证通过后)
|
||||
|
||||
同一轮次对通过 `_npu_test.py` 验收的归一化轮次 HEF
|
||||
(224×224,`hailo_infer.py` 帧归一化判定修复后)实测,预热 5 次取
|
||||
50 次均值(纯 NPU run 取 200 次):
|
||||
|
||||
| 方案 | 输入分辨率 | 延迟 | 备注 |
|
||||
|---|---|---|---|
|
||||
| Hailo-10 NPU 端到端 predict | 224×224 | 10.5 ms | 含预处理 + 量化/反量化 |
|
||||
| Hailo-10 NPU 纯 run(绑定后循环) | — | 4.4 ms | Python API/PCIe 开销主导,HEF 计算本身远小于此 |
|
||||
| Hailo-10 NPU 预处理(3 帧 resize+归一化) | 224×224 | 3.4 ms | PIL resize 为主 |
|
||||
| PyTorch CPU(同机复测) | 224×224 | 22.4 ms | 含预处理 |
|
||||
| ONNX Runtime CPU(同机复测) | 224×224 | 11.0 ms | 含预处理,重新导出的当前 checkpoint |
|
||||
|
||||
**结论**:
|
||||
|
||||
1. 224×224 下 NPU 端到端(10.5ms)与 ONNX Runtime(11.0ms)基本持平,
|
||||
纯 NPU run 的 4.4ms 是 Python 绑定/PCIe 同步开销主导(对照 §2
|
||||
resnet50 串行单帧同为 4.29ms,而 actpred 计算量仅其 ~1/55)。
|
||||
后续可用 `run_async` 流水线把预处理与 NPU 执行重叠。
|
||||
2. NPU 的价值不在绝对延迟而是**把推理从 CPU 卸载**,让 4×A76 专注
|
||||
pybot 主控循环;CPU 侧 ORT 仍是有效的纯软件方案。
|
||||
3. **1080p 采集是最大瓶颈**:预处理随输入像素量增长,实测
|
||||
640×480 输入 11.8ms、960×540 输入 15.7ms、1080p 输入 45.0ms
|
||||
(端到端分别 18.6 / 22.8 / 51.7ms)。部署建议摄像头以
|
||||
640×480 或更低分辨率采集(或 224×224 中心裁剪 ROI),
|
||||
而不是全分辨率后缩放。
|
||||
|
||||
## 参考
|
||||
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
# actpred 多头 one-hot 分类输出开发方案
|
||||
|
||||
> 文档类型:开发方案
|
||||
> 模块:actpred(联动 pybot 数据契约)
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-14
|
||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`、
|
||||
> `actpred_fixed_range_weighted_experiment_20260814.md`(模型 A-E 系列)
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景与动机
|
||||
|
||||
模型 A-E 系列实验定位到两个结构性难题:
|
||||
|
||||
1. **尺度放大**:fixed_range 线性 min-max 下 spk.F 的 1 个归一化单位 = 5000Hz 原始量,闭环反馈中误差滚雪球(模型 C/D 发散根因之一)
|
||||
2. **精度错配**:常用值(10Hz → 0.002)挤在线性尺度的低端,[0,1] 精度预算被浪费(log_range 正在验证的对数尺度解法)
|
||||
|
||||
多头 one-hot 是本问题的**结构性解法**:把 18 维连续回归换成 18 组独立分类头,每组输出一个 one-hot 分布。分类天然消除尺度放大(每箱等距)与物理界问题(箱只覆盖物理范围),且多列 softmax 是 NPU(Hailo)最擅长的计算形态。
|
||||
|
||||
## 2. 方案设计
|
||||
|
||||
### 2.1 分箱设计(每组 one-hot 的维度)
|
||||
|
||||
| 列组 | 分箱方式 | 箱数/列 | 依据 |
|
||||
|---|---|---|---|
|
||||
| 舵机 0-9(10 列) | 均匀 **1°/箱** | 180 | 匹配数据粒度:微动作插值 0.75°/步,5° 箱会把相邻步量化进同一箱、毁掉插值结构 |
|
||||
| spk.F(10) | 对数分箱 | 32 | 频率对数感知,低端聚集;log10(1+x) 均匀切 |
|
||||
| spk.delay(11) | 对数分箱 | 32 | 同左 |
|
||||
| 保留位 12-17 | 无头 | 0 | 恒 0,不参与 |
|
||||
|
||||
**输出总量**:10×180 + 32 + 32 = **1864 logits**(当前 18)。最后层 128→1864 ≈ 238K 参数,模型 409K → ~650K,仍属轻量。
|
||||
|
||||
### 2.2 各层改动
|
||||
|
||||
| 文件 | 改动 |
|
||||
|---|---|
|
||||
| `actpred/bins.py`(新增) | 分箱表 + `encode(x)→one-hot` / `decode(p)→连续值`(含对数箱);纯 numpy,供推理/部署复用 |
|
||||
| `model.py` | PredictionHead 输出 1864 维 → reshape (B, 18, bins) → 逐组 softmax |
|
||||
| `dataset.py` | target 从 (B,18) 连续值 → 每组 one-hot;**软标签**:真值高斯扩散到 ±1-2 箱 |
|
||||
| `train.py` | 损失 = Σ 列权重 × CE(组输出, one-hot 真值);列权重沿用 default_column_weights |
|
||||
| `predict.py` | 输出 decode:**期望值** `Σ pᵢ·箱中心`(推荐)或 argmax 箱中心 → 18 维原始量纲 |
|
||||
| `evaluate.py` | 无感——predict 返回仍是 18 维原始量纲 |
|
||||
| `export_deploy` / Hailo | 多组 softmax 为 NPU 友好形态 |
|
||||
|
||||
### 2.3 解码策略(执行时如何取连续值)
|
||||
|
||||
| 方式 | 行为 | 评估 |
|
||||
|---|---|---|
|
||||
| **期望值**(推荐) | `Σ pᵢ × 箱中心` | 平滑连续,分辨率超越箱宽(1° 箱可出 46.4°);单峰假设 |
|
||||
| argmax | 取最大概率箱中心 | 确定性,量化到 1° |
|
||||
| 按分布采样 | 随机抽箱 | 探索用,后续离线 RL 可复用 |
|
||||
|
||||
## 3. 关键决策点与理由
|
||||
|
||||
1. **舵机 1°/箱而非 5°/箱**:录制数据的微动作插值粒度为 0.75°,箱宽必须匹配数据粒度;5° 箱会把 45.0°/45.75° 压进同一箱,模型学不到特意录制的插值结构
|
||||
2. **软标签**:804 样本对 1864 logits 偏少,多数箱极少被访问。真值 one-hot 高斯扩散到邻近箱,未访问箱也能获得梯度
|
||||
3. **期望值解码补序数信息**:one-hot 丢失"45° 和 46° 相邻"先验,期望值解码在推理端部分恢复;训练端靠软标签缓解
|
||||
4. **pybot 契约零变化**:predict 输出仍是 18 维原始量纲,闭环与录制数据无感知
|
||||
5. **与现有归一化链路的兼容**:one-hot 方案独立于 zscore/fixed_range/log_range——分箱编码替代归一化,不叠加使用
|
||||
|
||||
## 4. 实施步骤
|
||||
|
||||
| 步骤 | 内容 | 产出 |
|
||||
|---|---|---|
|
||||
| 1 | `bins.py`:分箱表 + encode/decode + 对数箱 + 软标签工具 | 独立纯 numpy 模块 + 单测 |
|
||||
| 2 | `dataset.py`:目标 one-hot 化(软标签可开关) | 数据链路 + 单测 |
|
||||
| 3 | `model.py` + `train.py`:多头 softmax + 加权 CE | 模型与训练 + 单测 |
|
||||
| 4 | `predict.py`:期望值/argmax 解码集成 + 冒烟 | 推理链路 + 单测 |
|
||||
| 5 | 训练模型 F,与 A-E 闭环五方对比 | 对比报告 |
|
||||
|
||||
## 5. 与其他方案的定位
|
||||
|
||||
- **模型 E(log_range)正在训练**:若对数尺度已解决精度错配与闭环发散,one-hot 降级为可选优化(分类头天然有界、免归一化常量表的优势仍在)
|
||||
- 若 E 仍失败:one-hot 是下一个强候选(结构性消除尺度放大)
|
||||
- **互补关系**:one-hot 方案与逐列权重、R1-B rollout 正交,可叠加(rollout 治 exposure bias 的结论不变,任何损失形式都需要它)
|
||||
|
||||
## 6. 决策记录
|
||||
|
||||
- **箱宽**:初拟 5°/箱(保守,顾虑数据量),讨论后定为 1°/箱——数据插值粒度 0.75° 决定箱宽下限;spk 两列对数 32 箱
|
||||
- **单头多选 vs 多头单选**:softmax 单头在数学上是单选择分布;"多组 one-hot" = 18 个独立 softmax 头并行,每列一个选择
|
||||
- **多峰分布**:当前数据单峰为主,期望值解码够用;若未来出现双峰(如跨脚本姿态冲突),升级为逐箱 sigmoid 或 MDN
|
||||
@@ -0,0 +1,212 @@
|
||||
# actpred × pybot 研究方向报告:循环神经网络思想与机器人学习演进路线
|
||||
|
||||
日期:2026-08-14(Asia/Shanghai)
|
||||
范围:`actpred/`(动作预测,大脑)与 `pybot/`(机器人运行时,身体)
|
||||
类型:技术研究方向调研(`advisory`;涉及部署可行性的判断标注 `verify_pending`,需 Hailo 实测确认)
|
||||
|
||||
## 0. 报告说明
|
||||
|
||||
本报告源于两个递进的调研问题:(1)动作预测能否引入更多循环神经网络(RNN)思想;(2)在此之上的整体研究方向。报告先剖析 actpred 当前建模范式的本质与痛点,再给出 RNN 思想的四个切入点,最后组织成分层的演进路线图。所有判断基于当日代码基线(进程内闭环已落地、Hailo NPU 部署链已实机验证)与项目现实约束。
|
||||
|
||||
## 1. 现状建模范式与核心痛点
|
||||
|
||||
### 1.1 当前是"单步窗口预测",模型无跨步隐状态
|
||||
|
||||
actpred 的 `ActionPredictor`([model.py:89-147](../../actpred/model.py#L89-L147))每步独立地用
|
||||
`[F帧图像, H步动作, PRED]` token 序列做一次性 Transformer 注意力融合,取 PRED 位置回归下一动作。
|
||||
关键事实:**模型本身没有任何跨步隐状态**,跨步记忆完全由 H 步(默认 5)原始动作向量窗口承载。
|
||||
|
||||
训练侧是 teacher forcing([dataset.py:112-117](../../actpred/dataset.py#L112-L117)):
|
||||
`action_history` 取真实历史,`target` 取下一真实动作,单步监督。
|
||||
|
||||
### 1.2 由此产生的三个真实痛点
|
||||
|
||||
| 痛点 | 根因 | 后果 |
|
||||
|------|------|------|
|
||||
| 长程依赖被截断 | 记忆容量 = H 步原始动作 | 超过 5 步的运动模式(起步→加速→匀速的长惯性)在窗口外丢失 |
|
||||
| 动作不平滑 | 每步独立回归,无惯性先验 | 预测抖动,对物理舵机是磨损/损坏风险 |
|
||||
| exposure bias | 训练用真实历史、闭环推理用自预测历史 | 分布失配,误差随步数累积,闭环易发散 |
|
||||
|
||||
### 1.3 决定性约束:Hailo NPU 部署链
|
||||
|
||||
[model.py:94](../../actpred/model.py#L94) 明确"除 batch 外全部维度静态,保证 ONNX/NPU 部署友好";
|
||||
[export_deploy.py](../../actpred/export_deploy.py) 整套 DFC 改造(图像输入拆帧、AvgPool 替换、
|
||||
注意力手工展开)都围绕**静态图**。因此**循环算子(GRU/LSTM)是可行性的分水岭**——
|
||||
DFC 对循环算子的支持通常弱于静态卷积。这一约束直接决定下文各方案的取舍。
|
||||
|
||||
## 2. RNN 思想的四个切入点(按部署链冲击排序)
|
||||
|
||||
### 方案 C · 残差/惯性先验:预测 Δaction(零参数,立即见效)⭐
|
||||
|
||||
把预测目标从绝对动作改为增量:`â_t = a_{t-1} + Δ`,模型只学 `Δ`。
|
||||
|
||||
- **本质**:等价于对角权重固定为 1 的线性 RNN(积分器),是 RNN 思想的极简形式。
|
||||
- **收益**:输出被惯性锚定,天然平滑;训练更易收敛;舵机抖动大减。
|
||||
- **部署影响**:零(只改 dataset 的 target 构造 + predict 逆变换时加回 `a_{t-1}`)。
|
||||
- **代价**:极小。**强烈建议无论选哪条路线都先做这条**。
|
||||
|
||||
### 方案 B · 训练目标:多步展开 + scheduled sampling(治 exposure bias)
|
||||
|
||||
不改模型,只改训练:
|
||||
- dataset 一次取 K 步连续样本而非单步;
|
||||
- 第 1 步用真实历史,第 2 步起用模型自己的预测回传(scheduled sampling),loss 对 K 步累加。
|
||||
|
||||
- **收益**:直接消除训练/推理分布失配,闭环稳定性大幅提升。
|
||||
- **部署影响**:零(模型结构不变,已部署 Hailo 推理链不动)。
|
||||
- **代价**:只改 train.py 训练循环 + dataset 采样逻辑,风险最低。
|
||||
|
||||
### 方案 A · 混合架构:动作历史用 GRU 压缩(进阶)
|
||||
|
||||
`act_proj` 后不接 H 个独立 token,而是过一层 GRU,取最终隐状态压成 **1 个"动作记忆 token"**,
|
||||
与 F 个图像 token + PRED 一起进 Transformer:
|
||||
|
||||
```
|
||||
现状: H 个动作 token ──────────────────┐
|
||||
改进: H 步动作 → GRU → h_final (1 token) ├→ [img×F, mem×1, PRED] → Transformer
|
||||
```
|
||||
|
||||
- **收益**:跨步记忆容量与 H 解耦(GRU 隐状态压缩全部历史),长程惯性可建模。
|
||||
- **部署影响**:token 序列从 `F+H+1` 变 `F+2`,更短更静态;唯一新增算子是 GRU——
|
||||
**关键风险点**,需先验证 Hailo DFC 对 GRU 的支持(可用固定 H 步展开的 cell 堆叠规避真循环)。
|
||||
- **代价**:中等改动(model.py + export_deploy.py 的 GRU 展开)。
|
||||
|
||||
### 方案 D · 全 recurrent 编码器(表达力最强,部署风险最高)
|
||||
|
||||
图像帧也过 ConvGRU,跨步隐状态贯穿视觉与动作。**但这打破"除 batch 外全静态 shape"的
|
||||
NPU 部署前提**,需重新设计部署图。不推荐在当前 Hailo 目标下做。
|
||||
|
||||
### 切入点评审结论
|
||||
|
||||
```
|
||||
第 1 步(零风险): C 残差预测 ──→ 立即改善平滑性,验证收益
|
||||
第 2 步(零部署影响): B 多步展开训练 ──→ 治 exposure bias
|
||||
第 3 步(需验证): A GRU 记忆 token ──→ 先做 Hailo GRU 支持单测
|
||||
├─ 支持 → 展开 GRU 进部署链
|
||||
└─ 不支持 → 保持 Transformer,仅用 B+C
|
||||
```
|
||||
|
||||
**C+B 组合是性价比最高的第一步**:不改部署链、不增参数,直接解决"抖动 + exposure bias"
|
||||
两个最痛的问题。A 留作需要更强长程记忆时的进阶项。
|
||||
|
||||
## 3. 整体研究方向:分层演进路线
|
||||
|
||||
在 RNN 思想之外,从机器人学习前沿看,两模块的演进空间组织为四层。现实约束:
|
||||
**纯模仿学习(无奖励信号)、MPU6050 已装未用、数据靠实机 `--record`(稀缺)、Hailo NPU 部署硬约束**。
|
||||
|
||||
### Tier 0 · 立足现状(零部署风险,先做)
|
||||
|
||||
| 编号 | 方向 | 价值 | 改动 |
|
||||
|------|------|------|------|
|
||||
| R1 | 残差预测 + 多步展开训练(C+B) | 治抖动+失配,所有后续工作的地基 | dataset + train |
|
||||
| R2 | **MPU6050 姿态回流** ⭐ | 给模型"本体感觉",是机器人标准状态量 | pybot 读 IMU + actpred 加输入维 |
|
||||
| R3 | 数据增强扩充 | 稀缺数据的低成本放大 | dataset 增方法 |
|
||||
| R4 | **离线评估基准** | 没有基准,后续改进全是盲改 | 新增评估脚本 |
|
||||
|
||||
**R2 被低估**:pybot 已有 [mpu6050.py](../../pybot/pybot/hd/mpu6050.py) Kalman 姿态解算,但完全没进
|
||||
actpred 的状态表征——模型对自身俯仰/横滚/角速度是盲的。接入 IMU 6 轴作为额外 token 或并入
|
||||
动作向量,既补本体感觉,又为 R9 奖励设计提供信号源。
|
||||
|
||||
**R4 是基础设施**:现在只有开环 MSE。需要闭环代理指标——留出 episode 上开环滚动预测,
|
||||
测多步累积误差、动作平滑度(相邻动作差分方差)、物理可行性(越界率)。
|
||||
|
||||
**R3**:当前仅亮度/对比度/饱和度/色调增强([dataset.py:99-102](../../actpred/dataset.py#L99-L102)),
|
||||
可加时间扭曲(帧率扰动)、动作镜像(左右对称对)、episode 混合拼接。
|
||||
|
||||
### Tier 1 · 模型能力(需验证部署兼容)
|
||||
|
||||
| 编号 | 方向 | 说明 |
|
||||
|------|------|------|
|
||||
| R5 | GRU 记忆 token(方案 A) | 先验 Hailo DFC 循环算子支持 |
|
||||
| R6 | 多模态状态融合 | 视觉+IMU+本体感觉统一表征(R2 完整版) |
|
||||
| R7 | 世界模型雏形 | 学"动作→下一帧"预测头,支撑 MPC 规划与 model-based RL |
|
||||
|
||||
### Tier 2 · 学习范式跃迁(数据/奖励是瓶颈)
|
||||
|
||||
| 编号 | 方向 | 说明 |
|
||||
|------|------|------|
|
||||
| R8 | **离线强化学习(CQL/IQL)** | 关键契合点:已有大量录制 episode、无法实机安全探索(舵机会坏),离线 RL 从固定数据集学策略、无需在线交互 |
|
||||
| R9 | 奖励设计(MPU6050) | 姿态稳定性(姿态角方差小=奖励高)、能耗(动作差分平滑=省能)、任务进度 → 标量奖励;从"模仿"走向"优化"的关键一跃 |
|
||||
| R10 | 自监督视觉预训练 | MAE/对比学习在相机图像上预训练编码器,缓解数据稀缺(ImageNet 预训练与机器人第一视角差异大) |
|
||||
|
||||
**R8 的洞察**:离线 RL 正好契合"数据已有 + 不能实机探索"的现状。但前提是 R9 的奖励信号。
|
||||
|
||||
### R8 细化:预训练 → 监督微调 → 离线 RL 三层流水线
|
||||
|
||||
**流水线形态**(基于 2026-08 实验进展的定稿):
|
||||
|
||||
```
|
||||
阶段 0 预训练(.mt 动作动力学) → 已有(pretrain_mt)
|
||||
阶段 1 监督微调(录制 episode,MSE/CE)→ 已有(train.py)
|
||||
阶段 2 离线 RL(录制数据 + IMU 奖励) → 待建(本细化)
|
||||
```
|
||||
|
||||
**为什么离线而非在线**:在线探索在真机上有物理风险(舵机乱动、碰撞),且 R4 已证明闭环
|
||||
反馈敏感——RL 探索的随机扰动会放大不稳定。离线 RL 从固定数据集学策略,零硬件风险。
|
||||
|
||||
**算法选型**:
|
||||
|
||||
| 算法 | 形态 | 适配度 |
|
||||
|------|------|--------|
|
||||
| IQL | 隐式 Q 学习,数据高效 | 首选:单策略、录制数据规模小 |
|
||||
| CQL | 保守 Q 惩罚,抗分布外 | 备选:数据质量差时更稳 |
|
||||
| 决策 Transformer(DT) | 回报条件化,无价值网络 | 简替:实现最简,但回报需已知 |
|
||||
|
||||
**IMU 奖励工程(R9 落地)**——传感器从"输入"变"监督信号":
|
||||
|
||||
| 奖励分量 | 公式(示意) | 物理意义 |
|
||||
|---------|-------------|---------|
|
||||
| 姿态稳定 | r₁ = -α·(|roll| + |pitch|) | 越正越接近站直 |
|
||||
| 抖动惩罚 | r₂ = -β·var(角加速度) | 平滑运动,节能 |
|
||||
| 动作平滑 | r₃ = -γ·Σ‖aₜ - aₜ₋₁‖ | 与 R4 平滑度指标同源 |
|
||||
| 任务进度 | r₄ = 姿态变化率(完成检测) | 驱动任务完成 |
|
||||
|
||||
奖励为标量加权和:R = Σwᵢrᵢ。设计难点在于权重标定——需实机 IMU 统计
|
||||
(静止 vs 动作时的量级差)确定 α/β/γ。
|
||||
|
||||
**落地路径**:
|
||||
|
||||
```
|
||||
阶段 A(数据前提): 实机 --record 采集 + IMU 统计标定奖励权重
|
||||
阶段 B(离线 RL): IQL 在录制数据上训练,R4 闭环对比监督基线
|
||||
阶段 C(安全在线): 小幅度动作扰动微调 + 动作界/速率保护
|
||||
```
|
||||
|
||||
**与当前实验的衔接**:one-hot 方案(F 系列)的固定列机制与有界输出天然
|
||||
兼容 RL 的动作扰动保护;样本级验证集划分(20260815)为 RL 的价值函数
|
||||
评估提供更稳的口径。R8 的前提不变:奖励信号(R9)与真实数据缺一不可——
|
||||
当前 dry-run 数据(噪声帧 + 零 IMU)对 RL 无意义。
|
||||
|
||||
### Tier 3 · 远景(与 aibox 联动)
|
||||
|
||||
- **R11 语言条件策略**:aibox 高层语言指令("向前走""左转")作为条件,actpred 出低层动作——
|
||||
分层控制,打通项目"AI 智能体"与"机器人"两大域。
|
||||
- **R12 sim2real / 数字孪生**:pybot 的 dry_run 已具备合成帧 + Mock 总线,可演进为轻量仿真环境,
|
||||
先仿真预训练再迁移真机,解决数据稀缺与安全探索。
|
||||
- **R13 持续学习**:闭环运行中不断回流新 episode 增量微调,机器人越用越聪明。
|
||||
|
||||
## 4. 推荐切入点与关键决策点
|
||||
|
||||
**短期(本月)**:**R1 + R4**。先落地"残差+多步训练",同时建评估基准——没有基准,后续改进无法量化对比。
|
||||
|
||||
**中期**:**R2(MPU6050 回流)性价比之王**——硬件已装、代码已在、收益明确(本体感觉 + 奖励信号源),
|
||||
只缺"接一根线"。
|
||||
|
||||
**关键决策点**:R5 GRU 是否值得做,取决于先做 **"Hailo DFC 循环算子支持验证"** 这个小实验:
|
||||
|
||||
- 支持 → R5/R7 可行,循环路线打开;
|
||||
- 不支持 → Tier 1 循环类方案全部绕行(用更长静态窗口替代隐状态),精力转向
|
||||
**R2 多模态 + R8 离线 RL** 这条不依赖循环算子的路线。
|
||||
|
||||
## 5. 与本季其他工作的衔接
|
||||
|
||||
- 进程内闭环已落地(pybot `run_loop`/`--loop`),P1-2 帧源(守护线程队列)与 P1-1 overrun 告警已修复
|
||||
——见 [pybot_actpred_AUDIT_20260814.md](../reviews/pybot_actpred_AUDIT_20260814.md)。
|
||||
- 数据飞轮已通:`--record` 录制 → 训练 → 推理 → `--loop` 闭环,为 R8 离线 RL 与 R13 持续学习
|
||||
提供数据回路。
|
||||
- `--record` 微动作与录制设计见 [pybot_micro_action_and_recording_20260813.md](pybot_micro_action_and_recording_20260813.md);
|
||||
Hailo 部署链详见 [actpred_hailo_hat2_inference_research_20260813.md](actpred_hailo_hat2_inference_research_20260813.md)。
|
||||
|
||||
---
|
||||
|
||||
*本报告为研究方向调研(advisory),方案 A/R5/R7 的部署可行性标 verify_pending,
|
||||
落地前需做 Hailo DFC 循环算子支持实测。*
|
||||
@@ -0,0 +1,116 @@
|
||||
# actpred 训练时安全图像增强
|
||||
|
||||
> 文档类型:特性报告
|
||||
> 模块:actpred
|
||||
> 日期:2026-08-13
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景
|
||||
|
||||
pybot dry-run 录制的 episode 数据样本量有限(11 个脚本 → 12 个 episode → ~728 样本),训练容易过拟合。直接增大数据集需要更多 `.mt` 脚本或更多轮次录制,成本较高。
|
||||
|
||||
图像增强(data augmentation)是低成本扩充样本的标准手段,但机器人动作预测场景有特殊约束:**空间变换(翻转/旋转/裁剪)会改变图像的语义坐标,而动作标签未随之变换,导致训练对错误**。例如水平翻转后,左右舵机角度应互换,但标签不变,形成错误训练信号。
|
||||
|
||||
## 2. 方案
|
||||
|
||||
### 2.1 增强范围
|
||||
|
||||
仅启用**颜色空间变换**(ColorJitter),不涉及任何空间变换:
|
||||
|
||||
| 变换 | 参数 | 说明 |
|
||||
|------|------|------|
|
||||
| 亮度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同光照条件 |
|
||||
| 对比度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同摄像头增益 |
|
||||
| 饱和度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同色温环境 |
|
||||
| 色调 | `±5%`(factor ∈ [-0.05, 0.05]) | 轻微色偏,避免模型过依赖颜色 |
|
||||
|
||||
**不使用的变换**(因需同步变换动作标签):
|
||||
|
||||
| 变换 | 原因 |
|
||||
|------|------|
|
||||
| 水平/垂直翻转 | 左右舵机角度需互换,否则标签错误 |
|
||||
| 旋转 | 重力方向变化,姿态语义改变 |
|
||||
| 裁剪/平移 | 视野变化导致动作-画面对应关系改变 |
|
||||
| 透视变换 | 改变空间关系 |
|
||||
|
||||
### 2.2 同一样本内帧间参数一致性
|
||||
|
||||
一个训练样本包含 `frame_len` 帧连续图像(默认 3 帧,间隔 10ms)。真实场景中连续帧的光照条件不变,因此**同一样本内所有帧共享同一组增强参数**,而非每帧独立随机:
|
||||
|
||||
```python
|
||||
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) # 色调因子
|
||||
for p in image_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))
|
||||
```
|
||||
|
||||
不同样本(不同 `__getitem__` 调用)使用独立随机参数,保证整个数据集的增强多样性。
|
||||
|
||||
### 2.3 自适应 frame_len
|
||||
|
||||
增强逻辑在 `for p in image_paths` 循环前生成参数、循环内逐帧施加。`image_paths` 的长度由 `frame_len` 决定(`image_files[t - frame_len + 1 : t + 1]`),循环自动适配。**`frame_len` 为 1、2、3、5、10 等任意值时均无需修改代码**。
|
||||
|
||||
### 2.4 训练/验证隔离
|
||||
|
||||
| 数据集 | 增强 | 说明 |
|
||||
|--------|------|------|
|
||||
| 训练集 | `augment=cfg.train.augment` | CLI `--augment` 控制 |
|
||||
| 验证集 | `augment=False`(默认) | 验证 loss 反映真实分布表现 |
|
||||
|
||||
推理端(`predict.py` / `onnx_infer.py` / `hailo_infer.py`)不涉及增强,使用原始图像。
|
||||
|
||||
## 3. CLI 接口
|
||||
|
||||
```bash
|
||||
# 训练时启用图像增强
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --augment
|
||||
|
||||
# 显式关闭(默认关闭)
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --no-augment
|
||||
```
|
||||
|
||||
`--augment` 使用 `BooleanOptionalAction`,支持 `--augment` / `--no-augment` 双向开关。
|
||||
|
||||
## 4. 验证
|
||||
|
||||
### 4.1 功能验证
|
||||
|
||||
```
|
||||
dataset size: 728 samples
|
||||
augment diff (should be >0): 0.3767 # 增强前后图像有差异
|
||||
two augment calls diff (should be >0): 0.2626 # 两次增强参数不同(随机性)
|
||||
action_history diff (should be 0): 0.0 # 动作标签不受影响
|
||||
target diff (should be 0): 0.0
|
||||
```
|
||||
|
||||
### 4.2 端到端训练
|
||||
|
||||
`--augment --epochs 3`(64×64,batch=8)训练正常完成,loss 正常下降。
|
||||
|
||||
### 4.3 回归测试
|
||||
|
||||
全量单测 90/90 通过,无回归。
|
||||
|
||||
## 5. 改动文件
|
||||
|
||||
| 文件 | 改动 |
|
||||
|------|------|
|
||||
| `actpred/dataset.py` | `ActionDataset` 新增 `augment` 参数;`__getitem__` 在 Resize 前施加 ColorJitter(`torchvision.transforms.functional`) |
|
||||
| `actpred/config.py` | `TrainConfig` 新增 `augment: bool = False`;`build_argparser` 新增 `--augment`;`apply_overrides` 新增映射 |
|
||||
| `actpred/train.py` | 训练集构造传 `augment=cfg.train.augment`,验证集默认不增强 |
|
||||
|
||||
## 6. 后续方向
|
||||
|
||||
- **空间增强(需同步标签变换)**:水平翻转时自动交换左右舵机角度(s1↔s2, s3↔s4, s5↔s6, s7↔s8),可将训练数据量翻倍。需要了解舵机配对关系后实现。
|
||||
- **动作噪声增强**:`pretrain_mt.py` 已支持 `--noise-std`,`train.py` 可同样引入,对 action_history 和 target 施加高斯扰动。
|
||||
- **更多 `.mt` 脚本**:从根本上扩大动作多样性,是提升模型泛化能力的最有效手段。
|
||||
@@ -0,0 +1,59 @@
|
||||
# actpred 验证集划分方式实验报告:整集 vs 样本级随机
|
||||
|
||||
> 文档类型:实验报告
|
||||
> 模块:actpred
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-15
|
||||
> 前置:`actpred_zscore_baseline_20260815.md`(模型 B 基线)、
|
||||
> `actpred_onehot_multihead_dev_plan_20260814.md`(one-hot 方案)
|
||||
|
||||
---
|
||||
|
||||
## 1. 问题
|
||||
|
||||
模型 F/F2 系列(one-hot,v2 数据)出现**过拟合签名**:train loss 持续下降(2.40→1.57)而 val loss 平台后上行(1.73→2.09),且 val 波动剧烈(区间 1.74-2.0)。
|
||||
|
||||
排查后发现**选模信号本身有缺陷**:
|
||||
|
||||
- 旧划分方式(`split_episodes`)按 **episode 整集划分**:val_split=0.15 × 12 集 → 验证集只有 **1 个 episode、60 个样本**
|
||||
- 60 个样本全来自同一集(tilt 类),val loss 主要反映"这一集的难度"而非泛化能力
|
||||
- 波动大 + 分布偏 → early stopping 的 best 模型选择被噪声主导
|
||||
|
||||
## 2. 方案:样本级随机划分
|
||||
|
||||
新增 `--val-split-mode {episode, sample}`(**新默认 sample**):
|
||||
|
||||
- `dataset.py`:`ActionDataset` 新增 `sample_indices` 参数(构建后按索引过滤);**zscore 统计量仍只由训练子集计算**(经 stats 文件复用给验证集,统计口径正确性不变)
|
||||
- `train.py`:sample 模式下先建全量数据集取样本索引,按 seed 42 随机排列,取 val_ratio 为验证集
|
||||
- 效果:验证集从 60 → **142 样本**(810 train / 142 val),跨全部 12 个 episode 混合
|
||||
|
||||
## 3. 结果对比(同配置:onehot + spk 固定 + v2 数据)
|
||||
|
||||
| | F2(整集划分) | F2b(样本级划分) |
|
||||
|---|---|---|
|
||||
| 验证集 | 60 样本(1 集) | 142 样本(12 集混合) |
|
||||
| val 起点 | 1.914 | 2.056(更难、更有代表性) |
|
||||
| 最优 | 1.737(epoch 7) | **1.695(epoch 5,仍在下降)** |
|
||||
| 收敛形态 | 平台后上行(过拟合签名) | 平滑连续下降 |
|
||||
| val 波动 | ±0.15(1.74-2.0) | ±0.05(无跳变) |
|
||||
|
||||
F2b 在更难的代表性验证集上取得**更低的 val loss**,且曲线平滑无异常跳变——说明 F2 的"过拟合"判断部分被验证集噪声误导,真实模型能力被低估。
|
||||
|
||||
## 4. 已知代价(诚实记录)
|
||||
|
||||
样本级划分**略乐观**:微动作插值使相邻样本高度相关(0.75°/步),随机划分后每个验证样本的近邻大量在训练集,val 会略高估泛化。适用场景:
|
||||
|
||||
- **选模/早停**(本次用途):噪声小、跨集混合 → 更可靠的 best 选择信号 ✓
|
||||
- **能力评估**(需保守口径):仍建议整集划分或 leave-one-out,避免乐观偏差
|
||||
|
||||
## 5. 实施细节
|
||||
|
||||
- `dataset.py::ActionDataset(sample_indices=...)`:样本按索引过滤,位置在样本构建后、统计量计算前
|
||||
- 签名顺序调整:`sample_indices` 置于 `frame_len` 之后,避免破坏既有位置参数调用(曾因插入位置错误导致 frame_len 错位 TypeError)
|
||||
- `config.py`:`TrainConfig.val_split_mode`(默认 "sample")+ CLI + 校验
|
||||
- 测试:+2(样本过滤与统计口径、CLI 映射),全套 185 通过
|
||||
|
||||
## 6. 后续
|
||||
|
||||
- F2b 训练完成后跑闭环评估,验证 onehot + spk 固定 + 新划分的实际闭环收益
|
||||
- 若需严格能力评估,增加 `--val-split-mode episode` 复跑对照,或 leave-one-out 交叉验证
|
||||
@@ -0,0 +1,91 @@
|
||||
# actpred 模型 B 方案文档——zscore 等权稳定基线
|
||||
|
||||
> 文档类型:方案文档
|
||||
> 模块:actpred
|
||||
> 作者:自动生成
|
||||
> 日期:2026-08-15
|
||||
> 相关实验:`pybot_actpred_dataset_training_eval_20260814.md`(A/B 对照)、
|
||||
> `actpred_fixed_range_weighted_experiment_20260814.md`(C/D/E 失败系列)、
|
||||
> `actpred_onehot_multihead_dev_plan_20260814.md`(F 系列结构性方案)
|
||||
|
||||
---
|
||||
|
||||
## 1. 定位
|
||||
|
||||
模型 B 是六模型闭环评估(A/B/C/D/E/F)中**唯一稳定**的配置:overall_mse 2445、error_growth 0.082,远优于其余五个(6447~968459)。本文档固化其完整方案,作为:
|
||||
|
||||
1. 部署与实机测试的**首选基线**
|
||||
2. 后续所有实验的**对照基准**
|
||||
3. 归一化系列失败教训的**参照系**
|
||||
|
||||
## 2. 完整配置(checkpoint 实读)
|
||||
|
||||
来源:`actpred/checkpoints/recorded_run_norm/best_inference.pt`(epoch 26 早停,best val=0.0288)
|
||||
|
||||
### 2.1 数据侧(DataConfig)
|
||||
|
||||
| 项 | 值 |
|
||||
|---|---|
|
||||
| 分辨率 | 640×480 |
|
||||
| action_dim / history_len / frame_len | 18 / 5 / 3 |
|
||||
| normalize_actions | **True(z-score)** |
|
||||
| action_stats_path | `action_stats.npy`(272B,逐列 (2,18) mean/std,训练集统计) |
|
||||
| imu_dim | 0 |
|
||||
| column_weights | None(**等权**) |
|
||||
|
||||
### 2.2 训练侧(TrainConfig)
|
||||
|
||||
| 项 | 值 |
|
||||
|---|---|
|
||||
| 数据 | v1(`data/pybot_record`,804 样本,12/1 episode 划分,seed 42) |
|
||||
| 训练范式 | **单步 teacher forcing**(rollout_len=1, ss_prob=1.0) |
|
||||
| 增强 | ColorJitter(--augment) |
|
||||
| batch / lr / patience | 4 / 1e-3 / 15 |
|
||||
| predict_delta | False |
|
||||
|
||||
### 2.3 模型侧(ModelConfig)
|
||||
|
||||
SimpleCNN(128) → Transformer(d=128, 4 头, 2 层, pre-LN) → MLP 头。409K 参数。
|
||||
|
||||
### 2.4 闭环表现(v1 数据,R4)
|
||||
|
||||
- overall_mse **2445.4**,error_growth **0.082**(<1 无发散),oob_rate 28.1%
|
||||
- 逐 episode:stand 134 / jump 221 / hand 6850·8023 / tilt 60-752(唯一全部 episode 不爆表的模型)
|
||||
|
||||
## 3. 为什么 zscore 稳定(失败系列的对照结论)
|
||||
|
||||
| 方案 | 操作点 | 闭环行为 | 机理 |
|
||||
|---|---|---|---|
|
||||
| raw 等权(A) | 原始量纲 | 稳定但 spk.F 崩 | 舵机误差数值占优,spk.F 被忽略 |
|
||||
| **zscore(B)** | **数据均值 0** | **稳定** | **各列单位方差,归一化误差等量纲;误差收缩** |
|
||||
| fixed_range(C/D) | [0,1] 低端聚集 | 发散(growth 13.4/3.4) | 低端精度错配:小绝对误差=大相对误差 |
|
||||
| log_range(E) | log 空间 | 爆炸(growth 157) | 指数反变换在高端放大:0.81→1.0 = 1000→5000Hz |
|
||||
| onehot(F) | 箱(有界) | 部分稳定(oob=0,MSE 33K) | 有界生效,但 spk 无监督头闭环污染 |
|
||||
|
||||
**结论**:反馈回路需要**误差有界且收缩**的映射。z-score 的零均值单位方差恰好满足;固定范围类方案各自引入一端(低/高)的放大;one-hot 的有界性正确但需配 spk 列固定(见 F2)。
|
||||
|
||||
## 4. 已知局限
|
||||
|
||||
1. **依赖 stats 文件**:推理端需随 checkpoint 携带 action_stats.npy(与 fixed_range/onehot 的零状态部署相反)
|
||||
2. **训练于 v1 数据**:delay 语义混杂(首行 200ms 未分解),spk.delay 权重 0.2——两者已被后续修复(v2 均匀数据 + delay 权重 0)
|
||||
3. **spk.F 仍弱**:hand 类 episode 8023/6850 的 MSE 主导源是 spk.F(10→1000Hz 大跨度,等权下被均值化)
|
||||
|
||||
## 5. 演进建议:B2
|
||||
|
||||
模型 B 的两处已知缺陷(v1 数据、delay 权重)都已有修复就绪,B2 = **zscore + v2 数据 + 新权重(spk 列 0)**:
|
||||
|
||||
- 预期:闭环 MSE 进一步下降(delay 垃圾不再进历史)
|
||||
- 若 B2 稳定,作为部署基线导出 ONNX/HEF
|
||||
- 若 B2 需要 spk.F 频率跟踪(未来语音反馈需求),再叠加加权或 one-hot 方案
|
||||
|
||||
## 6. 复现命令
|
||||
|
||||
```bash
|
||||
# 模型 B(v1 数据,旧默认权重)
|
||||
./venv/bin/python train.py --data-dir /opt/micro_servers/data/pybot_record \
|
||||
--checkpoint-dir checkpoints/recorded_run_norm --augment --epochs 100
|
||||
|
||||
# B2(v2 数据,新权重:spk 两列 0)
|
||||
./venv/bin/python train.py --data-dir /opt/micro_servers/data/pybot_record_v2 \
|
||||
--checkpoint-dir checkpoints/recorded_run_norm_v2 --augment --epochs 100
|
||||
```
|
||||
@@ -0,0 +1,437 @@
|
||||
# aibox 模块安全与质量审计报告
|
||||
|
||||
- **日期**:2026-08-17(第一轮)/ 2026-08-18(第二轮复审)/ 2026-08-18(第三轮:遗留复核 + v3.4 专审)/ 2026-08-18(第四轮:3 路并行审读 + Batch H 修复)
|
||||
- **审计基线**:第一轮 HEAD `ce142fe`(Batch A–D 后)→ 第二轮 HEAD `177a9ac`(Batch E 修复后)→ 第三轮 HEAD `aa15841`(L4a/b/c 落地后)→ 第四轮 HEAD `d1f780e`(Batch G + SOCKS5 验证后)
|
||||
- **审计范围**:`aibox.py`(~3310 行)+ `llm_client.py` + `tools/`(7 文件)+ `subsystems/`(6 文件)+ `common/`(5 文件)
|
||||
- **方法**:第一轮逐文件代码审读 + PoC 动态验证;第二轮 Batch E 修复复审 + 全量重审 + 6 处关键发现实证抽查;第三轮遗留项复核 + v3.4 专审;第四轮 3 路并行审读 + 主审逐条复核 + Linux PoC 实证(6/6)
|
||||
|
||||
---
|
||||
|
||||
## 一、总体结论
|
||||
|
||||
**第一轮**(基线 ce142fe):架构设计(分层、参数化 SQL、密钥隔离、消息回滚)整体健康,但**命令执行链存在一条系统性断裂**:实际运行路径绕过了精心设计的执行器,导致 1MB 硬约束、超时、中断三项保护同时失效。
|
||||
|
||||
**第二轮**(基线 177a9ac):Batch E 五项修复(H1/H2/H3/M1/M2)中四项确认修复到位(PoC 5/5 验证),**M2 修复不完整**(超时路径跳过工具属性清理,僵尸输出问题仍存)。第二轮全量重审新发现 **1 高危 + 9 中危 + 若干低危**,其中三项与 Batch E 修复直接相关(reader 线程泄漏 / 子代理超时盲区 / M2 残留),其余为存量问题(代理密码明文落盘、drain 线程泄漏、事件循环阻塞等)。详见"八、第二轮复审"。
|
||||
|
||||
---
|
||||
|
||||
## 二、高危发现(3 项)
|
||||
|
||||
### H1:`run_command` 实际路径绕过注入的 executor——三重保护同时失效
|
||||
|
||||
证据链:
|
||||
- `aibox.py:1012` 将 `command_executor` 注入 `CommandLineTool`
|
||||
- 但 `tools/command_tool.py:86-219` 的 `execute_streaming` 自带一套 **150 行的重复 Popen 实现**,`self.executor` 全程未使用(死注入)
|
||||
- `tools/registry.py:71-86` 检测到 `execute_streaming` 即走该路径——**生产 100% 命中**
|
||||
|
||||
后果:
|
||||
|
||||
| 保护 | 设计位置 | 实际状态 |
|
||||
|---|---|---|
|
||||
| 1MB 输出上限(硬约束) | `executors.py:69` | **失效**——command_tool 的 `output_lines` 无限累积 |
|
||||
| `cancel()` 中断 | `executors.py:76-86` 只 kill `_current_process` | **失效**——实际进程引用只在 command_tool 局部变量,三层中断的第三层对 run_command 不生效 |
|
||||
| 超时/截断统一语义 | executors | 旁路(见 H2) |
|
||||
|
||||
附带问题:`msvcrt.kbhit()` 检查的是**当前进程**键盘而非子管道(`command_tool.py:123-125` 无意义代码);循环内重复 `import select`。
|
||||
|
||||
### H2:超时机制对无输出命令完全失效(PoC 已证)
|
||||
|
||||
`readline()` 是阻塞调用,超时检查只在**两次 readline 之间**执行。无输出命令会让 readline 阻塞到进程自然退出:
|
||||
|
||||
```
|
||||
PoC:timeout=3s,命令 powershell Start-Sleep 15(无 stdout 输出)
|
||||
结果:readline blocked 15.2s (timeout was 3s)
|
||||
VERDICT: timeout mechanism ineffective for silent commands: True
|
||||
```
|
||||
|
||||
两处同病:`command_tool.py:106-165` 与 `executors.py:97-113`。`--once` 模式有 `--once-timeout` 兜底,**WebSocket 模式无兜底**。
|
||||
|
||||
### H3:`stdin=subprocess.PIPE` 从不写入也不关闭——死锁放大器
|
||||
|
||||
`executors.py:156`、`command_tool.py:93`:等待 stdin 的命令(`cat` 无参、交互式 `python`、`ssh`)永不退出 → 与 H2 叠加:**readline 永久阻塞 + 超时检测永不执行 + cancel 杀不到进程** = 工具线程永久挂起。应为 `stdin=subprocess.DEVNULL`。
|
||||
|
||||
---
|
||||
|
||||
## 三、中危发现(4 项)
|
||||
|
||||
| # | 问题 | 位置 | 说明 |
|
||||
|---|---|---|---|
|
||||
| M1 | **heredoc 定界符注入** | `basic_tools.py:179` | `content` 含独立行 `AIBOX_EOF` 时 heredoc 提前终止,**后续行成为 shell 命令执行**。AI 写文件内容是不可信输入(可被网页内容/用户消息诱导)。修复:定界符唯一化(检测冲突则换 `AIBOX_EOF_2` 递增) |
|
||||
| M2 | 工具流式转发循环可永久挂起 | `registry.py:90-107` | `while thread.is_alive() or not empty()` + `get(0.1)`——H3 场景下永不退出;`join(5)` 超时后返回但 daemon 线程继续向输出队列写(僵尸输出混入后续对话流) |
|
||||
| M3 | web 工具无协议/目标校验 | `basic_tools.py:315-338` | `url` 未限制 scheme,`curl file:///etc/shadow`、内网地址均可达(SSRF)。本地单用户产品风险有限,建议白名单 http/https |
|
||||
| M4 | shlex.quote POSIX 转义在 Windows 失效 | `basic_tools.py:99-112` 等全部指引命令 | 单引号在 cmd.exe 无效,Windows 下 `head/sed/grep` 指引命令执行异常。Linux 部署(树莓派)不受影响 |
|
||||
|
||||
---
|
||||
|
||||
## 四、低危发现(3 项)
|
||||
|
||||
| # | 问题 | 位置 |
|
||||
|---|---|---|
|
||||
| L1 | `get_pending_memos` 的 `LIMIT {limit}` f-string 拼接 | `db.py:417`——当前调用链 limit 均为内部 int,无外部直达;防御性参数化即可 |
|
||||
| L2 | `check_path_safety` 仅查大小,无路径边界 | `base.py:96-110`——by design(AI 可全盘读写,本地单用户),已确认设计意图 |
|
||||
| L3 | `executors.execute()` finally kill 后未 `wait()` | `executors.py:205-207`——Windows 句柄短暂滞留 |
|
||||
|
||||
**做得好的**:SQL 全面参数化(LIKE 通配符正确转义 `db.py:200-209`)、API key 零日志泄漏、WAL 已启用、消息回滚/tool 配对修复、`_msg_lock`/`_process_lock` 使用正确、工具实例按连接隔离。
|
||||
|
||||
---
|
||||
|
||||
## 五、架构合规矩阵
|
||||
|
||||
| 硬约束 | 状态 |
|
||||
|---|---|
|
||||
| WebSocket :8765 | ✅ |
|
||||
| 依赖单向 `common←subsystems←tools←aibox` | ✅ |
|
||||
| Worker 仅 7 执行工具 | ✅ |
|
||||
| 工具输出 ≤1MB + 超时 + 中断 | ❌ **H1/H2/H3 三项全破** |
|
||||
| prompt 三级排序加载 / sqlite 线程安全 | ✅ |
|
||||
|
||||
---
|
||||
|
||||
## 六、修复建议(一批搞定 H1/H2/H3,约 ±120 行)
|
||||
|
||||
1. **删除 command_tool.py 的重复 Popen 实现**,`execute_streaming` 改为一行委托 `self.executor.execute_streaming(output_queue, ...)` → 1MB、cancel、统一语义全部自动恢复
|
||||
2. **executors.py 读取循环改线程化**:后台 reader 线程做阻塞 `readline` 入队列,主循环带超时轮询队列——跨平台解决 H2
|
||||
3. `stdin=subprocess.DEVNULL`(H3)
|
||||
4. M1 定界符唯一化 + M2 转发循环加总超时,顺手可修
|
||||
|
||||
---
|
||||
|
||||
## 七、修复落地(Batch E,2026-08-17)
|
||||
|
||||
五项修复(H1/H2/H3/M1/M2)已全部实施并通过验证:
|
||||
|
||||
| # | 修复 | 位置 | 状态 |
|
||||
|---|---|---|---|
|
||||
| H1 | `CommandLineTool.execute_streaming` 改为纯委托 `executor.execute_streaming`,删除 150 行重复 Popen(-165 行) | `tools/command_tool.py` | ✅ |
|
||||
| H2 | `_read_with_limit` 线程化:后台 reader 线程阻塞 `readline` 入队,主循环 0.1s 轮询 + 超时/1MB 检查;shell 已死但 reader 无 EOF 时 1s 宽限排空(Windows shell=True 孙进程持管道场景) | `subsystems/executors.py` | ✅ |
|
||||
| H3 | 两处 Popen 均改 `stdin=subprocess.DEVNULL` | `subsystems/executors.py` | ✅ |
|
||||
| M1 | heredoc 定界符唯一化:content 含独立行与 `AIBOX_EOF` 相同时递增 `_X` 后缀直到唯一 | `tools/basic_tools.py` | ✅ |
|
||||
| M2 | registry 转发循环加 3600s 总超时兜底,超时后放弃等待并返回错误信息 | `tools/registry.py` | ✅ |
|
||||
|
||||
### PoC 回归验证(`aibox/tests/_poc_audit_regression.py`,5/5 PASS)
|
||||
|
||||
```
|
||||
[PASS] PoC-1 无输出命令超时(H2): 3.0s 返回(修复前阻塞15s+),结果含超时: True
|
||||
[PASS] PoC-2 stdin等待命令不挂起(H3): 0.2s 返回(修复前永久挂起)
|
||||
[PASS] PoC-3 1MB输出截断: 0.2s,结果 1.00MB,含截断标记: True
|
||||
[PASS] PoC-4 cancel中断: cancel 后 3.0s 结束(修复前杀不到进程需等满120s)
|
||||
[PASS] PoC-5 工具层1MB贯通(H1): 0.2s,工具层结果 1.00MB(修复前无限累积)
|
||||
```
|
||||
|
||||
PoC-2 说明:`stdin=DEVNULL` 下等待输入的命令立即收到 EOF 快速退出(0.2s),优于按超时返回;判定标准为"不永久挂起"。
|
||||
|
||||
### 单元测试
|
||||
|
||||
aibox 全量单测 **264 passed**(零回归)。
|
||||
|
||||
### 遗留项(本轮未处理)
|
||||
|
||||
- M3(web 工具 scheme/目标校验)、M4(Windows 下 shlex.quote 转义失效)、L1/L3(防御性加固)——低优先级,待后续批次
|
||||
- 架构合规矩阵"工具输出 ≤1MB + 超时 + 中断"一项修复后为 ✅
|
||||
|
||||
---
|
||||
|
||||
## 八、第二轮复审(2026-08-18,HEAD `177a9ac`)
|
||||
|
||||
方法:Batch E 修复质量复审 + 全量重审(3 路并行审读 + 主审对 6 处关键发现逐一实证抽查)。
|
||||
|
||||
### 8.1 Batch E 修复复审结论
|
||||
|
||||
| # | 修复 | 复审结论 |
|
||||
|---|---|---|
|
||||
| H1 | command_tool 纯委托 executor | ✅ 到位(PoC-5 验证 1MB 贯通) |
|
||||
| H2 | reader 线程 + 超时轮询 | ✅ 到位(PoC-1 验证),**但有两个残余**:N4 reader 线程泄漏、N5 子代理超时盲区 |
|
||||
| H3 | stdin=DEVNULL | ✅ 到位(PoC-2 验证) |
|
||||
| M1 | heredoc 定界符唯一化 | ✅ 到位 |
|
||||
| M2 | registry 转发总超时 | ⚠️ **不完整**(见 N6) |
|
||||
|
||||
### 8.2 第一轮遗留项状态
|
||||
|
||||
| # | 问题 | 状态 |
|
||||
|---|---|---|
|
||||
| M3 | web 工具无 scheme 校验 | 未修复(实证:`basic_tools.py:319+` url 未限制,`file://` 等可达) |
|
||||
| M4 | shlex.quote POSIX 转义 Windows 失效 | 未修复(实证:`sed -r`/`grep -oP` 指引,Linux 部署不受影响) |
|
||||
| L1 | `LIMIT {limit}` f-string | 未修复,且**新发现同型第二处**(`db.py:417` 与 `db.py:459` `get_immediate_memos`) |
|
||||
| L2 | check_path_safety 无路径边界 | 维持 by design(本地单用户,AI 可全盘读写为设计意图) |
|
||||
| L3 | execute() finally kill 后未 wait() | 未修复(实证:`executors.py:236-241` 无 wait;`cancel()` 同样无 wait) |
|
||||
|
||||
### 8.3 新发现——高危(1 项)
|
||||
|
||||
#### N1:SOCKS5 代理密码明文持久化到配置文件
|
||||
|
||||
`common/config.py:238-251`(实证):`SetProxyTool.execute` 默认 `save=True` → `set_proxy` 将 `password` 明文写入 `_config["proxy"]["socks5"]` → `_save_config` 用 `json.dump` 全量落盘 `~/.aibox/config.json`。与 api_key 走环境变量的隔离设计形成反差,任何能读该文件的进程可获取代理凭据。
|
||||
|
||||
修复建议:密码不落盘(仅内存)或走系统凭据库;必须落盘时对称加密 + `os.chmod(0o600)`;`save` 默认改 `False`。
|
||||
|
||||
### 8.4 新发现——中危(9 项)
|
||||
|
||||
| # | 问题 | 位置 | 说明 |
|
||||
|---|---|---|---|
|
||||
| N2 | drain_thread 每次工具调用泄漏一个永不退出的线程 | `aibox.py:1232-1243`(实证) | 循环条件只查 cancel 事件不查队列空,注释与实现不符;闭包持 output_callback(含 websocket)阻碍 GC;长期运行累积线程/内存 |
|
||||
| N3 | 同步 threading.RLock 在 async 协程中阻塞事件循环 | `aibox.py:1855-1902, 2275, 2376` | `/new` 的 2s 轮询竞态窗口后仍可能 `with self._msg_lock:` 冻结整个事件循环 → 服务假死 |
|
||||
| N4 | reader 线程泄漏(H2 残余) | `subsystems/executors.py:130-144` | 宽限 1s break 后 reader 仍阻塞 readline(Windows shell=True 孙进程持管道),stdout 未 close、reader 未 join,常驻服务累积 fd 与线程 |
|
||||
| N5 | --once 超时对子代理完全无效(H2 盲区) | `aibox.py:3052-3058, 638-649` | Timer 到点只 `manager.cancel()` 父代理;子代理独立 `_cancel_event`/独立 executor/独立 HTTP client,cancel 不传播,进程挂起至子代理自行结束 |
|
||||
| N6 | registry 超时路径跳过工具属性清理(M2 残余) | `tools/registry.py:99, 116`(实证 L114-120) | 两个超时 return 均跳过 `streaming_output=False`/`output_queue=None`;daemon 线程继续写实例属性 `tool.output_queue`,该属性被后续调用覆盖到新会话队列 → 僵尸输出串流 |
|
||||
| N7 | FTS5 关键词注入 | `subsystems/db.py:179` | 关键词内 `"` 未转义(应翻倍),可破坏短语边界改变 AND/OR 语义;语法错误时回退 LIKE,影响有限 |
|
||||
| N8 | 流式输出旁路 TOOL_MAX_RESULT_LENGTH | `tools/db_tools.py:97-98` + `registry.py:104-110` | 超长记忆 content 经流式通道完整推到前端,100000 字符上限仅约束非流式返回值 |
|
||||
| N9 | prompt 注入存储 | `tools/prompt_tools.py:64, 127, 227` | Agent 可控内容原样写入 `00_IDENTITY.md`/`99_SOUL.md`/`100_STATUS.md` 并在后续会话作为系统提示词加载,形成持久化注入;update_soul 可自我强化 |
|
||||
| N10 | mark_as_triggered TOCTOU 竞态 | `subsystems/db.py:509-570` | 跨连接先读后写,并发触发 trigger_count 失真;建议 BEGIN IMMEDIATE 同事务读写 |
|
||||
|
||||
### 8.5 新发现——低危(7 项)
|
||||
|
||||
| # | 问题 | 位置 | 说明 |
|
||||
|---|---|---|---|
|
||||
| N11 | debug 日志打印工具敏感参数 | `tools/registry.py:66, 126` | debug=3 时 set_proxy 的 password 明文进 `~/.aibox/ai.log`(日志无 0600) |
|
||||
| N12 | _load_persona 路径遍历(限 .md) | `aibox.py:527-543` | persona 参数 LLM 可控未净化,可读 SKILLS_DIR 外 .md 文件并注入子代理 |
|
||||
| N13 | broadcast 裸 except + unregister KeyError 逃逸 | `aibox.py:2076-2082, 1971-1973` | 吞 CancelledError;并发注销时 `remove` 抛 KeyError;建议 `discard` + `except Exception` |
|
||||
| N14 | _processing_tasks 按 ws id 覆盖,旧任务失联 | `aibox.py:2728-2742` | 连接关闭只能 cancel 最后一个任务;handle_command 不受 is_processing 保护 |
|
||||
| N15 | llm_client 异常含完整响应体 | `llm_client.py:172-174` | 某些网关 4xx 页回显 Authorization 头 → 密钥进异常栈;建议截断 500 字符 |
|
||||
| N16 | 历史索引裸 except + 归档文件名撞名覆盖 | `subsystems/conversation_history.py:60-66, 135-137` | Ctrl+C 被吞后空索引覆盖原文件;时间戳粒度不足同秒覆盖 |
|
||||
| N17 | InstallSkillTool skill_path 无校验 | `tools/skill_tools.py:252-304` | 可引导复制任意敏感文件进技能目录,经 search_skills show_content 预览泄露 |
|
||||
|
||||
杂项(低危,随批处理):`executors.py:209/301/338` 裸 except 应收窄为 TimeoutExpired;`db.py` 多处 `with sqlite3.connect` 不显式 close(CPython 下无害,PyPy 下句柄延迟);`config.py:31-34` AI_CONFIG_PATH 可重定向配置加载(本地单用户风险有限);skill_manager 技能目录代码执行面维持 by design(建议文档化可信边界)。
|
||||
|
||||
### 8.6 更新后的架构合规矩阵
|
||||
|
||||
| 硬约束 | 第一轮 | 第二轮 |
|
||||
|---|---|---|
|
||||
| WebSocket :8765 | ✅ | ✅ |
|
||||
| 依赖单向 `common←subsystems←tools←aibox` | ✅ | ✅ |
|
||||
| Worker 仅 7 执行工具 | ✅ | ✅ |
|
||||
| 工具输出 ≤1MB + 超时 + 中断 | ❌ H1/H2/H3 全破 | ⚠️ **executor 链已恢复**;残余:流式 db_tools 旁路长度上限(N8)、子代理超时无效(N5)、reader 泄漏(N4) |
|
||||
| prompt 三级排序加载 / sqlite 线程安全 | ✅ | ⚠️ sqlite 主流程安全;mark_as_triggered 竞态(N10)、FTS5 注入(N7)待修 |
|
||||
|
||||
### 8.7 Batch F 修复建议(按优先级)
|
||||
|
||||
1. **F1(高)**:N1 代理密码脱敏——`set_proxy` 落盘前抹除 password 字段 + `save` 默认 False
|
||||
2. **F2(中)**:N6 registry 超时路径补清理(两个 return 前置 `tool.streaming_output=False; tool.output_queue=None`)+ N4 reader 线程收尾(break 后 `stdout.close()` + `reader.join(2)`)
|
||||
3. **F3(中)**:N5 子代理 cancel 传播——`run_full_subagent` 传递父 `_cancel_event`,每轮 think 前检查
|
||||
4. **F4(中)**:N2 drain_thread 加 sentinel 退出 + N3 reset_conversation 移出事件循环(run_in_executor)
|
||||
5. **F5(中)**:N7 FTS5 转义 + N8 流式截断 + N10 同事务读写 + 顺带修 L1(两处 LIMIT 参数化)+ L3(finally wait)
|
||||
6. **F6(低)**:N9 prompt 内容过滤/确认门槛 + N11 日志脱敏 + N12 persona 路径校验 + N13/N14/N15/N16/N17 择机处理 + M3/M4 视部署平台决定
|
||||
|
||||
### 8.8 Batch F 修复进展
|
||||
|
||||
**F1 ✅ 已落地(2026-08-18)**——N1 + N11 一并修复:
|
||||
|
||||
- `common/config.py`:`set_proxy` 新增 `password_env` 参数;`_save_config` 落盘前 deepcopy 脱敏(明文密码强制置 None,纵深防御,任何保存路径均拦截);`get_proxy` 按"内存明文优先 → password_env 查环境变量"解析,下游(llm_client / web 工具 `_proxy_curl_args`)无感知;加载时检测旧版明文自动迁移(内存保留供本会话 + 立即重写脱敏文件,一次性清除存量)
|
||||
- `tools/proxy_tools.py`:`SetProxyTool` 新增 `password_env` 参数,结果显示"密码来源"(环境变量 / 仅本次会话 / 无)
|
||||
- `tools/registry.py`(N11):debug 日志对 `password/api_key/token/secret/authorization` 参数打码
|
||||
- 测试:新增 `tests/test_config.py` 3 例(落盘无明文 / password_env 解析 / 旧配置迁移),全量 267 passed
|
||||
- 行为变化:已有明文密码的配置重启后代理需重设密码或改用环境变量(迁移日志有提示)——清除存量明文的必然代价
|
||||
|
||||
**F2 ✅ 已落地(2026-08-18)**——N6 + N4 一并修复:
|
||||
|
||||
- N6(registry 超时路径补清理):两个超时 return(deadline 命中、join(5) 超时)统一改为 `timed_out` 标记 + 出口统一执行 `tool.streaming_output=False` / `tool.output_queue=None`;僵尸 daemon 线程的 `send_line` 落到 None 上不再外写,消除僵尸输出串流路径。`stream_tool_timeout` 提为实例属性(默认 3600s,测试可调)
|
||||
- N4(reader 线程收尾,**含二次修正**):初版按审计建议"break 后 `stdout.close()` + `reader.join(2)`"实现,PoC 回归实证 **close 会阻塞**——CPython 流对象内部有锁,readline 持锁阻塞时 `close()` 一同挂起直到 readline 返回(PoC-1 从 3s 退化到 15.2s,恰为命令 Start-Sleep 时长)。修正后策略:reader 已正常退出(EOF)→ 显式 close 释放管道 fd;reader 仍阻塞 → 只做 `join(timeout=2)` 后放弃(daemon 兜底,孙进程退出收 EOF 后自然结束,fd 随管道消亡回收),主流程及时返回
|
||||
- **审计建议修正记录**:8.3 节 N4 与 8.7 节 F2 中"close stdout 可中断阻塞 readline(CPython 行为)"的表述**不成立**,已按实测修正——跨线程中断阻塞中的管道 read 无便携方案,正确做法是控制泄漏窗口而非强行中断
|
||||
- **测试**:新增 `tests/test_f2_fixes.py` 5 例(N6 超时/正常清理 ×2、N4 及时返回 ×2 + EOF 后 close ×1),PoC 回归 5/5(PoC-1/4 恢复 ~5s:超时 3s + join 宽限 2s),全量 272 passed
|
||||
|
||||
---
|
||||
|
||||
## 九、第三轮审计(2026-08-18,HEAD `aa15841`)
|
||||
|
||||
方法:前两轮遗留项逐条实证复核 + v3.4 新增代码(多模型档案 / L4a/b/c 失败转移,`llm_client.py`)全量审读 + 测试基线确认(全量 **355 passed**)。
|
||||
|
||||
### 9.1 遗留项修复状态复核
|
||||
|
||||
| # | 问题 | 状态 | 实证 |
|
||||
|---|---|---|---|
|
||||
| L3 | execute() kill 后未 wait() | ✅ **已修** | `executors.py:222-230`:kill/terminate → `wait(timeout=2)` → 超时再 kill 兜底 |
|
||||
| N17 | InstallSkillTool 任意文件复制 | ✅ **已消解(设计变更)** | 工具重构为纯指引(只返回安装说明文本,不再执行 shell/复制);`CreateSkillTool` 增加 skill_name 正则校验(`^[A-Za-z0-9_-]+$`) |
|
||||
| N1/N11 | 代理密码脱敏 / 日志打码 | ✅ 在位 | F1 修复无回归 |
|
||||
| N4/N6 | reader 收尾 / registry 超时清理 | ✅ 在位 | F2 修复无回归 |
|
||||
| N2 | drain 线程泄漏 | ❌ 未修 | `aibox.py:1280-1303`:循环条件仍只查 cancel 事件,注释"自动退出"与实现不符,每次工具调用泄漏一个常驻线程 |
|
||||
| N3 | `/new` 阻塞事件循环 | ❌ 未修 | `aibox.py:2325`:2s 轮询竞态窗口后仍 `with self._msg_lock:`(同步 RLock acquire 于事件循环线程) |
|
||||
| N5 | 子代理 cancel 不传播 | ❌ 未修 | `aibox.py:643-654`:子代理独立 `_cancel_event`/executor/client;`--once` Timer(3110-3113)只 cancel 父代理 |
|
||||
| N7 | FTS5 引号未翻倍 | ❌ 未修 | `db.py:178`:`f'"{kw}"'` 直接拼,关键词含 `"` 破坏短语边界 |
|
||||
| N8 | 流式旁路长度上限 | ❌ 未修 | `registry.py:120-129`:`("line", content)` 无限转发;`aibox.py:1326` 的 `MAX_RESULT_LENGTH` 只管回喂 LLM 的最终 result |
|
||||
| N9 | prompt 注入存储 | ❌ 未修 | `prompt_tools.py`:confirm 默认 True,Agent 可控内容仍直写 00/99/100 文件 |
|
||||
| N10 | mark_as_triggered TOCTOU | ❌ 未修 | `db.py:509-525`:跨连接先 `get_memo` 后 UPDATE,无 BEGIN IMMEDIATE |
|
||||
| N12 | _load_persona 路径遍历 | ❌ 未修 | `aibox.py:529-545`:persona LLM 可控,`os.path.join(SKILLS_DIR, name + ".md")` 无边界校验 |
|
||||
| N13 | broadcast 裸 except | ❌ 未修 | `aibox.py:2131`:`except:` 吞 CancelledError |
|
||||
| N14 | _processing_tasks 覆盖 | ❌ 未修 | `aibox.py:2788`:按 ws id 覆盖登记;先完成任务的 done_callback 会误 pop 后续任务的登记 |
|
||||
| N15 | 异常含完整响应体 | ⚠️ 部分改善 | 瞬时路径已截断 `[:200]`(163/213 行);**非瞬时路径仍全量**(`llm_client.py:215-216`) |
|
||||
| N16 | 历史索引裸 except / 同秒覆盖 | ❌ 未修 | `conversation_history.py:64` 裸 `except:`;136 行归档文件名秒级时间戳 |
|
||||
| M3/M4 | web scheme 校验 / shlex Windows | ❌ 未修 | 维持原判(Linux 部署不受影响) |
|
||||
| L1 | LIMIT f-string | ❌ 未修 | `db.py:417, 459` 两处在位(内部 int 直达,防御性待修) |
|
||||
|
||||
### 9.2 v3.4 新代码审计(多模型档案 + L4 失败转移)
|
||||
|
||||
**设计评价:整体高质量**。免密端点与密钥隔离(`allow_key_file=False` 防跨厂商凭据泄漏)、用户中断绝不转移(cancel 广播全链 + `_cancelled` 门控)、流式仅首事件前重放(`produced_any` 门控)、fallback 链一层不递归 + 自引用/悬空引用硬性报错、构造失败跳过不阻断主档案——关键决策均正确且有测试锚定(`test_models_profile.py` 约 60 个 L4 用例)。
|
||||
|
||||
新发现(P 系列):
|
||||
|
||||
| # | 级别 | 问题 | 位置 | 说明 |
|
||||
|---|---|---|---|---|
|
||||
| P1 | 中 | 非瞬时 HTTP 错误响应体无截断(N15 残留) | `llm_client.py:215-216` | `RuntimeError(f"HTTP {code}: {response.text}")` 全量响应体进异常栈;部分网关 4xx 页回显 Authorization 头 → 密钥进日志。与瞬时路径 `[:200]` 不一致,建议同样截断 |
|
||||
| P2 | 中 | Retry-After 无上限 | `llm_client.py:290` | `max(1.0, float(v))` 无上界——异常/恶意服务端 `Retry-After: 999999` 使档案冷却 11.5 天。建议 cap(如 `min(v, max(cooldown_seconds, 3600))`) |
|
||||
| P3 | 低中 | fallback 档案缺 model 不校验 | `llm_client.py:671-673` | 主档案经 `_require_model` 硬性报错,fallback 档案直接 `fb_profile.get("model")` → model=None → API 400 → 表现为神秘 None 而非清晰配置错误。建议构造链时同样校验 |
|
||||
| P4 | 低 | FailoverClient 非线程安全 | `llm_client.py:321-334` | `_cancelled`/`_call_started`/`_stats` 无锁。当前调用面无实际并发(compressor 在 `_msg_lock` 内同步调用、子代理独立 client),但 cancel(WebSocket 线程)与 generator 启动时 `_cancelled = False` 复位存在理论竞态窗口。建议注释声明"非并发契约"或加锁 |
|
||||
| P5 | 低 | `LLMClient._cancelled` 死状态 | `llm_client.py:52, 76` | 只写不读:cancel() 置位后无消费方也从不清除,真正判断靠 FailoverClient 自身标志。建议删除或在底层 client 消费 |
|
||||
| P6 | 低 | fallback 列表重复项不查重 | `llm_client.py:660-673` | `["kimi","kimi"]` 会在链内构造两个同名 client,冷却键同名冲突、重复撞同端点;影响有限 |
|
||||
| P7 | 信息 | timeout_budget 仅在尝试间检查 | `llm_client.py:363-366, 404, 443` | 预算不截断单档案进行中的阻塞读(单次仍可超时到 timeout)。与文档"限制整条链尝试的总秒数"有细微出入,可接受,建议文档注明"尝试间预算" |
|
||||
|
||||
### 9.3 架构合规矩阵(第三轮)
|
||||
|
||||
| 硬约束 | 状态 |
|
||||
|---|---|
|
||||
| WebSocket :8765 | ✅ |
|
||||
| 依赖单向 `common←subsystems←tools←aibox` | ✅ |
|
||||
| Worker 仅 7 执行工具 | ✅ |
|
||||
| 工具输出 ≤1MB + 超时 + 中断 | ⚠️ executor 链 ✅;残余:N8 流式旁路、N5 子代理超时、N2 drain 泄漏 |
|
||||
| prompt 三级排序 / sqlite 线程安全 | ⚠️ N7/N10 待修 |
|
||||
| 多模型档案凭据隔离(v3.4 新增维度) | ✅ 跨厂商 Key 不互发、免密端点显式声明 |
|
||||
|
||||
### 9.4 Batch G 修复建议(按优先级)
|
||||
|
||||
1. **G1(中)**:P1 非瞬时路径响应体截断 `[:200]` + P2 Retry-After 上限 cap——两处各一行改动
|
||||
2. **G2(中)**:P3 fallback 档案补 `_require_model` 校验(缺 model 的 fallback 跳过并告警,与缺 Key 行为对齐)
|
||||
3. **G3(中)**:N5 子代理 cancel 传播(`run_full_subagent` 传递父 `_cancel_event`,每轮 think 前检查)+ N2 drain 线程加 sentinel/空队列退出
|
||||
4. **G4(中低)**:N3 `/new` 的 `reset_conversation` 移入 `run_in_executor` + N8 registry 转发行截断 + N12 persona 路径边界校验(`os.path.realpath` 前缀检查)
|
||||
5. **G5(低)**:N7 FTS 引号翻倍 + N10 BEGIN IMMEDIATE + L1 LIMIT 参数化 + N13 `except Exception`/`discard` + N14 任务登记改按 task id 集合 + N16 收窄 except + 归档文件名加毫秒/序号
|
||||
6. **G6(低/择机)**:P4 并发契约注释或加锁 + P5 死状态清理 + P6 fallback 去重 + N9 prompt 更新确认门槛 + M3/M4 视部署平台决定
|
||||
|
||||
### 9.5 测试基线
|
||||
|
||||
全量 **355 passed**(21.8s,Windows 开发机 `py -3`),含 L4 失败转移约 60 用例、Batch E PoC 回归 5/5、Batch F1/F2 修复测试 8 例。零回归。
|
||||
|
||||
---
|
||||
|
||||
## 十、Batch G 修复落地(2026-08-18,第三轮审计修复)
|
||||
|
||||
按 §9.4 建议分 5 波实施,每波全量测试验证后进入下一波。
|
||||
|
||||
### 10.1 修复清单
|
||||
|
||||
| 波次 | # | 修复内容 | 文件/位置 |
|
||||
|---|---|---|---|
|
||||
| W1 | P1 | 非瞬时 HTTP 错误响应体截断 `[:200]`(与瞬时路径对齐,密钥/错误页不进异常栈) | `llm_client.py:216` |
|
||||
| W1 | P2 | Retry-After 解析上限 `min(max(1.0, v), 3600.0)`——恶意/异常头最多冷却 1h | `llm_client.py:284-296` |
|
||||
| W1 | P3 | fallback 档案构造前补 `_validate_profile` + `_require_model`,配置缺陷跳过并告警(对齐缺 Key 行为) | `llm_client.py:681-695` |
|
||||
| W1 | P5 | 删除 `LLMClient._cancelled` 死状态(只写不读;真实机制在 FailoverClient),cancel() 注释说明职责边界 | `llm_client.py:48-49, 71-76` |
|
||||
| W1 | P6 | fallback 列表 `seen` 去重(自引用硬性报错保持不变,测试锚定) | `llm_client.py:669-679` |
|
||||
| W1 | P4 | FailoverClient 类文档补并发契约声明(非线程安全 + 当前调用面为何安全) | `llm_client.py:324-327` |
|
||||
| W2 | N8 | registry 流式转发受 `TOOL_MAX_RESULT_LENGTH` 约束:累计预算 + 超限截断 + 停止转发提示(一次性) | `tools/registry.py:113-145` |
|
||||
| W2 | N12 | persona 路径边界校验:normcase+abspath 前缀检查,越界/绝对路径拒绝 | `aibox.py:536-544` |
|
||||
| W3 | N2 | drain 线程重构:批次级单线程 + 独立 stop 事件 + try/finally 统一回收(正常/异常/中断路径均 join≤2s);`set_output_queue(None)` 移入 finally | `aibox.py:1254-1386` |
|
||||
| W3 | N5 | 子代理取消传播:`run_full_subagent` 新增 `cancel_event` 参数→子代理共享父事件(think 前置/工具间/流式事件间既有检查点联动)+ 0.2s 观察者线程转发深度中断(强关在途 HTTP、kill 子进程),finally 通知退出;`SubAgentTool` 构造注入父 `_cancel_event` | `aibox.py:621-624, 669-684, 773-775, 786-790, 897-902, 1145-1147` |
|
||||
| W4 | N3 | `/new` 的 `reset_conversation` + `prompt_manager.reload()` 移入 `run_in_executor`(不再同步获取 `_msg_lock` 冻结事件循环) | `aibox.py:2386-2400` |
|
||||
| W4 | N13 | broadcast 裸 except → `except Exception`(不再吞 CancelledError);unregister `remove` → `discard`(幂等) | `aibox.py:2080-2084, 2185-2192` |
|
||||
| W4 | N14 | `_processing_tasks` 改按 websocket 键的任务集合登记:done_callback 从集合 discard;连接关闭取消该连接全部在途任务(此前仅最后一个) | `aibox.py:2022, 2847-2870` |
|
||||
| W5 | N7 | FTS5 关键词内双引号翻倍转义(`"` → `""`),MATCH 短语语法合法 | `subsystems/db.py:178-182` |
|
||||
| W5 | L1 | 两处 `LIMIT {limit}` f-string → `LIMIT ?` 参数化(`get_pending_memos`/`get_immediate_memos`) | `subsystems/db.py:419-421, 462-465` |
|
||||
| W5 | N10 | `mark_as_triggered` 读写收进单连接 `BEGIN IMMEDIATE` 事务(消灭 check-then-act 竞态:并发触发不再双计数/双推进 repeat) | `subsystems/db.py:515-590` |
|
||||
| W5 | N16 | 历史索引加载裸 except → `except Exception` + 告警;归档文件名冲突时加 `_N` 序号(同秒归档不覆盖) | `subsystems/conversation_history.py:58-69, 138-146` |
|
||||
|
||||
未实施(按 §9.4 G6 择机项):N9(prompt 更新确认门槛,涉及交互设计)、M3/M4(web scheme 校验 / shlex Windows,Linux 部署不受影响)、P7(timeout_budget 语义为文档注记,已在 FailoverClient 文档说明)。
|
||||
|
||||
### 10.2 新增测试
|
||||
|
||||
`tests/test_g_fixes.py` **20 例**(全部通过):
|
||||
- P1/P2 ×4:非瞬时错误体截断、Retry-After cap/正常值/非法值
|
||||
- P3/P5/P6 ×2:fallback 缺 model 跳过 + 重复去重(链 = [main, dupe, ok2])、死状态删除
|
||||
- N2 ×2:drain 线程批次结束/工具异常路径均回收
|
||||
- N5 ×3:子代理共享父取消事件 / 未传时保留自身 / SubAgentTool 注册注入链路
|
||||
- N7 ×2:含引号关键词不破坏检索、常规检索不回归
|
||||
- N8 ×1:120000 字符流式输出被截断至预算 + 提示标记
|
||||
- N10 ×2:缺失 id 事务路径返回 None、一次性任务完成语义
|
||||
- N12 ×3:`../` 越界/绝对路径拒绝、合法 persona 正常加载
|
||||
- N16 ×1:同秒两次归档文件名不冲突
|
||||
|
||||
N3/N13/N14 为 async WebSocket 路径,依赖事件循环与真实连接对象,由现有 WebSocket 测试面回归(全量通过)。
|
||||
|
||||
### 10.3 测试基线
|
||||
|
||||
全量 **375 passed**(22.5s,Windows 开发机 `py -3`)= 原 355 + Batch G 新增 20。零回归。
|
||||
|
||||
---
|
||||
|
||||
## 十一、第四轮审计(2026-08-18,HEAD `d1f780e`)
|
||||
|
||||
方法:3 路并行审读(aibox.py 全文 / llm_client+common+tools / subsystems+tests)+ 主审对全部候选发现逐条代码复核,关键项 Linux 动态实证(PoC 6/6)。
|
||||
|
||||
### 11.1 新发现与修复(Batch H,13 项全落地)
|
||||
|
||||
| # | 级别 | 问题 | 位置 | 修复 |
|
||||
|---|---|---|---|---|
|
||||
| H1 | **高危** | **N5 修复引入回归**:子代理共享父取消事件后,首轮 `_think_and_respond_impl` 无条件 `clear()` 吞掉父级已置位的取消信号(用户中断 / --once 超时恰落在子代理生成窗口时被清,观察者线程随之失效)——取消传播被绕过,子代理跑完全程 | `aibox.py:675, 1470` | 共享事件时置 `_shared_cancel_event` 跳过 clear;赋事件后立即检查 is_set 直接返回取消 |
|
||||
| M1 | 中 | `/reload_prompts` 在事件循环线程同步获取 `_msg_lock` 冻结整个事件循环(N3 同类遗漏,只修了 `/new`) | `aibox.py:2492` | 移入 `run_in_executor`(与 /new 一致) |
|
||||
| M2 | 中 | SkillManager 懒加载在事件循环线程串行执行 `--help` 子进程(每技能 ≤10s),首条消息/首个任务触发时冻结服务 | `aibox.py:2885-2891` | 启动前预热 WS 与调度器两个共享实例 |
|
||||
| M3 | 中 | 调度器任务无超时:AI think 卡死 → `await ai_future` 永久阻塞 → 调度循环停摆、后续任务全部丢失(已 mark_as_triggered 不重试) | `aibox.py:333` | `call_later` 任务级超时(`scheduler.task_timeout`,默认 600s)→ `ai.cancel()`,循环继续 |
|
||||
| M4 | 中 | `get_recent_memories(limit=None)` 无 LIMIT 全表加载(db_tools 分页传 None),记忆库增长后单次调用全量入内存 | `db.py:234` | None → 默认上限(对齐 search_memories 语义) |
|
||||
| M5 | 中 | `memo delete` 实为 `complete_memo(force)`——只标记完成不删行,重复任务被推进而非移除,stats 计数失真 | `db_tools.py:747` | db 层新增 `delete_memo` 物理删除,工具改调 |
|
||||
| M6 | 中 | `PromptManager.reload()` 先清空 `_prompts` 再逐文件重读:并发连接构建 system 消息时读到空/半截列表(N3 移线程池后放大) | `prompt_manager.py:298` | 本地构建新列表后原子替换 |
|
||||
| L1 | 低 | read_file 的 `start_line/end_line` 未强转直拼 sed 单引号模板——LLM 传含引号字符串可闭合引号注入生成的指引命令 | `basic_tools.py:97` | int() 强转 + 范围夹取 |
|
||||
| L2 | 低 | `_save_config` 吞异常仅 print:磁盘只读时 set_proxy 仍报"已保存",重启后配置丢失无感知 | `config.py:284` | `_save_config` 返回落盘结果,set_proxy 透传,工具如实报告 |
|
||||
| L3 | 低 | SSE 行 `decode("utf-8")` 遇非 UTF-8 字节抛 UnicodeDecodeError,整条流中断且不触发 failover | `llm_client.py:222` | `errors="replace"` |
|
||||
| L4 | 低 | persona 路径校验用 abspath 未解析符号链接:skills 目录内软链指向任意文件可绕过前缀检查被读取注入 | `aibox.py:536` | realpath 后再做前缀校验 |
|
||||
| L5 | 低 | 1MB 截断按字节切后 `errors='replace'` 解码,行尾出现半个字符 � | `executors.py:150` | 改 `errors='ignore'` |
|
||||
| L6 | 低 | compressor 压缩块全空内容时 `original_tokens=0` → 除零崩溃 | `conversation_compressor.py:517` | `max(original_tokens, 1)` |
|
||||
| L7 | 低 | InstallSkillTool 指引命令 `ls -l {src}`/`cp {src} {dst}` 未引用:LLM 可控路径含空格/元字符时 shell 拆词/注入 | `skill_tools.py:269-285` | shlex.quote 源与目标路径 |
|
||||
| L8 | 低 | ask_user 的 required 为空,LLM 不传 question 时 `str + None` 抛 TypeError | `misc_tools.py:42` | required 补 question + 缺参明确报错 |
|
||||
| L9 | 低 | Linux 孙进程持管道泄漏(N4 残余变体):`(sleep 300 &)` 形态 shell 已死但孙进程持写端,reader 阻塞无 EOF;**实证修复前孙进程残留** | `executors.py` | POSIX `start_new_session` + 进程组信号(组 ID 在 Popen 后立即捕获,shell 死亡后 getpgid 不可用);宽限窗口到期强杀进程组 |
|
||||
| L10 | 低 | web 工具 url 无 scheme 校验(M3 遗留):`file://` 等经 curl 指引可达(SSRF) | `basic_tools.py:319` | http/https 白名单 |
|
||||
| L11 | 低 | registry 日志打码精确匹配:apiKey/passwd 等变体漏网 | `registry.py:19` | 子串包含匹配(含 apiKey/passwd/credential) |
|
||||
| L12 | 低 | FailoverClient 异常消息经 `_log` 落 ai.log:网关错误页若在 200 字符内回显 Authorization 头则密钥进日志 | `llm_client.py` | `_log`/on_failover 前 Bearer 令牌正则脱敏 |
|
||||
|
||||
### 11.2 修复实证(Linux PoC 6/6 PASS)
|
||||
|
||||
```
|
||||
[PASS] L-1 无输出命令超时: 3.0s 返回,含超时(原 PoC-1 为 Windows powershell 命令)
|
||||
[PASS] L-3 1MB截断: 结果 1.00MB(Linux 等价命令)
|
||||
[PASS] L-4 cancel中断: cancel 后 2.0s 结束
|
||||
[PASS] L-5 工具层1MB贯通: 结果 1.00MB
|
||||
[PASS] L-9 孙进程持管道不悬挂: 1.1s 返回,reader 正常收尾
|
||||
[PASS] L-9 孙进程已被清理: 残留无(修复前 `(sleep 300 &)` 孙进程残留至自然退出)
|
||||
```
|
||||
|
||||
L-9 实证过程说明:首版实现按"break 后 killpg"处理,但 shell 已死时 `os.getpgid(pid)` 抛 ProcessLookupError,组 ID 无法反查,孙进程仍残留(实证 FAIL 抓出)。修正为 **Popen 成功后立即捕获组 ID 并保存**,killpg 用保存值——6/6 通过。
|
||||
|
||||
### 11.3 遗留项(维持原判 / 择机)
|
||||
|
||||
- **N9**(prompt 注入存储):confirm 默认 True、100_STATUS 直写无门槛——涉及交互设计(用户确认门槛),维持 G6 择机项
|
||||
- **M4/M3(旧编号)**:shlex Windows 转义(Linux 部署不受影响)、Q9 代理认证 curl 指引(修复需把密码嵌入命令串,反而引入密码进日志/会话史的新泄漏,维持现状)
|
||||
- **F5** express_emotion 不可中断、**F6** /new 占用默认线程池、**F8** WS 限速/空闲回收、**F9** --once 软超时(H1 修复后取消可传播,残余窗口受 timeout 有界)、**Q7a** stream_chat 取消空窗(受 timeout 有界)、Q4 busy_timeout 重试、Q5 提醒格式异常静默、Q7b load_conversation 无调用面、Q8 对话文件原子写、compressor 裸 print / token 双计、单例 __new__ 无锁、db.py `if limit:` limit=0 边界——均为低危/信息级,下批处理
|
||||
- **测试环境注记**:`test_executors.py::test_execute_timeout` 与 `test_models_profile.py::test_flat_api_still_reads_key_file` 为 Windows 向用例(ping -n 语法 / 依赖 ~/.cert/deepseek.key),Linux 环境预存失败,与本轮改动无关(git stash 实证)
|
||||
|
||||
### 11.4 测试基线(Linux,`aibox/venv`)
|
||||
|
||||
全量 **408 passed** = 原 382(Linux 环境)+ Batch H 新增 26。零回归;2 个预存环境性失败同前。PoC 回归脚本为 Windows 命令集,Linux 用等价命令实证(见 11.2)。
|
||||
|
||||
---
|
||||
|
||||
## 十二、Batch I 修复落地(2026-08-18,遗留项收尾)
|
||||
|
||||
按 §11.3 遗留清单分 3 波实施:安全项 → 稳定性项 → 信息级顺手项。
|
||||
|
||||
### 12.1 修复清单
|
||||
|
||||
| # | 级别 | 修复内容 | 文件/位置 |
|
||||
|---|---|---|---|
|
||||
| N9 | 中 | prompt 写入内容长度上限(持久化注入面约束):`PromptManager.MAX_CONTENT_LENGTH=10000`,update_identity/update_soul/update_status 超长拒绝(manager 兜底 + 工具友好提示);save_memory_and_end_conversation 状态内容超长拒绝 | `prompt_manager.py`、`prompt_tools.py` |
|
||||
| Q7b | 低 | `load_conversation` realpath 前缀校验,`../` 越界/软链逃逸返回 None(当前无 LLM 调用面,防御性加固) | `conversation_history.py:348` |
|
||||
| F5 | 低 | 情绪子代理取消传播:`run_light_subagent` 新增 `cancel_event`,改走 stream_chat(可被 cancel 强关在途 HTTP)+ 观察者线程;父已取消早退;`ExpressEmotionTool` 构造注入父 `_cancel_event` | `aibox.py:513, 608-660, 1177` |
|
||||
| Q4 | 低 | `mark_as_triggered` BEGIN IMMEDIATE 前 `PRAGMA busy_timeout = 10000`——多进程同库(CLI+服务)写竞争等待锁释放而非 5s 后直接报 locked | `db.py:527` |
|
||||
| Q5 | 低 | 提醒时间格式异常任务不再静默丢弃:去掉 SQL 层 `reminder_time <= ?` 预过滤(字符串比较会静默排除格式异常行),时间过滤改由 Python 二次验证权威判断;格式异常按"到期"返回并告警 | `db.py:405-411, 442-444` |
|
||||
| Q8 | 低 | 对话文件原子写:`_atomic_write_json`(tmp + os.replace)覆盖历史索引/当前对话/归档三处写路径,崩溃不留半截 JSON | `conversation_history.py` |
|
||||
| F6 | 低 | WS 命令专用线程池:`WebSocketHandler._cmd_executor`(max_workers=2),`/new`、`/reload_prompts` 不再占用默认池(繁忙 AI 时多连接并发 /new 不再耗尽线程池);服务退出时 shutdown | `aibox.py` |
|
||||
| F9 | 低 | --once 硬超时兜底:软超时 cancel 后若 AI 卡死在不可中断路径(HTTP 空窗、忽略信号的子进程),宽限一个超时周期后 `os._exit(4)` 强杀,保证硬性结束语义 | `aibox.py:3283-3299` |
|
||||
| — | 信息 | PromptManager 单例 `__new__` 加锁(并发首例化竞态防护) | `prompt_manager.py:49-53` |
|
||||
| — | 信息 | `get_pending_memos`/`get_immediate_memos` 的 `if limit:` → `if limit is not None`(limit=0 语义修正,None 才不加 LIMIT) | `db.py:420, 464` |
|
||||
|
||||
### 12.2 维持原判(未改动)
|
||||
|
||||
- **Q7a**(stream_chat 取消空窗):POST 返回前 `_current_response` 未赋值,cancel 无法中断——受 timeout 有界;根治需线程化 POST + 轮询取消,复杂度高,维持记录
|
||||
- **M4(旧编号)**(shlex Windows 转义):Linux 部署不受影响
|
||||
- **Q9(代理认证 curl 指引)**:修复会把密码嵌入命令串(进日志/会话史),反而引入新泄漏,维持现状
|
||||
- **F8**(WS 限速/空闲回收):本地单用户产品设计项
|
||||
- compressor 裸 print(终端过程展示,保留)、`estimate_tokens` 英文双计(测试锚定既有算法,仅影响压缩阈值)
|
||||
|
||||
### 12.3 新增测试(`tests/test_i_fixes.py` 16 例,全部通过)
|
||||
|
||||
N9 ×5(manager 兜底/身份/灵魂/状态超长拒绝 + 正常写入)、Q7b ×2(越界拒绝 + 正常加载)、F5 ×3(父已取消早退/流式中断/注册注入链路)、Q5 ×2(格式异常按到期返回 + 正常过滤不回归)、Q8 ×1(无 .tmp 残留 + JSON 合法)、F6 ×1(专用线程池存在与释放)、limit=0 ×2。
|
||||
|
||||
### 12.4 测试基线(Linux,`aibox/venv`)
|
||||
|
||||
全量 **424 passed** = 原 408 + Batch I 新增 16。零回归;2 个预存环境性失败同前。
|
||||
@@ -0,0 +1,358 @@
|
||||
# aibox CLI 单次调用模式 · 开发文档
|
||||
|
||||
> 文档类型:功能开发计划
|
||||
> 模块:aibox(单 AI 智能体系统)
|
||||
> 日期:2026-08-17
|
||||
> 状态:方案研究(待评审)
|
||||
> 影响范围:`aibox/aibox.py`(仅增量改动,WebSocket 模式零回归)
|
||||
|
||||
---
|
||||
|
||||
## 1. 概述
|
||||
|
||||
### 1.1 背景与动机
|
||||
|
||||
aibox 目前只有两种使用方式:
|
||||
|
||||
| 使用方式 | 入口 | 特点 |
|
||||
|----------|------|------|
|
||||
| WebSocket 服务器 | `python aibox.py` → `ws://localhost:8765` | 长连接、多轮对话、前端/BotOps 交互(核心场景) |
|
||||
| 备忘录 CLI | `--check-memo` / `--memo-add` / `--memo-list` / `--memo-complete` | 只做任务 CRUD,不触发 AI 对话 |
|
||||
|
||||
缺少一种**"问一次就退"**的 CLI 一次性对话模式。这导致以下场景无法覆盖:
|
||||
|
||||
1. **脚本化调用**:无法在 shell 脚本 / cron / CI 流水线中把 AI 当命令行工具用(如 `aibox --once "统计当前目录 py 文件行数" | grep ...`)
|
||||
2. **资源受限环境**:树莓派 / NPU 盒子上不想常驻一个 WebSocket 服务进程,只想按需拉起、用完即退
|
||||
3. **调试与测试**:复现某个 prompt 行为必须启动服务器 + 前端连接,链路太重
|
||||
4. **进程级隔离执行**:Worker 子代理当前是**同进程线程**(`run_full_subagent` → `DeepSeekChat`),子代理崩溃/死循环会影响主服务;有了 CLI 单次模式后,可演进为"独立子进程执行",OS 级隔离
|
||||
|
||||
### 1.2 设计目标
|
||||
|
||||
1. **新增 `--once QUERY` 参数**:执行一次性 AI 对话(含完整工具调用链),结果流式输出到 stdout,完成后进程退出
|
||||
2. **纯增量、零回归**:不改动 WebSocket 模式任何现有行为;不加参数时行为与现状完全一致
|
||||
3. **复用现有架构**:复用 `EnhancedAIManager` / `DeepSeekChat` / Orchestrator-Worker 工具分层 / 三层中断机制,不重写对话逻辑
|
||||
4. **管道友好**:stdout 只出 AI 回答正文,过程行(🔧 工具调用等)走 stderr,退出码语义化,可被 shell 脚本安全组合
|
||||
5. **状态隔离明确**:单次模式默认不加载历史对话、不写对话归档,记忆写入显式可控
|
||||
|
||||
### 1.3 非目标
|
||||
|
||||
- 不做多轮对话 CLI(交互式 REPL 已由 `client.py` 覆盖)
|
||||
- 不做 HTTP/REST 单次接口(如需可在后续批次评估,复用同一内核)
|
||||
- 不改变 Worker 子代理的线程内执行方式(进程化是后续可选演进,不在本期范围)
|
||||
|
||||
---
|
||||
|
||||
## 2. 现状分析
|
||||
|
||||
### 2.1 入口结构(aibox.py)
|
||||
|
||||
```
|
||||
parse_args() # L2786: argparse 定义
|
||||
main() # L2927: 主入口
|
||||
├── 语言/配置/调试覆盖 # L2931-2959
|
||||
├── 目录创建 + 日志初始化 # L2961-2973
|
||||
├── SubprocessCommandExecutor # L2976
|
||||
├── EnhancedAIManager # L2978-2981
|
||||
├── handle_check_commands() # L2983: --check-memo 处理后 return
|
||||
├── handle_memo_commands() # L2986: --memo-* 处理后 return
|
||||
└── manager.run() # L2990: WebSocket 服务器(唯一终点)
|
||||
```
|
||||
|
||||
关键观察:
|
||||
|
||||
- `handle_check_commands` / `handle_memo_commands` 已经是**"处理后即退出"的 CLI 模式先例**,单次对话模式可完全沿用该模式:在 `manager.run()` 之前插入 `handle_once_mode(manager)`,命中则执行并 return
|
||||
- `EnhancedAIManager.add_ai()`(L2584)已封装完整的 AI 实例组装(memory_db / memo_db / history_manager / prompt_manager / command_executor 注入),单次模式可直接复用
|
||||
- `DeepSeekChat.think_and_respond(input_text, output_callback)`(L1308)是同步阻塞函数,天然适合 CLI 直接调用——不需要 asyncio 事件循环
|
||||
|
||||
### 2.2 output_callback 消息类型(流式输出契约)
|
||||
|
||||
`think_and_respond` 通过 `output_callback(msg_type, content)` 实时推送:
|
||||
|
||||
| msg_type | 语义 | CLI 模式去向 |
|
||||
|----------|------|--------------|
|
||||
| `chunk` | AI 回答增量片段(**权威内容来源**) | **stdout**(流式 write + flush) |
|
||||
| `complete` | 收尾信号(内容与 chunk 累积重复) | 仅作收尾标记,不重复输出 |
|
||||
| `reasoning_start/reasoning/reasoning_end` | 思考模式流式 | stderr(可用 `--show-reasoning` 控制) |
|
||||
| `line` | 过程行:🔧 工具调用、工具输出、💬 前缀 | stderr |
|
||||
| `ai_break` | 工具调用分段信号 | 忽略(CLI 无 widget 概念) |
|
||||
| `error` | API/解析错误 | stderr + 退出码 2 |
|
||||
| `dialogue_ended` | AI 主动结束对话 | 正常收尾 |
|
||||
|
||||
### 2.3 需要处理的技术要点
|
||||
|
||||
| 要点 | 现状 | 风险与对策 |
|
||||
|------|------|-----------|
|
||||
| **压缩输出污染** | `_think_and_respond_impl` L1346 中 `print("📦 对话压缩...")` 直接打 stdout | 单次模式默认不加载历史 → 消息极少 → 几乎不触发压缩;另需确认该 print 走 stderr 或抑制(Batch C 处理) |
|
||||
| **记忆写入并发** | 多个 CLI 进程并发写 `memories.db`(SQLite) | 默认 `--once` 模式 AI 侧仍可调用 `memory` 工具(orchestrator profile);SQLite 有文件锁兜底,低频并发无碍;提供 `--no-save-memory` 开关时通过提示词引导(见 §3.4) |
|
||||
| **Ctrl+C 中断** | 三层取消机制(Event + HTTP close + 子进程 kill)依赖 `cancel()` 被调用 | CLI 模式注册 `signal.SIGINT` handler 调 `manager.cancel()`,与 WebSocket `handle_interrupt` 路径等价 |
|
||||
| **子代理输出转发** | `SubAgentTool._subagent_output` 经 `send_line` → `output_queue` → 回调链 | CLI 模式的 `output_callback` 会自然收到子代理的 `line`/`reasoning`,全部导流 stderr,无需特殊处理 |
|
||||
| **对话归档写入** | `global_exit_handler` 会调 `history_manager.end_current_conversation()` | 单次模式传 `exit_callback=None`(与 TaskScheduler 的任务实例一致),AI 主动 `save_memory_and_end_conversation` 时不触发全局归档 |
|
||||
|
||||
---
|
||||
|
||||
## 3. 设计方案
|
||||
|
||||
### 3.1 命令行接口
|
||||
|
||||
```bash
|
||||
# 基本用法:一次性对话,结果输出到 stdout
|
||||
python aibox.py --once "帮我统计当前目录下所有 py 文件的行数"
|
||||
|
||||
# 静默模式:只输出最终回答正文(工具过程完全静默,适合管道)
|
||||
python aibox.py --once "把 README.md 翻译成英文摘要" --quiet
|
||||
|
||||
# 显式控制
|
||||
python aibox.py --once "..." --no-save-memory # 不写入记忆数据库
|
||||
python aibox.py --once "..." --timeout 120 # 全局超时 120 秒(默认 300)
|
||||
python aibox.py --once "..." --show-reasoning # stderr 显示思考过程(默认不显示)
|
||||
```
|
||||
|
||||
新增参数(全部可选,`--once` 之外均有默认值):
|
||||
|
||||
| 参数 | 类型 | 默认 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `--once QUERY` | str | — | 一次性对话模式:执行 QUERY 后退出,不启动 WebSocket |
|
||||
| `--quiet` | flag | false | 仅 stdout 输出回答正文;过程行/thinking 全静默 |
|
||||
| `--no-save-memory` | flag | false | 禁止本次对话写入记忆数据库 |
|
||||
| `--timeout SEC` | int | 300 | 全局执行超时,超时 kill 并退出码 4 |
|
||||
| `--show-reasoning` | flag | false | 在 stderr 显示思考过程流 |
|
||||
|
||||
### 3.2 退出码语义
|
||||
|
||||
| 退出码 | 含义 |
|
||||
|--------|------|
|
||||
| 0 | 成功(AI 正常完成回答) |
|
||||
| 1 | 工具执行链路异常(如 run_command 失败且 AI 报告失败) |
|
||||
| 2 | API 请求失败(网络/HTTP/JSON 解析,服务端已回滚消息) |
|
||||
| 3 | 用户中断(Ctrl+C) |
|
||||
| 4 | 全局超时(--timeout 到期) |
|
||||
|
||||
> Batch A 只保证 0 / 2 / 3 / 4;退出码 1 的判定(从 AI 文本中识别失败)在 Batch B 实现,初版可统一归 0。
|
||||
|
||||
### 3.3 核心实现(Batch A,约 60~80 行增量)
|
||||
|
||||
**改动点 1:`parse_args()`(aibox.py L2786)** 追加 §3.1 的 5 个参数定义。
|
||||
|
||||
**改动点 2:新增 `handle_once_mode()`**,插入 `main()` 中 `handle_memo_commands()` 之后、`manager.run()` 之前:
|
||||
|
||||
```python
|
||||
def handle_once_mode(manager, args) -> bool:
|
||||
"""一次性对话模式:命中 --once 则执行并返回 True(进程退出)。"""
|
||||
if not args.once:
|
||||
return False
|
||||
|
||||
ai = manager.add_ai()
|
||||
ai.history_manager = None # 单次模式不加载/不写对话历史
|
||||
if args.no_save_memory:
|
||||
ai.memory_db = None # 工具注册时已注入;此处断引用防止运行期写入
|
||||
# 注:Batch A 中 memory 工具已在 _register_tools 注册;
|
||||
# 完整禁用见 §3.4 说明
|
||||
|
||||
exit_code = [0]
|
||||
|
||||
def cli_callback(msg_type, content):
|
||||
if msg_type == "chunk":
|
||||
sys.stdout.write(content)
|
||||
sys.stdout.flush()
|
||||
elif msg_type == "complete":
|
||||
sys.stdout.write("\n")
|
||||
sys.stdout.flush()
|
||||
elif msg_type == "error":
|
||||
print(content, file=sys.stderr)
|
||||
exit_code[0] = 2
|
||||
elif msg_type == "line" and not args.quiet:
|
||||
sys.stderr.write(content)
|
||||
sys.stderr.flush()
|
||||
elif msg_type.startswith("reasoning") and args.show_reasoning:
|
||||
sys.stderr.write(content)
|
||||
sys.stderr.flush()
|
||||
# ai_break / dialogue_ended 忽略
|
||||
|
||||
# Ctrl+C → 三层取消(与 WebSocket interrupt 路径等价)
|
||||
def _sigint(sig, frame):
|
||||
manager.cancel()
|
||||
signal.signal(signal.SIGINT, _sigint)
|
||||
|
||||
# 全局超时:守护线程到点 cancel
|
||||
timer = threading.Timer(args.timeout, manager.cancel)
|
||||
timer.daemon = True
|
||||
timer.start()
|
||||
try:
|
||||
response = ai.think_and_respond(args.once, cli_callback)
|
||||
if response is None or (isinstance(response, str)
|
||||
and response.startswith(f"\n{i18n.get('error_prefix')}")):
|
||||
exit_code[0] = max(exit_code[0], 2)
|
||||
# 被 cancel 打断时 think_and_respond 返回 interrupt 文案
|
||||
if ai._is_cancelled() and timer.is_alive() is False:
|
||||
exit_code[0] = 4
|
||||
finally:
|
||||
timer.cancel()
|
||||
ai.cancel() # 兜底清理子进程
|
||||
|
||||
sys.exit(exit_code[0])
|
||||
```
|
||||
|
||||
**改动点 3:`main()` 一行插入**:
|
||||
|
||||
```python
|
||||
if handle_memo_commands(manager):
|
||||
return
|
||||
|
||||
if handle_once_mode(manager, args): # 新增
|
||||
return
|
||||
|
||||
manager.run()
|
||||
```
|
||||
|
||||
> 注:`handle_once_mode` 内部 `sys.exit()`,return 分支实际不可达,保留 return 仅为与 `handle_check_commands` / `handle_memo_commands` 模式一致。
|
||||
|
||||
### 3.4 记忆写入的处理
|
||||
|
||||
`--no-save-memory` 的**完整实现**需要在工具注册层禁用 `MemoryTool`,但 `_register_tools` 依赖构造期注入的 `memory_db`。Batch A 采用轻量方案:
|
||||
|
||||
- **Batch A**:仅断 `ai.memory_db = None` 引用 + 文档声明"`--once` 场景不建议依赖记忆";AI 若调用 `memory` 工具会因注入的 db 为 None 由工具自身报错兜底(`MemoryTool.execute` 内部访问 db 返回错误字符串,不会崩溃)
|
||||
- **Batch B**:`DeepSeekChat.__init__` 增加 `enable_memory: bool = True` 形参,为 False 时 `_register_tools` 跳过 `MemoryTool` 注册,工具 schema 不出现在请求中,LLM 无从调用——这才是干净方案
|
||||
|
||||
### 3.5 与既有架构约束的对齐
|
||||
|
||||
| 项目硬约束 | 本方案如何满足 |
|
||||
|-----------|---------------|
|
||||
| WebSocket localhost:8765 通信 | 不受影响;`--once` 不启动 WebSocket,互斥使用 |
|
||||
| Worker 仅 7 种执行工具 | 不受影响;单次模式沿用 orchestrator profile,工具分层不变 |
|
||||
| 工具输出 ≤1MB + 超时中断 | 不受影响;复用 `SubprocessCommandExecutor` |
|
||||
| 单向依赖 common←subsystems←tools←aibox | 满足;改动全部在 aibox.py 聚合层 |
|
||||
| 新增/变更模块三处同步(根README/docs/MODULES/ARCHITECTURE) | 本功能非新模块,但落地时需同步更新 `aibox/README.md` 的"命令行参数"章节 |
|
||||
|
||||
---
|
||||
|
||||
## 4. 分批次实施计划
|
||||
|
||||
遵循项目"增量重构、低风险先行"惯例:
|
||||
|
||||
### Batch A:最小可用版(核心,建议先做)
|
||||
|
||||
- [ ] `parse_args()` 新增 `--once` / `--quiet` / `--timeout` 三个参数
|
||||
- [ ] 新增 `handle_once_mode()`(§3.3 实现)
|
||||
- [ ] `main()` 插入调用点
|
||||
- [ ] SIGINT handler + 全局超时 Timer
|
||||
- [ ] 退出码 0/2/3/4
|
||||
- [ ] **验收**:`python aibox.py --once "现在几点了"` 正常输出并退出;`echo $?` 为 0;WebSocket 模式回归(启动服务器 + client.py 对话一轮)
|
||||
|
||||
### Batch B:体验与边界补齐
|
||||
|
||||
- [ ] `--no-save-memory` / `--show-reasoning` 参数落地
|
||||
- [ ] `DeepSeekChat.__init__` 增加 `enable_memory` 开关,干净禁用 MemoryTool
|
||||
- [ ] `_think_and_respond_impl` L1346 压缩提示 `print` 改走 stderr(或注入 print 回调)
|
||||
- [ ] 退出码 1(工具失败识别)
|
||||
- [ ] 补充 `tests/test_once_mode.py`:mock LLM 客户端,验证 stdout/stderr 分流、退出码、超时路径
|
||||
- [ ] 更新 `aibox/README.md` 命令行参数章节
|
||||
|
||||
### Batch C:安全加固(可选,按需)
|
||||
|
||||
- [ ] SQLite WAL 模式(`PRAGMA journal_mode=WAL`),缓解多 CLI 进程并发写锁等待
|
||||
- [ ] `--worker-only` 开关:`tool_profile=worker` 纯执行面模式,供其他模块以子进程方式调用 AI 执行器(为"Worker 进程级隔离"演进铺路)
|
||||
- [ ] `--model` / `--temperature` 覆盖参数
|
||||
|
||||
### 明确不做(本期)
|
||||
|
||||
- 多轮交互式 CLI REPL(`client.py` 已覆盖)
|
||||
- HTTP 单次接口
|
||||
- Worker 子进程化(架构演进议题,单独立项)
|
||||
|
||||
---
|
||||
|
||||
## 5. 风险清单
|
||||
|
||||
| 风险 | 等级 | 缓解 |
|
||||
|------|------|------|
|
||||
| 冷启动开销(PromptManager 加载 + SQLite 初始化 + LLM 握手 ≈1~3s) | 低 | 可接受取舍;批量场景文档引导用 WebSocket 长连接 |
|
||||
| CLI 模式与 WebSocket 模式代码 drift | 中 | 复用 `EnhancedAIManager.add_ai()`,单次模式新增代码控制在 80 行内,核心对话逻辑零复制 |
|
||||
| `print` 污染 stdout 破坏管道 | 中 | Batch A 用 `--quiet` 兜底;Batch B 根治压缩提示 print |
|
||||
| 并发进程 SQLite 写锁 | 低 | 低频无碍;Batch C WAL 加固 |
|
||||
| Ctrl+C 时子进程泄漏 | 低 | 复用现成三层取消;`finally` 中 `ai.cancel()` 兜底 |
|
||||
|
||||
---
|
||||
|
||||
## 6. 测试方案
|
||||
|
||||
### 6.1 单元测试(tests/test_once_mode.py,Batch B)
|
||||
|
||||
| 用例 | 验证点 |
|
||||
|------|--------|
|
||||
| 正常问答 | mock LLM 返回纯文本 → stdout 含回答、退出码 0 |
|
||||
| 工具调用链 | mock LLM 返回 tool_calls → line 走 stderr、chunk 走 stdout |
|
||||
| API 失败 | mock LLM 抛 RequestException → 退出码 2、stderr 含错误 |
|
||||
| 超时 | mock LLM 阻塞 → Timer 触发 cancel → 退出码 4 |
|
||||
| SIGINT | 发送 SIGINT → 三层取消触发 → 退出码 3 |
|
||||
| `--quiet` | 过程行不输出到任何流 |
|
||||
| `--no-save-memory` | memories.db 无新增行 |
|
||||
|
||||
### 6.2 回归测试
|
||||
|
||||
- [ ] `py -m pytest aibox/tests -q` 现有 232 项全过
|
||||
- [ ] WebSocket 模式手测:启动服务器 + `client.py` 完成一轮对话 + Ctrl+C 中断
|
||||
|
||||
### 6.3 手动验收脚本
|
||||
|
||||
```bash
|
||||
# 基础问答(无工具)
|
||||
python aibox.py --once "用一句话介绍你自己" ; echo "exit=$?"
|
||||
|
||||
# 工具调用链(会派生 worker 子代理执行命令)
|
||||
python aibox.py --once "在 ~/.aibox/workspace 下创建一个 hello.txt 写入 hello" ; echo "exit=$?"
|
||||
|
||||
# 管道组合
|
||||
python aibox.py --once --quiet "列出当前目录 py 文件" | grep ".py" | wc -l
|
||||
|
||||
# 超时路径
|
||||
python aibox.py --once --timeout 5 "运行 sleep 60" ; echo "exit=$?" # 期望 4
|
||||
```
|
||||
|
||||
## 7. 实施状态(2026-08-17 落地记录)
|
||||
|
||||
> 本节为实施后回填,与第 3/4 节计划存在少量偏差,以本节为准。
|
||||
|
||||
### 7.1 提交与批次
|
||||
|
||||
| 批次 | 提交 | 内容 |
|
||||
|---|---|---|
|
||||
| Batch A | `ce3fb7f` | `--once` / `--quiet` / `--once-timeout`(300s 默认)+ `handle_once_mode` + SIGINT/超时复用三层取消 |
|
||||
| Batch B | `9e0d39c` | `--no-save-memory`(`enable_memory` 构造开关干净禁用 MemoryTool)、`--show-reasoning`(思考流默认静默走 stderr)、压缩提示改走 output_callback、退出码 1(orchestrator 层工具异常)+ `test_once_mode.py` 10 项 |
|
||||
| Batch C | `a5af4bb` | SQLite WAL、`--worker-only`(worker profile + 精简 system prompt 防人格污染)、`--model` / `--temperature`、子代理失败标记透传(`_SUBAGENT_TOOL_FAILED_MARKER` → orchestrator `stats.tool_failed`)+ 5 项测试 |
|
||||
| Batch D | `d03dbc7` | 多模型档案 `models` 配置节(本档案为实施中新增需求,超出原计划范围)+ `test_models_profile.py` 17 项 + README 更新 |
|
||||
|
||||
### 7.2 与计划的偏差
|
||||
|
||||
1. **超时参数名**:计划 `--timeout`,落地为 `--once-timeout`(避免与 `api.timeout` 语义混淆)
|
||||
2. **`models` 默认留空 `{}`**:计划未涉及;实施中发现默认档案硬编码 `DEEPSEEK_API_KEY` / 官方 URL 会覆盖用户自定义配置(本机为 `ARK_API_KEY` + 火山方舟端点,曾导致 401),故默认空 dict,未配置档案时零破坏回退旧 `api` 扁平节
|
||||
3. **`--worker-only` 需要精简 system prompt**:原计划仅切 tool_profile;实测 orchestrator 的人格/身份提示词会污染 worker 语义(自称"老板"、主动写文件、多轮不收敛),已对齐 `run_full_subagent` 的子代理语义修复
|
||||
4. **退出码 1 的边界**:仅覆盖 orchestrator 直接感知的工具异常;worker 子代理内的命令失败经 Batch C 的标记透传后同样计入。子命令失败但 AI 正常汇报时退出码为 0(AI 视角任务完成)
|
||||
|
||||
### 7.3 验证结果
|
||||
|
||||
- 单测 **264 项全过**(once 模式 15 + 多模型档案 17 新增)
|
||||
- 冒烟:基础问答 exit 0 / 工具调用链(子代理执行 echo)exit 0 / 超时路径 exit 4 / `--no-save-memory --quiet` stdout 干净 / `--show-reasoning` 思考流走 stderr / `--temperature` 透传 / 不存在的模型名 exit 2
|
||||
|
||||
### 7.4 明确不做(维持原判)
|
||||
|
||||
WebSocket 会话恢复、`--once` 结果缓存、批处理多 query——见第 4 节"明确不做"。
|
||||
|
||||
---
|
||||
|
||||
## 8. 后续演进方向(不在本期)
|
||||
|
||||
1. **Worker 进程级隔离**:`run_full_subagent` 的线程内 `DeepSeekChat` → `subprocess.run(["python", "aibox.py", "--once", "--worker-only", ...])`,子代理崩溃不波及主服务,可用 OS 用户/chroot 做权限隔离
|
||||
2. **HTTP 单次接口**:`POST /chat` 无状态接口,与 `--once` 共享内核
|
||||
3. **结构化输出**:`--json` 输出 `{answer, tool_calls, exit_reason}` JSON,供程序消费
|
||||
|
||||
---
|
||||
|
||||
## 9. 决策记录
|
||||
|
||||
| 决策点 | 结论 | 理由 |
|
||||
|--------|------|------|
|
||||
| 新入口文件还是复用 aibox.py | 复用 aibox.py + `--once` 参数 | 沿用 `--check-memo` 先例;避免维护两份入口 drift |
|
||||
| stdout 输出协议 | chunk 流式 + 末尾换行 | 与 WebSocket 的 chunk 权威语义一致;管道可逐段消费 |
|
||||
| 过程行去向 | stderr | Unix 管道惯例:stdout 是数据、stderr 是元信息 |
|
||||
| 单次模式是否加载历史 | 否 | "一次性"语义明确;避免历史污染导致行为不可复现 |
|
||||
| 是否默认写记忆 | 默认与 orchestrator 一致(可写),`--no-save-memory` 显式关闭 | 与 WebSocket 模式行为对齐,减少心智负担 |
|
||||
Loaded 100 of 149 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user