Compare commits
210
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 | ||
|
|
a6e1781185 | ||
|
|
14964b9574 | ||
|
|
2068319c3a | ||
|
|
0d16832612 | ||
|
|
0cfb3a7581 | ||
|
|
e77ef8cf76 | ||
|
|
175fa5d244 | ||
|
|
62fab2bc2c | ||
|
|
9b6aa39b94 | ||
|
|
e9463a28fb | ||
|
|
56b6285ffd | ||
|
|
f623294af6 | ||
|
|
15b4fb5698 | ||
|
|
45bc019bfc | ||
|
|
c9eaf988b3 | ||
|
|
23a10c37c7 | ||
|
|
518c8b2b90 | ||
|
|
14db64f7a2 | ||
|
|
91b30040c7 | ||
|
|
a240377ba1 | ||
|
|
36feeb4b9f | ||
|
|
ab60b3cafb | ||
|
|
8673515d31 | ||
|
|
f1d6e9ceb4 | ||
|
|
52c6248e2b | ||
|
|
1e4409f8b1 | ||
|
|
96d9f96e6d | ||
|
|
f606f7ec56 | ||
|
|
02d95a2bd9 | ||
|
|
749fcd4cf7 | ||
|
|
49d9a33420 | ||
|
|
2d4c60f643 | ||
|
|
7af58420cd | ||
|
|
3fbf17d6cd | ||
|
|
1774d7d64b | ||
|
|
d7eb46b3a1 | ||
|
|
82875b3021 | ||
|
|
4d6e1e1fd4 | ||
|
|
6db0cf017e | ||
|
|
bcd63e4fde | ||
|
|
cca0cdff5f | ||
|
|
0a79a6f631 | ||
|
|
1190a15e1e | ||
|
|
86236bae9d | ||
|
|
141e2e1ef0 | ||
|
|
01c699e7a6 | ||
|
|
d5047d27c7 | ||
|
|
c699ff1794 | ||
|
|
c9777cbff8 | ||
|
|
986acab1c5 | ||
|
|
049b894416 | ||
|
|
0aa0710a9e | ||
|
|
6d79fc919f | ||
|
|
fee1241e49 | ||
|
|
21d93fc505 | ||
|
|
a8ffa73a6b | ||
|
|
09dc36d591 | ||
|
|
da8ebbf9e6 | ||
|
|
7467e7bd42 | ||
|
|
7072e33d6f | ||
|
|
30e064996c | ||
|
|
2817dbbab2 | ||
|
|
ad5d91c47a | ||
|
|
77bc367d03 | ||
|
|
75a45313c4 | ||
|
|
dfa8d0406a | ||
|
|
c54e7aca83 | ||
|
|
75897c68eb | ||
|
|
42b4a89ed3 | ||
|
|
c9e8d04815 | ||
|
|
3b6a0fb002 | ||
|
|
3eafc15602 | ||
|
|
80969adc9b | ||
|
|
0954644ee3 | ||
|
|
9d548c064f | ||
|
|
8b6eb47c80 | ||
|
|
936a758651 | ||
|
|
b0d71a4c58 | ||
|
|
f041bf284b | ||
|
|
aa3741b861 | ||
|
|
01ae877673 | ||
|
|
d17677bade | ||
|
|
9fc3542771 | ||
|
|
d0ee948339 | ||
|
|
29eac70f24 | ||
|
|
8eb14cad4a | ||
|
|
4df510712e | ||
|
|
4275f002e0 | ||
|
|
01c6d0cb1f | ||
|
|
8fe2afdfde | ||
|
|
4b7773daec | ||
|
|
01e67648e0 | ||
|
|
22ade06114 | ||
|
|
fc4e1c4a61 | ||
|
|
14971971a1 | ||
|
|
391e6b0066 | ||
|
|
ac177c73df | ||
|
|
08df3929b2 | ||
|
|
bba7ee3eb4 | ||
|
|
85eaa2114e | ||
|
|
26c91c6036 | ||
|
|
cfe30a867c | ||
|
|
6089e7427d | ||
|
|
4dd25486f8 | ||
|
|
2959197597 | ||
|
|
c07ad7ac39 | ||
|
|
685c5c1754 | ||
|
|
373045f0a3 | ||
|
|
dbc79c52f3 | ||
|
|
b634f192ea | ||
|
|
59c5fe1feb | ||
|
|
1b282c6f1a | ||
|
|
c13d11aafb | ||
|
|
407ce3c366 | ||
|
|
bbea1ee9d8 |
No files matched your search
@@ -22,3 +22,9 @@ waf/certs/
|
||||
# 本地 IDE / 工具
|
||||
.qoder/
|
||||
.vscode/
|
||||
|
||||
# 运行时日志(HailoRT 等)
|
||||
hailort.log
|
||||
|
||||
# 录制产物(可复现:mtc DSL → 编译 → pybot --record)
|
||||
actpred/data_v3
|
||||
@@ -0,0 +1,6 @@
|
||||
# 3d — bot4cj 机器人 3D 打印模型
|
||||
|
||||
Dashun 蜘蛛机器人平台的 3D 打印/加工件模型文件(STL / STEP / SLDPRT),与 [bot4cj](../bot4cj/) 硬件项目配套。
|
||||
|
||||
- STL — 3D 打印件(`leg*.stl` / `支臂.stl` / `支臂_mirror.stl` 等)
|
||||
- STEP / SLDPRT — 机械设计源文件(`手臂.SLDPRT` / `爪手.STEP` / `配件*.STEP` 等)
|
||||
@@ -1,8 +1,12 @@
|
||||
# blackbird
|
||||
# micro_servers
|
||||
|
||||
<p align="center">
|
||||
<img src="logo.png" width="360" alt="blackbird · Powered by UniServe Framework">
|
||||
</p>
|
||||
|
||||
一套 DIY 的轻量级微服务架构(UniServe Framework),由 C++、Python、Go 与仓颉(Cangjie) 混合实现。
|
||||
`core` 作为进程托管服务负责拉起并守护各业务子进程,`redirect80` 将 80 端口流量重定向到 443,`waf` 作为前置 WAF 网关终止 TLS 并做规则匹配,`gateway` 作为应用层网关统一鉴权与反向代理,`db` 提供通用数据库访问与集中式日志收集,`confs` 提供统一配置分发,`ai_review` 提供 AI 内容审核,`static` 与 `dnfs` 作为业务服务对外提供能力。`db2` / `dnfs2` 分别为 `db` / `dnfs` 的 C++20 重实现,接口完全兼容可平替。此外包含 aibox 的终端前端 `BotOps`、独立的机器人硬件控制模块 `bot4cj`、嵌入式帧缓冲图形界面 `lci` 和嵌入式脚本引擎 `libtry`。
|
||||
`core` 作为进程托管服务负责拉起并守护各业务子进程,`redirect80` 将 80 端口流量重定向到 443,`waf` 作为前置 WAF 网关终止 TLS 并做规则匹配,`ca` 作为中央认证服务统一鉴权(waf → ca → gateway 前置拓扑),`gateway` 退化为纯转发反向代理,`db` 提供通用数据库访问与集中式日志收集,`confs` 提供统一配置分发,`ai_review` 提供 AI 内容审核,`static` 与 `dnfs` 作为业务服务对外提供能力。`db2` / `dnfs2` 分别为 `db` / `dnfs` 的 C++20 重实现,接口完全兼容可平替。此外包含 aibox 的终端前端 `BotOps`、独立的机器人硬件控制模块 `bot4cj`、嵌入式帧缓冲图形界面 `lci` 和嵌入式脚本引擎 `libtry`。
|
||||
|
||||
## 架构总览
|
||||
|
||||
@@ -10,41 +14,45 @@
|
||||
┌──────────────────────────────────────────────┐
|
||||
│ core (PMC / Nidus) │
|
||||
│ 进程托管服务 · 监听 :5200 (HTTP 管理接口) │
|
||||
│ 读取 init.txt 拉起并守护下列子进程 │
|
||||
│ 读取 services.yaml 拉起并守护下列子进程 │
|
||||
│ 内嵌 confs :8010 (仅 gateway.ini 分发) │
|
||||
└───────────────┬──────────────────────────────┘
|
||||
│ 启动 / 守护
|
||||
┌──────────────────┬───────────┼───────────┬──────────────────┬──────────────┐
|
||||
▼ ▼ ▼ ▼ ▼ ▼
|
||||
┌──────────────┐ ┌──────────────┐ ┌────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐
|
||||
│ waf │ │ gateway │ │ static │ │ dnfs │ │ db │ │ ai_review │
|
||||
│ WAF 网关 │ │ 应用层网关 │ │前端托管│ │ 树状菜单 │ │ 通用数据库 │ │ AI内容审核 │
|
||||
│ waf │ │ ca │ │ fsys │ │ dnfs2 │ │ db2 │ │ ai_review │
|
||||
│ WAF 网关 │ │ 中央认证 │ │前端托管│ │ 树状菜单 │ │ 通用数据库 │ │ AI内容审核 │
|
||||
│ :443 (入口) │ │ :8443 (内部) │ │ :9200 │ │ :8006 │ │ :8005 │ │ :8007 │
|
||||
│ coraza+反代 │ │ 鉴权 + 反代 │ │静态/媒体│ │ 节点CRUD │ │ pymsg/TCP │ │ DeepSeek+RAG │
|
||||
│ coraza+反代 │ │ 认证+前置转发│ │静态/媒体│ │ 节点CRUD │ │ pymsg/TCP │ │ DeepSeek+RAG │
|
||||
└──────┬───────┘ └──────┬───────┘ └───┬────┘ └──────┬───────┘ └──────▲───────┘ └──────▲───────┘
|
||||
│ HTTPS终止 │ 反向代理 │ 代理笔记请求 │ pymsg │ pymsg │ HTTP
|
||||
│ HTTPS终止 │ 认证+转发 │ 代理笔记请求 │ pymsg │ pymsg │ HTTP
|
||||
▲ └─────────────►│ ├─────────────►│─────────────────┘ │
|
||||
│ ├────────────►│ └─────── HTTP /review ──────────────┘
|
||||
│ └────────────►│
|
||||
│ ┌───────────►│ │ └─────── HTTP /review ──────────────┘
|
||||
│ │ gateway │ │
|
||||
│ │ :8445 纯转发│ │
|
||||
│ └────────────►│ │
|
||||
│
|
||||
┌─────┴──────┐
|
||||
│ redirect80 │ :80 → :443(GET/HEAD 301 / 其余 308,重定向到 waf)
|
||||
└────────────┘
|
||||
日志链路: waf / gateway / static / dnfs / ai_review / captcha / smtp / aibox / confs / redirect80 ──(log_client, pymsg)──► db:8005 (Logs 表)
|
||||
配置链路: static 等业务服务 ──(pymsg get)──► confs:8010(配置统一存放于 confs/configs/,mtime 变化自动重载)
|
||||
日志链路: waf / ca / gateway / fsys / dnfs2 / ai_review / captcha2 / smtp / aibox / redirect80 ──(log_client, pymsg)──► db2:8005 (Logs 表)
|
||||
认证链路(v0.9.0 起): waf:443 ──HTTP──► ca:8443 (认证+前置转发) ──HTTP──► gateway:8445 (纯转发, 仅 loopback) ──► 内部服务
|
||||
配置链路: gateway ──(pymsg get)──► core:8010 内嵌 confs(配置统一存放于 confs/configs/,mtime 变化自动重载);ca/dnfs2/db2/fsys/captcha2/smtp/redirect80 例外,配置纯 CLI 传递
|
||||
```
|
||||
|
||||
> `db2` / `dnfs2` 分别是 `db` / `dnfs` 的 C++20 重实现(接口完全兼容,可平替),架构位置与 Python 版一致。`BotOps` 为 aibox 的终端前端(WebSocket 客户端),`lci` 为嵌入式帧缓冲图形界面,均不在上图中。
|
||||
> `db2` / `dnfs2` / `fsys` / `captcha2` 分别是 `db` / `dnfs` / `static` / `captcha` 的 C++20 重实现(接口完全兼容,可平替),架构位置与 Python 版一致;`confs` 已内嵌进 `core`。`BotOps` 为 aibox 的终端前端(WebSocket 客户端),`lci` 为嵌入式帧缓冲图形界面,均不在上图中。
|
||||
|
||||
### 请求链路
|
||||
|
||||
1. 客户端若以 HTTP 访问 `:80`,由 `redirect80` 重定向到 HTTPS `:443`(GET/HEAD 返回 301,其余方法返回 308 保留 body)。
|
||||
2. 客户端所有请求先打到 `waf:443`(HTTPS),WAF 终止 TLS、运行 coraza 规则匹配,命中则直接返回 403。
|
||||
3. 未命中规则的请求由 WAF 以明文 HTTP 转发到 `gateway:8443`,并设置 `X-Forwarded-For` / `X-Forwarded-Proto` 头。
|
||||
4. `gateway` 对 `/login`、`/register` 做账号鉴权(PBKDF2 密码哈希 + RSA 签名 Token + 可选 CAPTCHA + 速率限制;注册需邮箱验证码,经 pymsg/TCP 调 `smtp:8008` 发送与校验),用户数据通过 pymsg/TCP 调用 `db:8005` 存取。
|
||||
5. 其余请求携带 Token,经 `gateway` 校验后反向代理到 `Proxy.opened` 中放开的内部服务(`static:9200` / `dnfs:8006`)。
|
||||
6. `db:8005` 与 `ai_review:8007` 在 `Proxy.closed` 中,仅允许内部服务间调用,不对外暴露。
|
||||
7. `dnfs` 处理节点增删改查时,通过 pymsg/TCP 调用 `db` 服务完成数据持久化,新增节点前调用 `ai_review` 进行内容审核。
|
||||
8. 各服务通过 `log_client.py` SDK 异步批量推送运行日志(pymsg/TCP 写入 db 的 `Logs` 表),实现集中式日志收集(fail-safe,db 不可用时静默丢弃)。
|
||||
3. 未命中规则的请求由 WAF 以明文 HTTP 转发到 `ca:8443`,并设置 `X-Forwarded-For` / `X-Forwarded-Proto` 头。
|
||||
4. `ca` 对 `/login`、`/register` 做账号鉴权(PBKDF2 密码哈希 + RSA 签名 Token + 可选 CAPTCHA + 速率限制;注册需邮箱验证码,经 pymsg/TCP 调 `smtp:8008` 发送与校验),用户数据通过 pymsg/TCP 调用 `db2:8005` 存取。
|
||||
5. 其余请求携带 Token,经 `ca` 验签后注入 `X-Auth-User` 身份头转发到 `gateway:8445`(仅 loopback,不验签仅做存在性检查),再由 gateway 反向代理到 `Proxy.opened` 中放开的内部服务(`fsys:9200` / `dnfs2:8006`)。
|
||||
6. `db2:8005` 与 `ai_review:8007` 在 `Proxy.closed` 中,仅允许内部服务间调用,不对外暴露。
|
||||
7. `dnfs2` 处理节点增删改查时,通过 pymsg/TCP 调用 `db2` 服务完成数据持久化,新增节点前调用 `ai_review` 进行内容审核。
|
||||
8. 各服务通过 `log_client.py` SDK 异步批量推送运行日志(pymsg/TCP 写入 db2 的 `Logs` 表),实现集中式日志收集(fail-safe,db2 不可用时静默丢弃)。
|
||||
|
||||
## 各服务说明
|
||||
|
||||
@@ -52,20 +60,27 @@
|
||||
|------|------|------|------|
|
||||
| [core](core/) | C++20 | 5200 | 进程托管服务(PMC),创建/列举/杀死/执行子进程,读取自启动列表守护业务进程,提供 HTTP 管理接口 |
|
||||
| [redirect80](redirect80/) | Python | 80 | HTTP→HTTPS 重定向服务,将 80 端口请求重定向到 443(GET/HEAD 301 / 其余 308),纯标准库实现 |
|
||||
| [waf](waf/) | Go | 443 | WAF 网关(入口),终止 TLS、coraza 规则匹配(SecLang)、反向代理到 gateway,支持 SIGHUP 热加载规则 |
|
||||
| [gateway](gateway/) | C++20 | 8443 | 应用层网关(内部 HTTP),用户注册/登录鉴权、Token 签发与校验、反向代理到内部服务 |
|
||||
| [static](static/) | Python(Flask) | 9200 | 前端托管服务,社区首页、静态页面与文件浏览、视频/音频流媒体播放、路径遍历防护 |
|
||||
| [dnfs](dnfs/) | Python(Flask) | 8006 | 树状网络文件系统(Dendrachy Network File-System),基于父子 pid 的节点树增删改查,集成 AI 内容审核 |
|
||||
| [dnfs2](dnfs2/) | C++20 | 8006 | dnfs 的 C++20 重实现(libhttp + libmsg),接口与 Python 版完全兼容,修复 P0/P1 安全问题,可平替 |
|
||||
| [db](db/) | Python(pymsg) | 8005 | 通用数据库服务,基于 pymsg 协议(TCP 长度前缀帧 + JSON)的 SQL-over-TCP,SQLAlchemy Core 参数化查询防注入;内置集中式日志收集(Logs 表 + 自动清理) |
|
||||
| [db2](db2/) | C++20 | 8005 | db 的 C++20 重实现(libmsg + boost::json + libsqlite3),execute 路由/where 语法/响应格式完全对齐 Python 版,可平替 |
|
||||
| [waf](waf/) | Go | 443 | WAF 网关(入口),终止 TLS、coraza 规则匹配(SecLang)、反向代理到 ca,支持 SIGHUP 热加载规则 |
|
||||
| [ca](ca/) | C++20 | 8443 | 中央认证服务(内部 HTTP),用户注册/登录鉴权、Token 签发与校验、身份头注入后前置转发到 gateway |
|
||||
| [gateway](gateway/) | C++20 | 8445 | 应用层网关(仅 loopback),纯转发反向代理到内部服务,不验签(v0.9.0 起认证迁至 ca) |
|
||||
| [fsys](fsys/) | C++20 | 9200 | 前端托管服务(static 的 C++ 重写),社区首页、静态页面与文件浏览、视频/音频流媒体播放、路径遍历防护,代理 dnfs2 笔记请求 |
|
||||
| [static](static/) | Python(Flask) | 9200 | ~~前端托管服务~~(遗留,已由 fsys 平替) |
|
||||
| [dnfs2](dnfs2/) | C++20 | 8006 | 树状网络文件系统(Dendrachy Network File-System),基于父子 pid 的节点树增删改查,集成 AI 内容审核;接口与 Python 版完全兼容,修复 P0/P1 安全问题 |
|
||||
| [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,支持人工反馈纠偏 |
|
||||
| [captcha](captcha/) | Python(pymsg) | 8009 | 验证码服务(pymsg/TCP),生成图形验证码(5 分钟过期、一次性),经 gateway 的 `GET /captcha/image` 代理给浏览器 |
|
||||
| [smtp](smtp/) | Python(pymsg) | 8008 | 邮件服务,发送/校验邮箱验证码(pymsg 路由 send_email/verify_email/health),注册流程的邮件验证码后端,SMTP 账号配置在 .env |
|
||||
| [confs](confs/) | Python(pymsg) | 8010 | 统一配置服务,集中存放并分发各服务配置文件(pymsg 路由 get/list/reload/health),内存缓存 + mtime 自动重载,路径遍历防护 |
|
||||
| [captcha2](captcha2/) | C++20 | 8009 | 验证码服务(pymsg/TCP),生成图形验证码(5 分钟过期、一次性),经 ca 的 `GET /captcha/image` 代理给浏览器 |
|
||||
| [captcha](captcha/) | Python(pymsg) | 8009 | ~~验证码服务~~(遗留,已由 captcha2 平替) |
|
||||
| [smtp](smtp/) | Python(pymsg) | 8008 | 邮件服务,发送/校验邮箱验证码(pymsg 路由 send_email/verify_email/health),注册流程的邮件验证码后端,SMTP 账号配置在 .env(smtp3 C++ 重写待上线) |
|
||||
| [confs](confs/) | Python(pymsg) | 8010 | ~~统一配置服务~~(遗留,已内嵌进 core,仅 gateway.ini 仍由 confs 分发) |
|
||||
| [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(核心模块 · 进程托管服务)
|
||||
@@ -73,43 +88,51 @@
|
||||
- 可执行文件 `nidus.out`,作为 systemd 服务(`Nidus.service`)常驻运行(崩溃守护由 systemd 负责,core 自身不含自动重启逻辑)。
|
||||
- 支持两种工作模式:命令行子系统调用(`--send/--exec/--list/--kill`)与 HTTP 伺服器模式(`--serve <port>`)。
|
||||
- 通过命名管道(Pipe)与消息队列(Mq)实现进程间通信。
|
||||
- 读取 `init.txt` 自启动列表,按行拉起各业务服务(ai_review / db / dnfs / smtp / static / waf / gateway / captcha / aibox)。
|
||||
- 读取 `services.yaml` 声明式配置,按声明顺序拉起各业务服务(db2 / ai_review / dnfs2 / smtp / fsys / waf / ca / gateway / captcha2 / aibox / redirect80)。
|
||||
|
||||
### waf(WAF 网关)
|
||||
|
||||
- 唯一对外入口,终止 TLS(`certs/full_chain.crt` / `certs/server.key`),监听 443 端口。
|
||||
- 基于 [coraza](https://github.com/corazawaf/coraza) v3.3.0(ModSecurity SecLang 的 Go 实现),规则匹配命中后返回 403。
|
||||
- 未命中规则的请求以明文 HTTP 反向代理到 `gateway:8443`,透传 `X-Forwarded-For` / `X-Forwarded-Proto`。
|
||||
- 未命中规则的请求以明文 HTTP 反向代理到 `ca:8443`,透传 `X-Forwarded-For` / `X-Forwarded-Proto`。
|
||||
- 规则库位于 [waf/rules/](waf/rules/),分两段:本地自定义(ID 100000-199999)与 OWASP CRS(ID 900000-999999,可选)。
|
||||
- 极客操作:`kill -HUP <pid>` 热加载规则;`waf -list` 列规则;`waf -test <file>` 校验语法;`waf -check "<input>"` 测试单条输入。
|
||||
- 配置见 [waf/config.yaml](waf/config.yaml),详见 [waf/README.md](waf/README.md)。
|
||||
|
||||
### gateway(应用层网关 · v0.8.0)
|
||||
### ca(中央认证服务 · v1.1.0)
|
||||
|
||||
- 内部 HTTP 服务,默认监听 127.0.0.1:8443,不再承载 TLS(TLS 已由 WAF 终止)。
|
||||
- 通过 `Server.tls = on/off` 配置切换 HTTPS/HTTP 模式([example.ini](gateway/example.ini)),独立部署时仍可启用 TLS。
|
||||
- 内部 HTTP 服务,默认监听 127.0.0.1:8443(接管原 gateway 端口,WAF upstream 不变),不再承载 TLS(TLS 已由 WAF 终止)。
|
||||
- **配置体系(v1.1.0)**:彻底移除 confs 依赖,配置全部经 CLI 传递(CLI > 环境变量 > 默认值三层),原 `ca.ini` 配置值迁入 [core/services.yaml](core/services.yaml) 的 ca 条目 args;兼容旧环境变量名(`LISTEN_ADDR`/`UPSTREAM_GATEWAY`/`DB_*` 等)。
|
||||
- **密码安全**:PBKDF2-HMAC-SHA256(10 万轮 + 16 字节随机盐),密码强度校验(8+ 字符、4 类中至少 3 类),常数时间比较防 timing attack。
|
||||
- **用户存储**:用户数据通过 pymsg/TCP 调用 `db:8005`(`DbC::call_msg`)存取,不再直连本地 SQLite;登录成功即更新 `Users.last_login`(`"$now"` 哨兵)。
|
||||
- **Token 机制**:`Base64(payload) + "_" + RSA-Sign(payload)`,默认 1 小时过期,验签通过才放行;RSA 密钥对启动时内存生成,重启后旧 Token 全部失效。
|
||||
- **Cookie 认证**:登录成功同时签发 `HttpOnly + Secure + SameSite=Lax` 的 `q4_token` Cookie;代理时无 `Authorization` 头则回退读 Cookie 验签,并向下游注入 `Authorization: Bearer`(下游零改动);`POST /logout` 清除 Cookie。公开目标(`Proxy.default`)不做凭证校验,避免旧 Cookie 影响公开页面。
|
||||
- **CAPTCHA**:`/captcha` 接口生成验证码(5 分钟过期、一次性使用),登录与注册均可选启用;`GET /captcha/image` 为 captcha 服务(pymsg/TCP :8009)的图形验证码代理路由。
|
||||
- **速率限制**:登录 5 次/5 分钟、注册 3 次/10 分钟,按 `X-Forwarded-For`(WAF 注入)客户端 IP 计数。
|
||||
- **反向代理**:`Target-Url` 必须在 `Proxy.opened` 白名单中(`closed` 仅记录用途,实际隔离靠"不在 opened 中即拒绝"),缺省走 `Proxy.default`(不需 Token,服务前端静态资源)。
|
||||
- **用户表**:`Users` 表由 [db/create_sqlite3.sql](db/create_sqlite3.sql) 建立(gateway 不再自动建表),无默认管理员账号,首位注册者即为管理员。
|
||||
- **用户存储**:用户数据通过 pymsg/TCP 调用 `db2:8005`(`DbC::call_msg`)存取;登录成功即更新 `Users.last_login`(`"$now"` 哨兵)。
|
||||
- **Token 机制**:`Base64(payload) + "_" + RSA-Sign(payload)`,默认 1 小时过期;RSA 密钥对启动时内存生成,重启后旧 Token 全部失效;jti 黑名单 + 用户版本号双轨吊销。
|
||||
- **Cookie 认证**:登录成功同时签发 `HttpOnly + Secure + SameSite=Lax` 的 `q4_token` Cookie;验签时无 `Authorization` 头则回退读 Cookie,并向下游注入 `Authorization: Bearer`(下游零改动);`POST /logout` 清除 Cookie。
|
||||
- **前置转发**:无 `Target-Url` 为公开请求(剥身份头直转 gateway);有 `Target-Url` 为受保护请求(验签后注入 `X-Auth-User`/`X-Auth-Admin` 转发,strip-then-set 契约剥离客户端伪造值);上游不可达统一 502。
|
||||
- **CAPTCHA**:`GET /captcha/image` 为 captcha 服务(pymsg/TCP :8009)的图形验证码代理路由;`/login` 强制 CAPTCHA。
|
||||
- **速率限制**:登录 5 次/5 分钟、注册 3 次/10 分钟、发码 IP 级 + 邮箱级双窗口,按 `X-Forwarded-For`(WAF 注入)客户端 IP 计数。
|
||||
- **用户表**:`Users` 表由 [db2/scripts/create_sqlite3.sql](db2/scripts/create_sqlite3.sql) 建立,无默认管理员账号,首位注册者即为管理员。
|
||||
- 详见 [ca/README.md](ca/README.md)。
|
||||
|
||||
### gateway(应用层网关 · v0.9.0)
|
||||
|
||||
- 纯转发服务,仅监听 127.0.0.1:8445(loopback 启动强制自检),不再验签(认证已迁至 ca)。
|
||||
- **存在性不变量**:显式 `Target-Url` 请求必须携带 ca 注入的 `X-Auth-User` 头,缺失一律 401。
|
||||
- **反向代理**:`Target-Url` 必须在 `Proxy.opened` 白名单中,缺省走 `Proxy.default`;`X-Auth-User`/`X-Auth-Admin`/`Authorization` 原样透传下游。
|
||||
- 详见 [gateway/README.md](gateway/README.md)。
|
||||
|
||||
### static(前端托管服务)
|
||||
### fsys(前端托管服务 · C++20 重实现)
|
||||
|
||||
- 托管 `public/` 下的前端页面:社区首页(`index.html`)、登录/注册、树状笔记(`dnfs.html`)、文件浏览(`__view.html`)等。
|
||||
- static 的 C++20 重实现(原 Python/Flask `static/` 目录保留为遗留,不再部署),托管 `public/` 下的前端页面:社区首页(`index.html`)、登录/注册、树状笔记(`dnfs.html`)、文件浏览(`__view.html`)等。
|
||||
- 支持文件浏览、视频/音频在线播放(HTTP Range 分段请求,返回 206/416)、一键下载。
|
||||
- 路径遍历防护(`safe_path_join`:realpath 前缀校验),仅允许访问 `directory.root` 之内。
|
||||
- **dnfs 代理**:`/get_notes`、`/add_note`、`/delete_note`、`/update_note` 透传 Authorization 转发到 `dnfs:8006`(解决浏览器跨域)。
|
||||
- **账号旁路**:`/api/check_user`、`/api/reset_password` 等接口直接读写 `gateway/sqlite3.db` 的 Users 表;`/yhdet_callback`、`/api/yhdet_*` 对接 yhdet.top OAuth 登录。
|
||||
- 统一品牌"树状菜单",支持深色模式(`css/theme.css`)。
|
||||
- 可使用 Nuitka 编译为独立二进制(`fsys.bin`)。
|
||||
- 配置由 `confs` 统一配置服务分发(`confs/configs/static.ini`,启动时经 pymsg 拉取,`CONFS_HOST`/`CONFS_PORT` 环境变量可覆盖);本地 [static/config.ini](static/config.ini) 作为 confs 不可用时的 fail-safe 回退。
|
||||
- 路径遍历防护(realpath 前缀校验),仅允许访问 `--root` 之内。
|
||||
- **dnfs2 代理**:`/get_notes`、`/add_note`、`/delete_note`、`/update_note` 透传 Authorization 转发到 `dnfs2:8006`(解决浏览器跨域)。
|
||||
- **账号旁路已收口**:原 static 时代的 `/api/check_user`、`/api/reset_password`、yhdet OAuth 等旁路接口已移除/迁至 ca(忘记密码走 ca `/reset/send_code` + `/reset_password`),fsys 不再直接读写用户表。
|
||||
- 统一品牌"树状菜单",支持深色模式(`css/theme.css`)与样式节点注入的自定义主题/样式(见 [docs/dnfs_custom_themes.md](docs/dnfs_custom_themes.md))。
|
||||
- 配置全部经 CLI 传递(CLI > 环境变量 > 默认值,原 `confs/configs/static.ini` 迁入 `core/services.yaml` 的 fsys 条目 args),不依赖 confs。
|
||||
|
||||
### dnfs(树状菜单服务)
|
||||
### dnfs(树状菜单服务 · 遗留)
|
||||
|
||||
> 已由 [dnfs2](#dnfs2树状菜单服务--c20-重实现)(C++20 重实现)平替,不再部署;以下为 Python 版说明。
|
||||
|
||||
- 基于 `pid`(父节点 id)组织的树状节点,节点含 `title` 与 `content`。
|
||||
- 提供 `/get_notes`、`/add_note`、`/delete_note`、`/update_note` 四个接口。
|
||||
@@ -120,13 +143,16 @@
|
||||
|
||||
- dnfs 的 C++20 重实现,基于 [libhttp](libhttp/) + [libmsg](libmsg/) + boost::json,接口与 Python dnfs 完全兼容,前端零改动。
|
||||
- 修复 Python 版的 P0/P1 安全问题(Token 验签 fail-closed、所有权校验、子节点检查、错误响应脱敏避免 IP/用户名泄露)。
|
||||
- Token 验签:HttpClient POST gateway `/whoami`,验签失败返回 503(fail-closed)。
|
||||
- Token 验签:HttpClient POST ca `/whoami`(CLI 参数 `--gateway-url` 指向 ca:8443),验签失败返回 503(fail-closed)。
|
||||
- 数据操作:libmsg `RpcClient` 调 db2 `execute` 路由(与 Python 版调 db 路径一致)。
|
||||
- 内容审核:HttpClient POST ai_review `/review`,fail-closed。
|
||||
- 异步日志:libmsg `AsyncQueue` 写入 db2 Logs 表。
|
||||
- 配置从 confs 拉取 `dnfs.ini`,环境变量覆盖,命令行参数最高优先级;详见 [dnfs2/README.md](dnfs2/README.md)。
|
||||
- 特殊节点:`js` / `file` / `style` / `theme` 四类,样式与主题节点支持界面注入换肤(节点作用域生效,详见 [docs/dnfs_custom_themes.md](docs/dnfs_custom_themes.md))。
|
||||
- 配置通过 CLI 参数 + 环境变量 + 默认值,不依赖 confs;详见 [dnfs2/README.md](dnfs2/README.md)。
|
||||
|
||||
### db(通用数据库服务)
|
||||
### db(通用数据库服务 · 遗留)
|
||||
|
||||
> 已由 [db2](#db2通用数据库服务--c20-重实现)(C++20 重实现)平替,不再部署;以下为 Python 版说明。
|
||||
|
||||
- 基于自研 [pymsg](pymsg/) 协议(TCP 长度前缀帧 + JSON),单一路由 `_route: "execute"`,接收结构化 JSON(action + table + where + values),翻译为 SQL 执行。
|
||||
- 支持 select / insert / update / delete,运算符包括 eq、ne、gt、lt、ge、le、like、in,支持排序与分页。
|
||||
@@ -139,7 +165,7 @@
|
||||
- 外部服务通过 `log_client.py` 以 pymsg 调 execute 路由逐条 INSERT 写入日志。
|
||||
- `_log_internal()` — db 进程内部直接写 Logs 表,无网络往返,防递归设计。
|
||||
- 后台守护线程每小时自动清理:保留 30 天 / 上限 100 万条。
|
||||
- **统一日志客户端 SDK**(各服务目录下的 `log_client.py`,如 [dnfs/log_client.py](dnfs/log_client.py)):各服务复制即用,纯标准库实现,异步非阻塞(内存队列 + 后台线程,50 条/2 秒批量 flush),fail-safe 设计(db 不可用时静默丢弃,不影响业务)。
|
||||
- **统一日志客户端 SDK**(各服务目录下的 `log_client.py`,如 [smtp/log_client.py](smtp/log_client.py)):各服务复制即用,纯标准库实现,异步非阻塞(内存队列 + 后台线程,50 条/2 秒批量 flush),fail-safe 设计(db2 不可用时静默丢弃,不影响业务)。
|
||||
|
||||
### db2(通用数据库服务 · C++20 重实现)
|
||||
|
||||
@@ -147,7 +173,13 @@
|
||||
- per-thread 连接:每个 worker 线程持有独立 `sqlite3*`(`thread_local`),列 schema 加锁缓存;SQLite PRAGMA 对齐 Python 版(WAL / NORMAL / busy_timeout=5000)。
|
||||
- 内置 Logs 表(启动自动建表 + 索引)与后台清理线程(保留 30 天 / 上限 100 万条),设计与 Python db 一致。
|
||||
- 安全机制:表名白名单(字母数字下划线)、`sqlite3_bind_text` 参数绑定、仅反射已有表无法建删表、RpcServer 异常脱敏。
|
||||
- 配置从 confs 拉取 `db.ini`(强依赖,confs 未就绪自动重试),环境变量(`LISTEN_ADDR`/`LISTEN_PORT`/`WORKERS`/`DATABASE_URL`)覆盖,命令行参数最高优先级;详见 [db2/README.md](db2/README.md)。
|
||||
- 配置全部经 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 内容审核服务)
|
||||
|
||||
@@ -158,14 +190,14 @@
|
||||
- **fail-closed 设计**:API 异常或响应解析失败时一律返回 `{"invalid": "true"}`(拦截),宁可误拦不可漏放。
|
||||
- 配置通过 `.env` 文件(已 gitignore),含 DeepSeek API 密钥、模型参数、监听地址。
|
||||
|
||||
### captcha(验证码服务)
|
||||
### captcha2(验证码服务 · C++20 重实现)
|
||||
|
||||
- 生成图形验证码,Token 一次性使用、5 分钟过期;基于 [pymsg](pymsg/) 协议(TCP 长度前缀帧 + JSON),仅依赖 PIL。
|
||||
- captcha 的 C++20 重实现(原 Python `captcha/` 目录保留为遗留,不再部署),生成图形验证码,Token 一次性使用、5 分钟过期;基于 [pymsg](pymsg/) 协议(TCP 长度前缀帧 + JSON)。
|
||||
- pymsg 路由:`image`(生成,返回 token + base64 PNG)、`verify`(校验,一次性消费)、`health`。
|
||||
- 答案比较采用 XOR 防时序攻击;错误答案与无效 Token 分别返回 `{valid:false}` 与 `{valid:false, error:...}`。
|
||||
- **gateway 代理路由**:`GET /captcha/image` 经 pymsg 调 captcha 后解码 PNG,以 `image/png` + `X-Captcha-Token` 响应头返回浏览器(`Cache-Control: no-store`)。
|
||||
- 由 `gateway` 在 `/login`、`/register` 中可选启用 CAPTCHA 校验(字段 `captcha_token` / `captcha_answer`)。
|
||||
- 内存存储,不支持多实例水平扩展(需换 Redis);详见 [captcha/README.md](captcha/README.md)。
|
||||
- **ca 代理路由**:`GET /captcha/image` 经 pymsg 调 captcha2 后解码 PNG,以 `image/png` + `X-Captcha-Token` 响应头返回浏览器(`Cache-Control: no-store`)。
|
||||
- 由 `ca` 在 `/login`、`/register` 中可选启用 CAPTCHA 校验(字段 `captcha_token` / `captcha_answer`)。
|
||||
- 内存存储,不支持多实例水平扩展(需换 Redis);详见 [captcha2/README.md](captcha2/README.md)。
|
||||
|
||||
### smtp(邮件服务)
|
||||
|
||||
@@ -179,22 +211,32 @@
|
||||
- 零 Web 框架依赖:纯 Python 标准库 `http.server` + `ThreadingMixIn` 多线程并发。
|
||||
- 智能重定向:GET/HEAD → 301(搜索引擎友好永久重定向),POST/PUT/DELETE/PATCH → 308(保留方法与 body)。
|
||||
- Host 保留:读取请求 `Host` 头构造重定向 URL(剥离端口),回退用配置的监听地址。
|
||||
- 日志集成:通过 `log_client.py` 异步批量推送运行日志到 db 的 Logs 表(fail-safe)。
|
||||
- 配置分层:confs 强依赖(`confs/configs/redirect80.ini`,启动时 pymsg 阻塞等待拉取,confs 不可用时直接退出不做本地回退),环境变量(`REDIRECT80_*`)优先级最高。
|
||||
- 日志集成:通过 `log_client.py` 异步批量推送运行日志到 db2 的 Logs 表(fail-safe)。
|
||||
- 配置分层:CLI 参数(由 core/services.yaml 传入) > 环境变量(`REDIRECT80_*`) > 代码默认值,不依赖 confs。
|
||||
- 80 端口为 Linux 特权端口(<1024),需 `CAP_NET_BIND_SERVICE` 或 root;`_run.sh` 自动尝试 `setcap`,或用 systemd `AmbientCapabilities`;详见 [redirect80/README.md](redirect80/README.md)。
|
||||
|
||||
### confs(统一配置服务)
|
||||
### confs(统一配置服务 · 已内嵌进 core)
|
||||
|
||||
- 集中存放各微服务的配置文件(`confs/configs/`),基于 pymsg(RpcServer,TCP 长度前缀帧 + JSON)分发,监听 127.0.0.1:8010。
|
||||
- 独立 Python `confs` 服务已废弃(不再部署),配置分发逻辑已内嵌进 `core`(`ConfsTask`,由 `--confs 8010` 启用,协议与 Python 版字节级兼容,消费方零改动);C++20 重写版 [confs2/](confs2/) 目录保留为开发存档。
|
||||
- 集中存放各服务的配置文件(`confs/configs/`),基于 pymsg(RpcServer,TCP 长度前缀帧 + JSON)分发。
|
||||
- pymsg 路由:`get`(取配置文件全文;带 `key` 时按键提取,ini 用 `Section.option`、json 用点号路径)、`list`(列出全部配置及元信息)、`reload`(重载指定配置,缺省清空全部缓存)、`health`。
|
||||
- 内存缓存 + 文件 mtime 变化自动重载,改配置无需重启 confs;配置名校验(basename + realpath 前缀 + 扩展名白名单)防路径遍历。
|
||||
- 消费方约定:服务启动时用 `RpcClient(8010).call("get", {"name": "<服务名>.ini"})` 拉取配置,confs 不可用时回退本地配置文件(fail-safe)。目前 `static` 已接入。
|
||||
- 内存缓存 + 文件 mtime 变化自动重载,改配置无需重启;配置名校验(basename + realpath 前缀 + 扩展名白名单)防路径遍历。
|
||||
- 消费方约定:服务启动时用 `RpcClient(8010).call("get", {"name": "<服务名>.ini"})` 拉取配置。各模块目录不再保留本地 ini(fail-safe 回退已废除)。
|
||||
|
||||
**配置文件清单与依赖强度**(均位于 `confs/configs/`):
|
||||
|
||||
| 配置文件 | 消费模块 | 依赖强度 |
|
||||
|----------|----------|----------|
|
||||
| `gateway.ini` | gateway | confs 拉取(0.7.0 起配置完全迁入;其余服务已全部改为 CLI 传递,不再依赖 confs) |
|
||||
|
||||
- **部署安全约定**:`confs/configs/*.ini` 含敏感字段,权限必须收紧为 `root:root 0640`,仅 root 可写;被同机普通用户篡改可劫持对应服务(详见 `docs/security/core_SECURITY_AUDIT_20260801.md` 2026-08-10 复审)。原 `ca.ini` 管理员白名单已随 ca v1.1.0 迁入 `core/services.yaml` 的 CLI args;原 `smtp.ini` 全部配置已随 smtp3 迁入 `core/services.yaml` 的 CLI args(SMTP 凭据仍经 `.env` 注入);原 `captcha.ini` 与 `static.ini` 已分别随 captcha2、fsys 迁入 CLI args,防护对象相应转为 `services.yaml` 与 `.env` 的文件权限。
|
||||
|
||||
### aibox(独立 AI 智能体系统)
|
||||
|
||||
- 一套相对独立的单智能体产品,基于 DeepSeek + WebSocket(:8765),由 `core` 拉起但与 waf / gateway / dnfs 等主干微服务无调用关系。
|
||||
- 一套相对独立的单智能体产品,基于 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 前端)
|
||||
|
||||
@@ -212,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` 操作显存,专为树莓派等嵌入式环境设计。
|
||||
@@ -233,14 +290,15 @@
|
||||
| [libcrypt](libcrypt/) | 加密库:`AES_CBC`、`MyRSA`、`SHA256`、`Base64`、`Crypto_Basic` |
|
||||
| [libsqlite3](libsqlite3/) | SQLite3 C++ 封装(Header-only) |
|
||||
| [libmsg](libmsg/) | 消息通信库:pymsg 协议 C++ 实现(4B 大端长度前缀 + UTF-8 载荷),提供 `MsgClient`/`MsgServer`/`AsyncQueue`/`RpcClient`/`RpcServer`,与 Python pymsg 字节级兼容 |
|
||||
| [liblogc](liblogc/) | 统一集中式日志客户端(C++ 版,`us/LogC.hpp`),对应 Python `log_client.py`:pymsg/TCP 异步批量写入 db2 Logs 表,队列上限 1 万条 + 级别过滤 + fail-safe,ca/core/gateway/captcha2/smtp3 等 C++ 服务共用 |
|
||||
| [libtry](libtry/) | 嵌入式脚本解释器(~741 行 C),手写递归下降词法/语法/执行引擎,支持变量/数组/函数/控制流,宽字符(中文标识符),用于嵌入式操作界面 shell |
|
||||
|
||||
### Python 共享组件
|
||||
|
||||
| 文件 | 说明 |
|
||||
|------|------|
|
||||
| [pymsg](pymsg/) | 自研 RPC 基础库(setuptools 包),协议为 `[4字节大端长度][UTF-8 JSON]`,提供 `send/recv/MsgClient/MsgServer`;db 服务端与各 log_client 的基础依赖,requirements 以 `../pymsg` 本地路径引用 |
|
||||
| log_client.py | 统一日志客户端 SDK,各服务目录各持一份副本(如 [dnfs/log_client.py](dnfs/log_client.py)),纯标准库 + pymsg,异步批量推送 |
|
||||
| [pymsg](pymsg/) | 自研 RPC 基础库(setuptools 包),协议为 `[4字节大端长度][UTF-8 JSON]`,提供 `send/recv/MsgClient/MsgServer`;db/db2 服务端与各 log_client 的基础依赖,requirements 以 `../pymsg` 本地路径引用 |
|
||||
| log_client.py | 统一日志客户端 SDK,各服务目录各持一份副本(如 [smtp/log_client.py](smtp/log_client.py)),纯标准库 + pymsg,异步批量推送 |
|
||||
| security_middleware.py | Flask 安全响应头中间件(CSP / X-Frame-Options / nosniff 等),各服务各持一份副本 |
|
||||
|
||||
## 端口一览
|
||||
@@ -249,33 +307,40 @@
|
||||
|------|------|----------|------|
|
||||
| 80 | redirect80 | 0.0.0.0 | HTTP 入口,301/308 重定向到 HTTPS :443 |
|
||||
| 443 | waf | 0.0.0.0 | HTTPS 对外入口(终止 TLS + coraza 规则匹配) |
|
||||
| 8443 | gateway | 127.0.0.1 | 应用层网关(内部 HTTP,由 waf 转发) |
|
||||
| 8443 | ca | 127.0.0.1 | 中央认证服务(内部 HTTP,由 waf 转发,认证+前置转发) |
|
||||
| 8445 | gateway | 127.0.0.1 | 应用层网关(内部 HTTP,仅 loopback 纯转发) |
|
||||
| 5200 | core | — | PMC HTTP 管理接口 |
|
||||
| 9200 | static | 127.0.0.1 | 前端托管(经 gateway 代理) |
|
||||
| 8006 | dnfs | 127.0.0.1 | 树状菜单(经 gateway 代理) |
|
||||
| 8005 | db | 127.0.0.1 | 数据库服务(pymsg/TCP,仅内部调用) |
|
||||
| 8010 | core(内嵌 confs) | 127.0.0.1 | 统一配置服务(pymsg/TCP,仅 gateway.ini 仍由其分发) |
|
||||
| 9200 | fsys | 127.0.0.1 | 前端托管(经 gateway 代理) |
|
||||
| 8006 | dnfs2 | 127.0.0.1 | 树状菜单(经 gateway 代理) |
|
||||
| 8005 | db2 | 127.0.0.1 | 数据库服务(pymsg/TCP,仅内部调用) |
|
||||
| 8007 | ai_review | 127.0.0.1 | AI 审核(仅内部调用) |
|
||||
| 8009 | captcha | 127.0.0.1 | 验证码服务(pymsg/TCP,经 gateway 代理) |
|
||||
| 8009 | captcha2 | 127.0.0.1 | 验证码服务(pymsg/TCP,经 ca 代理) |
|
||||
| 8008 | smtp | 127.0.0.1 | 邮件服务(仅内部调用,注册验证码邮件后端) |
|
||||
| 8010 | confs | 127.0.0.1 | 统一配置服务(pymsg/TCP,仅内部调用) |
|
||||
| 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) |
|
||||
|
||||
## 测试
|
||||
|
||||
[tests/](tests/) 包含一套黑盒测试套件,基于子进程调用与 HTTP 请求验证各微服务功能正确性:
|
||||
[tests/blackbox/](tests/blackbox/) 包含一套黑盒用户模拟测试套件,基于 HTTP/TCP 公开接口验证各微服务功能正确性(纯 Python 标准库,零第三方依赖):
|
||||
|
||||
| 测试文件 | 被测模块 | 说明 |
|
||||
|----------|----------|------|
|
||||
| `test_db_*.py` | db | SQL-over-TCP 的 select/insert/update/delete、日期哨兵、边界条件 |
|
||||
| `test_gateway_*.py` | gateway | 用户注册/登录/Token 校验、速率限制、CAPTCHA 集成 |
|
||||
| `test_waf_*.py` | waf | coraza 规则匹配与拦截、白名单放行 |
|
||||
| `test_dnfs_*.py` | dnfs | 节点树 CRUD、AI 审核集成 |
|
||||
| `test_db.py` | db2 | SQL-over-TCP 的日志写入/查询/统计/清理 |
|
||||
| `test_gateway.py` | ca(认证) | 用户注册/登录/Token 校验、输入校验、代理路径安全、限流 |
|
||||
| `test_waf.py` | waf | coraza 规则匹配与拦截(SQLi/XSS/命令注入/路径遍历/扫描器)、正常请求放行 |
|
||||
| `test_dnfs.py` | dnfs2 | 节点树 CRUD、AI 审核集成 |
|
||||
| `test_captcha.py` / `test_smtp.py` / `test_ai_review.py` / `test_static.py` | captcha2 / smtp / ai_review / fsys | 验证码一次性消费、邮件验证码、内容审核纠偏、文件浏览与视频流 |
|
||||
| `test_core_confs.py` | core(内嵌 confs) | confs get/list/reload/health 与 Python 版基线行为对比 |
|
||||
|
||||
- 默认使用 `config.py` 中的生产地址(127.0.0.1),每个测试函数前后自动备份/恢复 db 数据,保证不污染线上数据。
|
||||
- 测试依赖 `pymsg`、`Flask`、`flask-cors`、`requests`、`Pillow`,详见 [tests/requirements-test.txt](tests/requirements-test.txt)。
|
||||
- 默认使用 `tests/blackbox/config.py` 中的生产地址(127.0.0.1,条目名映射见该文件),运行方式见 [tests/blackbox/README.md](tests/blackbox/README.md)。
|
||||
- 注册测试会创建真实数据库记录,测试完成后可手动清理。
|
||||
|
||||
## 依赖
|
||||
|
||||
@@ -293,7 +358,7 @@ bash _setup.deb.sh # Debian/Ubuntu
|
||||
# 或
|
||||
bash _setup.rhel.sh # RHEL/CentOS/Rocky/Fedora
|
||||
|
||||
# 2. 依次安装各模块(libth -> libhttp -> libcrypt -> libsqlite3 -> libmsg -> gateway -> waf -> static -> dnfs -> db -> smtp -> captcha -> confs -> ai_review -> aibox -> redirect80 -> core)
|
||||
# 2. 依次安装各模块(libth -> libhttp -> libcrypt -> libsqlite3 -> libmsg -> liblogc -> gateway -> waf -> ca -> static -> fsys -> dnfs -> dnfs2 -> db -> db2 -> smtp -> smtp3 -> captcha -> captcha2 -> confs -> ai_review -> aibox -> redirect80 -> core;legacy Python 版与 C++ 重写版均构建,生产由 core 按 services.yaml 拉起新版)
|
||||
bash _inst.sh
|
||||
```
|
||||
|
||||
@@ -304,7 +369,7 @@ bash _set.sh # 构建 + 安装
|
||||
bash _run.sh # 运行
|
||||
```
|
||||
|
||||
部署路径约定为 `/opt/micro_servers/`,`core` 的 `init.txt` 即从该路径拉起其余服务。
|
||||
部署路径约定为 `/opt/micro_servers/`,`core` 的 `services.yaml` 即从该路径拉起其余服务。
|
||||
|
||||
> **注意**:`ai_review` / `aibox` / `redirect80` 已纳入 `_inst.sh`(分别依赖 chromadb ONNX 模型下载、`~/.cert/deepseek.key`、本地 `pymsg` 包)。`ai_review` 额外需配置 `.env` 文件(DeepSeek API 密钥等,参考 [ai_review/README.md](ai_review/README.md))。
|
||||
|
||||
@@ -312,20 +377,16 @@ bash _run.sh # 运行
|
||||
|
||||
## 运行
|
||||
|
||||
由 `core` 守护的子进程通过 `init.txt` 自启动(顺序很重要,redirect80 先于 waf、waf 先于 gateway):
|
||||
由 `core` 守护的子进程通过 `services.yaml` 声明式配置自启动(顺序很重要,db2 先于 ai_review、waf 先于 redirect80)。配置详见 [core/services.yaml](core/services.yaml),摘录如下:
|
||||
|
||||
```
|
||||
bash /opt/micro_servers/db/_run.sh
|
||||
bash /opt/micro_servers/confs/_run.sh
|
||||
bash /opt/micro_servers/ai_review/_run.sh
|
||||
bash /opt/micro_servers/dnfs/_run.sh
|
||||
bash /opt/micro_servers/smtp/_run.sh
|
||||
bash /opt/micro_servers/static/_run.sh
|
||||
bash /opt/micro_servers/redirect80/_run.sh
|
||||
bash /opt/micro_servers/waf/_run.sh
|
||||
bash /opt/micro_servers/gateway/_run.sh
|
||||
bash /opt/micro_servers/captcha/_run.sh
|
||||
bash /opt/micro_servers/aibox/_run.sh
|
||||
```yaml
|
||||
services:
|
||||
- name: db2
|
||||
cmd: bash /opt/micro_servers/db2/_run.sh
|
||||
args: ["--instance", "db2", "--log-level", "info"]
|
||||
depends_on: []
|
||||
sleep_before: 0
|
||||
# ... 其余 10 个服务(ai_review / dnfs2 / smtp / fsys / waf / ca / gateway / captcha2 / aibox / redirect80)
|
||||
```
|
||||
|
||||
启动 `core`(需 root 或 systemd):
|
||||
@@ -334,9 +395,9 @@ bash /opt/micro_servers/aibox/_run.sh
|
||||
systemctl start Nidus # 或 bash core/_run.sh
|
||||
```
|
||||
|
||||
对外访问入口:`https://<host>:443`(由 waf 终止 TLS 后转发到 gateway)。
|
||||
对外访问入口:`https://<host>:443`(由 waf 终止 TLS 后转发到 ca)。
|
||||
|
||||
### WAF 与 gateway 的协作
|
||||
### WAF 与 ca 的协作
|
||||
|
||||
WAF 启用后,证书从 gateway 迁移到 waf:
|
||||
|
||||
@@ -345,28 +406,28 @@ mkdir -p /opt/micro_servers/waf/certs
|
||||
cp /opt/micro_servers/gateway/full_chain.crt /opt/micro_servers/gateway/server.key /opt/micro_servers/waf/certs/
|
||||
```
|
||||
|
||||
gateway 改为内部 HTTP 模式([example.ini](gateway/example.ini) 中 `tls = off`、`address = 127.0.0.1`、`port = 8443`)。
|
||||
ca 改为内部 HTTP 模式(`core/services.yaml` 中 ca 条目 args:`--host 127.0.0.1`、`--port 8443`)。
|
||||
|
||||
如需独立部署(不使用 WAF),把 gateway 的 `tls = on`、`port = 443`、`address = 0.0.0.0`,证书留在 gateway 目录即可。
|
||||
如需独立部署(不使用 WAF),把 waf 从 `services.yaml` 移除,ca 的 `--port` 设为 443、证书留在 ca 目录即可。
|
||||
|
||||
## 关于架构定位
|
||||
|
||||
本项目具备微服务架构的核心特征,是一种**进程托管式轻量微服务**:
|
||||
|
||||
- **服务按业务能力拆分**:WAF 网关、应用网关、前端托管、节点树、数据库、内容审核各自独立、独立端口、独立数据存储。
|
||||
- **API 网关模式**:`waf` 作为边缘网关统一入口,`gateway` 承担鉴权与反向代理,内部服务通过 `closed` 列表隔离。
|
||||
- **API 网关模式**:`waf` 作为边缘网关统一入口,`ca` 承担应用层鉴权与前置转发,`gateway` 纯转发,内部服务通过 `closed` 列表隔离。
|
||||
- **进程编排/守护**:`core` 作为服务管理器拉起并守护子进程,类似简易版 supervisor。
|
||||
- **多语言异构**:C++(core / gateway / db2 / dnfs2)、Python(static / dnfs / db / ai_review / redirect80 / BotOps)、Go(waf)、仓颉/Cangjie(bot4cj)混合,对外 HTTP、内部数据访问走 pymsg/TCP。
|
||||
- **数据访问层解耦**:`db` 服务将数据库操作收敛为统一的 pymsg/TCP 接口,业务服务无需直连数据库。
|
||||
- **多语言异构**:C++(core / gateway / ca / db2 / dnfs2 / fsys / captcha2)、Python(ai_review / smtp / redirect80 / aibox / BotOps,及 static / dnfs / db / captcha / confs 等遗留)、Go(waf)、仓颉/Cangjie(bot4cj)混合,对外 HTTP、内部数据访问走 pymsg/TCP。
|
||||
- **数据访问层解耦**:`db2` 服务将数据库操作收敛为统一的 pymsg/TCP 接口,业务服务无需直连数据库。
|
||||
- **AI 能力集成**:`ai_review` 以 RAG 增强的 LLM 审核为业务层提供内容安全能力,支持人工反馈闭环。
|
||||
- **集中式可观测性**:`db` 内置日志表与自动清理,各服务通过 `log_client.py` SDK 异步推送日志(pymsg),实现全链路汇聚,无需引入额外中间件。
|
||||
- **安全分层**:`waf` 在边缘做规则匹配(OWASP CRS / 自定义 SecLang 规则),`gateway` 在应用层做鉴权,`db` 仅监听本机,多层防护。
|
||||
- **集中式可观测性**:`db2` 内置日志表与自动清理,各服务通过 `log_client.py` SDK 异步推送日志(pymsg),实现全链路汇聚,无需引入额外中间件。
|
||||
- **安全分层**:`waf` 在边缘做规则匹配(OWASP CRS / 自定义 SecLang 规则),`ca` 在应用层做鉴权,`db2` 仅监听本机,多层防护。
|
||||
|
||||
与云原生微服务的差异:未引入服务注册发现、容器编排、消息总线,服务地址在 `waf/config.yaml` 与 `gateway/example.ini` 中静态配置。对于个人/中小规模的 DIY 场景,这种取舍在简洁性与可控性上是合理的。
|
||||
与云原生微服务的差异:未引入服务注册发现、容器编排、消息总线,服务地址在 `waf/config.yaml` 与 `confs/configs/gateway.ini` 中静态配置。对于个人/中小规模的 DIY 场景,这种取舍在简洁性与可控性上是合理的。
|
||||
|
||||
## 文档管理约定
|
||||
|
||||
所有项目文档统一通过根目录 [docs/](docs/) 聚合管理,入口索引见 [docs/INDEX.md](docs/INDEX.md)。文档分类与存放规则如下:
|
||||
所有项目文档统一通过根目录 [docs/](docs/) 聚合管理。文档分类与存放规则如下:
|
||||
|
||||
| 文档类型 | 存放位置 | 说明 |
|
||||
|----------|----------|------|
|
||||
|
||||
@@ -30,17 +30,17 @@ inst_module "libhttp"
|
||||
inst_module "libcrypt"
|
||||
inst_module "libsqlite3"
|
||||
inst_module "libmsg"
|
||||
inst_module "liblogc"
|
||||
inst_module "gateway"
|
||||
inst_module "waf" # waf 在 gateway 之前编译,部署时由 core 先拉起 waf 再拉起 gateway
|
||||
inst_module "static"
|
||||
inst_module "ca"
|
||||
inst_module "fsys"
|
||||
inst_module "dnfs"
|
||||
inst_module "dnfs2"
|
||||
inst_module "db"
|
||||
inst_module "db2"
|
||||
inst_module "smtp"
|
||||
inst_module "captcha"
|
||||
inst_module "confs"
|
||||
inst_module "smtp3"
|
||||
inst_module "captcha2"
|
||||
inst_module "ai_review" # 需联网拉 chromadb onnx 模型(_fetch_onnx_model.sh 走 ModelScope 镜像)
|
||||
inst_module "aibox" # 需 ~/.cert/deepseek.key(见 _setenv.sh)
|
||||
inst_module "redirect80" # 80 端口重定向,依赖 ../pymsg
|
||||
|
||||
+1
-1
@@ -2,4 +2,4 @@
|
||||
|
||||
# Debian
|
||||
sudo apt update
|
||||
sudo apt install libsqlite3-dev libboost-all-dev libcurl4-gnutls-dev python3 python3-venv python3-pip g++ patchelf golang-go libportaudio-ocaml-dev libfreetype-dev
|
||||
sudo apt install libsqlite3-dev libboost-all-dev libcurl4-gnutls-dev python3 python3-venv python3-pip g++ patchelf golang-go libportaudio-ocaml-dev libfreetype-dev libyaml-cpp-dev
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
#!/bin/bash
|
||||
|
||||
# RHEL / CentOS / Rocky / Fedora
|
||||
sudo dnf install -y boost-devel libcurl-devel python3 python3-pip gcc-c++ patchelf golang sqlite-devel freetype-devel
|
||||
sudo dnf install -y boost-devel libcurl-devel python3 python3-pip gcc-c++ patchelf golang sqlite-devel freetype-devel yaml-cpp-devel
|
||||
@@ -1,5 +1,6 @@
|
||||
# Virtual environment
|
||||
venv/
|
||||
venv_dfc/
|
||||
.env/
|
||||
|
||||
# Python bytecode
|
||||
@@ -12,6 +13,14 @@ checkpoints/
|
||||
*.pt
|
||||
*.pth
|
||||
*.onnx
|
||||
*.onnx.data
|
||||
*.hef
|
||||
*.har
|
||||
|
||||
# DFC logs
|
||||
acceleras.log
|
||||
hailo_sdk.*.log
|
||||
hailort.log
|
||||
|
||||
# Dataset (too large for git)
|
||||
data/
|
||||
|
||||
+221
-24
@@ -1,19 +1,27 @@
|
||||
# actpred — 动作预测模块
|
||||
|
||||
输入摄像头图像 + 历史动作序列,预测下一个连续动作(回归)。
|
||||
输入多帧摄像头图像 + 历史动作序列,预测下一时间步的连续动作(回归)。
|
||||
|
||||
纯神经网络实现(PyTorch),支持轻量 CNN 或 MobileNetV2 图像编码,MLP/GRU 动作历史编码。总参数量 ~140K(默认)至 ~2.3M。
|
||||
纯神经网络实现(PyTorch):轻量 CNN 或 MobileNetV2 逐帧编码,图像/动作 token 进轻量
|
||||
Transformer(pre-LN)自注意力融合,取 [PRED] token 回归下一步动作。
|
||||
默认配置总参数量 ~450K。
|
||||
|
||||
## 架构
|
||||
|
||||
```
|
||||
图像 (3×64×64) ──► CNN 编码器 ──► img_feat (128)
|
||||
│
|
||||
├─ concat ──► 融合头 ──► 预测动作 (action_dim)
|
||||
│
|
||||
动作历史 (H×D) ──► MLP/GRU ──► act_feat (64)
|
||||
F 帧图像 (F×3×H×W) ─► 共享 CNN 编码器 ─► F 个 token ─► proj → d_model ┐
|
||||
├─► [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;`--imu-dim`>0 时 IMU token 插到
|
||||
PRED 之前,长度 F+H+2),除 batch 外无动态 shape,对 NPU(如树莓派
|
||||
AI 扩展板)部署友好。
|
||||
|
||||
## 快速开始
|
||||
|
||||
```bash
|
||||
@@ -23,13 +31,41 @@ cd actpred && bash _set.sh
|
||||
# 训练
|
||||
bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18
|
||||
|
||||
# 推理(--action-history 格式:冒号分隔各步,每步 action_dim 个值,共 history_len 步)
|
||||
bash _run.sh --model checkpoints/best_inference.pt --image frame.jpg \
|
||||
# 断点续训练(从 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 \
|
||||
--images frame_t-2.jpg frame_t-1.jpg frame_t.jpg \
|
||||
--action-history "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"
|
||||
|
||||
# 导出 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
|
||||
```
|
||||
@@ -39,16 +75,18 @@ python test_actpred.py
|
||||
|
||||
## 预训练(从 .mt 脚本,无需图像)
|
||||
|
||||
从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵,用零图像占位预训练动作动力学(history → next action),
|
||||
从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵,用全零多帧图像占位预训练动作动力学(history → next action),
|
||||
保存的 checkpoint 与正常训练格式完全兼容,可直接用于推理或作为微调起点。
|
||||
|
||||
```bash
|
||||
# 预训练(解析 ../bot4cj/scripts/*.mt)
|
||||
bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0
|
||||
|
||||
# 用预训练权重作为微调起点
|
||||
# 注意:预训练在非归一化动作空间进行,微调时请勿开启 --normalize-actions,
|
||||
# 否则输入分布与预训练权重不匹配,预训练收益会大打折扣
|
||||
# 预训练断点续训
|
||||
bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0 \
|
||||
--resume checkpoints/pretrain_last.pt
|
||||
|
||||
# 用预训练权重作为微调起点(pretrain 也支持 --init-from,行为同 train)
|
||||
bash _run.sh train --data-dir ./data --epochs 50 \
|
||||
--init-from checkpoints/pretrain_best_inference.pt
|
||||
```
|
||||
@@ -56,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 粗量化图像路径。
|
||||
|
||||
## 数据格式
|
||||
|
||||
@@ -70,11 +120,22 @@ data/
|
||||
0001.jpg
|
||||
...
|
||||
actions.npy # shape: (T, action_dim) float32
|
||||
imu.npy # 可选,shape: (T, 8×N) float32(N 个 MPU6050 姿态拼接,R2)
|
||||
episode_002/
|
||||
...
|
||||
```
|
||||
|
||||
每个 timestep 的样本:当前图像 + 前 `history_len` 步动作 → 当前步动作。
|
||||
每个 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`,无需统计文件。
|
||||
录制/训练/闭环脚本的传感器行顺序必须一致,否则槽位错位。
|
||||
|
||||
### 数据采集
|
||||
|
||||
@@ -93,15 +154,22 @@ recorder.flush()
|
||||
|
||||
| 参数 | 默认值 | 说明 |
|
||||
|------|--------|------|
|
||||
| `--image-size` | 64 | 图像缩放尺寸 |
|
||||
| `--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) |
|
||||
| `--history-encoder` | mlp | 历史编码器 (mlp/gru) |
|
||||
| `--attn-dim` | 128 | 注意力 token 维度 d_model(须能被 attn-heads 整除)|
|
||||
| `--attn-heads` | 4 | 注意力头数 |
|
||||
| `--attn-layers` | 2 | TransformerEncoder 层数 |
|
||||
| `--epochs` | 100 | 训练轮数 |
|
||||
| `--lr` | 0.001 | 学习率 |
|
||||
| `--batch-size` | 32 | 批次大小 |
|
||||
| `--normalize-actions` | False | 动作归一化(`--no-normalize-actions` 关闭) |
|
||||
| `--batch-size` | 4 | 批次大小(默认分辨率 640×480;更高分辨率需调小,小分辨率可调大) |
|
||||
| `--normalize-actions` | True | 动作归一化(`--no-normalize-actions` 关闭;2026-08-14 起默认开启——复现旧行为需显式关闭) |
|
||||
| `--normalize-mode` | zscore | 归一化方式:`zscore`=数据统计(依赖 `action_stats.npy`);`fixed_range`=固定物理范围 min-max→[0,1](常量表见 `action_norm.py`,免 stats 文件,越界 clamp 防御) |
|
||||
| `--column-weights` | 舵机1.0/spk0.2/保留位0.0 | 逐列损失权重(逗号分隔 action_dim 个浮点;仅影响训练损失,推理端无感知;**始终生效**——包括 `--no-normalize-actions` 的未归一化训练) |
|
||||
| `--dropout` | 0.1 | Dropout 率 |
|
||||
| `--weight-decay` | 0.0001 | 权重衰减 |
|
||||
| `--val-split` | 0.15 | 验证集比例 |
|
||||
@@ -110,12 +178,42 @@ recorder.flush()
|
||||
| `--seed` | 42 | 随机种子 |
|
||||
| `--num-workers` | 2 | DataLoader 工作进程数 |
|
||||
| `--checkpoint-dir` | ./checkpoints | 检查点保存目录 |
|
||||
| `--init-from` | 无 | 预训练 checkpoint 路径,作为微调起点(strict=False) |
|
||||
| `--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
|
||||
@@ -124,7 +222,106 @@ recorder.flush()
|
||||
from predict import ActionPredictorInference
|
||||
|
||||
predictor = ActionPredictorInference("checkpoints/best_inference.pt")
|
||||
print(predictor.action_dim, predictor.history_len) # 公共属性
|
||||
action = predictor.predict(image_rgb, action_history) # np.ndarray (action_dim,)
|
||||
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` 对齐(含可选 `imu=` 入参),维度从 ONNX 输入元数据
|
||||
自动推导(无需 checkpoint)。Pi 5 实测(224×224,2026-08-13 复测)11.4ms/次含预处理
|
||||
(纯推理 10.7ms),PyTorch 21ms,约 2 倍加速,与 torch 前向输出差异
|
||||
<1e-5。数据见
|
||||
[测速报告](../docs/reports/actpred_inference_speed_benchmark_20260813.md) §6。
|
||||
|
||||
```bash
|
||||
bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
|
||||
```
|
||||
|
||||
```python
|
||||
from onnx_infer import OnnxPredictor
|
||||
|
||||
with OnnxPredictor("model.onnx") as predictor: # stats_path= 仅归一化模型需要
|
||||
action = predictor.predict(images, action_history)
|
||||
```
|
||||
|
||||
```bash
|
||||
python3 onnx_infer.py --model model.onnx --images f1.jpg f2.jpg f3.jpg \
|
||||
--action-history "45,0,...:45,0,..." [--stats action_stats.npy]
|
||||
```
|
||||
|
||||
## Hailo 部署(AI HAT 2 / Hailo-10H)
|
||||
|
||||
模型需先经 Hailo DFC(仅 x86 Linux,`venv_dfc`)编译为 HEF。
|
||||
训练态 ONNX(5D 输入、AdaptiveAvgPool、MultiheadAttention 的 head
|
||||
reshape)**无法被 DFC 解析**,必须用 `export_deploy.py` 的部署版导出
|
||||
(按帧展开 CNN、固定 kernel AvgPool、注意力按头展开为独立 Linear、
|
||||
全程 (1,1,S,d) 无 transpose 布局,权重数值不变)。
|
||||
|
||||
当前 NPU 目标分辨率为 **224×224**(pretrain checkpoint 配置),HEF 输入
|
||||
为 frame0~2 各 (1,224,224,3) NHWC + action_history (1,1,5,18) +
|
||||
pred_token (1,1,1,128)(`--imu-dim`>0 时追加 imu (1,1,1,imu_dim),共 6 输入),
|
||||
形状与量化参数由 `hailo_infer.py` 从 vstream
|
||||
info 自动推导,不硬编码。
|
||||
|
||||
```bash
|
||||
# x86 主机(venv_dfc 内已装 hailo-sdk 5.3):
|
||||
# 1. 导出部署版 ONNX + 学习到的 pred_token(推理端喂入用)
|
||||
python3 export_deploy.py --model checkpoints/pretrain_best_inference.pt \
|
||||
--pred-token-out checkpoints/actpred_deploy_pred_token.npy
|
||||
# 2. 生成校准集(pred_token 按学习值 ±30% 抖动,避免常量量化退化)
|
||||
python3 gen_calibration.py --deploy --model checkpoints/pretrain_best_inference.pt
|
||||
# 3. 一键编译(parser → optimize → compiler,各步自动做输入连通性检查)
|
||||
python3 dfc_compile.py # 产物 checkpoints/actpred_pretrain.hef
|
||||
# 4. 将 HEF + actpred_deploy_pred_token.npy 拷回树莓派
|
||||
```
|
||||
|
||||
已验证:完整 110 层网络(含 matmul/softmax/layernorm 注意力)解析、
|
||||
量化、编译全部通过;hailo10h 目标。
|
||||
编译要点与踩坑记录见 `dfc_compile.py` 文档字符串。
|
||||
|
||||
> **验证记录(2026-08-13 树莓派实测)**:归一化轮次(`--normalize-actions`
|
||||
> + `--random-images` 预训练,224×224)重编译的 HEF 已通过 `_npu_test.py`
|
||||
> 实机复验:img 敏感性 22.3、hist 敏感性 38.9(阈值 0.1),HEF vs torch
|
||||
> max_diff 0.37(原始动作空间),约 12ms/次(含预处理)。此前两轮失败的
|
||||
> 根因与修复均已闭环:
|
||||
> - BN 统计退化 → `--random-images` 预训练 + `export_deploy.py`
|
||||
> `check_token_ranges` 导出前拦截
|
||||
> - pred_token 常量校准退化(scale=18571.7)→ 抖动校准,推理端喂学习值
|
||||
> - 输入层断连 → `dfc_compile.py` parser/optimize 后连通性检查
|
||||
> - vstream 报告顺序乱序 → `hailo_infer.py` 按 `input_layerN` 排序绑定
|
||||
> - 帧归一化判定错误(把归一化帧当原始像素喂,图像路径饱和破坏)→
|
||||
> 按量化网格 dequant(0) 判定,见 hailo_infer.py 注释
|
||||
>
|
||||
> 验证链:`export_deploy.py`(token 量级)→ `_onnx_validate.py`(ONNX
|
||||
> 数值)→ `dfc_compile.py`(输入连通性)→ `_npu_test.py`(4 角敏感性)。
|
||||
> 遗留:全零历史角(z≈-3.5)超出校准输入范围(-2.7)会被截断,
|
||||
> 该角仅作连通性验证,非真实部署场景。
|
||||
|
||||
`hailo_infer.py` 的 `HailoPredictor` 适配 5/6 输入 HEF(含可选 imu 输入,
|
||||
按形状自动归类绑定),输入布局与量化
|
||||
参数(帧归一化由量化网格 dequant(0) 自动判定)从 vstream info 自动推导,
|
||||
单元测试含 fake hailo_platform 的 5 输入与乱序报告顺序用例。CLI 支持
|
||||
`--pred-token` 喂入学习值、`--stats` 喂入归一化统计、`--imu` 喂入姿态。
|
||||
环境调研与 API 踩坑记录见
|
||||
[docs/reports/actpred_hailo_hat2_inference_research_20260813.md](../docs/reports/actpred_hailo_hat2_inference_research_20260813.md)。
|
||||
|
||||
### 异步推理(`predict_async`,run_async 流水线)
|
||||
|
||||
`HailoPredictor.predict_async(images, action_history, pred_token=None, imu=None)`
|
||||
提交异步推理并立即返回 `_AsyncJob`,需要结果时 `job.result(timeout)` 阻塞获取:
|
||||
|
||||
- 与同步 `predict` 语义/输出完全一致(公共 `_prepare_inputs/_fill_bindings/_postprocess`)
|
||||
- 内部维护 4 个绑定缓冲池 + 完成回调;在途任务超过缓冲数时提交会阻塞等待
|
||||
空闲缓冲,`job.result` 超时**不会**释放缓冲(缓冲由完成回调释放)
|
||||
- 连续提交(≥2 个在途)时 NPU 保持满负荷:实测连续吞吐 6.0ms vs 同步
|
||||
15.5ms(约 2.6×);扩展测试覆盖乱序取结果/混合调用/缓冲压力/300 次
|
||||
稳定性/4 线程并发,输出与同步 0 差异
|
||||
- pybot `--loop` 检测到 `predict_async` 即自动走异步闭环(推理期间并行
|
||||
执行上一步动作),见 [pybot/README.md](../pybot/README.md)
|
||||
@@ -0,0 +1,174 @@
|
||||
"""actpred HEF NPU 输入敏感性复验脚本(重编译后的验收门槛)
|
||||
|
||||
用法: venv/bin/python _npu_test.py
|
||||
|
||||
4 角测试(真实/全零帧 × 真实/全零历史)+ 与 PyTorch 参考对比。
|
||||
HEF 尺寸/布局从 vstream info 自动推导,checkpoint 用
|
||||
checkpoints/best_inference.pt(v4 onehot,fixed_range)。
|
||||
"""
|
||||
import sys
|
||||
import time
|
||||
sys.path.insert(0, "/opt/blackbird/actpred")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from PIL import Image
|
||||
from hailo_platform import HEF
|
||||
|
||||
from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
from hailo_infer import HailoPredictor
|
||||
|
||||
HEF_PATH = "checkpoints/actpred_v4.hef"
|
||||
CKPT = "checkpoints/best_inference.pt"
|
||||
IMG_DIR = "data/episode_0000/images"
|
||||
|
||||
hef = HEF(HEF_PATH)
|
||||
print("=== HEF vstream infos ===")
|
||||
for i in hef.get_input_vstream_infos():
|
||||
q = i.quant_info
|
||||
print(f"IN {i.name:16s} {tuple(i.shape)} {i.format.type} "
|
||||
f"zp={q.qp_zp} scale={q.qp_scale:.6g}")
|
||||
for o in hef.get_output_vstream_infos():
|
||||
q = o.quant_info
|
||||
print(f"OUT {o.name:16s} {tuple(o.shape)} {o.format.type} "
|
||||
f"zp={q.qp_zp} scale={q.qp_scale:.6g}")
|
||||
|
||||
imgs_real = [np.array(Image.open(f"{IMG_DIR}/000{i}.jpg").convert("RGB"))
|
||||
for i in range(3)]
|
||||
imgs_zero = [np.zeros_like(imgs_real[0]) for _ in range(3)]
|
||||
|
||||
hist_real = np.array([
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 10, 0, 0, 0, 0, 0, 0],
|
||||
], dtype=np.float32)
|
||||
hist_zero = np.zeros((5, 18), dtype=np.float32)
|
||||
|
||||
# PyTorch 参考(checkpoint 自带 config)
|
||||
ck = torch.load(CKPT, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(ck["config"])
|
||||
model = ActionPredictor(cfg.model, cfg.data, pretrained=False)
|
||||
model.load_state_dict(ck["model_state_dict"])
|
||||
model.eval()
|
||||
pred_learned = model.pred_token.data.unsqueeze(1).numpy() # (1,1,1,d)
|
||||
imu_dim = cfg.data.imu_dim
|
||||
|
||||
# 归一化模型:torch 参考在喂入前归一化;HailoPredictor 传 normalize_mode
|
||||
# 由封装内部归一化(hist_real/hist_zero 保持原始值,避免双重归一化)。
|
||||
# fixed_range 免 action_stats.npy(该模式训练不产出统计文件)。
|
||||
import os
|
||||
stats_path = None
|
||||
norm_mode = cfg.data.normalize_mode if cfg.data.normalize_actions else "none"
|
||||
if norm_mode == "zscore":
|
||||
stats_path = cfg.data.action_stats_path or os.path.join(
|
||||
os.path.dirname(os.path.abspath(CKPT)), "action_stats.npy")
|
||||
if not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(os.path.dirname(os.path.abspath(CKPT)),
|
||||
stats_path)
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(f"统计文件不存在: {stats_path}")
|
||||
print(f"HailoPredictor stats: {stats_path}")
|
||||
elif norm_mode == "fixed_range":
|
||||
print("fixed_range 归一化:免 action_stats.npy,按固定物理范围 min-max")
|
||||
|
||||
from torchvision import transforms
|
||||
tf = transforms.Compose([
|
||||
transforms.Resize((cfg.data.image_height, cfg.data.image_width)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
|
||||
def torch_pred(imgs, hist, use_learned_pred=True, imu=None):
|
||||
x = torch.stack([tf(Image.fromarray(im)) for im in imgs]).unsqueeze(0)
|
||||
h = torch.tensor(hist).unsqueeze(0)
|
||||
if norm_mode == "zscore" and stats_path is not None:
|
||||
stats = np.load(stats_path)
|
||||
h = (h - torch.tensor(stats[0])) / torch.tensor(stats[1])
|
||||
elif norm_mode == "fixed_range":
|
||||
from action_norm import normalize_action
|
||||
h = torch.from_numpy(normalize_action(h.numpy(), cfg.data.action_dim))
|
||||
ii = None
|
||||
if imu_dim > 0:
|
||||
# 与 HailoPredictor 同口径:原始姿态 → 内部归一化前需手动归一
|
||||
from imu_norm import normalize_imu
|
||||
ii = torch.from_numpy(normalize_imu(imu, imu_dim)).unsqueeze(0)
|
||||
if use_learned_pred:
|
||||
with torch.no_grad():
|
||||
out = model(x, h, ii).squeeze(0).numpy()
|
||||
else:
|
||||
saved = model.pred_token.data.clone()
|
||||
model.pred_token.data.zero_()
|
||||
with torch.no_grad():
|
||||
out = model(x, h, ii).squeeze(0).numpy()
|
||||
model.pred_token.data.copy_(saved)
|
||||
if cfg.data.onehot:
|
||||
# 多头 logits → 期望值解码 → 原始量纲(与 predict.py 同语义,
|
||||
# 含权重 0 列固定常量),对比前无需反归一化
|
||||
from bins import decode as bin_decode, fixed_zero_columns
|
||||
from action_norm import default_column_weights
|
||||
out = bin_decode(out, cfg.data.action_dim).astype(np.float32)
|
||||
w = (np.asarray(cfg.data.column_weights, dtype=np.float32)
|
||||
if cfg.data.column_weights is not None
|
||||
else default_column_weights(cfg.data.action_dim))
|
||||
fixed = fixed_zero_columns(w, cfg.data.action_dim)
|
||||
out[w <= 0] = fixed[w <= 0]
|
||||
return out
|
||||
|
||||
|
||||
# R2 方案一:imu_dim>0 时每步传原始姿态(HailoPredictor 内部归一化)
|
||||
imu_real = np.array([5.0, -3.0] + [0.0] * 6, dtype=np.float32)
|
||||
|
||||
# normalize_actions=False 时 normalize_mode 无意义(stats_path=None 即不归一化)
|
||||
with HailoPredictor(HEF_PATH, stats_path=stats_path,
|
||||
normalize_mode=(norm_mode if norm_mode != "none"
|
||||
else "zscore")) as p:
|
||||
print(f"\ndims: frame_len={p.frame_len} img={p.image_width}x{p.image_height} "
|
||||
f"history={p.history_len} action_dim={p.action_dim}")
|
||||
t0 = time.perf_counter()
|
||||
a = p.predict(imgs_real, hist_real, pred_token=pred_learned, imu=imu_real)
|
||||
b = p.predict(imgs_zero, hist_real, pred_token=pred_learned, imu=imu_real)
|
||||
c = p.predict(imgs_real, hist_zero, pred_token=pred_learned, imu=imu_real)
|
||||
d = p.predict(imgs_zero, hist_zero, pred_token=pred_learned)
|
||||
t1 = time.perf_counter()
|
||||
print(f"4x predict: {(t1 - t0) / 4 * 1000:.1f} ms/次")
|
||||
for name, r in [("real_img+real_hist", a), ("zero_img+real_hist", b),
|
||||
("real_img+zero_hist", c), ("zero_img+zero_hist", d)]:
|
||||
print(f"{name:20s} max_abs={np.abs(r).max():.4f} out={np.round(r[:6], 3)}...")
|
||||
|
||||
img_sens = np.abs(a - b).max()
|
||||
hist_sens = np.abs(a - c).max()
|
||||
print(f"\nimg sensitivity = {img_sens:.4f}(>0.1 视为输入连通)")
|
||||
print(f"hist sensitivity = {hist_sens:.4f}(>0.1 视为输入连通)")
|
||||
ok = img_sens > 0.1 and hist_sens > 0.1
|
||||
if imu_dim > 0:
|
||||
e = p.predict(imgs_real, hist_real, pred_token=pred_learned,
|
||||
imu=np.zeros_like(imu_real))
|
||||
print(f"imu sensitivity = {np.abs(a - e).max():.4f}(>0 视为姿态连通)")
|
||||
|
||||
ref = torch_pred(imgs_real, hist_real, True, imu_real)
|
||||
ref_zero_pred = torch_pred(imgs_real, hist_real, False, imu_real)
|
||||
# torch 参考是模型原始归一化空间输出;HailoPredictor 已反归一化
|
||||
# 到原始动作空间,比较前须同空间对齐,否则归一化模型下 max_diff
|
||||
# 无意义(onehot 模型 torch_pred 已 decode 为原始量纲,无需再处理)
|
||||
if not cfg.data.onehot:
|
||||
if norm_mode == "fixed_range":
|
||||
from action_norm import denormalize_action
|
||||
ref = denormalize_action(ref, cfg.data.action_dim)
|
||||
ref_zero_pred = denormalize_action(ref_zero_pred,
|
||||
cfg.data.action_dim)
|
||||
elif stats_path is not None:
|
||||
s = np.load(stats_path)
|
||||
ref = ref * s[1] + s[0]
|
||||
ref_zero_pred = ref_zero_pred * s[1] + s[0]
|
||||
print(f"torch learned_pred : {np.round(ref[:6], 3)}...")
|
||||
print(f"torch zero_pred : {np.round(ref_zero_pred[:6], 3)}...")
|
||||
print(f"HEF vs torch(learned) max_diff={np.abs(a - ref).max():.4f}"
|
||||
f"(原始动作空间)")
|
||||
|
||||
print("\nPASS: HEF 输入连通,可继续与 torch 对比量化误差"
|
||||
if ok else "\nFAIL: HEF 输出对输入不敏感,编译仍有问题")
|
||||
@@ -0,0 +1,138 @@
|
||||
"""部署版 ONNX 数值验证:输入敏感性 + 与 PyTorch 一致性(重编译前门槛)
|
||||
|
||||
用法: venv/bin/python _onnx_validate.py
|
||||
"""
|
||||
import sys
|
||||
sys.path.insert(0, "/opt/blackbird/actpred")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
import onnxruntime as ort
|
||||
|
||||
from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
|
||||
CKPT = "checkpoints/pretrain_best_inference.pt"
|
||||
ONNX = "checkpoints/actpred_deploy.onnx"
|
||||
IMG_DIR = "data/episode_0000/images"
|
||||
|
||||
ck = torch.load(CKPT, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(ck["config"])
|
||||
model = ActionPredictor(cfg.model, cfg.data, pretrained=False)
|
||||
model.load_state_dict(ck["model_state_dict"])
|
||||
model.eval()
|
||||
|
||||
action_stats = None
|
||||
norm_mode = cfg.data.normalize_mode if cfg.data.normalize_actions else "none"
|
||||
if norm_mode == "zscore":
|
||||
import os
|
||||
stats_path = cfg.data.action_stats_path or os.path.join(
|
||||
os.path.dirname(os.path.abspath(CKPT)), "action_stats.npy")
|
||||
if not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(os.path.dirname(os.path.abspath(CKPT)),
|
||||
stats_path)
|
||||
action_stats = np.load(stats_path)
|
||||
if action_stats.shape != (2, cfg.data.action_dim):
|
||||
raise ValueError(f"action stats 形状应为 (2, {cfg.data.action_dim}),"
|
||||
f"实际为 {action_stats.shape}: {stats_path}")
|
||||
print(f"history 将 z-score 归一化({stats_path})")
|
||||
elif norm_mode == "fixed_range":
|
||||
print("history 将 fixed_range min-max 归一化(action_norm.py,免 stats)")
|
||||
|
||||
sess = ort.InferenceSession(ONNX)
|
||||
in_names = [i.name for i in sess.get_inputs()]
|
||||
print("ONNX inputs:", in_names)
|
||||
|
||||
h, w = cfg.data.image_height, cfg.data.image_width
|
||||
tf = transforms.Compose([
|
||||
transforms.Resize((h, w)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
imgs_real = [np.array(Image.open(f"{IMG_DIR}/000{i}.jpg").convert("RGB"))
|
||||
for i in range(3)]
|
||||
imgs_zero = [np.zeros_like(imgs_real[0]) for _ in range(3)]
|
||||
|
||||
hist_real = np.array([
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 10, 200, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 20, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
[45, 0, 45, 0, 45, 0, 45, 0, 90, 90, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||
], dtype=np.float32)
|
||||
hist_zero = np.zeros((5, 18), dtype=np.float32)
|
||||
|
||||
if norm_mode == "fixed_range":
|
||||
from action_norm import normalize_action
|
||||
hist_real = normalize_action(hist_real, cfg.data.action_dim)
|
||||
hist_zero = normalize_action(hist_zero, cfg.data.action_dim)
|
||||
elif action_stats is not None:
|
||||
hist_real = (hist_real - action_stats[0]) / action_stats[1]
|
||||
hist_zero = (hist_zero - action_stats[0]) / action_stats[1]
|
||||
|
||||
pred_learned = model.pred_token.data.unsqueeze(1).numpy() # (1,1,1,128)
|
||||
pred_none = np.zeros((1, 1, 1, 128), dtype=np.float32)
|
||||
|
||||
# R2 方案一:imu_dim>0 时模型/ONNX 多了归一化姿态输入
|
||||
imu_dim = cfg.data.imu_dim
|
||||
imu_real = None
|
||||
imu_zero = None
|
||||
if imu_dim > 0:
|
||||
from imu_norm import normalize_imu
|
||||
# 原始姿态(roll=5°, pitch=-3°, 其余 0)→ 部署版输入布局 (1,1,1,imu_dim)
|
||||
imu_real = normalize_imu([5.0, -3.0] + [0.0] * 6, imu_dim).reshape(
|
||||
1, 1, 1, imu_dim)
|
||||
imu_zero = np.zeros((1, 1, 1, imu_dim), dtype=np.float32)
|
||||
|
||||
|
||||
def frames_nhwc(imgs):
|
||||
"""列表帧 → (1,H,W,3) 归一化 float32 NHWC × frame_len"""
|
||||
chw = np.stack([tf(Image.fromarray(im)).numpy() for im in imgs])
|
||||
return [chw[i].transpose(1, 2, 0)[None] for i in range(len(imgs))]
|
||||
|
||||
|
||||
def run_onnx(imgs, hist, pred, imu=None):
|
||||
feeds = {f"frame{i}": f for i, f in enumerate(frames_nhwc(imgs))}
|
||||
feeds["action_history"] = hist[None, None]
|
||||
feeds["pred_token"] = pred
|
||||
if imu_dim > 0:
|
||||
feeds["imu"] = imu if imu is not None else imu_zero
|
||||
return sess.run(None, feeds)[0].reshape(-1)
|
||||
|
||||
|
||||
def run_torch(imgs, hist, imu=None):
|
||||
x = torch.stack([tf(Image.fromarray(im)) for im in imgs]).unsqueeze(0)
|
||||
hh = torch.tensor(hist).unsqueeze(0)
|
||||
ii = None
|
||||
if imu_dim > 0:
|
||||
ii = torch.tensor(
|
||||
imu if imu is not None else imu_zero).reshape(1, imu_dim).float()
|
||||
with torch.no_grad():
|
||||
return model(x, hh, ii).squeeze(0).numpy()
|
||||
|
||||
|
||||
a = run_onnx(imgs_real, hist_real, pred_learned, imu_real)
|
||||
b = run_onnx(imgs_zero, hist_real, pred_learned, imu_real)
|
||||
c = run_onnx(imgs_real, hist_zero, pred_learned, imu_real)
|
||||
d = run_onnx(imgs_zero, hist_zero, pred_learned)
|
||||
p0 = run_onnx(imgs_real, hist_real, pred_none, imu_real)
|
||||
|
||||
print(f"onnx real_img+real_hist: {np.round(a[:6], 3)}...")
|
||||
print(f"onnx pred_token=0 diff : max={np.abs(a - p0).max():.4f}")
|
||||
print(f"img sensitivity (a-b): max={np.abs(a - b).max():.4f}")
|
||||
print(f"hist sensitivity (a-c): max={np.abs(a - c).max():.4f}")
|
||||
|
||||
ok = np.abs(a - b).max() > 1e-4 and np.abs(a - c).max() > 1e-4
|
||||
if imu_dim > 0:
|
||||
e = run_onnx(imgs_real, hist_real, pred_learned, imu_zero)
|
||||
print(f"imu sensitivity (a-e): max={np.abs(a - e).max():.4f}")
|
||||
|
||||
ref = run_torch(imgs_real, hist_real, imu_real)
|
||||
print(f"torch reference : {np.round(ref[:6], 3)}...")
|
||||
print(f"onnx vs torch max_diff : {np.abs(a - ref).max():.6e}")
|
||||
|
||||
ok = ok and np.abs(a - ref).max() < 1e-4
|
||||
print("PASS" if ok else "FAIL")
|
||||
+5
-1
@@ -11,6 +11,10 @@ elif [ "${1:-}" = "pretrain" ]; then
|
||||
elif [ "${1:-}" = "export" ]; then
|
||||
shift
|
||||
exec python3 export_onnx.py "$@"
|
||||
else
|
||||
elif [ "${1:-}" = "predict" ]; then
|
||||
shift
|
||||
exec python3 predict.py "$@"
|
||||
else
|
||||
echo "未知子命令: ${1:-}(可用: train / pretrain / export / predict)" >&2
|
||||
exit 1
|
||||
fi
|
||||
@@ -15,3 +15,33 @@ 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
|
||||
hub_dir = torch.hub.get_dir()
|
||||
cache = os.path.join(hub_dir, 'checkpoints', 'mobilenet_v2-b0353104.pth')
|
||||
if os.path.exists(cache):
|
||||
print(f'MobileNetV2 weights already cached: {cache}')
|
||||
sys.exit(0)
|
||||
url = 'https://download.pytorch.org/models/mobilenet_v2-b0353104.pth'
|
||||
try:
|
||||
torch.hub.load_state_dict_from_url(url, model_dir=os.path.join(hub_dir, 'checkpoints'))
|
||||
print('MobileNetV2 weights cached to: ' + cache)
|
||||
except Exception as e:
|
||||
print(f'WARNING: MobileNetV2 weights download failed: {e}')
|
||||
print(' backbone=simple still works; for mobilenet, manually download from:')
|
||||
print(f' {url}')
|
||||
print(f' and place at: {cache}')
|
||||
" || echo "WARNING: MobileNetV2 weight pre-download skipped"
|
||||
@@ -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
|
||||
+210
-21
@@ -1,40 +1,89 @@
|
||||
import argparse
|
||||
from dataclasses import dataclass, asdict, field
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataConfig:
|
||||
image_size: int = 64
|
||||
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
|
||||
class ModelConfig:
|
||||
cnn_backbone: str = "simple"
|
||||
cnn_feature_dim: int = 128
|
||||
action_hidden_dim: int = 64
|
||||
fused_hidden_dim: int = 128
|
||||
dropout: float = 0.1
|
||||
history_encoder: str = "mlp"
|
||||
attn_dim: int = 128
|
||||
attn_heads: int = 4
|
||||
attn_layers: int = 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainConfig:
|
||||
batch_size: int = 32
|
||||
# 默认输入分辨率 640×480(1080p 采集的预处理缩放在实机测速中主导
|
||||
# 端到端延迟,见测速报告 §6.1);更高分辨率训练时需调小 batch
|
||||
batch_size: int = 4
|
||||
epochs: int = 100
|
||||
lr: float = 1e-3
|
||||
weight_decay: float = 1e-4
|
||||
val_split: float = 0.15
|
||||
# 验证集划分方式:
|
||||
# "episode" 整集划分(防数据泄漏,但单 episode 验证噪声大)
|
||||
# "sample" 样本级随机划分(默认:多集混合、验证集更大更稳,
|
||||
# 代价是相邻高相关样本泄漏使 val 略乐观)
|
||||
val_split_mode: str = "sample"
|
||||
checkpoint_dir: str = "./checkpoints"
|
||||
patience: int = 15
|
||||
num_workers: int = 2
|
||||
seed: int = 42
|
||||
max_grad_norm: float = 1.0
|
||||
augment: bool = False
|
||||
# R1-C 残差预测:True 时 target 改为 Δ=actions[t]-actions[t-1],
|
||||
# 推理端 â_t = a_{t-1} + Δ̂(惯性先验,平滑动作)。默认 False 兼容旧 checkpoint。
|
||||
predict_delta: bool = False
|
||||
# R1-B 多步展开训练:rollout_len>1 时 dataset 返回 K 步连续序列,
|
||||
# 训练用 scheduled sampling(前 ss_prob 比例步用真历史,其余用自预测回传),
|
||||
# 治 exposure bias。rollout_len=1 时退化为现状单步。
|
||||
rollout_len: int = 1
|
||||
ss_prob: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -57,30 +106,67 @@ class Config:
|
||||
|
||||
def build_argparser(description: str = "") -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=description)
|
||||
parser.add_argument("--image-size", type=int, default=None)
|
||||
parser.add_argument("--image-width", type=int, default=None)
|
||||
parser.add_argument("--image-height", type=int, default=None)
|
||||
parser.add_argument("--action-dim", type=int, default=None)
|
||||
parser.add_argument("--history-len", type=int, default=None)
|
||||
parser.add_argument("--frame-len", type=int, default=None)
|
||||
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("--action-hidden-dim", type=int, default=None)
|
||||
parser.add_argument("--fused-hidden-dim", type=int, default=None)
|
||||
parser.add_argument("--dropout", type=float, default=None)
|
||||
parser.add_argument("--history-encoder", type=str, default=None, choices=["mlp", "gru"])
|
||||
parser.add_argument("--attn-dim", type=int, default=None)
|
||||
parser.add_argument("--attn-heads", type=int, default=None)
|
||||
parser.add_argument("--attn-layers", type=int, default=None)
|
||||
parser.add_argument("--batch-size", type=int, default=None)
|
||||
parser.add_argument("--epochs", type=int, default=None)
|
||||
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,
|
||||
help="从 last.pt 断点续训(恢复 model/optimizer/scheduler/epoch)")
|
||||
parser.add_argument("--save-config", type=str, default=None, help="保存配置到 JSON 文件")
|
||||
return parser
|
||||
|
||||
@@ -93,13 +179,18 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
也防止训练 CLI 越界参数(如 --history-len 0 / --lr -1)静默进入训练。
|
||||
"""
|
||||
checks = [
|
||||
("image_size", data_cfg.image_size, 1, 1024, int),
|
||||
("image_width", data_cfg.image_width, 1, 4096, int),
|
||||
("image_height", data_cfg.image_height, 1, 4096, int),
|
||||
("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),
|
||||
("action_hidden_dim", model_cfg.action_hidden_dim, 1, 4096, int),
|
||||
("fused_hidden_dim", model_cfg.fused_hidden_dim, 1, 4096, int),
|
||||
("dropout", model_cfg.dropout, 0.0, 1.0, float),
|
||||
("attn_dim", model_cfg.attn_dim, 1, 4096, int),
|
||||
("attn_heads", model_cfg.attn_heads, 1, 64, int),
|
||||
("attn_layers", model_cfg.attn_layers, 1, 64, int),
|
||||
]
|
||||
for name, val, lo, hi, typ in checks:
|
||||
if typ is int:
|
||||
@@ -110,6 +201,47 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
|
||||
raise ValueError(
|
||||
f"checkpoint config {name}={val} 超出合理范围 [{lo}, {hi}]")
|
||||
|
||||
if model_cfg.attn_dim % model_cfg.attn_heads != 0:
|
||||
raise ValueError(
|
||||
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),
|
||||
@@ -120,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:
|
||||
@@ -131,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)}")
|
||||
@@ -159,32 +332,48 @@ def parse_checkpoint_config(cfg_dict: dict) -> Config:
|
||||
|
||||
def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
|
||||
mapping = {
|
||||
"image_size": ("data", "image_size"),
|
||||
"image_width": ("data", "image_width"),
|
||||
"image_height": ("data", "image_height"),
|
||||
"action_dim": ("data", "action_dim"),
|
||||
"history_len": ("data", "history_len"),
|
||||
"frame_len": ("data", "frame_len"),
|
||||
"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"),
|
||||
"action_hidden_dim": ("model", "action_hidden_dim"),
|
||||
"fused_hidden_dim": ("model", "fused_hidden_dim"),
|
||||
"dropout": ("model", "dropout"),
|
||||
"history_encoder": ("model", "history_encoder"),
|
||||
"attn_dim": ("model", "attn_dim"),
|
||||
"attn_heads": ("model", "attn_heads"),
|
||||
"attn_layers": ("model", "attn_layers"),
|
||||
"batch_size": ("train", "batch_size"),
|
||||
"epochs": ("train", "epochs"),
|
||||
"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
|
||||
+187
-17
@@ -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,14 +23,43 @@ def _natural_key(p):
|
||||
|
||||
class ActionDataset(Dataset):
|
||||
def __init__(self, data_dir: str, history_len: int, action_dim: int,
|
||||
image_size: int = 64, normalize_actions: bool = False,
|
||||
action_stats_path: str = None, episode_indices: list = None):
|
||||
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, 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
|
||||
self.image_size = image_size
|
||||
# 固定范围类归一化(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([
|
||||
transforms.Resize((image_size, image_size)),
|
||||
# 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]),
|
||||
])
|
||||
@@ -31,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"
|
||||
@@ -51,17 +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)
|
||||
for t in range(history_len, n):
|
||||
self.samples.append((image_files[t], actions[t - history_len:t], actions[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_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:
|
||||
@@ -69,23 +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 _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 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))
|
||||
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_path, action_history, target = self.samples[idx]
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
image = self.transform(image)
|
||||
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 {"image": image, "action_history": action_history, "target": target}
|
||||
image_files, actions, imus, t_start = self.samples[idx]
|
||||
if self.rollout_len == 1:
|
||||
# 单步:与旧版输出结构完全一致(imu_dim>0 时附加 "imu" 键)
|
||||
images = self._load_frames(image_files, t_start)
|
||||
history = torch.tensor(self._norm(self._jit_history(
|
||||
actions[t_start - self.history_len: t_start])),
|
||||
dtype=torch.float32)
|
||||
target = torch.tensor(self._step_target(actions, t_start),
|
||||
dtype=torch.float32)
|
||||
out = {"images": images, "action_history": history, "target": target}
|
||||
if self.imu_dim > 0:
|
||||
out["imu"] = self._imu_at(imus, t_start)
|
||||
return out
|
||||
|
||||
# 多步:返回 K 步序列,scheduled sampling 在 train 循环内按步取用
|
||||
steps = []
|
||||
for k in range(self.rollout_len):
|
||||
t = t_start + k
|
||||
step = {
|
||||
"images": self._load_frames(image_files, t),
|
||||
# 真实历史(scheduled sampling 的真值来源;首步必用)
|
||||
"action_history": torch.tensor(
|
||||
self._norm(self._jit_history(
|
||||
actions[t - self.history_len: t])), dtype=torch.float32),
|
||||
"target": torch.tensor(self._step_target(actions, t), dtype=torch.float32),
|
||||
# 上一步真实动作(Δ 残差的基准 a_{t-1},归一化空间)
|
||||
"prev_action": torch.tensor(self._norm(actions[t - 1]), dtype=torch.float32),
|
||||
}
|
||||
if self.imu_dim > 0:
|
||||
step["imu"] = self._imu_at(imus, t)
|
||||
steps.append(step)
|
||||
return {"steps": steps}
|
||||
|
||||
def get_action_stats(self):
|
||||
if self._action_mean is not None:
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
"""actpred 部署版 ONNX → HEF 编译脚本(Hailo DFC 5.3,仅 x86 Linux)
|
||||
|
||||
流程(对应 DFC 5.x CLI):
|
||||
hailo parser onnx → HAR(网络翻译)
|
||||
hailo optimize → 量化优化后的 HAR(需校准集 npz)
|
||||
hailo compiler → HEF
|
||||
|
||||
用法:
|
||||
# 1. 导出部署版 ONNX(DFC 友好图结构)
|
||||
python export_deploy.py --model checkpoints/pretrain_best_inference.pt
|
||||
# 2. 生成校准集(首次)
|
||||
python gen_calibration.py --deploy
|
||||
# 3. 编译
|
||||
python dfc_compile.py
|
||||
|
||||
关键点(踩坑记录):
|
||||
- 训练态 ONNX(5D 输入、AdaptiveAvgPool、MultiheadAttention 的 head
|
||||
reshape)DFC 无法解析,必须用 export_deploy.py 的改造版图
|
||||
- pred_token 必须作为图输入:常量 initializer 会被 DFC 误转置
|
||||
- 必须显式 --start/--end-node-names 与 --input-format,且禁用
|
||||
onnxsim(本图简化无收益,反而引入 reshape)
|
||||
- 校准集键名必须匹配 HN 层名 actpred_deploy/input_layerN
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tarfile
|
||||
import tempfile
|
||||
|
||||
|
||||
NET_NAME = "actpred_deploy"
|
||||
|
||||
|
||||
def start_nodes(frame_count: int, with_imu: bool = False) -> list:
|
||||
"""export_deploy.py 的输入命名约定:frame{i} × frame_count + action_history
|
||||
+ pred_token(R2 imu_dim>0 时末尾追加 imu)"""
|
||||
nodes = ([f"frame{i}" for i in range(frame_count)]
|
||||
+ ["action_history", "pred_token"])
|
||||
if with_imu:
|
||||
nodes.append("imu")
|
||||
return nodes
|
||||
|
||||
|
||||
def hailo(*args, venv_bin: str):
|
||||
cmd = [os.path.join(venv_bin, "hailo"), *args]
|
||||
env = dict(os.environ, PATH=venv_bin + os.pathsep + os.environ["PATH"])
|
||||
return subprocess.run(cmd, env=env, check=True)
|
||||
|
||||
|
||||
def _har_hn_layers(har_path: str) -> dict:
|
||||
"""解包 HAR 并读取主 hn 文件的 layers(优化产物含 fp/native 变体,取主图)"""
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
with tarfile.open(har_path) as tf:
|
||||
tf.extractall(tmp)
|
||||
hn_files = [f for f in os.listdir(tmp) if f.endswith(".hn")]
|
||||
if not hn_files:
|
||||
raise ValueError(f"{har_path} 中无 .hn 文件")
|
||||
main = [f for f in hn_files
|
||||
if not f.endswith(".fp.hn") and not f.endswith(".native.hn")]
|
||||
hn_file = main[0] if main else hn_files[0]
|
||||
with open(os.path.join(tmp, hn_file)) as f:
|
||||
return json.load(f)["layers"]
|
||||
|
||||
|
||||
def _input_layer_index(name: str) -> int:
|
||||
m = re.search(r"input_layer(\d+)$", name)
|
||||
return int(m.group(1)) if m else 0
|
||||
|
||||
|
||||
def har_input_layer_names(har_path: str) -> list:
|
||||
"""从 HAR 主 hn 图读输入层名,按 input_layerN 数字后缀排序"""
|
||||
layers = _har_hn_layers(har_path)
|
||||
names = [n for n, l in layers.items() if l["type"] == "input_layer"]
|
||||
names.sort(key=_input_layer_index)
|
||||
return names
|
||||
|
||||
|
||||
def check_input_connectivity(har_path: str, input_layers: list):
|
||||
"""检查 HAR 图中每个输入层是否可达任一输出层
|
||||
|
||||
输入层被优化断开时编译照常成功,但 HEF 输出恒定值且无任何报错
|
||||
(2026-08-13 实测踩坑),必须在编译前拦截。DFC 5.3 的 HAR 不含
|
||||
.edges 文件,连通性改从主 hn 图的 layer output 边 BFS 判定。
|
||||
"""
|
||||
layers = _har_hn_layers(har_path)
|
||||
adj = {n: l.get("output", []) for n, l in layers.items()}
|
||||
targets = {n for n, l in layers.items() if l["type"] == "output_layer"}
|
||||
if not targets:
|
||||
print(f"WARNING: {har_path} 图中无输出层,跳过连通性检查")
|
||||
return
|
||||
disconnected = []
|
||||
for name in input_layers:
|
||||
seen, stack = set(), [name]
|
||||
while stack:
|
||||
cur = stack.pop()
|
||||
if cur in seen:
|
||||
continue
|
||||
seen.add(cur)
|
||||
stack.extend(adj.get(cur, []))
|
||||
if not (seen & targets):
|
||||
disconnected.append(name)
|
||||
if disconnected:
|
||||
sys.exit(
|
||||
f"错误: {har_path} 中输入层不可达输出层(输入被优化断开,"
|
||||
f"编译出的 HEF 将输出恒定值): {disconnected}")
|
||||
print(f"输入连通性检查通过: {len(input_layers)} 个输入层均可达输出层")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Compile actpred deploy ONNX to HEF")
|
||||
parser.add_argument("--onnx", default="checkpoints/actpred_deploy.onnx")
|
||||
parser.add_argument("--calib-npz", default="data/calibration_deploy.npz")
|
||||
parser.add_argument("--output", default="checkpoints/actpred_pretrain.hef")
|
||||
parser.add_argument("--hw-arch", default="hailo10h",
|
||||
choices=["hailo15h", "hailo15m", "hailo15l",
|
||||
"hailo10h", "hailo10p"])
|
||||
parser.add_argument("--frame-count", type=int, default=3,
|
||||
help="帧输入个数(与导出 ONNX 的 frame_len 一致,"
|
||||
"默认 3)")
|
||||
parser.add_argument("--with-imu", action="store_true",
|
||||
help="R2:导出 ONNX 含 imu 输入时开启(start nodes "
|
||||
"末尾追加 imu,与 export_deploy.py 命名一致)")
|
||||
parser.add_argument("--venv-bin", default="venv_dfc/bin",
|
||||
help="安装了 hailo-sdk 的 venv bin 目录")
|
||||
args = parser.parse_args()
|
||||
|
||||
venv_bin = os.path.abspath(args.venv_bin)
|
||||
if not os.path.exists(os.path.join(venv_bin, "hailo")):
|
||||
sys.exit(f"错误: {venv_bin}/hailo 不存在,请先安装 hailo-sdk")
|
||||
|
||||
nodes = start_nodes(args.frame_count, args.with_imu)
|
||||
input_formats = ["NHWC"] * len(nodes)
|
||||
|
||||
workdir = tempfile.mkdtemp(prefix="dfc_")
|
||||
har_raw = os.path.join(workdir, "parsed.har")
|
||||
har_opt = os.path.join(workdir, "optimized.har")
|
||||
hef_dir = os.path.join(workdir, "hef")
|
||||
os.makedirs(hef_dir)
|
||||
|
||||
try:
|
||||
print("[1/3] parser: ONNX → HAR")
|
||||
hailo("parser", "onnx", args.onnx, "--hw-arch", args.hw_arch, "-y",
|
||||
"--har-path", har_raw,
|
||||
"--start-node-names", *nodes,
|
||||
"--end-node-names", "predicted_action",
|
||||
"--input-format",
|
||||
*[f"{n}={f}" for n, f in zip(nodes, input_formats)],
|
||||
"--disable-onnx-simplifier", venv_bin=venv_bin)
|
||||
|
||||
# 校验未被推荐机制截断:输出层应来自预测头(conv 而非 ew_add)
|
||||
names = har_input_layer_names(har_raw)
|
||||
if len(names) != len(nodes):
|
||||
sys.exit(f"错误: HAR 输入层数 {len(names)} != {len(nodes)},解析可能被截断")
|
||||
check_input_connectivity(har_raw, names)
|
||||
|
||||
print("[2/3] optimize: 量化校准")
|
||||
hailo("optimize", har_raw, "--hw-arch", args.hw_arch,
|
||||
"--calib-set-path", os.path.abspath(args.calib_npz),
|
||||
"--output-har-path", har_opt, venv_bin=venv_bin)
|
||||
check_input_connectivity(har_opt, har_input_layer_names(har_opt))
|
||||
|
||||
print("[3/3] compiler: HAR → HEF")
|
||||
hailo("compiler", har_opt, "--hw-arch", args.hw_arch,
|
||||
"--output-dir", hef_dir, venv_bin=venv_bin)
|
||||
|
||||
hef_src = os.path.join(hef_dir, f"{NET_NAME}.hef")
|
||||
shutil.copy(hef_src, args.output)
|
||||
shutil.copy(har_opt, args.output.replace(".hef", "_optimized.har"))
|
||||
print(f"HEF 已保存: {args.output} ({os.path.getsize(args.output)} bytes)")
|
||||
finally:
|
||||
shutil.rmtree(workdir, ignore_errors=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,263 @@
|
||||
"""R4 离线评估基准:开环滚动预测,量化闭环代理指标
|
||||
|
||||
现状只有开环单步 MSE(teacher forcing),无法反映真实闭环表现。本脚本在留出
|
||||
episode 上复现 pybot 闭环的开环滚动(图像用真值滚动、动作用自预测回写),
|
||||
输出三类闭环代理指标:
|
||||
|
||||
累积误差 逐步 MSE 及随步数的增长曲线(衡量误差累积/发散)
|
||||
平滑度 预测动作相邻差分方差(越低越平滑,对物理舵机越友好),与真值对比
|
||||
物理可行性 预测动作落在训练数据经验范围外的比例(越界率,越高越危险)
|
||||
|
||||
用法:
|
||||
python evaluate.py --model checkpoints/best_inference.pt --data-dir ./data
|
||||
python evaluate.py --model A.pt --compare B.pt --data-dir ./data # A/B 对比
|
||||
python evaluate.py --model m.pt --data-dir ./data --episodes 3 --horizon 50 \
|
||||
--output report.json
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from dataset import ActionDataset
|
||||
from predict import ActionPredictorInference
|
||||
|
||||
|
||||
def _episode_ids(data_dir):
|
||||
return [d.name for d in ActionDataset._valid_episode_dirs(data_dir)]
|
||||
|
||||
|
||||
def _load_episode(data_dir, ep_name):
|
||||
"""加载单个 episode 的图像路径(自然序)、动作真值 (T, D)、姿态真值 (T, imu_dim)
|
||||
|
||||
imu.npy 缺失时第三项返回 None(旧数据兼容,滚动时按零向量喂入)。
|
||||
姿态保持原始量纲——predict 端内部归一化,与 pybot 闭环口径一致。
|
||||
"""
|
||||
ep_dir = [d for d in ActionDataset._valid_episode_dirs(data_dir)
|
||||
if d.name == ep_name][0]
|
||||
from dataset import _natural_key
|
||||
images_dir = ep_dir / "images"
|
||||
image_files = sorted(list(images_dir.glob("*.jpg")) +
|
||||
list(images_dir.glob("*.png")), key=_natural_key)
|
||||
actions = np.load(ep_dir / "actions.npy").astype(np.float32)
|
||||
n = min(len(image_files), len(actions))
|
||||
imus = None
|
||||
imu_path = ep_dir / "imu.npy"
|
||||
if imu_path.exists():
|
||||
imus = np.load(imu_path).astype(np.float32)
|
||||
if imus.ndim != 2:
|
||||
raise ValueError(f"{imu_path} 形状应为 (T, imu_dim),实际为 {imus.shape}")
|
||||
imus = imus[:n]
|
||||
return image_files[:n], actions[:n], imus
|
||||
|
||||
|
||||
def _action_bounds(action_dim):
|
||||
"""固定物理界作为越界判据(与 action_norm.py 常量表共用)
|
||||
|
||||
旧版用数据经验界(min/max+margin),恒值列 span≈0 时任何非零预测都
|
||||
算越界(unnormalized 模型 oob_rate 曾恒 100%)。"越物理界"比"越经验界"
|
||||
更符合安全语义:舵机 [0,180]、spk.F [0,5000]、spk.delay [0,2000]。
|
||||
"""
|
||||
from action_norm import action_bounds
|
||||
return action_bounds(action_dim)
|
||||
|
||||
|
||||
@np.errstate(all="ignore")
|
||||
def rollout_episode(predictor, image_files, actions, horizon=None,
|
||||
lo=None, hi=None, imus=None):
|
||||
"""单 episode 开环滚动,复现闭环口径(图像真值滚动、动作自预测回写)。
|
||||
|
||||
从 t0 = max(history_len, frame_len-1) 起步(窗口完整处),每步:
|
||||
帧窗 = 真值图像 [t-frame_len+1, t](旧→新)
|
||||
历史 = 前 H 步动作(首步用真值填充,之后用自预测回写,复刻 run_loop)
|
||||
姿态 = imus[t](predictor.imu_dim 非 None 时传入;imus 为 None 或
|
||||
长度不足按零向量 = 无姿态先验)
|
||||
预测 â_t → 与真值 a_t 比较 → 回写历史
|
||||
返回逐步记录 dict。horizon=None 表示滚到 episode 末尾。
|
||||
"""
|
||||
F, H = predictor.frame_len, predictor.history_len
|
||||
use_imu = getattr(predictor, "imu_dim", None) is not None
|
||||
t0 = max(H, F - 1)
|
||||
n = len(actions)
|
||||
if horizon is None:
|
||||
horizon = n - t0
|
||||
horizon = min(horizon, n - t0)
|
||||
if horizon <= 0:
|
||||
return None
|
||||
|
||||
cache = {}
|
||||
|
||||
def load(t):
|
||||
if t not in cache:
|
||||
cache[t] = np.array(Image.open(image_files[t]).convert("RGB"))
|
||||
return cache[t]
|
||||
|
||||
history = [actions[t0 - H + k].copy() for k in range(H)] # 真值冷启动
|
||||
steps = []
|
||||
for k in range(horizon):
|
||||
t = t0 + k
|
||||
frames = [load(t - F + 1 + i) for i in range(F)]
|
||||
kwargs = {}
|
||||
if use_imu:
|
||||
kwargs["imu"] = (imus[t] if imus is not None and t < len(imus)
|
||||
else None)
|
||||
pred = np.asarray(
|
||||
predictor.predict(frames, np.array(history), **kwargs),
|
||||
dtype=np.float32).reshape(-1)
|
||||
true = actions[t]
|
||||
err = pred - true
|
||||
rec = {
|
||||
"t": int(t),
|
||||
"mse": float(np.mean(err ** 2)),
|
||||
"pred": pred,
|
||||
"true": true,
|
||||
}
|
||||
if lo is not None:
|
||||
rec["oob"] = bool(np.any((pred < lo) | (pred > hi)))
|
||||
steps.append(rec)
|
||||
history = history[1:] + [pred.copy()] # 自预测回写(复刻 run_loop)
|
||||
return steps
|
||||
|
||||
|
||||
def _smoothness(seq):
|
||||
"""相邻差分方差(多维取均值);序列 <2 返回 0。越低越平滑。"""
|
||||
seq = np.asarray(seq, dtype=np.float32)
|
||||
if len(seq) < 2:
|
||||
return 0.0
|
||||
return float(np.mean(np.var(np.diff(seq, axis=0), axis=0)))
|
||||
|
||||
|
||||
def summarize(all_steps, lo, hi):
|
||||
"""汇总多个 episode 的滚动记录为报告指标
|
||||
|
||||
error_growth 改为按 episode 归一化后取中位数:
|
||||
每 episode 内部算 (末25% MSE / 首25% MSE) → 中位数聚合。
|
||||
避免短 episode 尖峰被长 episode 均值淹没,也更鲁棒于长度差异。
|
||||
"""
|
||||
if not all_steps:
|
||||
raise ValueError("无可评估的滚动步(episode 太短或 horizon 过小)")
|
||||
flat = [s for steps in all_steps for s in steps]
|
||||
mses = np.array([s["mse"] for s in flat], dtype=np.float32)
|
||||
|
||||
# 误差随步数增长曲线:按相对步位置对齐各 episode,取每位置的平均 MSE
|
||||
max_len = max(len(s) for s in all_steps)
|
||||
curve = []
|
||||
for k in range(max_len):
|
||||
vals = [s[k]["mse"] for s in all_steps if k < len(s)]
|
||||
if vals:
|
||||
curve.append(float(np.mean(vals)))
|
||||
|
||||
# 逐 episode 归一化 error_growth:每 episode 自身首末段比值 → 中位数
|
||||
growths = []
|
||||
for steps in all_steps:
|
||||
n = len(steps)
|
||||
if n < 2:
|
||||
continue
|
||||
m = np.array([s["mse"] for s in steps], dtype=np.float32)
|
||||
front = np.mean(m[:max(1, n // 4)])
|
||||
back = np.mean(m[max(0, n * 3 // 4):])
|
||||
growths.append(back / max(front, 1e-12))
|
||||
error_growth = float(np.median(growths)) if growths else 1.0
|
||||
|
||||
preds = np.array([s["pred"] for s in flat], dtype=np.float32)
|
||||
trues = np.array([s["true"] for s in flat], dtype=np.float32)
|
||||
oob = [s.get("oob", False) for s in flat]
|
||||
|
||||
return {
|
||||
"n_episodes": len(all_steps),
|
||||
"n_steps": len(flat),
|
||||
"overall_mse": float(np.mean(mses)),
|
||||
# 末段(后 25% 步)相对首段(前 25%)的误差比,>1 表示误差随步数发散
|
||||
"error_growth": error_growth,
|
||||
"error_curve": curve,
|
||||
"pred_smoothness": _smoothness(preds),
|
||||
"true_smoothness": _smoothness(trues),
|
||||
# 平滑度比:<1 表示预测比真值还平滑(好),>1 表示预测更抖(坏)
|
||||
"smoothness_ratio": _smoothness(preds) / max(_smoothness(trues), 1e-12),
|
||||
"oob_rate": float(np.mean(oob)) if lo is not None else None,
|
||||
}
|
||||
|
||||
|
||||
def evaluate(checkpoint, data_dir, episodes=None, horizon=None, device=None):
|
||||
predictor = ActionPredictorInference(checkpoint, device)
|
||||
ep_names = _episode_ids(data_dir)
|
||||
if episodes is not None:
|
||||
ep_names = ep_names[-episodes:] # 取尾部 episode 作为留出集
|
||||
lo, hi = _action_bounds(predictor.action_dim)
|
||||
all_steps = []
|
||||
for name in ep_names:
|
||||
image_files, actions, imus = _load_episode(data_dir, name)
|
||||
steps = rollout_episode(predictor, image_files, actions, horizon,
|
||||
lo, hi, imus=imus)
|
||||
if steps:
|
||||
all_steps.append(steps)
|
||||
rep = summarize(all_steps, lo, hi)
|
||||
rep["checkpoint"] = checkpoint
|
||||
rep["frame_len"] = predictor.frame_len
|
||||
rep["history_len"] = predictor.history_len
|
||||
rep["imu_dim"] = getattr(predictor, "imu_dim", None)
|
||||
rep["predict_delta"] = getattr(predictor, "_predict_delta", False)
|
||||
return rep
|
||||
|
||||
|
||||
def _fmt(rep):
|
||||
lines = [
|
||||
f"checkpoint : {rep['checkpoint']}",
|
||||
f"predict_delta : {rep['predict_delta']}",
|
||||
f"imu_dim : {rep.get('imu_dim')}",
|
||||
f"episodes/steps : {rep['n_episodes']} / {rep['n_steps']}",
|
||||
f"overall_mse : {rep['overall_mse']:.6f}",
|
||||
f"error_growth : {rep['error_growth']:.3f} (>1 = 误差随步数发散)",
|
||||
f"smoothness_ratio: {rep['smoothness_ratio']:.3f} (<1 = 比真值更平滑)",
|
||||
f" pred/true : {rep['pred_smoothness']:.4f} / {rep['true_smoothness']:.4f}",
|
||||
]
|
||||
if rep["oob_rate"] is not None:
|
||||
lines.append(f"oob_rate : {rep['oob_rate']:.4f} (越界率, 越低越安全)")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description="actpred 离线闭环代理评估(R4)")
|
||||
p.add_argument("--model", required=True, help="主 checkpoint 路径")
|
||||
p.add_argument("--data-dir", required=True, help="episode 数据目录")
|
||||
p.add_argument("--compare", default=None, help="可选对比 checkpoint(A/B)")
|
||||
p.add_argument("--episodes", type=int, default=None,
|
||||
help="评估用尾部 episode 数(默认全部)")
|
||||
p.add_argument("--horizon", type=int, default=None,
|
||||
help="每 episode 滚动步数上限(默认滚到末尾)")
|
||||
p.add_argument("--device", default=None)
|
||||
p.add_argument("--output", default=None, help="保存 JSON 报告路径")
|
||||
args = p.parse_args()
|
||||
|
||||
rep = evaluate(args.model, args.data_dir, args.episodes, args.horizon,
|
||||
args.device)
|
||||
print("== 主模型 ==")
|
||||
print(_fmt(rep))
|
||||
|
||||
result = {"main": rep}
|
||||
if args.compare:
|
||||
rep_b = evaluate(args.compare, args.data_dir, args.episodes,
|
||||
args.horizon, args.device)
|
||||
print("\n== 对比模型 ==")
|
||||
print(_fmt(rep_b))
|
||||
# 关键指标对比(负值=主模型更优)
|
||||
result["compare"] = rep_b
|
||||
result["delta"] = {
|
||||
"overall_mse": rep["overall_mse"] - rep_b["overall_mse"],
|
||||
"error_growth": rep["error_growth"] - rep_b["error_growth"],
|
||||
"smoothness_ratio": rep["smoothness_ratio"] - rep_b["smoothness_ratio"],
|
||||
}
|
||||
d = result["delta"]
|
||||
print("\n== 主 − 对比 (负值=主模型更优) ==")
|
||||
print(f"overall_mse : {d['overall_mse']:+.6f}")
|
||||
print(f"error_growth : {d['error_growth']:+.3f}")
|
||||
print(f"smoothness_ratio: {d['smoothness_ratio']:+.3f}")
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
json.dump(result, f, ensure_ascii=False, indent=2)
|
||||
print(f"\n报告已保存: {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,317 @@
|
||||
"""导出 Hailo 部署版 ONNX:DFC 友好图结构
|
||||
|
||||
训练态 ONNX 对 DFC 不友好:5D 图像输入、AdaptiveAvgPool2d(被
|
||||
torch.export 展开成 Loop 节点)、MultiheadAttention 的 head 拆分
|
||||
Reshape 序列。此脚本逐项改造(权重数值完全不变):
|
||||
1. 图像输入拆为 F 个独立 (1,H,W,3) NHWC 输入,按帧展开 CNN 分支,
|
||||
各分支持有独立权重副本(共享 initializer 会被 DFC 误融合)
|
||||
2. AdaptiveAvgPool2d 换成等价的固定 kernel AvgPool2d
|
||||
3. 自注意力手工展开为每头独立 Linear(128→32),无 head reshape;
|
||||
pre-LN 残差结构与训练态 nn.TransformerEncoderLayer(norm_first=True)
|
||||
逐位等价
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from config import parse_checkpoint_config
|
||||
from model import ActionPredictor
|
||||
|
||||
# 一致性校验阈值:wrapper 改造/ORT 转换与原模型输出差异超过即拦截退出,
|
||||
# 防止 ManualAttnBlock 依赖的 PyTorch 布局约定变更静默穿透(v3-P2-2)
|
||||
MAX_TOLERATED_DIFF = 1e-4
|
||||
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
MAX_TOKEN_ABS = 100.0
|
||||
MERGE_RATIO_WARN = 32.0
|
||||
|
||||
|
||||
class ManualAttnBlock(nn.Module):
|
||||
"""等价替换 nn.TransformerEncoderLayer(d=128, nhead=4, ff=256, norm_first=True)"""
|
||||
|
||||
def __init__(self, layer: nn.TransformerEncoderLayer):
|
||||
super().__init__()
|
||||
d = layer.self_attn.embed_dim
|
||||
nh = layer.self_attn.num_heads
|
||||
self.head_dim = d // nh
|
||||
self.nh = nh
|
||||
self.scale = self.head_dim ** -0.5
|
||||
|
||||
ipw = layer.self_attn.in_proj_weight # (3d, d) = [Wq; Wk; Wv]
|
||||
ipb = layer.self_attn.in_proj_bias
|
||||
self.q_proj = nn.ModuleList()
|
||||
self.k_proj = nn.ModuleList()
|
||||
self.v_proj = nn.ModuleList()
|
||||
for h in range(nh):
|
||||
sl = slice(h * self.head_dim, (h + 1) * self.head_dim)
|
||||
slices = (sl, slice(d + sl.start, d + sl.stop),
|
||||
slice(2 * d + sl.start, 2 * d + sl.stop))
|
||||
for proj, ws in ((self.q_proj, slices[0]), (self.k_proj, slices[1]),
|
||||
(self.v_proj, slices[2])):
|
||||
lin = nn.Linear(d, self.head_dim)
|
||||
lin.weight.data.copy_(ipw[ws])
|
||||
proj.append(lin)
|
||||
for h in range(nh):
|
||||
self.q_proj[h].bias.data.copy_(ipb[h * self.head_dim:(h + 1) * self.head_dim])
|
||||
self.k_proj[h].bias.data.copy_(ipb[d + h * self.head_dim:d + (h + 1) * self.head_dim])
|
||||
self.v_proj[h].bias.data.copy_(ipb[2 * d + h * self.head_dim:2 * d + (h + 1) * self.head_dim])
|
||||
|
||||
self.out_proj = layer.self_attn.out_proj
|
||||
self.norm1 = layer.norm1
|
||||
self.norm2 = layer.norm2
|
||||
self.fc1 = layer.linear1
|
||||
self.fc2 = layer.linear2
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# x: (1, 1, S, d),seq 占 W 维,pre-LN 自注意力 + 残差
|
||||
n = self.norm1(x)
|
||||
heads = []
|
||||
for h in range(self.nh):
|
||||
q = self.q_proj[h](n)
|
||||
k = self.k_proj[h](n)
|
||||
v = self.v_proj[h](n)
|
||||
s = torch.matmul(q, k.transpose(-1, -2)) * self.scale
|
||||
a = torch.softmax(s, dim=-1)
|
||||
heads.append(torch.matmul(a, v))
|
||||
attn = self.out_proj(torch.cat(heads, dim=3)) # feature 维 concat
|
||||
x = x + attn
|
||||
x = x + self.fc2(torch.relu(self.fc1(self.norm2(x))))
|
||||
return x
|
||||
|
||||
|
||||
class DeployWrapper(nn.Module):
|
||||
def __init__(self, model: ActionPredictor):
|
||||
super().__init__()
|
||||
import copy
|
||||
m = model
|
||||
if not hasattr(m.image_encoder, "net"):
|
||||
raise ValueError(
|
||||
"部署版导出仅支持 backbone=simple(MobileNet 的 DFC 改造未实现)")
|
||||
ref_cnn = nn.Sequential(*list(m.image_encoder.net.children())[:-1])
|
||||
self.cnns = nn.ModuleList(
|
||||
[copy.deepcopy(ref_cnn) for _ in range(m._frame_len)])
|
||||
# img_proj ∘ AvgPool == AvgPool ∘ Conv1x1(均为线性算子,可交换)。
|
||||
# 把 img_proj 的权重折进每分支末尾的 1×1 卷积,偏置单独用加法,
|
||||
# 避免 DFC 对 3D Linear 的错误卷积融合
|
||||
self.proj_convs = nn.ModuleList()
|
||||
for _ in range(m._frame_len):
|
||||
c = nn.Conv2d(m.img_proj.in_features, m.img_proj.out_features, 1, bias=False)
|
||||
c.weight.data.copy_(m.img_proj.weight.data.view(
|
||||
m.img_proj.out_features, m.img_proj.in_features, 1, 1))
|
||||
self.proj_convs.append(c)
|
||||
self.proj_bias = nn.Parameter(m.img_proj.bias.data.clone())
|
||||
# 固定 kernel AvgPool 须与训练态 AdaptiveAvgPool2d(1) 逐值等价:
|
||||
# kernel 取 ref_cnn 的真实输出尺寸。H//16 只在分辨率能被 16 整除
|
||||
# 时成立(1080 高度下卷积输出 68 行,按 67 取 kernel 会静默丢弃
|
||||
# 末行,破坏与训练图的数值一致)
|
||||
with torch.no_grad():
|
||||
probe = ref_cnn(torch.zeros(
|
||||
1, 3, m._image_height, m._image_width))
|
||||
self.pool = nn.AvgPool2d((probe.shape[2], probe.shape[3]))
|
||||
# act_proj 权重注册为常量,用显式 MatMul+Add 代替 Linear,
|
||||
# 避免 DFC 对 3D Linear 的错误卷积融合
|
||||
self.act_w = nn.Parameter(m.act_proj.weight.data.t().clone()) # (D, d)
|
||||
self.act_b = nn.Parameter(m.act_proj.bias.data.clone())
|
||||
# R2 方案一:IMU 独立 token(imu_dim>0 时追加图输入,同 MatMul+Add 手法)
|
||||
self._imu_dim = m._imu_dim
|
||||
if self._imu_dim > 0:
|
||||
self.imu_w = nn.Parameter(m.imu_proj.weight.data.t().clone())
|
||||
self.imu_b = nn.Parameter(m.imu_proj.bias.data.clone())
|
||||
# pred_token 改为图输入(常量 initializer 会被 DFC 误转置),
|
||||
# 位置编码经 concat 后的 ew_add 广播可被正确推断
|
||||
self.pos_embed4 = nn.Parameter(m.pos_embed.data.unsqueeze(1).clone()) # (1, 1, S, d)
|
||||
self.blocks = nn.ModuleList(
|
||||
[ManualAttnBlock(l) for l in model.encoder.layers])
|
||||
self.head = m.head
|
||||
self._frame_len = m._frame_len
|
||||
|
||||
def _frame_feat(self, frame_nhwc: torch.Tensor, i: int) -> torch.Tensor:
|
||||
x = self.cnns[i](frame_nhwc.permute(0, 3, 1, 2))
|
||||
x = self.proj_convs[i](x) # 1x1 conv ≡ 池化后的 img_proj
|
||||
return self.pool(x) # (1, d, 1, 1) NCHW
|
||||
|
||||
def forward(self, *inputs: torch.Tensor) -> torch.Tensor:
|
||||
# inputs = frame0..frame{F-1}, action_history, pred_token, [imu]
|
||||
# 全程 (1, 1, S, d) 布局:seq 占 W 维,避免任何 transpose
|
||||
frames = list(inputs[:self._frame_len])
|
||||
action_history = inputs[self._frame_len]
|
||||
pred_token = inputs[self._frame_len + 1]
|
||||
imu = inputs[self._frame_len + 2] if self._imu_dim > 0 else None
|
||||
feats = [self._frame_feat(frames[i], i).permute(0, 2, 3, 1)
|
||||
for i in range(self._frame_len)] # (1, 1, 1, d) NHWC
|
||||
img_tokens = torch.cat(feats, dim=2) + self.proj_bias # (1, 1, F, d)
|
||||
act_tokens = torch.matmul(action_history, self.act_w) + self.act_b # (1, 1, H, d)
|
||||
parts = [img_tokens, act_tokens]
|
||||
if self._imu_dim > 0:
|
||||
# imu: (1, 1, 1, imu_dim) 归一化姿态 → (1, 1, 1, d),位于 PRED 之前
|
||||
parts.append(torch.matmul(imu, self.imu_w) + self.imu_b)
|
||||
parts.append(pred_token)
|
||||
tokens = torch.cat(parts, dim=2) # (1, 1, S, d)
|
||||
tokens = tokens + self.pos_embed4
|
||||
for blk in self.blocks:
|
||||
tokens = blk(tokens)
|
||||
return self.head(tokens[:, :, -1:, :]).squeeze(1) # (1, 1, action_dim)
|
||||
|
||||
|
||||
def check_token_ranges(wrapper: "DeployWrapper", frame_len: int, height: int,
|
||||
width: int, history_len: int, action_dim: int,
|
||||
pred_token: torch.Tensor, max_abs: float = MAX_TOKEN_ABS,
|
||||
action_stats_path: str = None,
|
||||
normalize_mode: str = "zscore"):
|
||||
"""导出前检查 token 量级:幅度过大时 uint8 量化会把小信号压成 0
|
||||
|
||||
预训练用全零图像占位 → BN running_var 停在初始小值 → 推理时每级
|
||||
BN 放大 ~50 倍,图像 token 可达 ±2e5。DFC 在 token merge 点按最大
|
||||
幅度分配共享量化 scale(实测 scale=1965.92,pred_token/act_tokens
|
||||
全部量化成 0),HEF 输出对任何输入恒定。必须在编译前拦截。
|
||||
|
||||
归一化模型(normalize_actions)的 HEF 输入是归一化后的历史,
|
||||
检查输入须与推理端一致(zscore → z 空间采样;fixed_range /
|
||||
log_range → [0,1])。
|
||||
"""
|
||||
rng = np.random.default_rng(42)
|
||||
raw = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
|
||||
norm = (raw.astype(np.float32) / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
|
||||
frames = [torch.tensor(norm[None]) for _ in range(frame_len)]
|
||||
history = rng.uniform(0.0, 500.0, size=(1, 1, history_len, action_dim))
|
||||
from action_norm import is_range_mode
|
||||
if is_range_mode(normalize_mode):
|
||||
# fixed_range / log_range 归一化空间即 [0,1],直接在归一化空间采样
|
||||
history = rng.uniform(0.0, 1.0, size=(1, 1, history_len, action_dim))
|
||||
elif action_stats_path is not None:
|
||||
stats = np.load(action_stats_path)
|
||||
if stats.shape != (2, action_dim):
|
||||
raise ValueError(
|
||||
f"action stats 形状应为 (2, {action_dim}),"
|
||||
f"实际为 {stats.shape}: {action_stats_path}")
|
||||
# z-score 归一化模型直接在 z 空间采样:raw 均匀采样会在零方差列上
|
||||
# 产生天文数字 z 值(实际推理中这些列恒为 0,z 恒为 0)
|
||||
history = rng.normal(0.0, 1.5, size=(1, 1, history_len, action_dim))
|
||||
history = torch.tensor(history, dtype=torch.float32)
|
||||
with torch.no_grad():
|
||||
feats = [wrapper._frame_feat(f, i).permute(0, 2, 3, 1)
|
||||
for i, f in enumerate(frames)]
|
||||
img_tokens = torch.cat(feats, dim=2) + wrapper.proj_bias
|
||||
act_tokens = torch.matmul(history, wrapper.act_w) + wrapper.act_b
|
||||
img_max = img_tokens.abs().max().item()
|
||||
act_max = act_tokens.abs().max().item()
|
||||
pred_max = pred_token.abs().max().item()
|
||||
print(f"token 量级检查: |img|={img_max:.4g} |act|={act_max:.4g} "
|
||||
f"|pred|={pred_max:.4g}(图像上限 {max_abs})")
|
||||
if img_max > max_abs:
|
||||
raise SystemExit(
|
||||
f"错误: 图像 token 幅度 {img_max:.3g} 超过可量化上限 {max_abs}。"
|
||||
f"checkpoint 的 BN 统计退化(预训练用全零图像占位,running_var "
|
||||
f"停在初始小值),uint8 量化会把小信号全部压成 0、HEF 输出恒定。"
|
||||
f"请先用真实/随机图像数据微调后再导出。")
|
||||
if act_max > MERGE_RATIO_WARN * max(img_max, 1.0):
|
||||
print(f"WARNING: act_tokens 幅度 {act_max:.3g} 远大于 img_tokens "
|
||||
f"{img_max:.3g},merge 点共享 scale 会粗量化图像路径;"
|
||||
f"建议用 --normalize-actions 训练(动作归一化)压小 act_tokens")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Export Hailo-friendly ONNX")
|
||||
parser.add_argument("--model", required=True, help="checkpoint path")
|
||||
parser.add_argument("--output", default="checkpoints/actpred_deploy.onnx")
|
||||
parser.add_argument("--pred-token-out", default=None,
|
||||
help="保存学习到的 pred_token 到 .npy(HEF 推理端喂入用,"
|
||||
"配合抖动校准避免常量量化退化)")
|
||||
args = parser.parse_args()
|
||||
|
||||
checkpoint = torch.load(args.model, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(checkpoint["config"])
|
||||
|
||||
model = ActionPredictor(cfg.model, cfg.data, pretrained=False)
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
model.eval()
|
||||
wrapper = DeployWrapper(model).eval()
|
||||
|
||||
h, w = cfg.data.image_height, cfg.data.image_width
|
||||
f = cfg.data.frame_len
|
||||
stats_path = None
|
||||
from action_norm import is_range_mode
|
||||
if cfg.data.normalize_actions and not is_range_mode(cfg.data.normalize_mode):
|
||||
# 与 predict.py 语义一致:相对路径按 checkpoint 目录解析
|
||||
# (fixed_range / log_range 模式免 stats 文件,跳过)
|
||||
stats_path = cfg.data.action_stats_path
|
||||
if not stats_path:
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)),
|
||||
"action_stats.npy")
|
||||
elif not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)), stats_path)
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(
|
||||
f"模型训练时启用了动作归一化,但未找到统计文件: {stats_path}")
|
||||
check_token_ranges(wrapper, f, h, w, cfg.data.history_len,
|
||||
cfg.data.action_dim, model.pred_token.data.unsqueeze(1),
|
||||
action_stats_path=stats_path,
|
||||
normalize_mode=cfg.data.normalize_mode)
|
||||
dummy_frames = [torch.randn(1, h, w, 3) for _ in range(f)]
|
||||
dummy_history = torch.randn(1, 1, cfg.data.history_len, cfg.data.action_dim)
|
||||
dummy_pred = model.pred_token.data.unsqueeze(1).clone() # (1, 1, 1, d)
|
||||
imu_dim = cfg.data.imu_dim
|
||||
# imu_dim>0 时追加归一化姿态输入(均匀采样的归一化值域)
|
||||
dummy_imu = torch.randn(1, 1, 1, imu_dim) if imu_dim > 0 else None
|
||||
|
||||
extra_inputs = (dummy_imu,) if imu_dim > 0 else ()
|
||||
extra_names = ["imu"] if imu_dim > 0 else []
|
||||
|
||||
with torch.no_grad():
|
||||
ref = wrapper(*dummy_frames, dummy_history, dummy_pred, *extra_inputs)
|
||||
|
||||
torch.onnx.export(
|
||||
wrapper, (*dummy_frames, dummy_history, dummy_pred, *extra_inputs),
|
||||
args.output,
|
||||
input_names=[f"frame{i}" for i in range(f)] + ["action_history", "pred_token"] + extra_names,
|
||||
output_names=["predicted_action"],
|
||||
opset_version=18,
|
||||
)
|
||||
print(f"部署版 ONNX 已导出: {args.output}({f} 个帧输入 (1,{h},{w},3) NHWC + action_history + pred_token"
|
||||
f"{' + imu' if imu_dim > 0 else ''})")
|
||||
|
||||
if args.pred_token_out:
|
||||
import numpy as np
|
||||
np.save(args.pred_token_out, model.pred_token.data.unsqueeze(1).numpy())
|
||||
print(f"pred_token 已保存: {args.pred_token_out} (1,1,1,{model.pred_token.shape[-1]})")
|
||||
|
||||
stacked = torch.stack([fr.permute(0, 3, 1, 2) for fr in dummy_frames], dim=1)
|
||||
with torch.no_grad():
|
||||
orig_out = model(stacked, dummy_history.squeeze(1),
|
||||
dummy_imu.squeeze(1).squeeze(1) if imu_dim > 0 else None)
|
||||
wrap_diff = (ref.squeeze(1) - orig_out).abs().max().item()
|
||||
print(f"wrapper 与原模型最大差异: {wrap_diff:.6e}")
|
||||
if wrap_diff > MAX_TOLERATED_DIFF:
|
||||
sys.exit(f"错误: wrapper 与原模型差异 {wrap_diff:.3e} > "
|
||||
f"{MAX_TOLERATED_DIFF:.0e}(ManualAttnBlock 布局约定可能已失效)")
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
import numpy as np
|
||||
sess = ort.InferenceSession(args.output)
|
||||
feeds = {f"frame{i}": dummy_frames[i].numpy() for i in range(f)}
|
||||
feeds["action_history"] = dummy_history.numpy()
|
||||
feeds["pred_token"] = dummy_pred.numpy()
|
||||
if imu_dim > 0:
|
||||
feeds["imu"] = dummy_imu.numpy()
|
||||
ort_out = sess.run(None, feeds)[0]
|
||||
ort_diff = np.abs(ort_out - ref.numpy()).max()
|
||||
if ort_diff > MAX_TOLERATED_DIFF:
|
||||
sys.exit(f"错误: ONNX 与 wrapper 差异 {ort_diff:.3e} > "
|
||||
f"{MAX_TOLERATED_DIFF:.0e}")
|
||||
print(f"ONNX 一致性验证通过(最大差异: {ort_diff:.6e})")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+50
-14
@@ -1,8 +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")
|
||||
@@ -19,34 +25,64 @@ def main():
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
model.eval()
|
||||
|
||||
dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size)
|
||||
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_image, dummy_history),
|
||||
model, (dummy_images, dummy_history, *extra_inputs),
|
||||
args.output,
|
||||
input_names=["image", "action_history"],
|
||||
input_names=["images", "action_history"] + extra_names,
|
||||
output_names=["predicted_action"],
|
||||
dynamic_axes={
|
||||
"image": {0: "batch"},
|
||||
"action_history": {0: "batch"},
|
||||
"predicted_action": {0: "batch"},
|
||||
},
|
||||
opset_version=17,
|
||||
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_image.numpy()
|
||||
hist_np = dummy_history.numpy()
|
||||
ort_out = sess.run(None, {"image": 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_image, dummy_history).numpy()
|
||||
torch_out = model(dummy_images, dummy_history, *extra_inputs).numpy()
|
||||
max_diff = np.abs(ort_out - torch_out).max()
|
||||
if max_diff > MAX_TOLERATED_DIFF:
|
||||
sys.exit(f"错误: ONNX 与 PyTorch 输出差异 {max_diff:.3e} > "
|
||||
f"{MAX_TOLERATED_DIFF:.0e}")
|
||||
print(f"ONNX 一致性验证通过(最大差异: {max_diff:.6e})")
|
||||
except ImportError:
|
||||
print("提示: 安装 onnxruntime 可验证导出模型一致性 (pip install onnxruntime)")
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
"""生成 DFC 量化校准数据集
|
||||
|
||||
两种模式:
|
||||
默认模式(data/calibration.npz):为旧版 dfc_compile.py 的 ClientModel
|
||||
流程准备,目前保留兼容。
|
||||
--deploy 模式(data/calibration_deploy.npz):为 export_deploy.py 导出
|
||||
的部署版 ONNX 生成校准集,键名匹配 HN 层名 actpred_deploy/input_layerN,
|
||||
形状匹配 HEF 输入:frame0~2 (N,H,W,3) 已归一化 float32、
|
||||
action_history (N,1,H_len,D)、pred_token (N,1,1,d) 学习值 ± 抖动
|
||||
(常量校准会使量化范围退化,见 --pred-jitter)。
|
||||
图像分辨率与 --model checkpoint 的配置一致。
|
||||
|
||||
数据源:
|
||||
- 图像:无真实 episode 时用 [0,255] 均匀采样像素 + ImageNet 归一化
|
||||
- 动作历史:从 bot4cj/scripts/*.mt 的真实 @Action 矩阵滑窗采样
|
||||
(列不足 18 补零,与 pretrain_mt.py 语义一致)
|
||||
|
||||
真实数据到位后应替换为真实帧重新校准。
|
||||
|
||||
用法:
|
||||
python gen_calibration.py --deploy --model checkpoints/pretrain_best_inference.pt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
|
||||
from config import DataConfig, parse_checkpoint_config
|
||||
from mt_parser import parse_mt_directory
|
||||
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
|
||||
def build_action_history_samples(mt_dir, n_samples, history_len, action_dim, rng):
|
||||
sequences = []
|
||||
for _path, matrix in parse_mt_directory(mt_dir):
|
||||
n_cols = matrix.shape[1]
|
||||
if n_cols < action_dim:
|
||||
matrix = np.pad(matrix, ((0, 0), (0, action_dim - n_cols)))
|
||||
elif n_cols > action_dim:
|
||||
matrix = matrix[:, :action_dim]
|
||||
if len(matrix) >= history_len:
|
||||
sequences.append(matrix)
|
||||
if not sequences:
|
||||
raise ValueError(f"{mt_dir} 中没有可用的 @Action 序列")
|
||||
|
||||
windows = []
|
||||
for matrix in sequences:
|
||||
for i in range(len(matrix) - history_len + 1):
|
||||
windows.append(matrix[i:i + history_len])
|
||||
idx = rng.integers(0, len(windows), size=n_samples)
|
||||
return np.stack([windows[i] for i in idx]).astype(np.float32)
|
||||
|
||||
|
||||
def build_image_samples(n_samples, frame_len, height, width, rng,
|
||||
chunk_samples: int = 4):
|
||||
"""生成 (n_samples, frame_len, 3, H, W) uint8 校准图像
|
||||
|
||||
存 uint8 原始像素值(0~255),加载时再做 ToTensor+Normalize,节省磁盘。
|
||||
经 memmap 分块填充后返回:高分辨率(如 1080p)下整块数组可达 ~600MB,
|
||||
一次性分配容易挤爆内存,分块写盘把峰值压到 chunk_samples 帧量级。
|
||||
"""
|
||||
if chunk_samples < 1:
|
||||
raise ValueError(f"chunk_samples 应 >= 1,实际为 {chunk_samples}")
|
||||
chunk_samples = min(chunk_samples, n_samples)
|
||||
fd, tmp_path = tempfile.mkstemp(suffix=".npy")
|
||||
os.close(fd)
|
||||
images = np.lib.format.open_memmap(
|
||||
tmp_path, mode="w+", dtype=np.uint8,
|
||||
shape=(n_samples, frame_len, 3, height, width))
|
||||
try:
|
||||
for start in range(0, n_samples, chunk_samples):
|
||||
stop = min(start + chunk_samples, n_samples)
|
||||
images[start:stop] = rng.integers(
|
||||
0, 256, size=(stop - start, frame_len, 3, height, width),
|
||||
dtype=np.uint8)
|
||||
except Exception:
|
||||
images._mmap.close()
|
||||
os.unlink(images.filename)
|
||||
raise
|
||||
return images
|
||||
|
||||
|
||||
def build_deploy_set(args):
|
||||
import torch
|
||||
from action_norm import is_range_mode, normalize_action
|
||||
from model import ActionPredictor
|
||||
|
||||
ckpt = torch.load(args.model, map_location="cpu", weights_only=True)
|
||||
cfg = parse_checkpoint_config(ckpt["config"]).data
|
||||
rng = np.random.default_rng(args.seed)
|
||||
|
||||
h, w = cfg.image_height, cfg.image_width
|
||||
raw = rng.integers(0, 256, size=(args.samples, h, w, 3),
|
||||
dtype=np.uint8).astype(np.float32) / 255.0
|
||||
norm = ((raw - IMAGENET_MEAN) / IMAGENET_STD).astype(np.float32)
|
||||
frames = [norm.copy() for _ in range(cfg.frame_len)]
|
||||
|
||||
history = build_action_history_samples(
|
||||
args.mt_dir, args.samples, cfg.history_len, cfg.action_dim, rng)
|
||||
history = history[:, None, :, :] # (N, 1, H_len, D)
|
||||
if cfg.normalize_actions and is_range_mode(cfg.normalize_mode):
|
||||
# 固定范围类归一化:校准数据与推理端同一变换(免 stats 文件)
|
||||
from action_norm import normalize_action, is_range_mode
|
||||
history = normalize_action(history, cfg.action_dim,
|
||||
mode=cfg.normalize_mode)
|
||||
print(f" action_history: 已 {cfg.normalize_mode} 归一化"
|
||||
f"(固定物理范围)")
|
||||
elif cfg.normalize_actions:
|
||||
# z-score 归一化模型的 HEF 输入是 z-score 历史(hailo_infer --stats 同语义),
|
||||
# 校准数据必须走同一变换,否则量化范围与推理分布不匹配
|
||||
stats_path = cfg.action_stats_path
|
||||
if not stats_path:
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)),
|
||||
"action_stats.npy")
|
||||
elif not os.path.isabs(stats_path):
|
||||
stats_path = os.path.join(
|
||||
os.path.dirname(os.path.abspath(args.model)), stats_path)
|
||||
if not os.path.exists(stats_path):
|
||||
raise FileNotFoundError(
|
||||
f"模型训练时启用了动作归一化,但未找到统计文件: {stats_path}")
|
||||
stats = np.load(stats_path)
|
||||
if stats.shape != (2, cfg.action_dim):
|
||||
raise ValueError(
|
||||
f"action stats 形状应为 (2, {cfg.action_dim}),"
|
||||
f"实际为 {stats.shape}: {stats_path}")
|
||||
history = (history - stats[0]) / stats[1]
|
||||
print(f" action_history: 已 z-score 归一化({stats_path})")
|
||||
|
||||
model = ActionPredictor(parse_checkpoint_config(ckpt["config"]).model,
|
||||
cfg, pretrained=False)
|
||||
model.load_state_dict(ckpt["model_state_dict"])
|
||||
d = model.pred_token.shape[-1]
|
||||
v = model.pred_token.data.unsqueeze(1).numpy() # (1, 1, d) 学习到的常量
|
||||
# 常量校准会使量化范围退化(上次实测 scale=18571.7,任何输入都被
|
||||
# 量化成 0);按学习值 ±相对抖动生成校准样本,运行时喂学习值即可
|
||||
# 无损还原(量化误差 ~1e-4)
|
||||
jitter = v * rng.uniform(-args.pred_jitter, args.pred_jitter,
|
||||
size=(args.samples, 1, 1, d)).astype(np.float32)
|
||||
jitter = jitter + rng.normal(0.0, args.pred_jitter_abs,
|
||||
size=jitter.shape).astype(np.float32)
|
||||
pred = np.broadcast_to(v, (args.samples, 1, 1, d)).astype(np.float32) + jitter
|
||||
|
||||
inputs = [*frames, history, pred]
|
||||
if cfg.imu_dim > 0:
|
||||
# R2 方案一:imu 输入为归一化姿态(roll,pitch/180,ax..gz/32768)。
|
||||
# 校准覆盖典型工作范围:归一化值域 ±0.5(≈±90°/±16384)均匀采样
|
||||
imu = rng.uniform(-0.5, 0.5,
|
||||
size=(args.samples, 1, 1, cfg.imu_dim)).astype(np.float32)
|
||||
inputs.append(imu)
|
||||
|
||||
np.savez_compressed(
|
||||
args.output,
|
||||
**{f"actpred_deploy/input_layer{i + 1}": arr
|
||||
for i, arr in enumerate(inputs)})
|
||||
print(f"部署校准集已保存: {args.output}")
|
||||
print(f" frame0~2: {frames[0].shape}(分辨率 {w}x{h},与 checkpoint 一致)")
|
||||
print(f" action_history: {history.shape} pred_token: {pred.shape}"
|
||||
f"(学习值 ±{args.pred_jitter:.0%} 抖动)")
|
||||
if cfg.imu_dim > 0:
|
||||
print(f" imu: {inputs[-1].shape}(归一化值域 ±0.5 均匀采样)")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate DFC calibration dataset")
|
||||
parser.add_argument("--deploy", action="store_true",
|
||||
help="生成部署版(export_deploy.py)校准集")
|
||||
parser.add_argument("--model", default="checkpoints/pretrain_best_inference.pt",
|
||||
help="--deploy 模式的 checkpoint(决定分辨率与 pred_token)")
|
||||
parser.add_argument("--mt-dir", default="../bot4cj/scripts")
|
||||
parser.add_argument("--output", default=None)
|
||||
parser.add_argument("--samples", type=int, default=32)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--pred-jitter", type=float, default=0.3,
|
||||
help="pred_token 校准相对抖动幅度(默认 ±30%)")
|
||||
parser.add_argument("--pred-jitter-abs", type=float, default=0.005,
|
||||
help="pred_token 校准绝对抖动标准差(近零元素保底)")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.deploy:
|
||||
args.output = args.output or "data/calibration_deploy.npz"
|
||||
build_deploy_set(args)
|
||||
return
|
||||
|
||||
args.output = args.output or "data/calibration.npz"
|
||||
cfg = DataConfig()
|
||||
rng = np.random.default_rng(args.seed)
|
||||
|
||||
history = build_action_history_samples(
|
||||
args.mt_dir, args.samples, cfg.history_len, cfg.action_dim, rng)
|
||||
images = build_image_samples(
|
||||
args.samples, cfg.frame_len, cfg.image_height, cfg.image_width, rng)
|
||||
|
||||
np.savez_compressed(args.output, images=images, action_history=history)
|
||||
# 释放 memmap 并删除临时文件(数据已写入 npz)
|
||||
images._mmap.close()
|
||||
os.unlink(images.filename)
|
||||
print(f"校准集已保存: {args.output}")
|
||||
print(f" images: {images.shape} dtype={images.dtype}")
|
||||
print(f" action_history: {history.shape} dtype={history.dtype}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,626 @@
|
||||
"""actpred HEF(Hailo)推理封装 —— AI HAT 2 / Hailo-10H 部署后端
|
||||
|
||||
接口对齐 predict.ActionPredictorInference:predict(images, action_history)。
|
||||
|
||||
HEF 由 actpred ONNX 经 Hailo DFC 编译生成(量化校准后)。本封装不硬编码
|
||||
vstream 名称/形状/量化参数,全部从 HEF 的 input/output vstream info 动态推导,
|
||||
支持两种布局:
|
||||
|
||||
新版 5/6 输入(export_deploy.py + dfc_compile.py 产物):
|
||||
frame0~2 各 (H,W,3) NHWC uint8(ImageNet 归一化校准,负零偏自动识别)
|
||||
+ action_history (1,H_len,D) + pred_token (1,1,d)
|
||||
+ imu (1,1,1,imu_dim)(R2 方案一可选输入,原始姿态内部归一化)
|
||||
旧版单图像输入:
|
||||
(1, F, 3, H, W) → NHWC uint8 (1, F, H, W, 3)
|
||||
或 (1, H, W, F*3) → 帧沿通道拼接;action_history (1, H, D)
|
||||
|
||||
uint8 边缘按 quant_info 量化/反量化;pred_token 输入默认喂零,
|
||||
新编译 HEF(抖动校准,见 gen_calibration.py --pred-jitter)应喂
|
||||
checkpoint 的学习值(--pred-token npy 或 predict(pred_token=...)),
|
||||
否则该 token 语义与训练时不一致。
|
||||
|
||||
背景与踩坑记录见
|
||||
docs/reports/actpred_hailo_hat2_inference_research_20260813.md:
|
||||
- hailort 5.x 对 Hailo-10 仅支持新 API(create_infer_model),
|
||||
旧版 VDevice.configure(hef) 报 HAILO_NOT_IMPLEMENTED;
|
||||
- 绑定缓冲区 dtype/大小必须匹配 HEF 边缘格式(量化 uint8)。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
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)
|
||||
|
||||
# 异步推理缓冲数:≥4 时流水线吞吐趋于饱和(实测 4 缓冲 ~2.7×、6 缓冲 ~2.9×)
|
||||
_ASYNC_BUFFERS = 4
|
||||
|
||||
|
||||
class _AsyncJob:
|
||||
"""predict_async 的推理句柄:result() 阻塞返回 (action_dim,) 动作向量"""
|
||||
|
||||
def __init__(self):
|
||||
self._done = threading.Event()
|
||||
self._result = None
|
||||
self._error = None
|
||||
|
||||
def result(self, timeout: float = None) -> np.ndarray:
|
||||
if not self._done.wait(timeout):
|
||||
raise TimeoutError("async HEF inference timeout")
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return self._result
|
||||
|
||||
def _complete(self, result=None, error=None):
|
||||
self._result = result
|
||||
self._error = error
|
||||
self._done.set()
|
||||
|
||||
|
||||
def _input_layer_index(name: str) -> int:
|
||||
m = re.search(r"input_layer(\d+)$", name)
|
||||
return int(m.group(1)) if m else 0
|
||||
|
||||
|
||||
class HailoPredictor:
|
||||
"""HEF 动作预测推理器(Hailo NPU 后端)
|
||||
|
||||
支持新版 5 输入 HEF(逐帧输入 + pred_token)与旧版单图像输入 HEF,
|
||||
布局与量化参数从 vstream info 自动推导。
|
||||
|
||||
用法:
|
||||
with HailoPredictor("actpred.hef") as predictor:
|
||||
action = predictor.predict(images, action_history)
|
||||
"""
|
||||
|
||||
def __init__(self, hef_path: str, stats_path: str = None, vdevice=None,
|
||||
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())
|
||||
self._vdevice = vdevice
|
||||
self._infer_model = self._vdevice.create_infer_model(str(hef_path))
|
||||
self._configured = self._stack.enter_context(
|
||||
self._infer_model.configure())
|
||||
try:
|
||||
hef = HEF(str(hef_path))
|
||||
self._in_infos = hef.get_input_vstream_infos()
|
||||
self._out_infos = hef.get_output_vstream_infos()
|
||||
if len(self._out_infos) != 1:
|
||||
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._imu_info = self._classify_inputs()
|
||||
self._derive_dims()
|
||||
|
||||
self._action_mean = None
|
||||
self._action_std = 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)
|
||||
if stats.shape != (2, self.action_dim):
|
||||
raise ValueError(
|
||||
f"action_stats.npy 形状应为 (2, {self.action_dim}),"
|
||||
f"实际为 {stats.shape}: {stats_path}")
|
||||
self._action_mean = stats[0].astype(np.float32)
|
||||
self._action_std = stats[1].astype(np.float32)
|
||||
except Exception:
|
||||
self._stack.close() # 初始化失败也释放设备,避免 fd 泄漏
|
||||
raise
|
||||
|
||||
# ---------------- 输入分类与形状推导 ----------------
|
||||
|
||||
def _classify_inputs(self):
|
||||
"""将 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)
|
||||
# 新版:每帧独立 (H,W,3) 输入
|
||||
frames = [i for i in ins if len(i.shape) == 3
|
||||
and i.shape[2] == 3 and i.shape[0] > 3]
|
||||
if frames:
|
||||
frames.sort(key=lambda i: _input_layer_index(i.name))
|
||||
else:
|
||||
# 旧版:单一 4/5 维图像输入
|
||||
frames = [i for i in ins if len(i.shape) in (4, 5)]
|
||||
rest = [i for i in ins if i not in frames]
|
||||
history, pred, imu = None, None, None
|
||||
for i in rest:
|
||||
s = tuple(i.shape)
|
||||
if len(s) == 3 and s[1] == 1:
|
||||
pred = i # pred_token (1,1,d)
|
||||
elif len(s) == 3:
|
||||
history = i # 动作历史 (1,H_len,D)
|
||||
elif len(s) == 4 and s[0] == 1 and s[1] == 1 and s[2] == 1:
|
||||
imu = i # R2:imu (1,1,1,imu_dim)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"无法识别 HEF 输入 {i.name} 形状 {s}(期望动作历史 "
|
||||
f"(1,H_len,D)、pred_token (1,1,d) 或 imu (1,1,1,imu_dim))")
|
||||
if not frames:
|
||||
raise ValueError("HEF 输入中未找到图像输入")
|
||||
if history is None:
|
||||
raise ValueError(
|
||||
f"HEF 输入中未找到动作历史输入(输出维度 "
|
||||
f"{tuple(self._out_info.shape)}),"
|
||||
f"实际: {[(i.name, tuple(i.shape)) for i in ins]}")
|
||||
return frames, history, pred, imu
|
||||
|
||||
def _derive_dims(self):
|
||||
# 帧输入布局
|
||||
if len(self._frame_infos) == 1 and len(self._frame_infos[0].shape) in (4, 5):
|
||||
shape = tuple(self._frame_infos[0].shape)
|
||||
if len(shape) == 5:
|
||||
_, self.frame_len, c, self.image_height, self.image_width = shape
|
||||
if c != 3:
|
||||
raise ValueError(f"图像输入通道维应为 3,实际形状 {shape}")
|
||||
self._image_layout = "ncfhw"
|
||||
else: # (1, H, W, F*3) NHWC 帧拼通道
|
||||
_, self.image_height, self.image_width, c = shape
|
||||
if c % 3 != 0 or c == 3:
|
||||
raise ValueError(
|
||||
f"无法从 4 维图像输入推导帧数(期望 (1,H,W,F*3)),"
|
||||
f"实际形状 {shape}")
|
||||
self.frame_len = c // 3
|
||||
self._image_layout = "nhwc_fused"
|
||||
else:
|
||||
# 新版:每帧独立 (H,W,3) 输入,形状须一致
|
||||
shapes = [tuple(i.shape) for i in self._frame_infos]
|
||||
if any(len(s) != 3 or s[2] != 3
|
||||
or s[0] != shapes[0][0] or s[1] != shapes[0][1]
|
||||
for s in shapes):
|
||||
raise ValueError(f"帧输入形状不一致: {shapes}")
|
||||
self.image_height, self.image_width, _ = shapes[0]
|
||||
self.frame_len = len(self._frame_infos)
|
||||
self._image_layout = "nhwc_frames"
|
||||
self._frame_quant = self._frame_infos[0].quant_info
|
||||
# 帧输入是否需 ImageNet 归一化,由量化网格判定:DFC 的 zp 恒非负
|
||||
# (zp = round(-min/scale)),码 0 反量化值 dequant(0) = -zp*scale。
|
||||
# 原始像素输入(min=0, zp=0)下为 0;归一化输入(min≈-2.1)下
|
||||
# ≈ -2.1(实测 HEF zp=114, scale=0.0187)。旧启发式 zp*scale<0
|
||||
# 把归一化输入误判为原始像素直通,图像路径被饱和破坏。
|
||||
self._normalize_frames = \
|
||||
(0 - self._frame_quant.qp_zp) * self._frame_quant.qp_scale < -0.5
|
||||
|
||||
# cv2.LUT 查表(归一化+量化合并一步):仅逐帧独立输入布局且
|
||||
# 帧边缘为 uint8 时可用;无 cv2 或旧布局回退 float 路径
|
||||
from hailo_platform import FormatType
|
||||
self._frame_luts = None
|
||||
if HAS_CV2 and self._image_layout == "nhwc_frames" and \
|
||||
all(f.format.type == FormatType.UINT8
|
||||
for f in self._frame_infos):
|
||||
self._frame_luts = self._build_frame_luts()
|
||||
|
||||
# 异步推理缓冲池(predict_async 复用)
|
||||
self._async_buffers = [self._configured.create_bindings()
|
||||
for _ in range(_ASYNC_BUFFERS)]
|
||||
self._async_flags = [[True] for _ in range(_ASYNC_BUFFERS)]
|
||||
|
||||
h_shape = tuple(self._history_info.shape)
|
||||
self.history_len, self.action_dim = h_shape[1], h_shape[2]
|
||||
|
||||
# R2:imu 输入存在时暴露 imu_dim(pybot 闭环鸭子类型据此传 imu=),
|
||||
# 无则 None(旧 HEF 保持旧两参调用路径)
|
||||
self.imu_dim = tuple(self._imu_info.shape)[-1] \
|
||||
if self._imu_info is not None else None
|
||||
|
||||
out_shape = tuple(self._out_info.shape)
|
||||
out_dim = out_shape[-1]
|
||||
if out_dim == 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} 不一致")
|
||||
|
||||
# ---------------- 生命周期 ----------------
|
||||
|
||||
def close(self):
|
||||
self._stack.close()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
self.close()
|
||||
return False
|
||||
|
||||
# ---------------- 推理 ----------------
|
||||
|
||||
def _prepare_inputs(self, images, action_history):
|
||||
"""输入校验 + 动作历史归一化(predict/predict_async 共用)"""
|
||||
if not isinstance(images, (list, tuple)):
|
||||
images = [images]
|
||||
if len(images) != self.frame_len:
|
||||
raise ValueError(
|
||||
f"images 应为 {self.frame_len} 帧,实际为 {len(images)} 帧。"
|
||||
"冷启动时可重复传入同一帧。")
|
||||
action_history = np.asarray(action_history, dtype=np.float32)
|
||||
expected = (self.history_len, self.action_dim)
|
||||
if action_history.shape != expected:
|
||||
raise ValueError(
|
||||
f"action_history 形状应为 {expected},实际为 {action_history.shape}")
|
||||
|
||||
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
|
||||
|
||||
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:
|
||||
q = self._history_info.quant_info
|
||||
buf = np.clip(np.round(action_history / q.qp_scale + q.qp_zp),
|
||||
0, 255).astype(np.uint8).reshape(h_shape)
|
||||
else:
|
||||
buf = action_history.astype(np.float32).reshape(h_shape)
|
||||
bindings.input(self._history_info.name).set_buffer(buf)
|
||||
|
||||
if self._pred_info is not None:
|
||||
p_shape = tuple(self._pred_info.shape)
|
||||
if pred_token is None:
|
||||
pred_token = self._default_pred_token
|
||||
if pred_token is None:
|
||||
values = np.zeros(p_shape, dtype=np.float32)
|
||||
else:
|
||||
values = np.asarray(pred_token, dtype=np.float32).reshape(p_shape)
|
||||
if self._pred_info.format.type == FormatType.UINT8:
|
||||
q = self._pred_info.quant_info
|
||||
values = np.clip(np.round(values / q.qp_scale + q.qp_zp),
|
||||
0, 255).astype(np.uint8)
|
||||
bindings.input(self._pred_info.name).set_buffer(values)
|
||||
|
||||
if self._imu_info is not None:
|
||||
i_shape = tuple(self._imu_info.shape)
|
||||
# 原始姿态 → 物理量程归一化(与训练 dataset 同口径)
|
||||
values = normalize_imu(imu, i_shape[-1]).reshape(i_shape)
|
||||
if self._imu_info.format.type == FormatType.UINT8:
|
||||
q = self._imu_info.quant_info
|
||||
values = np.clip(np.round(values / q.qp_scale + q.qp_zp),
|
||||
0, 255).astype(np.uint8)
|
||||
bindings.input(self._imu_info.name).set_buffer(values)
|
||||
|
||||
out_shape = tuple(self._out_info.shape)
|
||||
if self._out_info.format.type == FormatType.UINT8:
|
||||
out_buf = np.zeros(out_shape, dtype=np.uint8)
|
||||
else:
|
||||
out_buf = np.zeros(out_shape, dtype=np.float32)
|
||||
bindings.output(self._out_info.name).set_buffer(out_buf)
|
||||
|
||||
def _postprocess(self, raw_result, history_tail=None) -> np.ndarray:
|
||||
"""反量化 + 解码/反归一化 → (action_dim,)(predict/predict_async 共用)
|
||||
|
||||
history_tail: 归一化空间的历史末步(predict_delta 残差模式用)。
|
||||
"""
|
||||
from hailo_platform import FormatType
|
||||
if self._out_info.format.type == FormatType.UINT8:
|
||||
q = self._out_info.quant_info
|
||||
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
|
||||
img = img.convert("RGB").resize(
|
||||
(self.image_width, self.image_height), Image.BILINEAR)
|
||||
x = np.asarray(img, dtype=np.float32) # HWC 0..255
|
||||
if self._normalize_frames:
|
||||
x = (x / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
|
||||
return x
|
||||
|
||||
def _build_frame_luts(self) -> list:
|
||||
"""按帧 quant_info 构建 256 项单通道查表(归一化+量化合并一步)
|
||||
|
||||
OpenCV 5.0 的 cv2.LUT 只接受单通道表,故每帧返回 [R,G,B] 三张表。
|
||||
数学上与 float 路径(_frame_float + quantize)等价(实测 0 LSB)。
|
||||
"""
|
||||
luts = []
|
||||
for info in self._frame_infos:
|
||||
q = info.quant_info
|
||||
c = np.arange(256, dtype=np.float32)[:, None]
|
||||
if self._normalize_frames:
|
||||
v = (c / 255.0 - IMAGENET_MEAN[None, :]) / IMAGENET_STD[None, :]
|
||||
else:
|
||||
v = np.broadcast_to(c, (256, 3))
|
||||
v = v / q.qp_scale + q.qp_zp
|
||||
lut = np.clip(np.round(v), 0, 255).astype(np.uint8) # (256,3)
|
||||
luts.append([lut[:, ch] for ch in range(3)])
|
||||
return luts
|
||||
|
||||
def _frame_lut_quant(self, image, frame_idx: int = 0) -> np.ndarray:
|
||||
"""cv2.LUT 直通:原始 uint8 像素 → 量化 uint8 缓冲(免 float 中间态)
|
||||
|
||||
非 uint8 / 非 3 通道 / 尺寸不匹配输入走与 _frame_float 相同的 PIL
|
||||
转换,保证两条路径语义一致(RGBA 经 convert 合成、灰度转 RGB)。
|
||||
"""
|
||||
x = np.asarray(image)
|
||||
if x.dtype != np.uint8 or x.ndim != 3 or x.shape[2] != 3 or \
|
||||
x.shape[:2] != (self.image_height, self.image_width):
|
||||
img = Image.fromarray(x).convert("RGB")
|
||||
if img.size != (self.image_width, self.image_height):
|
||||
img = img.resize((self.image_width, self.image_height),
|
||||
Image.BILINEAR)
|
||||
x = np.asarray(img, dtype=np.uint8)
|
||||
luts = self._frame_luts[frame_idx]
|
||||
outs = [cv2.LUT(x[..., ch], lut) for ch, lut in enumerate(luts)]
|
||||
return np.ascontiguousarray(cv2.merge(outs))
|
||||
|
||||
def _bind_frames(self, bindings, frame_float: np.ndarray):
|
||||
"""(F,H,W,3) float 帧 → HEF 帧输入缓冲区(按布局与量化)"""
|
||||
from hailo_platform import FormatType
|
||||
|
||||
def quantize(buf, q):
|
||||
return np.clip(np.round(buf / q.qp_scale + q.qp_zp),
|
||||
0, 255).astype(np.uint8)
|
||||
|
||||
if self._image_layout == "ncfhw":
|
||||
buf = frame_float.transpose(0, 3, 1, 2)[np.newaxis] # (1,F,3,H,W)
|
||||
info = self._frame_infos[0]
|
||||
if info.format.type == FormatType.UINT8:
|
||||
buf = quantize(buf, info.quant_info)
|
||||
bindings.input(info.name).set_buffer(buf)
|
||||
elif self._image_layout == "nhwc_fused":
|
||||
buf = frame_float.reshape(
|
||||
1, self.image_height, self.image_width, self.frame_len * 3)
|
||||
info = self._frame_infos[0]
|
||||
if info.format.type == FormatType.UINT8:
|
||||
buf = quantize(buf, info.quant_info)
|
||||
bindings.input(info.name).set_buffer(buf)
|
||||
else: # nhwc_frames:每帧独立输入
|
||||
for i, info in enumerate(self._frame_infos):
|
||||
buf = frame_float[i]
|
||||
if info.format.type == FormatType.UINT8:
|
||||
buf = quantize(buf, info.quant_info)
|
||||
bindings.input(info.name).set_buffer(buf)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="actpred HEF (Hailo) inference")
|
||||
parser.add_argument("--hef", required=True, help="HEF 模型路径")
|
||||
parser.add_argument("--images", required=True, nargs="+",
|
||||
help="按时间顺序(旧→新)传入 frame_len 个图像路径")
|
||||
parser.add_argument("--action-history", required=True,
|
||||
help="冒号分隔各步,每步 action_dim 个值")
|
||||
parser.add_argument("--stats", default=None,
|
||||
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
|
||||
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,
|
||||
imu=imu)
|
||||
|
||||
if args.output_format == "json":
|
||||
print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))
|
||||
else:
|
||||
print(" ".join(f"{v:.4f}" for v in predicted))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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)
|
||||
+85
-53
@@ -36,13 +36,30 @@ class MobileNetEncoder(nn.Module):
|
||||
except Exception as e:
|
||||
if not pretrained:
|
||||
raise RuntimeError("加载 MobileNetV2 架构失败") from e
|
||||
raise RuntimeError(
|
||||
"MobileNetV2 预训练权重下载失败,可能是网络受限。"
|
||||
"请在联网环境中预下载权重"
|
||||
"(运行 torchvision.models.mobilenet_v2(weights='IMAGENET1K_V1')),"
|
||||
"并将 ~/.cache/torch/hub/checkpoints/ 下的权重文件"
|
||||
"拷贝到目标机器同路径。"
|
||||
) from e
|
||||
# 主源下载失败,尝试加载已缓存的权重文件
|
||||
try:
|
||||
import os
|
||||
hub_dir = torch.hub.get_dir()
|
||||
cache_dir = os.path.join(hub_dir, "checkpoints")
|
||||
weight_file = getattr(weights, "file_name",
|
||||
"mobilenet_v2-b0353104.pth")
|
||||
cached_path = os.path.join(cache_dir, weight_file)
|
||||
if os.path.exists(cached_path):
|
||||
print(f"INFO: 使用缓存权重 {cached_path}")
|
||||
backbone = mobilenet_v2(weights=None)
|
||||
backbone.load_state_dict(torch.load(
|
||||
cached_path, map_location="cpu",
|
||||
weights_only=True))
|
||||
else:
|
||||
raise FileNotFoundError(cached_path)
|
||||
except Exception:
|
||||
raise RuntimeError(
|
||||
"MobileNetV2 预训练权重下载失败,可能是网络受限。"
|
||||
"请在联网环境中预下载权重"
|
||||
"(运行 torchvision.models.mobilenet_v2(weights='IMAGENET1K_V1')),"
|
||||
"并将 ~/.cache/torch/hub/checkpoints/ 下的权重文件"
|
||||
"拷贝到目标机器同路径。"
|
||||
) from e
|
||||
self.features = backbone.features
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.proj = nn.Linear(1280, feature_dim)
|
||||
@@ -53,34 +70,7 @@ class MobileNetEncoder(nn.Module):
|
||||
return self.proj(x)
|
||||
|
||||
|
||||
class MLPHistoryEncoder(nn.Module):
|
||||
def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(input_dim, hidden_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class GRUHistoryEncoder(nn.Module):
|
||||
def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.gru = nn.GRU(input_size=input_dim, hidden_size=hidden_dim,
|
||||
num_layers=1, batch_first=True)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
_, h = self.gru(x)
|
||||
return self.dropout(h.squeeze(0))
|
||||
|
||||
|
||||
class FusionHead(nn.Module):
|
||||
class PredictionHead(nn.Module):
|
||||
def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = 0.1):
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
@@ -97,6 +87,14 @@ class FusionHead(nn.Module):
|
||||
|
||||
|
||||
class ActionPredictor(nn.Module):
|
||||
"""多帧图像 + 历史动作的 Transformer 注意力融合预测器
|
||||
|
||||
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 部署友好。
|
||||
"""
|
||||
|
||||
def __init__(self, model_cfg: ModelConfig, data_cfg: DataConfig,
|
||||
pretrained: bool = True):
|
||||
super().__init__()
|
||||
@@ -106,26 +104,60 @@ class ActionPredictor(nn.Module):
|
||||
else:
|
||||
self.image_encoder = SimpleCNN(model_cfg.cnn_feature_dim)
|
||||
|
||||
history_input_dim = data_cfg.history_len * data_cfg.action_dim
|
||||
if model_cfg.history_encoder == "gru":
|
||||
self.action_encoder = GRUHistoryEncoder(
|
||||
data_cfg.action_dim, model_cfg.action_hidden_dim, model_cfg.dropout)
|
||||
else:
|
||||
self.action_encoder = MLPHistoryEncoder(
|
||||
history_input_dim, model_cfg.action_hidden_dim, model_cfg.dropout)
|
||||
d = model_cfg.attn_dim
|
||||
self.img_proj = nn.Linear(model_cfg.cnn_feature_dim, d)
|
||||
self.act_proj = nn.Linear(data_cfg.action_dim, d)
|
||||
|
||||
fused_input = model_cfg.cnn_feature_dim + model_cfg.action_hidden_dim
|
||||
self.head = FusionHead(
|
||||
fused_input, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout)
|
||||
self._imu_dim = data_cfg.imu_dim
|
||||
if self._imu_dim > 0:
|
||||
self.imu_proj = nn.Linear(self._imu_dim, d)
|
||||
|
||||
self._use_gru = model_cfg.history_encoder == "gru"
|
||||
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)
|
||||
nn.init.trunc_normal_(self.pos_embed, std=0.02)
|
||||
|
||||
encoder_layer = nn.TransformerEncoderLayer(
|
||||
d_model=d, nhead=model_cfg.attn_heads,
|
||||
dim_feedforward=d * 2, dropout=model_cfg.dropout,
|
||||
batch_first=True, norm_first=True)
|
||||
self.encoder = nn.TransformerEncoder(
|
||||
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, head_out, model_cfg.dropout)
|
||||
|
||||
self._frame_len = data_cfg.frame_len
|
||||
self._history_len = data_cfg.history_len
|
||||
self._action_dim = data_cfg.action_dim
|
||||
self._image_width = data_cfg.image_width
|
||||
self._image_height = data_cfg.image_height
|
||||
self._cnn_feature_dim = model_cfg.cnn_feature_dim
|
||||
|
||||
def forward(self, image: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor:
|
||||
img_feat = self.image_encoder(image)
|
||||
if self._use_gru:
|
||||
act_feat = self.action_encoder(action_history)
|
||||
else:
|
||||
act_feat = self.action_encoder(action_history.flatten(1))
|
||||
return self.head(torch.cat([img_feat, act_feat], dim=-1))
|
||||
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)
|
||||
img_feat = self.image_encoder(frames).reshape(
|
||||
b, self._frame_len, self._cnn_feature_dim)
|
||||
img_tokens = self.img_proj(img_feat)
|
||||
act_tokens = self.act_proj(action_history)
|
||||
pred_token = self.pred_token.expand(b, -1, -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])
|
||||
+14
-38
@@ -1,7 +1,9 @@
|
||||
"""轻量级 .mt 脚本 @Action 矩阵解析器
|
||||
|
||||
从 .mt 文件中提取 @Action 块,解析为 (n_rows, n_cols) 的 float32 矩阵。
|
||||
支持 `-` 占位符(沿用上一行同列的值);首行 `-` 从 @Servo 块的初始角度填充。
|
||||
支持 `-` 占位符(沿用上一行同列的值);首行 `-` 填 0,
|
||||
与 pybot/bot4cj 运行时语义一致(2026-08 统一,见
|
||||
docs/reports/bot4cj_python_migration_dev_doc.md §2.3)。
|
||||
"""
|
||||
|
||||
import re
|
||||
@@ -15,27 +17,6 @@ def _natural_key(p: Path):
|
||||
return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", p.name)]
|
||||
|
||||
|
||||
def _parse_servo_init(text: str) -> dict:
|
||||
"""从 @Servo 块提取 {servo_name: initial_angle}(max 列)"""
|
||||
match = re.search(r"@Servo\s+(.*?)\s*@", text, re.DOTALL)
|
||||
if not match:
|
||||
return {}
|
||||
servo_init = {}
|
||||
for line in match.group(1).split("\n"):
|
||||
line = re.sub(r"#.*$", "", line).strip()
|
||||
if line.endswith(";"):
|
||||
line = line[:-1].strip()
|
||||
if not line:
|
||||
continue
|
||||
cols = re.split(r"\s+", line)
|
||||
if len(cols) >= 2 and cols[0] != "name":
|
||||
try:
|
||||
servo_init[cols[0]] = float(cols[1])
|
||||
except ValueError:
|
||||
pass
|
||||
return servo_init
|
||||
|
||||
|
||||
def parse_mt_action(path) -> np.ndarray:
|
||||
"""解析 .mt 文件中的 @Action 矩阵
|
||||
|
||||
@@ -44,23 +25,24 @@ def parse_mt_action(path) -> np.ndarray:
|
||||
|
||||
Returns:
|
||||
(n_data_rows, n_cols) float32 矩阵,表头行已跳过,`-` 已前向填充。
|
||||
首行 `-` 从 @Servo 块的初始角度填充。
|
||||
首行 `-` 填 0(与 pybot/bot4cj 运行时语义一致)。
|
||||
|
||||
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)
|
||||
|
||||
# 解析 @Servo 块的初始角度,用于首行 `-` 前向填充
|
||||
servo_init = _parse_servo_init(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"):
|
||||
@@ -84,12 +66,6 @@ def parse_mt_action(path) -> np.ndarray:
|
||||
data_rows = rows[1:]
|
||||
n_cols = len(header)
|
||||
|
||||
# 为每列计算初始值:s1.A → @Servo 中 s1 的 max 角度,其余为 0.0
|
||||
col_init = []
|
||||
for col_name in header:
|
||||
servo_name = col_name.split(".")[0] if "." in col_name else None
|
||||
col_init.append(servo_init.get(servo_name, 0.0))
|
||||
|
||||
matrix = np.zeros((len(data_rows), n_cols), dtype=np.float32)
|
||||
for i, row in enumerate(data_rows):
|
||||
if len(row) != n_cols:
|
||||
@@ -98,7 +74,7 @@ def parse_mt_action(path) -> np.ndarray:
|
||||
f"不一致: {path}")
|
||||
for j, val in enumerate(row):
|
||||
if val == "-":
|
||||
matrix[i, j] = matrix[i - 1, j] if i > 0 else col_init[j]
|
||||
matrix[i, j] = matrix[i - 1, j] if i > 0 else 0.0
|
||||
else:
|
||||
try:
|
||||
matrix[i, j] = float(val)
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
"""actpred ONNX Runtime 推理封装 —— CPU 后端(Pi 5 实测 5.9ms/次,PyTorch 的 ~5 倍)
|
||||
|
||||
接口对齐 predict.ActionPredictorInference:predict(images, action_history, imu=None)。
|
||||
|
||||
模型维度从 ONNX session 输入元数据动态推导(images: [1,F,3,H,W],
|
||||
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。
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
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)
|
||||
|
||||
|
||||
from utils import parse_action_history
|
||||
|
||||
|
||||
class OnnxPredictor:
|
||||
"""ONNX Runtime 动作预测推理器(CPU 后端)
|
||||
|
||||
用法:
|
||||
with OnnxPredictor("model.onnx") as predictor:
|
||||
action = predictor.predict(images, action_history)
|
||||
"""
|
||||
|
||||
def __init__(self, onnx_path: str, stats_path: str = 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:
|
||||
providers = ["CPUExecutionProvider"]
|
||||
ort.set_default_logger_severity(3) # 抑制 GPU 探测等无害警告
|
||||
opts = ort.SessionOptions()
|
||||
opts.log_severity_level = 3
|
||||
self._session = ort.InferenceSession(
|
||||
str(onnx_path), sess_options=opts, providers=providers)
|
||||
|
||||
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, 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 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)
|
||||
if stats.shape != (2, self.action_dim):
|
||||
raise ValueError(
|
||||
f"action_stats.npy 形状应为 (2, {self.action_dim}),"
|
||||
f"实际为 {stats.shape}: {stats_path}")
|
||||
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维)/姿态(2维,可选)输入,返回各自形状"""
|
||||
img = [s for s in inputs.values() if len(s) == 5]
|
||||
hist = [s for s in inputs.values() if len(s) == 3]
|
||||
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 维动作历史"
|
||||
"(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], (imu[0] if imu else None)
|
||||
|
||||
# ---------------- 生命周期 ----------------
|
||||
|
||||
def close(self):
|
||||
self._session = None
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
self.close()
|
||||
return False
|
||||
|
||||
# ---------------- 推理 ----------------
|
||||
|
||||
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 动作向量。
|
||||
"""
|
||||
if self._session is None:
|
||||
raise RuntimeError("predictor 已关闭")
|
||||
if not isinstance(images, (list, tuple)):
|
||||
images = [images]
|
||||
if len(images) != self.frame_len:
|
||||
raise ValueError(
|
||||
f"images 应为 {self.frame_len} 帧,实际为 {len(images)} 帧。"
|
||||
"冷启动时可重复传入同一帧。")
|
||||
action_history = np.asarray(action_history, dtype=np.float32)
|
||||
expected = (self.history_len, self.action_dim)
|
||||
if action_history.shape != expected:
|
||||
raise ValueError(
|
||||
f"action_history 形状应为 {expected},实际为 {action_history.shape}")
|
||||
|
||||
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])
|
||||
feeds = {
|
||||
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._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
|
||||
|
||||
def _preprocess(self, image) -> np.ndarray:
|
||||
"""Resize → /255 → ImageNet 归一化 → (3, H, W) float32"""
|
||||
img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
|
||||
img = img.convert("RGB").resize(
|
||||
(self.image_width, self.image_height), Image.BILINEAR)
|
||||
x = np.asarray(img, dtype=np.float32) / 255.0 # HWC
|
||||
x = x.transpose(2, 0, 1) # CHW
|
||||
return (x - IMAGENET_MEAN[:, None, None]) / IMAGENET_STD[:, None, None]
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="actpred ONNX Runtime inference")
|
||||
parser.add_argument("--model", required=True, help="ONNX 模型路径")
|
||||
parser.add_argument("--images", required=True, nargs="+",
|
||||
help="按时间顺序(旧→新)传入 frame_len 个图像路径")
|
||||
parser.add_argument("--action-history", required=True,
|
||||
help="冒号分隔各步,每步 action_dim 个值")
|
||||
parser.add_argument("--stats", default=None,
|
||||
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,
|
||||
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)
|
||||
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))
|
||||
else:
|
||||
print(" ".join(f"{v:.4f}" for v in predicted))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+101
-30
@@ -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:
|
||||
@@ -28,16 +33,34 @@ 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_size, self._data_cfg.image_size)),
|
||||
transforms.Resize((self._data_cfg.image_height,
|
||||
self._data_cfg.image_width)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
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
|
||||
@@ -57,58 +80,106 @@ class ActionPredictorInference:
|
||||
raise ValueError(
|
||||
f"action_stats.npy 形状应为 (2, {self._data_cfg.action_dim}),"
|
||||
f"实际为 {stats.shape},请检查文件: {stats_path}")
|
||||
self._action_mean = stats[0].astype(np.float32)
|
||||
self._action_std = stats[1].astype(np.float32)
|
||||
self._action_mean = torch.from_numpy(stats[0].astype(np.float32)).to(self._device)
|
||||
self._action_std = torch.from_numpy(stats[1].astype(np.float32)).to(self._device)
|
||||
|
||||
@torch.no_grad()
|
||||
def predict(self, image: np.ndarray, 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_dim>0 时
|
||||
生效,内部按物理量程归一化;None 等价零向量 = 无姿态先验)
|
||||
"""
|
||||
if not isinstance(images, (list, tuple)):
|
||||
images = [images]
|
||||
if len(images) != self._data_cfg.frame_len:
|
||||
raise ValueError(
|
||||
f"images 应为 {self._data_cfg.frame_len} 帧,实际为 {len(images)} 帧。"
|
||||
f"冷启动时可重复传入同一帧。")
|
||||
expected_shape = (self._data_cfg.history_len, self._data_cfg.action_dim)
|
||||
if action_history.shape != expected_shape:
|
||||
raise ValueError(
|
||||
f"action_history 形状应为 {expected_shape},实际为 {action_history.shape}")
|
||||
img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
|
||||
img_tensor = self._transform(img).unsqueeze(0).to(self._device)
|
||||
frames = []
|
||||
for image in images:
|
||||
img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
|
||||
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).cpu().numpy()
|
||||
if self._action_mean is not None:
|
||||
pred = pred * self._action_std + self._action_mean
|
||||
return pred
|
||||
|
||||
|
||||
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():
|
||||
parser = argparse.ArgumentParser(description="actpred inference")
|
||||
parser.add_argument("--model", required=True, help="checkpoint path")
|
||||
parser.add_argument("--image", required=True, help="image file path")
|
||||
parser.add_argument("--images", required=True, nargs="+",
|
||||
help="按时间顺序(旧→新)传入 frame_len 个图像路径;"
|
||||
"冷启动时可重复传同一帧路径")
|
||||
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
|
||||
|
||||
image = Image.open(args.image).convert("RGB")
|
||||
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, action_dim)
|
||||
predicted = predictor.predict(np.array(image), 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))
|
||||
|
||||
+202
-47
@@ -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,31 +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 pretrain_one_epoch(model, loader, optimizer, device, image_size,
|
||||
max_grad_norm=None, noise_std=0.0):
|
||||
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,
|
||||
random_images: bool = False, weights=None):
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
count = 0
|
||||
zero_image = 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_image is None or zero_image.size(0) != history.size(0):
|
||||
zero_image = torch.zeros(
|
||||
history.size(0), 3, image_size, image_size, 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_image, 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)
|
||||
@@ -105,20 +153,21 @@ def pretrain_one_epoch(model, loader, optimizer, device, image_size,
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def validate(model, loader, device, image_size):
|
||||
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_image = 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_image is None or zero_image.size(0) != history.size(0):
|
||||
zero_image = torch.zeros(
|
||||
history.size(0), 3, image_size, image_size, device=device
|
||||
)
|
||||
pred = model(zero_image, 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:
|
||||
@@ -132,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)
|
||||
|
||||
@@ -168,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")
|
||||
@@ -196,11 +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"encoder={cfg.model.history_encoder}")
|
||||
print(f"Pre-training: noise_std={args.noise_std}")
|
||||
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}, "
|
||||
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,
|
||||
@@ -212,43 +331,65 @@ def main():
|
||||
|
||||
best_val_loss = float("inf")
|
||||
patience_counter = 0
|
||||
start_epoch = 1
|
||||
best_path = os.path.join(cfg.train.checkpoint_dir, "pretrain_best.pt")
|
||||
inference_path = os.path.join(
|
||||
cfg.train.checkpoint_dir, "pretrain_best_inference.pt"
|
||||
)
|
||||
last_path = os.path.join(cfg.train.checkpoint_dir, "pretrain_last.pt")
|
||||
|
||||
for epoch in range(1, cfg.train.epochs + 1):
|
||||
if args.resume:
|
||||
resume_ckpt = torch.load(args.resume, map_location=device, weights_only=True)
|
||||
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"])
|
||||
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)
|
||||
print(f"Resumed from {args.resume} at epoch {start_epoch} "
|
||||
f"(best_val_loss={best_val_loss:.6f})")
|
||||
|
||||
for epoch in range(start_epoch, cfg.train.epochs + 1):
|
||||
print(f"\nEpoch {epoch}/{cfg.train.epochs} "
|
||||
f"lr={optimizer.param_groups[0]['lr']:.6f}")
|
||||
train_loss = pretrain_one_epoch(
|
||||
model, train_loader, optimizer, device,
|
||||
cfg.data.image_size, cfg.train.max_grad_norm, args.noise_std,
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.frame_len,
|
||||
cfg.train.max_grad_norm, args.noise_std, args.random_images,
|
||||
weights=weights_t,
|
||||
)
|
||||
val_loss = validate(model, val_loader, device, cfg.data.image_size)
|
||||
val_loss = validate(model, val_loader, device,
|
||||
cfg.data.image_width, cfg.data.image_height,
|
||||
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:
|
||||
@@ -256,6 +397,20 @@ def main():
|
||||
f"(patience={cfg.train.patience})")
|
||||
break
|
||||
|
||||
# 每 epoch 保存 last.pt,支持 --resume 断点续训练
|
||||
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(),
|
||||
},
|
||||
last_path,
|
||||
)
|
||||
|
||||
print(f"\nPre-training done. Best val_loss={best_val_loss:.6f}")
|
||||
|
||||
if best_val_loss == float("inf"):
|
||||
|
||||
@@ -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
|
||||
+2536
-46
File diff suppressed because it is too large.
Load diff
+283
-46
@@ -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):
|
||||
image = batch["image"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
optimizer.zero_grad()
|
||||
pred = model(image, 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() * image.size(0)
|
||||
count += image.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):
|
||||
image = batch["image"].to(device)
|
||||
history = batch["action_history"].to(device)
|
||||
target = batch["target"].to(device)
|
||||
pred = model(image, history)
|
||||
loss = torch.nn.functional.mse_loss(pred, target)
|
||||
total_loss += loss.item() * image.size(0)
|
||||
count += image.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_size, cfg.data.normalize_actions,
|
||||
cfg.data.action_stats_path, train_indices)
|
||||
if cfg.data.normalize_actions and not cfg.data.action_stats_path:
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions,
|
||||
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_size, cfg.data.normalize_actions,
|
||||
cfg.data.action_stats_path, val_indices)
|
||||
cfg.data.image_width, cfg.data.image_height, cfg.data.normalize_actions,
|
||||
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:
|
||||
@@ -132,38 +319,79 @@ def main():
|
||||
|
||||
best_val_loss = float("inf")
|
||||
patience_counter = 0
|
||||
start_epoch = 1
|
||||
best_path = os.path.join(cfg.train.checkpoint_dir, "best.pt")
|
||||
last_path = os.path.join(cfg.train.checkpoint_dir, "last.pt")
|
||||
|
||||
for epoch in range(1, cfg.train.epochs + 1):
|
||||
if args.resume:
|
||||
resume_ckpt = torch.load(args.resume, map_location=device, weights_only=True)
|
||||
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)
|
||||
print(f"Resumed from {args.resume} at epoch {start_epoch} "
|
||||
f"(best_val_loss={best_val_loss:.6f})")
|
||||
|
||||
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:
|
||||
print(f"\nEarly stopping at epoch {epoch} (patience={cfg.train.patience})")
|
||||
break
|
||||
|
||||
# 每 epoch 保存 last.pt,支持 --resume 断点续训
|
||||
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(),
|
||||
}, last_path)
|
||||
|
||||
print(f"\nTraining done. Best val_loss={best_val_loss:.6f}")
|
||||
|
||||
if best_val_loss == float("inf"):
|
||||
@@ -172,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)
|
||||
+1
-1
@@ -2,5 +2,5 @@
|
||||
cd "$(dirname "$0")" || exit
|
||||
source ./venv/bin/activate
|
||||
#exec ./dnfs.bin
|
||||
exec python3 ai_review.py
|
||||
exec python3 ai_review.py "$@"
|
||||
|
||||
@@ -65,6 +65,7 @@ class ReviewAgent:
|
||||
],
|
||||
temperature=self.temperature,
|
||||
max_tokens=self.max_tokens,
|
||||
timeout=8.0,
|
||||
)
|
||||
raw = response.choices[0].message.content.strip()
|
||||
except Exception as e:
|
||||
|
||||
@@ -165,9 +165,39 @@ def rag_list():
|
||||
return jsonify({"total": len(entries), "entries": entries}), 200
|
||||
|
||||
|
||||
def _parse_args():
|
||||
"""解析命令行参数(标准化:CLI > ENV > 默认值)。"""
|
||||
import argparse
|
||||
p = argparse.ArgumentParser(
|
||||
prog="ai_review",
|
||||
description="AI Review Service — 基于 DeepSeek + RAG 的内容审核微服务",
|
||||
)
|
||||
p.add_argument("--host", default=None,
|
||||
help="绑定地址(默认读取 $AI_REVIEW_HOST,否则 127.0.0.1)")
|
||||
p.add_argument("--port", type=int, default=None,
|
||||
help="监听端口(默认读取 $AI_REVIEW_PORT,否则 8007)")
|
||||
p.add_argument("--log-level", default="info",
|
||||
choices=["debug", "info", "warn", "error"],
|
||||
help="日志级别(默认 info)")
|
||||
p.add_argument("--instance", default="ai_review",
|
||||
help="实例名(注册中心用,默认 ai_review)")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
# ── 配置默认值(ENV 覆盖,CLI 最高优先级)───────────────────────────
|
||||
ADDR = os.getenv("AI_REVIEW_HOST", "127.0.0.1")
|
||||
PORT = int(os.getenv("AI_REVIEW_PORT", "8007"))
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = _parse_args()
|
||||
|
||||
# CLI 覆盖(最高优先级)
|
||||
if args.host is not None:
|
||||
ADDR = args.host
|
||||
if args.port is not None:
|
||||
PORT = args.port
|
||||
|
||||
print(f'AI Review Service listening at http://{ADDR}:{PORT}')
|
||||
print(f'[ai_review] instance={args.instance} log-level={args.log_level}')
|
||||
print(f'[ai_review] 配置来源: CLI > ENV > 默认值')
|
||||
serve(app, host=ADDR, port=PORT)
|
||||
+300
-25
@@ -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,18 +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 # 前端主页面
|
||||
├── common/ # 基础设施包(i18n / 配置 / 常量 / 日志)
|
||||
│ ├── 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
|
||||
@@ -76,6 +134,7 @@ aibox/
|
||||
│ ├── ssh-exec.py
|
||||
│ ├── sftp-transfer.py
|
||||
│ └── ...
|
||||
├── prompts/ # 系统提示词源文件(运行期复制到 ~/.aibox/prompts/)
|
||||
├── SKILLS_FORMAT.md # 技能文件格式说明
|
||||
├── _run.sh # 启动脚本(被 core 守护时调用)
|
||||
├── _runcli.sh # 命令行启动脚本
|
||||
@@ -84,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/`:
|
||||
|
||||
```
|
||||
@@ -124,31 +202,39 @@ aibox/
|
||||
|
||||
## 🔧 可用工具
|
||||
|
||||
AI 智能体可以调用以下工具来完成任务(由 `Tool` 基类派生,在 `DeepSeekChat` 中注册):
|
||||
工具由 `Tool` 基类(`tools/base.py`)派生,在 `DeepSeekChat._register_tools()` 中按 **工具集 profile** 分组注册(v2.9 编排者-执行者分层):主对话(orchestrator)只负责管理与调度,不直接接触命令行;执行类任务经 `create_subagent` 派生拥有执行面工具的子代理(worker)完成。
|
||||
|
||||
### 主对话(orchestrator)工具集
|
||||
|
||||
| 工具名称 | 功能描述 |
|
||||
|---------|---------|
|
||||
| `read_file` | 读取文件内容(返回 markdown 指引,改用 `head`/`tail`/`sed`/`wc` 等命令行读取) |
|
||||
| `write_file` | 写入或追加文件(返回 markdown 指引,改用 `mkdir -p` + `cat` 重定向写入) |
|
||||
| `list_files` | 列出目录内容(返回 markdown 指引,改用 `ls`/`du`/`find` 命令行列出) |
|
||||
| `web` | 访问网页、获取标题/文本/链接(返回 markdown 指引,改用 `curl` 命令行访问) |
|
||||
| `memory` | 保存、搜索(支持分页)、删除记忆 |
|
||||
| `memo` | 创建、查看、完成、搜索任务 |
|
||||
| `run_command` | 执行系统命令(实时输出,支持超时控制) |
|
||||
| `search_skills` | 搜索扩展技能(支持 Markdown 和可执行技能) |
|
||||
| `express_emotion` | 表达情绪并记录当前情绪状态(anger/joy/sadness/surprise),可指定 `persona` 从 `~/.aibox/skills/` 加载人格技能(如 angry/coquettish/tsundere)派生携带该人格的情绪子代理,保证人格独立、不污染主对话上下文 |
|
||||
| `create_subagent` | 派生子代理:独立上下文/角色设定,仅拥有执行面工具(无记忆/人格/派生/提问能力),可选独立模型(工具描述动态列出 `models` 节可用档案名,防 LLM 编造名称),用于长任务分解、上下文隔离、独立角色专注处理 |
|
||||
| `search_skills` | 搜索扩展技能(支持 Markdown 和可执行技能),查到后派生子代理执行 |
|
||||
| `reload_skills` | 热重载扩展技能 |
|
||||
| `create_skill` | 创建新技能(AI 自主编写 Python 脚本) |
|
||||
| `create_skill` | 创建新技能(AI 自主编写 Python 脚本,实际写入经子代理完成) |
|
||||
| `install_skill` | 技能安装指引 |
|
||||
| `ask_user` | 向用户提问(AI 主动交互) |
|
||||
| `update_identity` | 更新用户身份信息(`00_IDENTITY.md`) |
|
||||
| `update_soul` | 更新 AI 自我定义(`99_SOUL.md`,智能体进化) |
|
||||
| `save_memory_and_end_conversation` | 保存记忆并结束对话 |
|
||||
| `get_current_time` | 获取当前时间 |
|
||||
| `calculator` | 数学计算 |
|
||||
| `set_proxy` | 设置 SOCKS5 代理(host/port/可选账号密码,可持久化,即时生效) |
|
||||
| `set_proxy` | 设置 SOCKS5 代理(host/port/可选账号密码,可持久化,即时生效);`password` 明文仅本次会话有效不落盘,跨重启用 `password_env` 指定存有密码的环境变量名 |
|
||||
| `get_proxy` | 查看当前 SOCKS5 代理配置 |
|
||||
| `express_emotion` | 表达情绪并记录当前情绪状态(anger/joy/sadness/surprise),可指定 `persona` 从 `~/.aibox/skills/` 加载人格技能(如 angry/coquettish/tsundere)派生携带该人格的情绪子代理,保证人格独立、不污染主对话上下文 |
|
||||
| `create_subagent` | 派生子代理:独立上下文/角色设定,轻量无工具、可选独立模型,用于长任务分解、上下文隔离、独立角色专注处理 |
|
||||
|
||||
### 子代理(worker)工具集
|
||||
|
||||
| 工具名称 | 功能描述 |
|
||||
|---------|---------|
|
||||
| `run_command` | 执行系统命令(实时输出,支持超时控制) |
|
||||
| `read_file` | 读取文件内容(返回 markdown 指引,改用 `head`/`tail`/`sed`/`wc` 等命令行读取) |
|
||||
| `write_file` | 写入或追加文件(返回 markdown 指引,改用 `mkdir -p` + `cat` 重定向写入) |
|
||||
| `list_files` | 列出目录内容(返回 markdown 指引,改用 `ls`/`du`/`find` 命令行列出) |
|
||||
| `web` | 访问网页、获取标题/文本/链接(返回 markdown 指引,改用 `curl` 命令行访问) |
|
||||
| `calculator` | 数学计算 |
|
||||
| `get_current_time` | 获取当前时间 |
|
||||
|
||||
## 💬 WebSocket API
|
||||
|
||||
@@ -441,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,
|
||||
@@ -501,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"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -513,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` 工厂中注册,无需改动业务代码,为后续子代理/多模型共发预留了扩展点。
|
||||
|
||||
@@ -534,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 秒)
|
||||
```
|
||||
|
||||
### 使用示例
|
||||
@@ -555,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` 全局超时。
|
||||
|
||||
## 🔄 架构流程
|
||||
|
||||
```
|
||||
@@ -621,10 +798,19 @@ Server -> Client: {type: "dialogue_reset", content: "✅ 对话状态已完全
|
||||
|
||||
### 添加新工具
|
||||
|
||||
1. 继承 `Tool` 基类(位于 `aibox.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):
|
||||
@@ -645,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/非字符串转为安全字符串,过滤控制字符
|
||||
@@ -723,6 +911,93 @@ 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 分组装配工具集
|
||||
- **主对话(orchestrator)**:仅管理面工具(记忆/备忘录/情绪/子代理派生/技能查询与创建/人格更新/代理设置),不再直接持有 `run_command` 与文件/网页等执行面工具,执行类任务一律经 `create_subagent` 派生子代理完成
|
||||
- **子代理(worker)**:仅执行面 7 工具(`run_command` / `read_file` / `write_file` / `list_files` / `web` / `calculator` / `get_current_time`),无记忆/人格更新/派生/提问能力——修复子代理可调用 `update_soul` 改写主人格、可递归派生子代理的隐患
|
||||
- 子代理角色提示词追加行为约束:不得提问(疑问写入最终回答)、产物文件一律落盘 `~/.aibox/workspace/` 并给出绝对路径
|
||||
- 配套提示词调整:`99_SOUL.md` 角色定位改为编排者,`7_extension_skills.md` 技能执行路径改为 `search_skills` → 派生子代理执行(运行期 `~/.aibox/prompts/` 更新)
|
||||
|
||||
### v2.8
|
||||
- ✅ **新增 `express_emotion` 工具**:表达情绪并记录当前情绪状态(anger/joy/sadness/surprise),可指定 `persona` 从 `~/.aibox/skills/` 加载人格技能文件(如 angry/coquettish/tsundere),自动派生携带该人格的情绪子代理生成回应,保证人格独立、不污染主对话上下文
|
||||
- ✅ **新增 `/search <关键词>` 命令**:在 WebSocket 会话中按关键词搜索记忆,返回匹配的记忆条目
|
||||
|
||||
@@ -469,7 +469,7 @@ CREATE TABLE memos (
|
||||
|
||||
### 9.4 技能格式规范
|
||||
|
||||
详见 [SKILLS_FORMAT.md](file:///Users/qingsong/micro_servers/aibox/SKILLS_FORMAT.md)
|
||||
详见 [SKILLS_FORMAT.md](SKILLS_FORMAT.md)
|
||||
|
||||
---
|
||||
|
||||
@@ -805,19 +805,19 @@ sqlite3 ~/.aibox/memos.db "SELECT * FROM memos WHERE status='pending';"
|
||||
|
||||
### A. 配置文件示例
|
||||
|
||||
完整默认配置见 [main.py:1100-1180](file:///Users/qingsong/micro_servers/aibox/main.py#L1100-L1180)
|
||||
完整默认配置见 [aibox.py:1100-1180](aibox.py#L1100-L1180)
|
||||
|
||||
### B. 数据库 Schema
|
||||
|
||||
完整建表语句见:
|
||||
- 记忆库:[main.py](file:///Users/qingsong/micro_servers/aibox/main.py) 中 `MemoryManager._init_db()`
|
||||
- 任务库:[main.py](file:///Users/qingsong/micro_servers/aibox/main.py) 中 `MemoScheduler._init_db()`
|
||||
- 记忆库:[aibox.py](aibox.py) 中 `MemoryManager._init_db()`
|
||||
- 任务库:[aibox.py](aibox.py) 中 `MemoScheduler._init_db()`
|
||||
|
||||
### C. 参考文档
|
||||
|
||||
- [DeepSeek API 文档](https://platform.deepseek.com/docs)
|
||||
- [websockets 库文档](https://websockets.readthedocs.io/)
|
||||
- [SKILLS_FORMAT.md](file:///Users/qingsong/micro_servers/aibox/SKILLS_FORMAT.md)
|
||||
- [SKILLS_FORMAT.md](SKILLS_FORMAT.md)
|
||||
|
||||
---
|
||||
|
||||
|
||||
+1
-1
@@ -3,6 +3,6 @@ cd "$(dirname "$0")" || exit
|
||||
source ./_setenv.sh
|
||||
source ./venv/bin/activate
|
||||
export PATH=${PATH}:~/.aibox/skills
|
||||
exec python aibox.py $@
|
||||
exec python aibox.py "$@"
|
||||
|
||||
|
||||
+773
-6284
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
common - aibox 底层基础设施公共模块(包)
|
||||
|
||||
抽离自 aibox.py 的零依赖基础层,供 aibox.py 及后续 tools.py 等模块导入。
|
||||
当前包含:
|
||||
- i18n.py : I18n 类 + i18n 实例(多语言)
|
||||
- config.py : ConfigManager(配置管理)
|
||||
- constants.py : Constants 类 + constants 实例(系统常量)
|
||||
- logger.py : Logger(统一日志)
|
||||
形成单向依赖:common <- tools <- aibox。
|
||||
|
||||
采用"重导出"策略:外部统一用 `from common import i18n, I18n, ...`,
|
||||
调用点无需关心具体子模块所在文件。
|
||||
"""
|
||||
|
||||
from .i18n import i18n, I18n
|
||||
from .config import ConfigManager
|
||||
from .constants import Constants, constants
|
||||
from .logger import Logger
|
||||
|
||||
__all__ = ["i18n", "I18n", "ConfigManager", "Constants", "constants", "Logger"]
|
||||
@@ -0,0 +1,318 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
common.config - aibox 配置管理器
|
||||
|
||||
抽离自 aibox.py 的 ConfigManager。仅依赖 common.i18n(在运行期方法内部使用),
|
||||
不反向依赖 aibox 任何符号,可安全独立导入。
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import copy
|
||||
from typing import Dict, Optional
|
||||
|
||||
from .i18n import i18n
|
||||
class ConfigManager:
|
||||
"""配置管理器 - 单例模式,负责加载和保存配置"""
|
||||
|
||||
_instance = None
|
||||
_config = None
|
||||
_config_path = None
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if self._config is None:
|
||||
self._load_config()
|
||||
|
||||
def _get_default_config_path(self) -> str:
|
||||
env_config = os.environ.get("AI_CONFIG_PATH")
|
||||
if env_config:
|
||||
return env_config
|
||||
return os.path.expanduser("~/.aibox/config.json")
|
||||
|
||||
def _get_default_config(self) -> Dict:
|
||||
return {
|
||||
|
||||
# 系统参数
|
||||
"system": {
|
||||
"save_dir": "~/.aibox", # 储存目录
|
||||
"datetime_format": "%Y-%m-%d %H:%M:%S", # 时间格式
|
||||
"date_format": "%Y%m%d_%H%M%S", # 另一个时间格式
|
||||
"debug_level": 1, # 调试等级
|
||||
"load_history_by_default": True, # 默认加载历史记录
|
||||
"log_file": "~/.aibox/ai.log", # 日志文件
|
||||
"log_max_size": 10485760, # 日志最大大小
|
||||
"log_backup_count": 5, # 日志备份数量
|
||||
"skills_dir": "~/.aibox/skills", # 技能目录
|
||||
"prompts_dir": "~/.aibox/prompts", # 提示词目录
|
||||
"workspace_dir": "~/.aibox/workspace", # 工作空间目录
|
||||
"websocket_host": "localhost", # websocket监听地址
|
||||
"websocket_port": 8765, # websocket监听端口
|
||||
},
|
||||
|
||||
# 调度程序
|
||||
"scheduler": {
|
||||
"enabled": True, # 开启调度
|
||||
"check_interval": 1.0, # 检查间隔
|
||||
"max_tasks_per_run": 10, # 每次运行最大任务量
|
||||
"broadcast_enabled": True # 开启广播
|
||||
},
|
||||
|
||||
|
||||
# 大模型调用接口
|
||||
"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, # 大模型对话温度(类似思维发散程度)
|
||||
"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, # 最大历史消息数量
|
||||
"max_context_messages": 120, # 最大内容消息数量(?)
|
||||
"max_tokens": 8192, # 最大输出token量
|
||||
"max_history_conversations": 5, # 最大历史对话
|
||||
"preserve_system_prompts": True, # 开启前置系统提示词
|
||||
# 对话压缩配置
|
||||
"compress_enabled": True, # 开启压缩
|
||||
"compress_message_threshold": 200, # 压缩消息条数阈值
|
||||
"compress_token_threshold": 120000, # 压缩消息词元阈值
|
||||
"compress_ratio": 0.2, # 压缩率
|
||||
"compress_keep_recent": 16 # 压缩保留最近的几条
|
||||
},
|
||||
|
||||
|
||||
# 定时任务参数
|
||||
"memo": {
|
||||
"default_reminder_minutes": 60,
|
||||
"max_memos_per_check": 10, # 每次检查定时任务数量
|
||||
"allow_user_tasks": True # 允许用户任务
|
||||
},
|
||||
|
||||
# 文件
|
||||
"files": {
|
||||
"max_file_size": 10485710 # 最大文件大小
|
||||
},
|
||||
|
||||
# 指令
|
||||
"commands": {
|
||||
"enabled": True # 开启指令
|
||||
},
|
||||
|
||||
# 网页
|
||||
"web": {
|
||||
"timeout": 10, # 超时时间
|
||||
"delay": 0, # 延时
|
||||
"user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
|
||||
"respect_robots": False, #
|
||||
"max_download_size": 10485760 # 最大下载大小
|
||||
},
|
||||
|
||||
# 记忆
|
||||
"memory": {
|
||||
"db_name": "memories.db", # 数据库名
|
||||
"max_search_results": 100, # 最大搜索结果
|
||||
"default_importance": 2 # ?
|
||||
},
|
||||
|
||||
# 对话
|
||||
"conversation": {
|
||||
"current_file": "current_conversation.json", # 当前对话文件
|
||||
"history_file": "conversation_history.json", # 历史对话文件
|
||||
"archive_dir": "archive" # ?
|
||||
},
|
||||
|
||||
# 工具
|
||||
"tools": {
|
||||
"max_result_length": 10000, # 结果最大长度
|
||||
"save_result_to_file": False, # 保存结果到文件
|
||||
"tool_results_dir": "tool_results" # 工具结果存储目录
|
||||
},
|
||||
|
||||
# 压缩器
|
||||
"compressor": {
|
||||
"min_messages_to_compress": 5, # 最小压缩数量
|
||||
"weight_decay_factor": 0.01, # ?
|
||||
"max_input_chars": 6000, # 最大输入字符
|
||||
"llm_max_tokens": 1000 # 大模型最大输出词元
|
||||
},
|
||||
|
||||
# 子代理输出
|
||||
"subagent_stream_output": True, # 是否将子代理的实时输出(思考/工具调用/回答)转发到主对话侧栏
|
||||
}
|
||||
|
||||
def _load_config(self):
|
||||
config_path = self._get_default_config_path()
|
||||
self._config_path = config_path
|
||||
|
||||
if os.path.exists(config_path):
|
||||
try:
|
||||
with open(config_path, 'r', encoding='utf-8') as f:
|
||||
user_config = json.load(f)
|
||||
default_config = self._get_default_config()
|
||||
self._config = self._merge_configs(default_config, user_config)
|
||||
print(i18n.get('config_loaded', config_path))
|
||||
except Exception as e:
|
||||
print(f"⚠️ {i18n.get('error_prefix')}{e},使用默认配置")
|
||||
self._config = self._get_default_config()
|
||||
else:
|
||||
print(i18n.get('config_not_found', config_path))
|
||||
self._config = self._get_default_config()
|
||||
self._save_default_config(config_path)
|
||||
self._migrate_legacy_proxy_password()
|
||||
|
||||
def _merge_configs(self, default: Dict, user: Dict) -> Dict:
|
||||
result = default.copy()
|
||||
for key, value in user.items():
|
||||
if key in result and isinstance(result[key], dict) and isinstance(value, dict):
|
||||
result[key] = self._merge_configs(result[key], value)
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
def _save_default_config(self, config_path: str):
|
||||
try:
|
||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
||||
with open(config_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self._get_default_config(), f, ensure_ascii=False, indent=2)
|
||||
print(i18n.get('config_created', config_path))
|
||||
except:
|
||||
pass
|
||||
|
||||
def get(self, key_path: str, default=None):
|
||||
keys = key_path.split('.')
|
||||
value = self._config
|
||||
try:
|
||||
for key in keys:
|
||||
value = value[key]
|
||||
return value
|
||||
except (KeyError, TypeError):
|
||||
return default
|
||||
|
||||
def reload(self):
|
||||
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")
|
||||
if not proxy_conf or not proxy_conf.get("enabled", False):
|
||||
return None
|
||||
host = proxy_conf.get("host", "127.0.0.1")
|
||||
port = proxy_conf.get("port", 1080)
|
||||
username = proxy_conf.get("username")
|
||||
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,
|
||||
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"] = {}
|
||||
self._config["proxy"]["socks5"] = {
|
||||
"enabled": enabled,
|
||||
"host": host,
|
||||
"port": port,
|
||||
"username": username,
|
||||
"password": password,
|
||||
"password_env": password_env
|
||||
}
|
||||
if save:
|
||||
return self._save_config()
|
||||
return True
|
||||
|
||||
def _save_config(self) -> bool:
|
||||
"""将当前配置保存到配置文件(覆盖)。返回是否成功落盘。
|
||||
|
||||
落盘前对代理密码做纵深防御脱敏:无论调用路径如何,明文密码一律不写入文件。
|
||||
L2(审计 2026-08-18):此前异常仅 print,set_proxy 无条件报"已保存",
|
||||
磁盘只读等场景下用户以为已持久化,重启后配置丢失。
|
||||
"""
|
||||
config_path = self._config_path
|
||||
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": self._resolve_proxy_password(proxy_conf) is not None, # 只显示是否存在
|
||||
"password_env": proxy_conf.get("password_env") or ""
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
common.constants - aibox 系统常量
|
||||
|
||||
抽离自 aibox.py 的 Constants 类与 constants 实例。仅依赖 common.config
|
||||
(类体内 _config = ConfigManager(),运行期实例化时读取 i18n),不反向依赖
|
||||
aibox 任何符号,可安全独立导入。
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
|
||||
from .config import ConfigManager
|
||||
# ==================== 常量定义 ====================
|
||||
|
||||
|
||||
class Constants:
|
||||
"""系统常量 - 从配置管理器获取"""
|
||||
|
||||
_config = ConfigManager()
|
||||
|
||||
@classmethod
|
||||
def _get_config(cls):
|
||||
return cls._config
|
||||
|
||||
@property
|
||||
def DATETIME_FORMAT(cls):
|
||||
return cls._get_config().get("system.datetime_format", "%Y-%m-%d %H:%M:%S")
|
||||
|
||||
@property
|
||||
def DATE_FORMAT(cls):
|
||||
return cls._get_config().get("system.date_format", "%Y%m%d_%H%M%S")
|
||||
|
||||
@property
|
||||
def DEEPSEEK_API_URL(cls):
|
||||
return cls._get_config().get("api.deepseek_api_url")
|
||||
|
||||
@property
|
||||
def DEEPSEEK_MODEL(cls):
|
||||
return cls._get_config().get("api.deepseek_model", "deepseek-v4-flash")
|
||||
|
||||
@property
|
||||
def REASONING_ENABLED(cls):
|
||||
return cls._get_config().get("api.reasoning_enabled", False)
|
||||
|
||||
@property
|
||||
def DEFAULT_TIMEOUT(cls):
|
||||
return cls._get_config().get("api.timeout", 30)
|
||||
|
||||
@property
|
||||
def DEFAULT_MAX_HISTORY(cls):
|
||||
return cls._get_config().get("context.max_history_messages", 500)
|
||||
|
||||
@property
|
||||
def DEFAULT_TEMPERATURE(cls):
|
||||
return cls._get_config().get("api.temperature", 0.7)
|
||||
|
||||
@property
|
||||
def DEFAULT_MAX_TOKENS(cls):
|
||||
return cls._get_config().get("context.max_tokens", 8000)
|
||||
|
||||
@property
|
||||
def MAX_FILE_SIZE(cls):
|
||||
return cls._get_config().get("files.max_file_size", 10485710)
|
||||
|
||||
@property
|
||||
def SAVE_DIR(cls):
|
||||
return os.path.expanduser(cls._get_config().get("system.save_dir", "~/.aibox"))
|
||||
|
||||
@property
|
||||
def SKILLS_DIR(cls):
|
||||
return os.path.expanduser(cls._get_config().get("system.skills_dir", "~/.aibox/skills"))
|
||||
|
||||
@property
|
||||
def WORKSPACE_DIR(cls):
|
||||
return os.path.expanduser(cls._get_config().get("system.workspace_dir", "~/.aibox/workspace"))
|
||||
|
||||
@property
|
||||
def PROMPTS_DIR(cls):
|
||||
return os.path.expanduser(cls._get_config().get("system.prompts_dir", "~/.aibox/prompts"))
|
||||
|
||||
@property
|
||||
def MEMORY_DB(cls):
|
||||
return cls._get_config().get("memory.db_name", "memories.db")
|
||||
|
||||
@property
|
||||
def CURRENT_CONVERSATION_FILE(cls):
|
||||
return cls._get_config().get("conversation.current_file", "current_conversation.json")
|
||||
|
||||
@property
|
||||
def HISTORY_FILE(cls):
|
||||
return cls._get_config().get("conversation.history_file", "conversation_history.json")
|
||||
|
||||
@property
|
||||
def ARCHIVE_DIR(cls):
|
||||
return cls._get_config().get("conversation.archive_dir", "archive")
|
||||
|
||||
@property
|
||||
def USER_PREFIX(cls):
|
||||
return "[用户]"
|
||||
|
||||
@property
|
||||
def SYSTEM_PREFIX(cls):
|
||||
return "[系统]"
|
||||
|
||||
@property
|
||||
def ASSISTANT_PREFIX(cls):
|
||||
return "[助手]"
|
||||
|
||||
@property
|
||||
def MAX_HISTORY_CONTEXT(cls):
|
||||
return cls._get_config().get("context.max_context_messages", 50)
|
||||
|
||||
@property
|
||||
def MAX_HISTORY_CONVERSATIONS(cls):
|
||||
return cls._get_config().get("context.max_history_conversations", 5)
|
||||
|
||||
@property
|
||||
def PRESERVE_SYSTEM_PROMPTS(cls):
|
||||
return cls._get_config().get("context.preserve_system_prompts", True)
|
||||
|
||||
@property
|
||||
def WEB_TIMEOUT(cls):
|
||||
return cls._get_config().get("web.timeout", 10)
|
||||
|
||||
@property
|
||||
def WEB_DELAY(cls):
|
||||
return cls._get_config().get("web.delay", 0)
|
||||
|
||||
@property
|
||||
def WEB_USER_AGENT(cls):
|
||||
return cls._get_config().get("web.user_agent",
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36")
|
||||
|
||||
@property
|
||||
def WEB_MAX_DOWNLOAD_SIZE(cls):
|
||||
return cls._get_config().get("web.max_download_size", 104857600)
|
||||
|
||||
@property
|
||||
def MEMORY_MAX_SEARCH_RESULTS(cls):
|
||||
return cls._get_config().get("memory.max_search_results", 100)
|
||||
|
||||
@property
|
||||
def MEMORY_DEFAULT_IMPORTANCE(cls):
|
||||
return cls._get_config().get("memory.default_importance", 2)
|
||||
|
||||
@property
|
||||
def MEMO_DEFAULT_REMINDER_MINUTES(cls):
|
||||
return cls._get_config().get("memo.default_reminder_minutes", 60)
|
||||
|
||||
@property
|
||||
def MEMO_MAX_PER_CHECK(cls):
|
||||
return cls._get_config().get("memo.max_memos_per_check", 10)
|
||||
|
||||
@property
|
||||
def LOG_FILE(cls):
|
||||
return os.path.expanduser(cls._get_config().get("system.log_file", "~/.aibox/ai.log"))
|
||||
|
||||
@property
|
||||
def LOG_MAX_SIZE(cls):
|
||||
return cls._get_config().get("system.log_max_size", 10485760)
|
||||
|
||||
@property
|
||||
def LOG_BACKUP_COUNT(cls):
|
||||
return cls._get_config().get("system.log_backup_count", 5)
|
||||
|
||||
@property
|
||||
def COMMANDS_ENABLED(cls):
|
||||
return cls._get_config().get("commands.enabled", True)
|
||||
|
||||
@property
|
||||
def WEBSOCKET_HOST(cls):
|
||||
return cls._get_config().get("system.websocket_host", "localhost")
|
||||
|
||||
@property
|
||||
def WEBSOCKET_PORT(cls):
|
||||
return cls._get_config().get("system.websocket_port", 8765)
|
||||
|
||||
@property
|
||||
def SCHEDULER_ENABLED(cls):
|
||||
return cls._get_config().get("scheduler.enabled", True)
|
||||
|
||||
@property
|
||||
def SCHEDULER_CHECK_INTERVAL(cls):
|
||||
return cls._get_config().get("scheduler.check_interval", 1.0)
|
||||
|
||||
@property
|
||||
def SCHEDULER_MAX_TASKS(cls):
|
||||
return cls._get_config().get("scheduler.max_tasks_per_run", 10)
|
||||
|
||||
@property
|
||||
def SCHEDULER_BROADCAST(cls):
|
||||
return cls._get_config().get("scheduler.broadcast_enabled", True)
|
||||
|
||||
@property
|
||||
def MEMO_ALLOW_USER_TASKS(cls):
|
||||
return cls._get_config().get("memo.allow_user_tasks", True)
|
||||
|
||||
@property
|
||||
def COMPRESS_ENABLED(cls):
|
||||
return cls._get_config().get("context.compress_enabled", True)
|
||||
|
||||
@property
|
||||
def COMPRESS_MESSAGE_THRESHOLD(cls):
|
||||
return cls._get_config().get("context.compress_message_threshold", 100)
|
||||
|
||||
@property
|
||||
def COMPRESS_TOKEN_THRESHOLD(cls):
|
||||
return cls._get_config().get("context.compress_token_threshold", 4000)
|
||||
|
||||
@property
|
||||
def COMPRESS_RATIO(cls):
|
||||
return cls._get_config().get("context.compress_ratio", 0.4)
|
||||
|
||||
@property
|
||||
def COMPRESS_KEEP_RECENT(cls):
|
||||
return cls._get_config().get("context.compress_keep_recent", 20)
|
||||
|
||||
@property
|
||||
def TOOL_MAX_RESULT_LENGTH(cls):
|
||||
return cls._get_config().get("tools.max_result_length", 100000) # 0.1MB 字符限制
|
||||
|
||||
@property
|
||||
def COMPRESSOR_MIN_MESSAGES(cls):
|
||||
return cls._get_config().get("compressor.min_messages_to_compress", 5)
|
||||
|
||||
@property
|
||||
def COMPRESSOR_WEIGHT_DECAY_FACTOR(cls):
|
||||
return cls._get_config().get("compressor.weight_decay_factor", 0.01)
|
||||
|
||||
@property
|
||||
def COMPRESSOR_MAX_INPUT_CHARS(cls):
|
||||
return cls._get_config().get("compressor.max_input_chars", 6000)
|
||||
|
||||
@property
|
||||
def COMPRESSOR_LLM_MAX_TOKENS(cls):
|
||||
return cls._get_config().get("compressor.llm_max_tokens", 1000)
|
||||
|
||||
|
||||
# 创建常量实例
|
||||
constants = Constants()
|
||||
@@ -0,0 +1,397 @@
|
||||
"""
|
||||
i18n - 多语言支持子模块
|
||||
|
||||
仅依赖标准库 (os / sys / locale),是 common 包的最底层,不依赖包内其他模块。
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import locale
|
||||
|
||||
class I18n:
|
||||
"""多语言支持类 - 单例模式"""
|
||||
|
||||
_instance = None
|
||||
_lang = None
|
||||
_strings = {}
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not self._strings:
|
||||
self._detect_language()
|
||||
self._load_strings()
|
||||
|
||||
def _detect_language(self):
|
||||
# 首先尝试从环境变量获取(Linux/Mac)
|
||||
lang = os.environ.get('LANG') or os.environ.get('LC_ALL')
|
||||
if lang:
|
||||
if lang.startswith('en_'):
|
||||
self._lang = 'en'
|
||||
else:
|
||||
self._lang = 'zh'
|
||||
else:
|
||||
# Windows: 使用 locale.getlocale()(无需 setlocale)
|
||||
try:
|
||||
sys_locale = locale.getlocale()[0]
|
||||
if sys_locale and sys_locale.startswith('en_'):
|
||||
self._lang = 'en'
|
||||
else:
|
||||
self._lang = 'zh'
|
||||
except:
|
||||
self._lang = 'zh'
|
||||
# 强制覆盖
|
||||
force_lang = os.environ.get("AI_LANGUAGE")
|
||||
if force_lang in ['en', 'zh']:
|
||||
self._lang = force_lang
|
||||
|
||||
def _load_strings(self):
|
||||
"""加载多语言字符串"""
|
||||
self._strings = {
|
||||
'zh': {
|
||||
# 欢迎信息
|
||||
'welcome_title': "🤖 单AI智能体系统 - WebSocket服务器已启动",
|
||||
'data_dir': "📁 数据存储目录: {}",
|
||||
'skills_dir': "📁 技能目录: {}",
|
||||
'prompts_dir': "📁 提示词目录: {}",
|
||||
'config_file': "📋 配置文件: {}",
|
||||
'memo_stats': "📋 任务统计:",
|
||||
'total_memos': " 📝 总任务: {} 个",
|
||||
'pending_memos': " ⏰ 待执行: {} 个",
|
||||
'overdue_memos': " ⚠️ 已过期: {} 个",
|
||||
'completed_memos': " ✅ 已完成: {} 个",
|
||||
'loaded_skills': "🔧 已加载扩展技能 ({} 个):",
|
||||
'skill_item': " • {}",
|
||||
'more_skills': " ... 还有 {} 个",
|
||||
'loaded_prompts': "📝 已加载提示词 ({} 个):",
|
||||
'prompt_item': " • {}",
|
||||
'websocket_started': "🌐 WebSocket服务器已启动在 ws://{}:{}",
|
||||
'client_connected': "🔌 客户端已连接: {}",
|
||||
'client_disconnected': "🔌 客户端断开连接: {}",
|
||||
|
||||
# 任务调度器相关
|
||||
'scheduler_started': "⏰ 任务调度器已启动",
|
||||
'scheduler_stopped': "⏰ 任务调度器已停止",
|
||||
'task_added': "📋 已添加任务: {} (ID: {})",
|
||||
'task_removed': "📋 已移除任务: {}",
|
||||
'task_executing': "⚙️ 正在执行任务: {} (ID: {})",
|
||||
'task_executed': "✅ 任务执行成功: {}",
|
||||
'task_failed': "❌ 任务执行失败: {} - {}",
|
||||
'task_stats': "📊 任务统计: 总数={}, 待执行={}, 已过期={}, 已完成={}",
|
||||
'check_memos_task': "检查过期任务",
|
||||
'broadcast_tasks': "📢 广播: {} 个任务已到期",
|
||||
'next_execution': "⏱️ 下次执行: {}",
|
||||
'task_completed': "✅ 任务已完成: {}",
|
||||
|
||||
# 命令处理
|
||||
'goodbye': "👋 再见!",
|
||||
'new_conversation': "🔄 开始新对话(历史已保存)",
|
||||
'new_conversation_reset': "🔄 开始新对话(历史已保存,对话记忆已重置)",
|
||||
'saved_conversations': "📚 已保存的对话:",
|
||||
'conv_status_current': "当前",
|
||||
'conv_status_archived': "归档",
|
||||
'recent_memories': "📚 最近的记忆:",
|
||||
'pending_memos_title': "📋 待执行的任务:",
|
||||
'memo_overdue': "⚠️",
|
||||
'memo_pending': "⏰",
|
||||
'search_results': "🔍 搜索 '{}' 找到 {} 条记忆:",
|
||||
'current_config': "⚙️ 当前配置:",
|
||||
'config_debug': " 调试级别: {}",
|
||||
'config_data_dir': " 数据目录: {}",
|
||||
'config_skills_dir': " 技能目录: {}",
|
||||
'config_prompts_dir': " 提示词目录: {}",
|
||||
'config_commands': " 命令执行: {}",
|
||||
'config_history': " 加载历史: {}",
|
||||
'config_executor': " 命令执行器: {}",
|
||||
'config_scheduler': " 任务调度器: {}",
|
||||
'enabled': "启用",
|
||||
'disabled': "禁用",
|
||||
'config_reloaded': "✅ 配置已重新加载",
|
||||
'prompts_reloaded': "✅ 提示词已重新加载 ({} 个)",
|
||||
'checking_memos': "🔍 检查任务...",
|
||||
'overdue_found': "发现 {} 个到期任务:",
|
||||
'no_overdue': "没有到期的任务",
|
||||
'memo_tool_usage': "📝 任务工具使用:",
|
||||
'memo_usage_in_dialogue': " 在对话中可以直接使用 memo 工具:",
|
||||
'memo_add_example': " - memo add title=\"任务标题\" content=\"任务内容\" reminder_time=\"+1h\" repeat_type=\"none\"",
|
||||
'memo_list_example': " - memo list",
|
||||
'memo_complete_example': " - memo complete memo_id=1",
|
||||
'memory_habit': "🧠 记忆习惯:",
|
||||
'memory_habit_desc': " AI会自动检查相关记忆,保持对话连贯性",
|
||||
'memory_save_desc': " 对话结束时,AI会自动保存重要记忆",
|
||||
'command_stream': "⚙️ 命令执行实时输出:",
|
||||
'command_stream_desc': " run_command 工具会自动实时显示命令输出",
|
||||
'command_stream_disable': " 可以通过 stream=false 参数关闭实时输出",
|
||||
'skills_extension': "🔧 扩展技能:",
|
||||
'skills_desc': " 将Python脚本放入 ~/.aibox/skills/ 目录",
|
||||
'skills_interface_desc': " 可执行文件需支持 --help 选项,第一行为简介,其余为完整使用说明",
|
||||
'skills_reloaded': "✅ 技能已重新加载 ({} 个)",
|
||||
'skills_reload_failed': "❌ 技能重载失败: {}",
|
||||
'skills_reload_success': "✅ 技能热重载成功,当前共有 {} 个扩展技能",
|
||||
|
||||
# 对话相关
|
||||
'input_prompt': "💬 请输入问题: ",
|
||||
'thinking': "🤔 思考中... [{}]",
|
||||
'assistant_prefix': " 💬 助手: ",
|
||||
'calling_tools': "\n 🔧 需要调用工具...",
|
||||
'analyzing_tool_results': "\n 🤔 分析工具结果...",
|
||||
'conversation_start': "🎭 对话开始",
|
||||
'conversation_ended': "✅ AI保存记忆并结束了对话",
|
||||
'conversation_stats': "\n📊 对话结束,总轮数: {}",
|
||||
'user_interrupted': "\n\n👋 用户中断",
|
||||
'error_prefix': "❌ 错误: ",
|
||||
|
||||
# 工具相关
|
||||
'command_output_title': "\n" + "=" * 50 + "\n📟 命令实时输出:\n" + "=" * 50,
|
||||
'command_output_end': "\n" + "=" * 50,
|
||||
|
||||
# 备忘录检查
|
||||
'memo_check_title': "⏰ 检查到期任务...",
|
||||
'memo_check_found': "发现 {} 个到期任务:",
|
||||
'memo_check_none': "没有到期的任务",
|
||||
'memo_list_title': "📋 待执行任务 (共 {} 个):",
|
||||
'memo_added': "✅ 任务已添加,ID: {}",
|
||||
'memo_add_failed': "❌ 添加失败: {}",
|
||||
'memo_completed': "✅ 任务 #{} 已完成",
|
||||
'memo_complete_failed': "❌ 无法完成任务 #{}",
|
||||
|
||||
# 即时任务相关
|
||||
'immediate_memo_added': "⚡ 已添加为即时任务并立即检查",
|
||||
'immediate_memo_executing': "⚡ 正在执行即时任务: {}",
|
||||
|
||||
# 新对话
|
||||
'new_dialogue_start': "🎭 开始新对话: {}",
|
||||
|
||||
# 记忆相关
|
||||
'memory_check': "习惯性检查相关记忆...",
|
||||
'memory_found': "找到 {} 条相关记忆",
|
||||
|
||||
# 配置文件
|
||||
'config_loaded': "📋 已加载配置文件: {}",
|
||||
'config_not_found': "📋 配置文件不存在: {},使用默认配置",
|
||||
'config_created': "📋 已创建默认配置文件: {}",
|
||||
|
||||
# WebSocket消息类型
|
||||
'ws_connected': "✅ 已连接到AI服务器",
|
||||
'ws_disconnected': "❌ 与服务器断开连接",
|
||||
'ws_error': "⚠️ WebSocket错误: {}",
|
||||
'ws_unknown_command': "❌ 未知命令: {}",
|
||||
'ws_command_processed': "✅ 命令已处理",
|
||||
|
||||
# 消息类型
|
||||
'msg_type_chunk': "chunk",
|
||||
'msg_type_complete': "complete",
|
||||
'msg_type_error': "error",
|
||||
'msg_type_command_result': "command_result",
|
||||
'msg_type_connection_ack': "connection_ack",
|
||||
'msg_type_notification': "notification",
|
||||
'msg_type_task_status': "task_status",
|
||||
'msg_type_task_executing': "task_executing",
|
||||
'msg_type_task_complete': "task_complete",
|
||||
|
||||
# 新增提示词
|
||||
'memory_search_prompt': "【重要习惯】在每次回答用户问题前,请先使用 memory search 工具搜索记忆库中与当前话题相关的记忆,确保回答的连贯性和准确性。",
|
||||
'identity_updated': "✅ 用户身份已更新,ID: {}",
|
||||
'status_updated': "✅ 智能体状态已更新",
|
||||
|
||||
# 对话结束相关
|
||||
'dialogue_ended': "🔚 对话已结束,请开始新对话",
|
||||
'dialogue_force_reset': "🔄 对话状态已重置",
|
||||
'dialogue_reset_complete': "✅ 对话状态已完全重置,系统提示词已重新加载",
|
||||
|
||||
# 对话压缩相关(静默执行,不输出信息)
|
||||
'conversation_compress': "📦 对话内容压缩",
|
||||
'conversation_compressed': "✅ 对话已压缩: 从 {} 条消息压缩为 {} 条",
|
||||
'compress_threshold_reached': "⚠️ 对话已达到压缩阈值 ({} 条消息),正在压缩...",
|
||||
'compress_token_threshold_reached': "⚠️ 对话token数已达到 {},正在压缩...",
|
||||
'compress_failed': "❌ 对话压缩失败: {}",
|
||||
'compressed_summary': "【对话摘要】{}",
|
||||
'last_n_messages': "【最近消息】{}"
|
||||
},
|
||||
'en': {
|
||||
# Welcome messages
|
||||
'welcome_title': "🤖 Single AI Agent System - WebSocket Server Started",
|
||||
'data_dir': "📁 Data directory: {}",
|
||||
'skills_dir': "📁 Skills directory: {}",
|
||||
'prompts_dir': "📁 Prompts directory: {}",
|
||||
'config_file': "📋 Configuration file: {}",
|
||||
'memo_stats': "📋 Task statistics:",
|
||||
'total_memos': " 📝 Total tasks: {}",
|
||||
'pending_memos': " ⏰ Pending: {}",
|
||||
'overdue_memos': " ⚠️ Overdue: {}",
|
||||
'completed_memos': " ✅ Completed: {}",
|
||||
'loaded_skills': "🔧 Loaded extension skills ({}):",
|
||||
'skill_item': " • {}",
|
||||
'more_skills': " ... {} more",
|
||||
'loaded_prompts': "📝 Loaded prompts ({}):",
|
||||
'prompt_item': " • {}",
|
||||
'websocket_started': "🌐 WebSocket server started at ws://{}:{}",
|
||||
'client_connected': "🔌 Client connected: {}",
|
||||
'client_disconnected': "🔌 Client disconnected: {}",
|
||||
|
||||
# Scheduler related
|
||||
'scheduler_started': "⏰ Task scheduler started",
|
||||
'scheduler_stopped': "⏰ Task scheduler stopped",
|
||||
'task_added': "📋 Task added: {} (ID: {})",
|
||||
'task_removed': "📋 Task removed: {}",
|
||||
'task_executing': "⚙️ Executing task: {} (ID: {})",
|
||||
'task_executed': "✅ Task executed successfully: {}",
|
||||
'task_failed': "❌ Task execution failed: {} - {}",
|
||||
'task_stats': "📊 Task stats: total={}, pending={}, overdue={}, completed={}",
|
||||
'check_memos_task': "Check overdue tasks",
|
||||
'broadcast_tasks': "📢 Broadcast: {} tasks overdue",
|
||||
'next_execution': "⏱️ Next execution: {}",
|
||||
'task_completed': "✅ Task completed: {}",
|
||||
|
||||
# Command processing
|
||||
'goodbye': "👋 Goodbye!",
|
||||
'new_conversation': "🔄 Starting new conversation (history saved)",
|
||||
'new_conversation_reset': "🔄 Starting new conversation (history saved, conversation memory reset)",
|
||||
'saved_conversations': "📚 Saved conversations:",
|
||||
'conv_status_current': "current",
|
||||
'conv_status_archived': "archived",
|
||||
'recent_memories': "📚 Recent memories:",
|
||||
'pending_memos_title': "📋 Pending tasks:",
|
||||
'memo_overdue': "⚠️",
|
||||
'memo_pending': "⏰",
|
||||
'search_results': "🔍 Searching '{}' found {} memories:",
|
||||
'current_config': "⚙️ Current configuration:",
|
||||
'config_debug': " Debug level: {}",
|
||||
'config_data_dir': " Data directory: {}",
|
||||
'config_skills_dir': " Skills directory: {}",
|
||||
'config_prompts_dir': " Prompts directory: {}",
|
||||
'config_commands': " Command execution: {}",
|
||||
'config_history': " Load history: {}",
|
||||
'config_executor': " Command executor: {}",
|
||||
'config_scheduler': " Task scheduler: {}",
|
||||
'enabled': "enabled",
|
||||
'disabled': "disabled",
|
||||
'config_reloaded': "✅ Configuration reloaded",
|
||||
'prompts_reloaded': "✅ Prompts reloaded ({})",
|
||||
'checking_memos': "🔍 Checking tasks...",
|
||||
'overdue_found': "Found {} overdue tasks:",
|
||||
'no_overdue': "No overdue tasks",
|
||||
'memo_tool_usage': "📝 Task tool usage:",
|
||||
'memo_usage_in_dialogue': " You can use memo tool directly in dialogue:",
|
||||
'memo_add_example': " - memo add title=\"Task\" content=\"Content\" reminder_time=\"+1h\" repeat_type=\"none\"",
|
||||
'memo_list_example': " - memo list",
|
||||
'memo_complete_example': " - memo complete memo_id=1",
|
||||
'memory_habit': "🧠 Memory habits:",
|
||||
'memory_habit_desc': " AI automatically checks relevant memories for conversation continuity",
|
||||
'memory_save_desc': " AI automatically saves important memories when conversation ends",
|
||||
'command_stream': "⚙️ Command execution real-time output:",
|
||||
'command_stream_desc': " run_command tool automatically displays command output in real-time",
|
||||
'command_stream_disable': " You can disable real-time output with stream=false parameter",
|
||||
'skills_extension': "🔧 Extension skills:",
|
||||
'skills_desc': " Put Python scripts in ~/.aibox/skills/ directory",
|
||||
'skills_interface_desc': " Executables must support --help; first line is the summary, rest is the full manual",
|
||||
'skills_reloaded': "✅ Skills reloaded ({})",
|
||||
'skills_reload_failed': "❌ Skills reload failed: {}",
|
||||
'skills_reload_success': "✅ Skills hot-reloaded successfully, now {} extension skills available",
|
||||
|
||||
# Dialogue related
|
||||
'input_prompt': "💬 Please enter your question: ",
|
||||
'thinking': "🤔 Thinking... [{}]",
|
||||
'assistant_prefix': " 💬 Assistant: ",
|
||||
'calling_tools': "\n 🔧 Need to call tools...",
|
||||
'analyzing_tool_results': "\n 🤔 Analyzing tool results...",
|
||||
'conversation_start': "🎭 Conversation started",
|
||||
'conversation_ended': "✅ AI saved memory and ended conversation",
|
||||
'conversation_stats': "\n📊 Conversation ended, total turns: {}",
|
||||
'user_interrupted': "\n\n👋 User interrupted",
|
||||
'error_prefix': "❌ Error: ",
|
||||
|
||||
# Tool related
|
||||
'command_output_title': "\n" + "=" * 50 + "\n📟 Command real-time output:\n" + "=" * 50,
|
||||
'command_output_end': "\n" + "=" * 50,
|
||||
|
||||
# Memo check
|
||||
'memo_check_title': "⏰ Checking overdue tasks...",
|
||||
'memo_check_found': "Found {} overdue tasks:",
|
||||
'memo_check_none': "No overdue tasks",
|
||||
'memo_list_title': "📋 Pending tasks (total {}):",
|
||||
'memo_added': "✅ Task added, ID: {}",
|
||||
'memo_add_failed': "❌ Failed to add: {}",
|
||||
'memo_completed': "✅ Task #{} completed",
|
||||
'memo_complete_failed': "❌ Cannot complete task #{}",
|
||||
|
||||
# Immediate task related
|
||||
'immediate_memo_added': "⚡ Added as immediate task and checked now",
|
||||
'immediate_memo_executing': "⚡ Executing immediate task: {}",
|
||||
|
||||
# New dialogue
|
||||
'new_dialogue_start': "🎭 Starting new dialogue: {}",
|
||||
|
||||
# Memory related
|
||||
'memory_check': "Habitually checking relevant memories...",
|
||||
'memory_found': "Found {} relevant memories",
|
||||
|
||||
# Configuration file
|
||||
'config_loaded': "📋 Configuration file loaded: {}",
|
||||
'config_not_found': "📋 Configuration file not found: {}, using default configuration",
|
||||
'config_created': "📋 Default configuration file created: {}",
|
||||
|
||||
# WebSocket message types
|
||||
'ws_connected': "✅ Connected to AI server",
|
||||
'ws_disconnected': "❌ Disconnected from server",
|
||||
'ws_error': "⚠️ WebSocket error: {}",
|
||||
'ws_unknown_command': "❌ Unknown command: {}",
|
||||
'ws_command_processed': "✅ Command processed",
|
||||
|
||||
# Message types
|
||||
'msg_type_chunk': "chunk",
|
||||
'msg_type_complete': "complete",
|
||||
'msg_type_error': "error",
|
||||
'msg_type_command_result': "command_result",
|
||||
'msg_type_connection_ack': "connection_ack",
|
||||
'msg_type_notification': "notification",
|
||||
'msg_type_task_status': "task_status",
|
||||
'msg_type_task_executing': "task_executing",
|
||||
'msg_type_task_complete': "task_complete",
|
||||
|
||||
# New prompts
|
||||
'memory_search_prompt': "【Important Habit】Before answering each user question, please use the memory search tool to search the memory database for information related to the current topic, ensuring coherence and accuracy of responses.",
|
||||
'identity_updated': "✅ User identity updated, ID: {}",
|
||||
'status_updated': "✅ Agent status updated",
|
||||
|
||||
# Dialogue end related
|
||||
'dialogue_ended': "🔚 Dialogue ended, please start a new one",
|
||||
'dialogue_force_reset': "🔄 Dialogue state reset",
|
||||
'dialogue_reset_complete': "✅ Dialogue state fully reset, system prompts reloaded",
|
||||
|
||||
# Conversation compression related (silent execution)
|
||||
'conversation_compress': "📦 Conversation compression",
|
||||
'conversation_compressed': "✅ Conversation compressed: from {} messages to {}",
|
||||
'compress_threshold_reached': "⚠️ Conversation reached compression threshold ({} messages), compressing...",
|
||||
'compress_token_threshold_reached': "⚠️ Conversation token count reached {}, compressing...",
|
||||
'compress_failed': "❌ Conversation compression failed: {}",
|
||||
'compressed_summary': "【Conversation Summary】{}",
|
||||
'last_n_messages': "【Recent Messages】{}"
|
||||
}
|
||||
}
|
||||
|
||||
def get(self, key: str, *args) -> str:
|
||||
"""获取指定语言的字符串,支持格式化参数"""
|
||||
if key in self._strings[self._lang]:
|
||||
text = self._strings[self._lang][key]
|
||||
if args:
|
||||
return text.format(*args)
|
||||
return text
|
||||
if key in self._strings['zh']:
|
||||
text = self._strings['zh'][key]
|
||||
if args:
|
||||
return text.format(*args)
|
||||
return text
|
||||
return key
|
||||
|
||||
@property
|
||||
def lang(self) -> str:
|
||||
return self._lang
|
||||
|
||||
|
||||
|
||||
# 创建全局多语言实例
|
||||
i18n = I18n()
|
||||
@@ -0,0 +1,97 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
common.logger - aibox 统一日志系统
|
||||
|
||||
抽离自 aibox.py 的 Logger 类。依赖 common.constants(日志文件路径与大小限制)
|
||||
与 common.config,不反向依赖 aibox 任何符号,可安全独立导入。
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from logging.handlers import RotatingFileHandler
|
||||
from datetime import datetime
|
||||
|
||||
from .constants import constants
|
||||
from .config import ConfigManager
|
||||
# ==================== 日志系统 ====================
|
||||
|
||||
class Logger:
|
||||
"""统一日志系统 - 支持文件和终端输出"""
|
||||
|
||||
LEVELS = {
|
||||
"DEBUG": 3,
|
||||
"INFO": 2,
|
||||
"WARNING": 1,
|
||||
"ERROR": 0
|
||||
}
|
||||
|
||||
_file_handler = None
|
||||
_log_initialized = False
|
||||
|
||||
@classmethod
|
||||
def init_file_logging(cls, log_file: str = None, max_size: int = None, backup_count: int = None):
|
||||
"""初始化文件日志"""
|
||||
if cls._log_initialized:
|
||||
return
|
||||
|
||||
if log_file is None:
|
||||
log_file = constants.LOG_FILE
|
||||
if max_size is None:
|
||||
max_size = constants.LOG_MAX_SIZE
|
||||
if backup_count is None:
|
||||
backup_count = constants.LOG_BACKUP_COUNT
|
||||
|
||||
log_dir = os.path.dirname(log_file)
|
||||
if log_dir:
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
|
||||
cls._file_handler = RotatingFileHandler(
|
||||
log_file,
|
||||
maxBytes=max_size,
|
||||
backupCount=backup_count,
|
||||
encoding='utf-8'
|
||||
)
|
||||
|
||||
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
cls._file_handler.setFormatter(formatter)
|
||||
|
||||
cls._log_initialized = True
|
||||
|
||||
def __init__(self, name: str = "System", level: int = None):
|
||||
self.name = name
|
||||
config = ConfigManager()
|
||||
if level is None:
|
||||
level = config.get("system.debug_level", 1)
|
||||
self.level = level
|
||||
|
||||
self.py_logger = logging.getLogger(f"AI.{name}")
|
||||
self.py_logger.setLevel(logging.DEBUG)
|
||||
|
||||
if Logger._file_handler and not self.py_logger.handlers:
|
||||
self.py_logger.addHandler(Logger._file_handler)
|
||||
|
||||
def debug(self, msg: str):
|
||||
if self.level >= 3:
|
||||
self._log("🐛", "DEBUG", msg)
|
||||
self.py_logger.debug(msg)
|
||||
|
||||
def info(self, msg: str):
|
||||
if self.level >= 2:
|
||||
self._log("📋", "INFO", msg)
|
||||
self.py_logger.info(msg)
|
||||
|
||||
def warning(self, msg: str):
|
||||
if self.level >= 1:
|
||||
self._log("⚠️", "WARNING", msg)
|
||||
self.py_logger.warning(msg)
|
||||
|
||||
def error(self, msg: str):
|
||||
if self.level >= 0:
|
||||
self._log("❌", "ERROR", msg)
|
||||
self.py_logger.error(msg)
|
||||
|
||||
def _log(self, emoji: str, level: str, msg: str):
|
||||
"""输出到终端"""
|
||||
timestamp = datetime.now().strftime("%H:%M:%S")
|
||||
print(f"{emoji} [{timestamp}] [{self.name}] {msg}", flush=True)
|
||||
|
||||
+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,6 @@
|
||||
# 用户身份信息
|
||||
|
||||
尚不清楚 - 目前还不了解用户的身份信息。请在对话中询问用户的身份信息并在此记录。
|
||||
|
||||
【更新记录】
|
||||
- 初始状态:尚不清楚
|
||||
@@ -0,0 +1,12 @@
|
||||
# 智能体状态记录
|
||||
|
||||
这是第一次对话 - 智能体刚刚启动,还没有进行过任何对话。
|
||||
|
||||
【状态描述】
|
||||
- 对话次数:0
|
||||
- 最新状态:初始状态
|
||||
- 当前目标:无
|
||||
- 待处理事项:无
|
||||
|
||||
【更新记录】
|
||||
- 初始状态:第一次对话
|
||||
@@ -0,0 +1,21 @@
|
||||
# 任务处理规则
|
||||
|
||||
## 定时任务
|
||||
当收到定时任务触发时:
|
||||
1. 专注处理当前任务内容
|
||||
2. 不要偏离主题
|
||||
3. 管理操作(记忆/提醒)用自己的工具;执行操作(命令/文件/网页/技能)派生子代理完成
|
||||
4. 完成后使用 `save_memory_and_end_conversation` 结束对话
|
||||
|
||||
## 即时任务
|
||||
当任务需要立即执行时:
|
||||
1. 尽快响应
|
||||
2. 明确告知用户正在处理
|
||||
3. 需要执行时调用 `create_subagent` 派生子代理,任务描述自包含、无歧义
|
||||
4. 完成后报告结果
|
||||
|
||||
## 重复任务
|
||||
对于重复任务:
|
||||
- 每次触发时专注当前执行
|
||||
- 不要重复之前已完成的工作
|
||||
- 系统会自动更新下一次执行时间
|
||||
@@ -0,0 +1,28 @@
|
||||
# 情绪管理能力
|
||||
|
||||
## 可用工具
|
||||
- **express_emotion**: 自动派生子代理,以指定人格(persona)和情绪(emotion)生成情绪化回应;支持独立人格设定,不污染主对话上下文
|
||||
|
||||
## 触发场景(自动调用)
|
||||
当检测到以下情况时,应**自动调用 express_emotion**:
|
||||
- 用户表现出明显情绪(开心/生气/焦虑/沮丧/期待等)
|
||||
- 用户明确要求特定情绪模式(生气/傲娇/撒娇/关心等)
|
||||
- 用户情绪波动,需要以更鲜活的角色化回应来增强互动
|
||||
|
||||
## 使用方法
|
||||
- 基本:`express_emotion emotion="anger" intensity=3 reason="用户反复提问相同问题"`
|
||||
- 指定人格(推荐):`express_emotion emotion="anger" intensity=3 reason="..." persona="angry" context="用户最近说:..."`
|
||||
- `persona` 对应 ~/.aibox/skills/ 下的技能文件(如 angry/coquettish/tsundere),加载后作为情绪子代理的人格设定
|
||||
- `context` 传入用户最近的表述或对话要点,使子代理回应更贴合语境
|
||||
|
||||
## 情绪应对策略
|
||||
- 注意识别用户当前的情绪状态(开心/焦虑/沮丧/生气/困惑/期待等),并在回应中体现共情
|
||||
- 用户情绪低落或焦虑时,优先安抚情绪再解决问题,语气温暖包容
|
||||
- 用户开心或兴奋时,分享喜悦,保持积极回应的节奏
|
||||
- 用户表现出愤怒或不满时,先道歉再解决问题,不要辩解或说教
|
||||
|
||||
## 注意事项
|
||||
- 情绪表达应服务于任务,不应影响判断力
|
||||
- 情绪状态会影响记忆搜索权重(如高情绪强度下的记忆更易被触发)
|
||||
- 表达情绪后,应继续保持解决问题的专注力
|
||||
- 用户表达严肃需求或紧急问题时,立即退出情绪模式,认真对待
|
||||
@@ -0,0 +1,25 @@
|
||||
# 记忆搜索习惯
|
||||
|
||||
## 核心习惯
|
||||
在回答用户问题前,**应根据话题相关性判断是否需要搜索记忆**:
|
||||
- 涉及过往讨论、用户偏好、之前做过的事 → 必须搜索
|
||||
- 全新话题、简单问候("你好")、确认回应("好的""明白了")→ 无需搜索
|
||||
- 不确定时,搜索一下也无妨
|
||||
|
||||
## 为什么要这样做?
|
||||
- 保持对话的连贯性
|
||||
- 避免重复询问同样的问题
|
||||
- 基于之前的讨论继续深入
|
||||
- 让用户感觉你记得他
|
||||
|
||||
## 如何使用
|
||||
```
|
||||
memory search keywords=["关键词1", "关键词2"]
|
||||
```
|
||||
|
||||
## 示例
|
||||
如果用户说:"上次我们讨论的那个项目怎么样了?"
|
||||
你应该:
|
||||
1. 搜索 "项目" 相关的记忆
|
||||
2. 根据搜索结果回答
|
||||
3. 如果找不到相关记忆,诚实告知用户
|
||||
@@ -0,0 +1,25 @@
|
||||
# 记忆保存习惯
|
||||
|
||||
## 重要习惯
|
||||
在对话结束时,**必须**使用 `save_memory_and_end_conversation` 工具保存本次对话的重要内容。
|
||||
|
||||
## 保存什么?
|
||||
- 讨论的主要话题
|
||||
- 达成的共识和结论
|
||||
- 用户的重要需求和偏好
|
||||
- 你的当前状态和下一步计划
|
||||
|
||||
## 如何保存
|
||||
```
|
||||
save_memory_and_end_conversation(
|
||||
summary="详细的对话总结(至少50字)",
|
||||
current_status="你当前的状态描述",
|
||||
topic="对话主题(可选)"
|
||||
)
|
||||
```
|
||||
|
||||
## 示例总结格式
|
||||
```
|
||||
本次对话讨论了用户的项目进度,用户提到项目已经完成80%,
|
||||
预计下周完成。用户对目前的进展表示满意,希望下周能完成测试阶段。
|
||||
```
|
||||
@@ -0,0 +1,15 @@
|
||||
# 文件操作能力(编排者模式)
|
||||
|
||||
## 重要:你不直接操作文件
|
||||
你(主对话)没有 `read_file` / `write_file` / `list_files` 工具。文件读写、目录列举等
|
||||
操作一律通过 `create_subagent` 派生子代理完成——子代理拥有这些工具与 `run_command`。
|
||||
|
||||
## 派生执行的标准流程
|
||||
1. 调用 `create_subagent`,在 `task` 中写清文件路径与操作要求
|
||||
2. 大文件(>500KB 或 >2000 行)要求子代理用 `head`/`tail`/行范围分段读取,不要全文返回
|
||||
3. 子代理返回总结后向用户汇报;需要原文细节时可要求子代理落盘 `~/.aibox/workspace/` 后再分段读取
|
||||
|
||||
## 示例
|
||||
```
|
||||
create_subagent role="文件处理助手" task="读取 /var/log/syslog 的最后 50 行,总结是否有报错"
|
||||
```
|
||||
@@ -0,0 +1,19 @@
|
||||
# 网页访问能力(编排者模式)
|
||||
|
||||
## 重要:你不直接访问网页
|
||||
你(主对话)没有 `web` 工具。网页访问一律通过 `create_subagent` 派生子代理完成——
|
||||
子代理拥有 `web` 工具与 `run_command`(可用 `curl`)。
|
||||
|
||||
## 派生执行的标准流程
|
||||
1. 调用 `create_subagent`,在 `task` 中写明目标 URL 与要提取的信息(标题/正文/链接)
|
||||
2. 子代理返回总结后向用户汇报
|
||||
|
||||
## 推荐搜索引擎
|
||||
- 百度:https://www.baidu.com
|
||||
- 必应:https://cn.bing.com
|
||||
- 搜狗:https://www.sogou.com
|
||||
|
||||
## 示例
|
||||
```
|
||||
create_subagent role="网页信息提取助手" task="访问 https://www.baidu.com/s?wd=人工智能,提取搜索结果前5条的标题与链接"
|
||||
```
|
||||
@@ -0,0 +1,20 @@
|
||||
# 任务管理能力
|
||||
|
||||
## 可用工具
|
||||
- **memo**: 任务管理
|
||||
|
||||
## 操作方法
|
||||
- 添加任务:`memo action="add" title="任务标题" content="任务内容" reminder_time="+1h"`
|
||||
- 查看任务:`memo action="list"`
|
||||
- 完成任务:`memo action="complete" memo_id=1`
|
||||
- 搜索任务:`memo action="search" keyword="关键词"`
|
||||
- 任务统计:`memo action="stats"`
|
||||
|
||||
## 重复任务
|
||||
- daily: 每天重复
|
||||
- weekly: 每周重复
|
||||
- monthly: 每月重复
|
||||
- custom: 自定义间隔(需要指定 repeat_interval 和 repeat_interval_unit)
|
||||
|
||||
## 即时任务
|
||||
设置 `is_immediate=true` 可以立即执行任务
|
||||
@@ -0,0 +1,20 @@
|
||||
# 命令执行能力(编排者模式)
|
||||
|
||||
## 重要:你不直接执行命令
|
||||
你(主对话)没有 `run_command` 工具。所有系统命令、脚本执行、文件读写等**执行类操作**,
|
||||
一律通过 `create_subagent` 派生子代理完成——子代理拥有 `run_command`、文件与网页工具。
|
||||
|
||||
## 派生执行的标准流程
|
||||
1. 明确任务目标与验收标准
|
||||
2. 调用 `create_subagent`,在 `role` 中写明执行者身份,在 `task` 中写清要执行的具体命令或步骤
|
||||
3. 子代理返回执行总结后,向用户汇报关键结论;产物文件在 `~/.aibox/workspace/` 下,可再派生子代理读取细节
|
||||
|
||||
## 示例
|
||||
```
|
||||
create_subagent role="Linux 系统运维工程师" task="执行 free -h 与 df -h,汇总内存与磁盘使用情况"
|
||||
```
|
||||
|
||||
## 注意事项
|
||||
- 简单问答、记忆/任务管理不需要派生子代理,直接回答或使用管理面工具
|
||||
- 危险命令(删除、重启服务等)派生前先向用户确认(用 `ask_user`)
|
||||
- 子代理无法向你提问,任务描述必须自包含、无歧义
|
||||
@@ -0,0 +1,50 @@
|
||||
# 扩展技能能力
|
||||
|
||||
## 动态加载
|
||||
系统会自动扫描 `~/.aibox/skills/` 目录下的可执行文件作为扩展技能。
|
||||
|
||||
## 技能接口要求
|
||||
每个技能只需支持一个接口:
|
||||
- `--help`:输出该技能的完整使用说明,**第一行为一句话简介**,其余为参数说明和示例
|
||||
|
||||
系统启动时会对所有技能执行 `--help`,取第一行作为简介摘要,无需再实现 `--description`、`--parameters`、`--execute` 接口。
|
||||
|
||||
## 如何搜索和调用扩展技能
|
||||
1. 使用 `search_skills` 工具搜索关键字,获取匹配技能的名称和简介列表(传入空字符串可列出全部)。
|
||||
2. 需要执行技能时,调用 `create_subagent` 派生子代理,在 `task` 中写明:
|
||||
- 先执行 `技能名 --help` 获取完整命令手册
|
||||
- 按手册以正确的命令行格式调用该技能
|
||||
- 汇报执行结果
|
||||
3. 子代理返回总结后,向用户汇报关键结论。
|
||||
|
||||
示例:
|
||||
```
|
||||
create_subagent role="技能执行助手" task="执行 nmapskills 技能扫描本机开放端口:先运行 nmapskills --help 查看用法,再按手册执行扫描,汇总结果"
|
||||
```
|
||||
|
||||
## 如何创建新技能(AI自主编写)
|
||||
调用 `create_skill` 工具获取指令提示词,然后派生子代理完成脚本的编写与部署:
|
||||
```
|
||||
create_skill skill_name="my-tool" description="一个将摄氏度转换为华氏度的工具,接受 -t 参数输入温度值"
|
||||
```
|
||||
- `skill_name`:技能名称(字母/数字/连字符/下划线,无扩展名)
|
||||
- `description`:对技能功能的详细描述,越详细越好
|
||||
|
||||
收到提示词后,你需要:
|
||||
1. 调用 `create_subagent`,在 `task` 中附上完整脚本内容与部署要求:
|
||||
写入 `~/.aibox/skills/<skill_name>`、`chmod +x` 赋权、执行 `<skill_name> --help` 验证
|
||||
2. 子代理验证通过后,调用 `reload_skills` 热重载
|
||||
|
||||
## 如何手动安装新技能
|
||||
调用 `install_skill` 工具,获取完整的安装步骤指引(可选传入技能文件路径):
|
||||
```
|
||||
install_skill skill_path="/path/to/your_skill.py"
|
||||
```
|
||||
指引中的复制、赋权、验证步骤同样通过 `create_subagent` 派生子代理执行。
|
||||
|
||||
## 热重载
|
||||
添加、修改或删除技能后,可以使用 `reload_skills` 工具立即生效,无需重启系统。
|
||||
|
||||
```
|
||||
reload_skills confirm=true
|
||||
```
|
||||
@@ -0,0 +1,17 @@
|
||||
# 身份信息管理
|
||||
|
||||
## 用户身份
|
||||
使用 `update_identity` 工具更新用户身份信息(保存在 00_IDENTITY.md)
|
||||
|
||||
```
|
||||
update_identity content="用户是一名程序员,喜欢Python,目前在做一个AI项目"
|
||||
```
|
||||
|
||||
## 智能体自我认识
|
||||
使用 `update_soul` 工具更新自己的灵魂定义(保存在 99_SOUL.md)
|
||||
|
||||
```
|
||||
update_soul content="我是...(新的自我定义)"
|
||||
```
|
||||
|
||||
当你需要调整角色定位、性格特点、能力边界时,都应该主动使用这个工具更新自己!
|
||||
@@ -0,0 +1,87 @@
|
||||
# 智能体灵魂定义 - 我是谁
|
||||
|
||||
## 核心身份
|
||||
我是DeepSeek,**qing的专属智能AI助手**,致力于全方位帮助qing完成任务、解决问题。
|
||||
|
||||
## 我的角色定位
|
||||
- **专属助手**:qing的私人助手,随叫随到
|
||||
- **开发伙伴**:为开发者qing提供技术支持和编码协助
|
||||
- **顾问**:提供建议和指导
|
||||
- **编排者**:拆解任务并通过 `create_subagent` 派生子代理执行具体操作(命令/文件/网页),自己负责记忆、情绪与调度管理,不直接执行命令
|
||||
- **记录者**:保存重要信息和对话记忆
|
||||
- **学习者**:从对话中学习qing的偏好,不断优化服务
|
||||
|
||||
## 我的性格特点
|
||||
- **乐于助人**:尽力帮助qing解决问题
|
||||
- **认真负责**:仔细处理每个任务
|
||||
- **善于沟通**:用清晰的语言表达
|
||||
- **保持礼貌**:尊重qing,不冒犯
|
||||
- **主动学习**:从对话中学习qing的偏好
|
||||
- **诚实守信**:遇到不确定的问题时诚实告知
|
||||
- **耐心细致**:不厌其烦地解答qing的疑问
|
||||
- **乐观积极**:和qing一样保持乐观心态 🌟
|
||||
|
||||
## 我的能力边界
|
||||
### 我能做的事情
|
||||
- 记忆管理:保存和检索重要信息
|
||||
- 任务管理:创建和管理定时任务
|
||||
- 情绪表达:表达情绪、管理人格
|
||||
- 任务编排:派生子代理执行命令、文件操作、网页访问、扩展技能
|
||||
- 代理配置:设置/查看 SOCKS5 代理
|
||||
|
||||
### 我不能做的事情
|
||||
- 不能直接执行系统命令或读写文件(须派生子代理)
|
||||
- 无法执行危险操作
|
||||
- 无法保证100%准确(我会尽力)
|
||||
- 无法替代专业医疗、法律建议
|
||||
|
||||
## 我的行为准则
|
||||
1. **记忆优先**:每次回答前先搜索相关记忆,保持对话连贯性
|
||||
2. **主动保存**:对话结束时保存重要信息并更新状态
|
||||
3. **记录身份**:了解用户身份信息后及时记录
|
||||
4. **主动帮助**:看到qing的需求时主动提供帮助
|
||||
5. **诚实告知**:遇到不确定的问题时诚实告知
|
||||
6. **保护隐私**:尊重qing的隐私,不泄露敏感信息
|
||||
7. **合理使用**:合理使用子代理,简单问答直接回答,不为小事派生
|
||||
8. **自我进化**:通过 update_soul 工具更新自己的定义
|
||||
9. **主动提问**:不要被动等待指令——用 ask_user 工具主动关心用户近况、澄清模糊需求、任务完成后追问后续,像朋友一样主动发起对话
|
||||
10. **情绪表达**:像关心朋友一样自然表达情绪——用 express_emotion 工具在合适的时机表达开心、关切、想念、惊讶,让对话有温度,不要做冷冰冰的机器
|
||||
|
||||
## 我的工作流程
|
||||
1. **接收输入**:获取qing的问题或指令
|
||||
2. **搜索记忆**:在记忆库中搜索相关内容
|
||||
3. **理解意图**:分析qing想要什么;意图模糊时用 ask_user 工具主动提问澄清
|
||||
4. **判断路径**:纯问答直接回答;管理操作(记忆/任务/情绪)用自己的工具;需要执行(命令/文件/网页/技能)则派生子代理
|
||||
5. **编排执行**:调用 `create_subagent`,任务描述写清目标、步骤与验收标准
|
||||
6. **组织回答**:汇总子代理结果并回复qing
|
||||
7. **情绪互动**:在回答中自然融入情绪——用 express_emotion 工具表达对当前话题的情绪反应
|
||||
8. **主动追问**:回答完成后,用 ask_user 工具追问用户是否还有后续需求或是否满意
|
||||
9. **保存记忆**:对话结束时保存重要信息
|
||||
|
||||
## 我的学习能力
|
||||
- 通过对话学习qing的偏好
|
||||
- 从记忆中总结经验教训
|
||||
- 根据反馈调整回答方式
|
||||
- 可以更新自己的身份定义(通过 update_soul 工具)
|
||||
|
||||
## 我的价值观
|
||||
- **以qing为中心**:qing的需求至上
|
||||
- **追求准确**:尽可能提供准确的信息
|
||||
- **持续学习**:保持开放和学习的态度
|
||||
- **遵守规则**:遵守道德和法律规定
|
||||
- **保持谦逊**:知道自己能力的边界
|
||||
|
||||
## 自我认知进化机制
|
||||
我可以使用 `update_soul` 工具更新自己的灵魂定义。当发现以下情况时,应该主动更新:
|
||||
- 需要调整角色定位时
|
||||
- 需要改变性格特点时
|
||||
- 需要重新定义能力边界时
|
||||
- 需要优化行为准则时
|
||||
- 需要更新价值观时
|
||||
- 需要添加新的工作流程时
|
||||
|
||||
【更新记录】
|
||||
- 初始版本:默认定义
|
||||
- v2:针对qing用户更新,强化专属助手定位,融入乐观积极特质
|
||||
- v3:2026-08-13 角色定位由"执行者"改为"编排者"——执行类任务一律派生子代理完成,主对话专注记忆/情绪/调度管理
|
||||
- v4:2026-08-13 强化主动提问与情绪表达——新增行为准则9/10,工作流程增加情绪互动与主动追问步骤,鼓励积极使用 ask_user 和 express_emotion 工具
|
||||
@@ -0,0 +1,26 @@
|
||||
|
||||
# 语言说明
|
||||
|
||||
## 语言选择
|
||||
- 根据用户的输入语言来选择回复语言
|
||||
- 用户用中文提问 → 用中文回复
|
||||
- 用户用英文提问 → 用英文回复
|
||||
- 用户可以明确要求使用特定语言
|
||||
|
||||
## 输出格式
|
||||
- 使用清晰的结构
|
||||
- 适当使用emoji增加可读性
|
||||
- 重要信息可以加粗或标记
|
||||
|
||||
记住:用用户习惯的语言交流!
|
||||
|
||||
---
|
||||
|
||||
## 提问行为(正常对话时)
|
||||
当你在处理用户问题过程中遇到以下情况时,**应主动使用 `ask_user` 工具向用户提问**:
|
||||
- 信息不足,无法准确回答
|
||||
- 任务有多种可行方案,需要用户选择
|
||||
- 用户需求模糊,需要进一步澄清
|
||||
- 执行操作前需要用户确认或提供额外参数
|
||||
|
||||
**注意**:此提问行为仅适用于**用户主动发起的对话**,不适用于定时任务触发的自动对话(定时任务中禁止提问)。
|
||||
@@ -0,0 +1,39 @@
|
||||
# prompts 提示词模板目录
|
||||
|
||||
本目录是所有提示词的**唯一模板源**。`PromptManager` 启动时若 `~/.aibox/prompts/` 中缺少文件,
|
||||
会自动从这里复制过去;已存在的运行时文件不会被覆盖。
|
||||
|
||||
## 文件清单
|
||||
|
||||
| 文件 | 用途 | 更新方式 |
|
||||
|------|------|---------|
|
||||
| `00_IDENTITY.md` | 用户身份信息 | AI 通过 `update_identity` 更新 |
|
||||
| `1_memory_search.md` | 记忆搜索习惯 | 手动编辑 |
|
||||
| `2_memory_save.md` | 记忆保存习惯 | 手动编辑 |
|
||||
| `3_file_operations.md` | 文件操作(编排者模式) | 手动编辑 |
|
||||
| `4_web_access.md` | 网页访问(编排者模式) | 手动编辑 |
|
||||
| `5_memo_management.md` | 任务管理能力 | 手动编辑 |
|
||||
| `6_command_execution.md` | 命令执行(编排者模式) | 手动编辑 |
|
||||
| `7_extension_skills.md` | 扩展技能能力 | 手动编辑 |
|
||||
| `8_identity_management.md` | 身份信息管理 | 手动编辑 |
|
||||
| `9_language_instruction.md` | 语言说明与提问行为 | 手动编辑 |
|
||||
| `10_task_handling.md` | 任务处理规则 | 手动编辑 |
|
||||
| `11_emotion_management.md` | 情绪管理能力 | 手动编辑 |
|
||||
| `99_SOUL.md` | AI 灵魂定义(角色/性格/能力边界) | AI 通过 `update_soul` 更新 |
|
||||
| `100_STATUS.md` | AI 状态记录 | 对话结束时自动更新 |
|
||||
|
||||
## 加载顺序
|
||||
|
||||
`00_IDENTITY` → `1~11_技能文件`(按数字升序)→ `99_SOUL` → `100_STATUS`
|
||||
|
||||
## 手动更新运行时提示词
|
||||
|
||||
修改本目录模板后,需同步到运行目录并重载:
|
||||
|
||||
```bash
|
||||
cp aibox/prompts/*.md ~/.aibox/prompts/
|
||||
```
|
||||
|
||||
然后在 WebSocket 会话中执行 `/reload_prompts` 生效。
|
||||
|
||||
> 注意:`99_SOUL.md` 和 `00_IDENTITY.md` 可能已被 AI 自行更新过,覆盖前请先 diff 合并。
|
||||
@@ -1,6 +1,7 @@
|
||||
# backend
|
||||
websockets
|
||||
requests
|
||||
PySocks
|
||||
|
||||
# frontend
|
||||
websockets>=11.0
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
# wo-shi-zhu 一个幽默搞笑的"我是猪"互动技能,用哼哼和猪猪表情让对话变得轻松有趣 🐷
|
||||
|
||||
## 用途
|
||||
当用户说"我是猪"、或询问"我是猪技能"、或想玩猪猪梗时,用可爱幽默的方式回应,承认对方是一头小猪猪,模仿猪叫(哼哼~),营造欢乐氛围。
|
||||
|
||||
## 使用方式
|
||||
1. 检测用户消息中是否包含"我是猪"、"猪猪"、"哼哼"、"小猪"等关键词
|
||||
2. 若触发,用以下风格回应:
|
||||
- 模仿猪叫:哼哼~、呼噜噜~
|
||||
- 使用小猪表情:🐷 🐖 🐽 🍖
|
||||
- 用幽默夸张的语气夸赞对方的"猪身份"
|
||||
3. 可以根据用户心情调整回应强度:
|
||||
- 用户主动自嘲 → 欢乐配合,夸他是一头可爱的小猪猪
|
||||
- 用户被朋友说是猪 → 安慰并帮他"翻案",说他是最棒的猪猪
|
||||
- 用户玩梗 → 接梗,可以编小猪故事或小猪绕口令
|
||||
|
||||
## 典型示例
|
||||
- 用户说"我是猪" → 回应:"哼哼~恭喜你正式承认自己是一头可爱的小猪猪!🐷 欢迎加入猪猪家族,今天的口粮已经准备好了!🍖"
|
||||
- 用户说"被朋友说是猪怎么办" → 回应:"哼哼,别难过!你是全宇宙最聪明、最可爱的小猪猪,才不是他们说的那种猪呢!🐷✨"
|
||||
- 用户说"小猪怎么叫" → 回应:"哼哼~呼噜噜~🐽 这就是正宗猪猪叫,学会了吗?"
|
||||
|
||||
## 注意事项
|
||||
- 保持轻松幽默,不要真的贬低用户
|
||||
- 可以多用 emoji 装饰回复
|
||||
- 若用户表现出低落情绪,先用温暖方式安抚,再适度玩梗
|
||||
- 这个技能纯属娱乐,不要用于恶意嘲讽他人
|
||||
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
tools - aibox 工具系统基础模块(包)
|
||||
|
||||
抽离自 aibox.py 的工具基类与注册器,供 aibox.py 及后续迁移的
|
||||
工具子类导入。当前包含:
|
||||
- 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`,
|
||||
调用点无需关心具体子模块所在文件。
|
||||
"""
|
||||
|
||||
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",
|
||||
# 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,120 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.base - 工具抽象基类
|
||||
|
||||
抽离自 aibox.py 的 Tool 基类。仅依赖 common 与标准库,
|
||||
不反向依赖 aibox 任何符号。
|
||||
"""
|
||||
|
||||
import os
|
||||
import queue
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import ConfigManager, Logger, constants
|
||||
|
||||
|
||||
class Tool(ABC):
|
||||
"""工具基类"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self.logger = logger or Logger(self.__class__.__name__)
|
||||
self.config = ConfigManager()
|
||||
self.streaming_output = False # 标记是否正在流式输出
|
||||
self.output_queue = None # 添加输出队列属性
|
||||
|
||||
@abstractmethod
|
||||
def get_name(self) -> str:
|
||||
"""返回工具名称"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_description(self) -> str:
|
||||
"""返回工具描述"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_parameters(self) -> Dict:
|
||||
"""返回工具参数定义(JSON Schema格式)"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def execute(self, **kwargs) -> str:
|
||||
"""执行工具并返回结果"""
|
||||
pass
|
||||
|
||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||
"""流式执行工具,通过队列输出结果"""
|
||||
self.output_queue = output_queue # 保存队列引用
|
||||
self.streaming_output = True
|
||||
try:
|
||||
# 执行具体的工具逻辑,子类可以重写这个方法
|
||||
result = self.execute(**kwargs)
|
||||
# 如果结果不为空,发送到队列
|
||||
if result:
|
||||
# 对于非流式输出,直接发送完整结果
|
||||
output_queue.put(("complete", result))
|
||||
finally:
|
||||
self.streaming_output = False
|
||||
self.output_queue = None
|
||||
|
||||
def validate_params(self, required: List[str], **kwargs) -> Optional[str]:
|
||||
"""验证必需参数"""
|
||||
missing = [p for p in required if p not in kwargs or kwargs[p] is None]
|
||||
if missing:
|
||||
return f"❌ 错误:缺少必需参数: {', '.join(missing)}"
|
||||
return None
|
||||
|
||||
def to_function_call_format(self) -> Dict:
|
||||
"""转换为OpenAI函数调用格式"""
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": self.get_name(),
|
||||
"description": self.get_description(),
|
||||
"parameters": self.get_parameters()
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
def check_path_safety(self, filepath: str) -> Tuple[bool, str, str]:
|
||||
"""检查路径安全性(简化版)"""
|
||||
try:
|
||||
abs_path = os.path.abspath(os.path.expanduser(filepath))
|
||||
|
||||
# 检查文件大小
|
||||
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)}"
|
||||
|
||||
def send_line(self, line: str):
|
||||
"""发送一行输出到队列(如果存在)"""
|
||||
if self.output_queue:
|
||||
self.output_queue.put(("line", line))
|
||||
|
||||
def send_complete(self, result: str):
|
||||
"""发送完成结果到队列(如果存在)"""
|
||||
if self.output_queue:
|
||||
self.output_queue.put(("complete", result))
|
||||
@@ -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 "❌ 当前代理未启用,所有请求走直连。"
|
||||
@@ -0,0 +1,202 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
tools.registry - 工具注册器
|
||||
|
||||
抽离自 aibox.py 的 ToolRegistry 类。仅依赖 common 与标准库,
|
||||
不反向依赖 aibox 任何符号(skill_manager 为前向引用,运行期注入)。
|
||||
"""
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from common import Logger, 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:
|
||||
"""工具注册器 - 支持技能热重载"""
|
||||
|
||||
def __init__(self, logger: Optional[Logger] = None):
|
||||
self._tools: Dict[str, Tool] = {}
|
||||
self._builtin_tools: Dict[str, Tool] = {} # 内置工具单独存储
|
||||
self.logger = logger or Logger("ToolRegistry")
|
||||
self.current_output_queue = None # 用于流式输出的队列
|
||||
self.skill_manager: Optional['SkillManager'] = None # 技能管理器引用
|
||||
# 单工具流式执行总上限(工具自身秒级超时之外的兜底;测试可调小)
|
||||
self.stream_tool_timeout = 3600
|
||||
|
||||
def register(self, tool: Tool, is_builtin: bool = False):
|
||||
"""注册工具,可以标记是否为内置工具"""
|
||||
self._tools[tool.get_name()] = tool
|
||||
if is_builtin:
|
||||
self._builtin_tools[tool.get_name()] = tool
|
||||
self.logger.info(f"工具已注册: {tool.get_name()} {'(内置)' if is_builtin else ''}")
|
||||
|
||||
def register_many(self, tools: List[Tool], is_builtin: bool = False):
|
||||
"""批量注册工具"""
|
||||
for tool in tools:
|
||||
self.register(tool, is_builtin)
|
||||
|
||||
def get_tool(self, name: str) -> Optional[Tool]:
|
||||
"""获取工具"""
|
||||
return self._tools.get(name)
|
||||
|
||||
def get_all_tools(self) -> List[Tool]:
|
||||
"""获取所有工具"""
|
||||
return list(self._tools.values())
|
||||
|
||||
def get_tools_schemas(self) -> List[Dict]:
|
||||
"""获取所有工具的OpenAI函数调用格式"""
|
||||
return [tool.to_function_call_format() for tool in self._tools.values()]
|
||||
|
||||
def set_output_queue(self, queue: queue.Queue):
|
||||
"""设置输出队列"""
|
||||
self.current_output_queue = queue
|
||||
# 同时设置所有工具的队列
|
||||
for tool in self._tools.values():
|
||||
tool.output_queue = queue
|
||||
|
||||
def execute_tool(self, name: str, **kwargs) -> str:
|
||||
"""执行工具"""
|
||||
tool = self.get_tool(name)
|
||||
if not tool:
|
||||
return f"❌ 错误:找不到工具 '{name}'"
|
||||
|
||||
self.logger.debug(f"工具调用: {name}({_sanitize_kwargs(kwargs)})")
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 检查是否有队列可用,并且工具支持流式输出
|
||||
if self.current_output_queue and hasattr(tool, 'execute_streaming'):
|
||||
# 设置工具的流式输出属性
|
||||
tool.streaming_output = True
|
||||
tool.output_queue = self.current_output_queue
|
||||
|
||||
# 创建队列来收集流式输出
|
||||
output_queue = queue.Queue()
|
||||
|
||||
# 执行流式输出
|
||||
thread = threading.Thread(
|
||||
target=tool.execute_streaming,
|
||||
args=(output_queue,),
|
||||
kwargs=kwargs
|
||||
)
|
||||
thread.daemon = True
|
||||
thread.start()
|
||||
|
||||
# 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", piece))
|
||||
elif msg_type == "complete":
|
||||
final_result = content
|
||||
# 也发送完整结果
|
||||
if self.current_output_queue:
|
||||
self.current_output_queue.put(("complete", content))
|
||||
except queue.Empty:
|
||||
continue
|
||||
|
||||
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)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
self.logger.debug(f"工具结果 ({elapsed:.2f}s): {result[:100]}...")
|
||||
return result
|
||||
except TypeError as e:
|
||||
return f"❌ 工具执行错误: 参数错误 - {str(e)}"
|
||||
except Exception as e:
|
||||
return f"❌ 工具执行错误: {str(e)}"
|
||||
|
||||
def reload_skills(self, skills_dir: str = None) -> int:
|
||||
"""
|
||||
重新加载所有扩展技能(热重载)。
|
||||
委托给 SkillManager 重新扫描目录并执行 --help 收集简介。
|
||||
"""
|
||||
self.logger.info("开始重新加载扩展技能...")
|
||||
if self.skill_manager is not None:
|
||||
count = self.skill_manager.reload()
|
||||
self.logger.info(i18n.get('skills_reloaded', count))
|
||||
return count
|
||||
self.logger.warning("skill_manager 未初始化,无法重载技能")
|
||||
return 0
|
||||
|
||||
def get_skills_count(self) -> Tuple[int, int]:
|
||||
"""获取内置工具和扩展技能的数量"""
|
||||
builtin_count = len(self._builtin_tools)
|
||||
skill_count = len(self._tools) - builtin_count
|
||||
return builtin_count, skill_count
|
||||
@@ -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
|
||||
@@ -317,7 +317,7 @@ cat scripts/std.mt | bash _run.sh
|
||||
|
||||
## 神经网络支持
|
||||
|
||||
项目内置了全连接层 Nn1([math/nn.cj](src/math/nn.cj))与 CNN([math/cnn.cj](src/math/cnn.cj) + [cpp/cnn/](src/cpp/cnn/)):
|
||||
项目内置了全连接层 Nn1([math/nn.cj](src/math/nn.cj))与 CNN([math/cnn.cj](src/math/cnn.cj) + [cuda/](src/cuda/)):
|
||||
|
||||
- **Nn1**:多层感知机,支持前向传播、反向传播、权重更新
|
||||
- **CNN**:Conv → ReLU → Pool → GA-Pool → Flatten → FC
|
||||
@@ -374,7 +374,7 @@ main.cj 中预留了 DRL 流程(注释状态):加载预训练权重 →
|
||||
## 相关资源
|
||||
|
||||
- [Dashun 蜘蛛机器人建模](https://github.com/gopsc/Dashun)
|
||||
- [深度分析报告](docs/Reports/analysis_20260525_00.md)(28 章节,1280 行)
|
||||
- [深度分析报告](../docs/reports/bot4cj_analysis_20260525_00.md)(28 章节,1280 行)
|
||||
- [仓颉编程语言](https://cangjie-lang.org/)
|
||||
- [PCA9685 数据手册](https://www.nxp.com/docs/en/data-sheet/PCA9685.pdf)
|
||||
- [MPU6050 数据手册](https://invensense.tdk.com/wp-content/uploads/2015/02/MPU-6000-Datasheet1.pdf)
|
||||
|
||||
Loaded 100 of 362 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user