Compare commits
67
Commits
a5af4bb404
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2832566e87 | ||
|
|
7d5d7b9e36 | ||
|
|
4497d80a94 | ||
|
|
04f1375b45 | ||
|
|
d9b2359d25 | ||
|
|
cc72248e26 | ||
|
|
c4154d4db9 | ||
|
|
0ed03ac095 | ||
|
|
2ea3fa14d8 | ||
|
|
dbfb925f65 | ||
|
|
67acae708c | ||
|
|
7e44025671 | ||
|
|
2dc79f0637 | ||
|
|
5894e64287 | ||
|
|
3e860c16e0 | ||
|
|
68a8e49896 | ||
|
|
5cb35945a4 | ||
|
|
f3bb09e6b5 | ||
|
|
6f1f2bd296 | ||
|
|
38be729509 | ||
|
|
d4e627fbb2 | ||
|
|
3e6e171831 | ||
|
|
d1357ae483 | ||
|
|
9db2bc112d | ||
|
|
9064750f33 | ||
|
|
5dc5ec9af2 | ||
|
|
e43aeeaf44 | ||
|
|
7dc6796f3a | ||
|
|
f7fc06c0ff | ||
|
|
04e3357dd0 | ||
|
|
2c6d1e8ce6 | ||
|
|
457bb900fb | ||
|
|
f4fcf175ec | ||
|
|
d513c5f892 | ||
|
|
4cb97318b5 | ||
|
|
31b2a4d6ca | ||
|
|
db54921066 | ||
|
|
9fee9a1478 | ||
|
|
d3acccd146 | ||
|
|
d1f780e381 | ||
|
|
acb6ddbb4c | ||
|
|
cccc946660 | ||
|
|
aa15841d73 | ||
|
|
0246943165 | ||
|
|
d53eba77f0 | ||
|
|
3d7f5ca3cb | ||
|
|
d29bb5504a | ||
|
|
a583f44d5d | ||
|
|
5f3ca9a246 | ||
|
|
77262f8c3a | ||
|
|
84a3f6dd5e | ||
|
|
0e9efe4895 | ||
|
|
4480bce684 | ||
|
|
58725c8a12 | ||
|
|
90f5544673 | ||
|
|
5bcf22f9a7 | ||
|
|
6ef5d11d9d | ||
|
|
eb88a77238 | ||
|
|
a833cd0d5f | ||
|
|
207137c0dd | ||
|
|
b2ecb8e12b | ||
|
|
910ab8aa38 | ||
|
|
8de7f6696c | ||
|
|
1348660cd8 | ||
|
|
177a9ac9a6 | ||
|
|
ce142fe32f | ||
|
|
d03dbc7f91 |
No files matched your search
@@ -28,3 +28,6 @@ hailort.log
|
|||||||
|
|
||||||
# 录制产物(可复现:mtc DSL → 编译 → pybot --record)
|
# 录制产物(可复现:mtc DSL → 编译 → pybot --record)
|
||||||
actpred/data_v3
|
actpred/data_v3
|
||||||
|
|
||||||
|
# pybot 训练/评估数据(pybot_record_* 动作录制 + eval_recorded*.json)
|
||||||
|
data/
|
||||||
@@ -78,7 +78,7 @@
|
|||||||
| [aibox](aibox/) | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度,由 core 拉起但与主干微服务无调用关系 |
|
| [aibox](aibox/) | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度,由 core 拉起但与主干微服务无调用关系 |
|
||||||
| [BotOps](BotOps/) | Python | — | aibox 智能体 TUI 前端(Textual),WebSocket 连接 aibox,支持 Markdown 增量流式渲染、思考模式展示、任务独立侧栏、自动重连 |
|
| [BotOps](BotOps/) | Python | — | aibox 智能体 TUI 前端(Textual),WebSocket 连接 aibox,支持 Markdown 增量流式渲染、思考模式展示、任务独立侧栏、自动重连 |
|
||||||
| [bot4cj](bot4cj/) | 仓颉(Cangjie) | — | 机器人硬件脚本解析器,解析 .mt 矩阵脚本通过 I2C 驱动 PCA9685 舵机 / MPU6050 陀螺仪,面向树莓派 Dashun 蜘蛛机器人平台(独立硬件项目) |
|
| [bot4cj](bot4cj/) | 仓颉(Cangjie) | — | 机器人硬件脚本解析器,解析 .mt 矩阵脚本通过 I2C 驱动 PCA9685 舵机 / MPU6050 陀螺仪,面向树莓派 Dashun 蜘蛛机器人平台(独立硬件项目) |
|
||||||
| [pybot](pybot/) | Python | — | bot4cj 的 Python 迁移版:.mt 动作脚本解析器(零第三方依赖,与仓颉版逐语义对齐);支持执行 / `--dry-run` / `--record` 录制 actpred episode(含 imu.npy)/ `--loop` 实时预测闭环 |
|
| [pybot](pybot/) | Python | — | bot4cj 的 Python 迁移版:.mt 动作脚本解析器(零第三方依赖,与仓颉版逐语义对齐);支持执行 / `--dry-run` / `--record` 录制 actpred episode(含 imu.npy)/ `--loop` 实时预测闭环(HEF 后端自动异步,推理与执行并行) |
|
||||||
| [mtc](mtc/) | Python | — | 命令式动作脚本(Python DSL)→ .mt 编译器:非阻塞时间轴 + 10ms 栅格展开,输出可由 pybot/bot4cj 直接执行;零第三方依赖 |
|
| [mtc](mtc/) | Python | — | 命令式动作脚本(Python DSL)→ .mt 编译器:非阻塞时间轴 + 10ms 栅格展开,输出可由 pybot/bot4cj 直接执行;零第三方依赖 |
|
||||||
| [actpred](actpred/) | Python(PyTorch) | — | 动作预测模型(图像 + 动作历史 → 下一动作),bot4cj/pybot 的「大脑」配套;两阶段训练(.mt 预训练 → episode 微调),ONNX/Hailo HEF 部署(实验) |
|
| [actpred](actpred/) | Python(PyTorch) | — | 动作预测模型(图像 + 动作历史 → 下一动作),bot4cj/pybot 的「大脑」配套;两阶段训练(.mt 预训练 → episode 微调),ONNX/Hailo HEF 部署(实验) |
|
||||||
| [lci](lci/) | C++ | — | 嵌入式帧缓冲图形界面(LCI),基于 `/dev/fb0` 直接操作显存,无需 X11/Wayland,矢量字体、图片渲染、实时数据可视化,面向树莓派等嵌入式环境 |
|
| [lci](lci/) | C++ | — | 嵌入式帧缓冲图形界面(LCI),基于 `/dev/fb0` 直接操作显存,无需 X11/Wayland,矢量字体、图片渲染、实时数据可视化,面向树莓派等嵌入式环境 |
|
||||||
@@ -235,7 +235,8 @@
|
|||||||
|
|
||||||
- 一套相对独立的单智能体产品,基于 DeepSeek + WebSocket(:8765),由 `core` 拉起但与 waf / gateway / dnfs2 等主干微服务无调用关系。
|
- 一套相对独立的单智能体产品,基于 DeepSeek + WebSocket(:8765),由 `core` 拉起但与 waf / gateway / dnfs2 等主干微服务无调用关系。
|
||||||
- 自带技能系统(`skills/`)、任务调度器、记忆管理、前端页面与语音客户端(`voice_client.py`)。
|
- 自带技能系统(`skills/`)、任务调度器、记忆管理、前端页面与语音客户端(`voice_client.py`)。
|
||||||
- 默认模型 `deepseek-v4-flash`,配置在 `config.json`;详见 [aibox/README.md](aibox/README.md)。
|
- 默认模型 `deepseek-v4-flash`,配置在 `config.json`;`models` 配置节可挂多个 OpenAI 兼容端点按档案名切换(档案级 reasoning/温度/展示名/`fallback` 失败转移),`model_config.py` 提供交互式接入点管理(列表/添加/删除/连通性验证,Key 仅记环境变量名不落盘),`--once` 支持无 WebSocket 的 CLI 一次性对话;详见 [aibox/README.md](aibox/README.md)。
|
||||||
|
- **六轮安全审计(2026-08,Batch E/F/G/H/I/K/L)**:命令执行保护链(1MB/超时/中断)、凭据脱敏(代理密码不落盘、日志打码、Bearer 脱敏)、取消传播(子代理/情绪子代理/调度器超时)、DB 事务与 FTS5 转义、POSIX 进程组终止、任务系统入参校验(畸形时间/数字字段防调度器热循环)、前端 XSS 净化(DOMPurify + SRI)等;第六轮交叉验证 Batch K 零回归,R1~R3 低危已随 Batch L 落地(72 项发现全部修复);单测 490 项全过(Linux);审计报告见 [docs/reports/aibox_audit_20260817.md](docs/reports/aibox_audit_20260817.md)。
|
||||||
|
|
||||||
### BotOps(aibox 智能体 TUI 前端)
|
### BotOps(aibox 智能体 TUI 前端)
|
||||||
|
|
||||||
@@ -256,7 +257,7 @@
|
|||||||
### pybot(机器人动作脚本解析器 · Python 迁移版)
|
### pybot(机器人动作脚本解析器 · Python 迁移版)
|
||||||
|
|
||||||
- [bot4cj](#bot4cj机器人硬件脚本解析器) 的 Python 迁移版,零第三方依赖(纯标准库),与仓颉版逐语义对齐(含占位符 `-` 递归继承、微步分解、块序依赖)。
|
- [bot4cj](#bot4cj机器人硬件脚本解析器) 的 Python 迁移版,零第三方依赖(纯标准库),与仓颉版逐语义对齐(含占位符 `-` 递归继承、微步分解、块序依赖)。
|
||||||
- 三种运行模式:脚本执行 `run()`、`--record` 录制 actpred episode(images + actions.npy + imu.npy,≤4 个 MPU6050 按块行序拼接)、`--loop` 实时预测闭环(daemon 线程队列非阻塞取帧,overrun 告警)。
|
- 三种运行模式:脚本执行 `run()`、`--record` 录制 actpred episode(images + actions.npy + imu.npy,≤4 个 MPU6050 按块行序拼接)、`--loop` 实时预测闭环(daemon 线程队列非阻塞取帧,overrun 告警;HEF 后端自动走异步推理——执行上一步动作与 NPU 推理并行,滞后一拍)。
|
||||||
- 扬声器静音语义:`frequency <= 0` 不出声仅保持节拍(对齐仓颉版 sleep 占位语义)。
|
- 扬声器静音语义:`frequency <= 0` 不出声仅保持节拍(对齐仓颉版 sleep 占位语义)。
|
||||||
- Windows 开发机可跑测试(fcntl mock + 符号链接 fallback);详见 [pybot/README.md](pybot/README.md)。
|
- Windows 开发机可跑测试(fcntl mock + 符号链接 fallback);详见 [pybot/README.md](pybot/README.md)。
|
||||||
|
|
||||||
|
|||||||
@@ -311,3 +311,17 @@ python3 dfc_compile.py # 产物 checkpoints/actpred_pretrain.hef
|
|||||||
`--pred-token` 喂入学习值、`--stats` 喂入归一化统计、`--imu` 喂入姿态。
|
`--pred-token` 喂入学习值、`--stats` 喂入归一化统计、`--imu` 喂入姿态。
|
||||||
环境调研与 API 踩坑记录见
|
环境调研与 API 踩坑记录见
|
||||||
[docs/reports/actpred_hailo_hat2_inference_research_20260813.md](../docs/reports/actpred_hailo_hat2_inference_research_20260813.md)。
|
[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)
|
||||||
+1
-1
@@ -1,6 +1,6 @@
|
|||||||
"""多头 one-hot 分箱 —— 18 维连续动作向量与多组 one-hot 分类目标互转
|
"""多头 one-hot 分箱 —— 18 维连续动作向量与多组 one-hot 分类目标互转
|
||||||
|
|
||||||
分箱设计(actpred_onehot_multihead_dev_plan_20260814.md):
|
分箱设计:
|
||||||
0-9 舵机角:均匀 1°/箱 → 180 箱(匹配 0.75° 微动作插值数据粒度)
|
0-9 舵机角:均匀 1°/箱 → 180 箱(匹配 0.75° 微动作插值数据粒度)
|
||||||
10 spk.F:对数分箱 log10(1+x) ∈ [0, log10(5001)] → 32 箱
|
10 spk.F:对数分箱 log10(1+x) ∈ [0, log10(5001)] → 32 箱
|
||||||
11 spk.delay:对数分箱 log10(1+x) ∈ [0, log10(2001)] → 32 箱
|
11 spk.delay:对数分箱 log10(1+x) ∈ [0, log10(2001)] → 32 箱
|
||||||
|
|||||||
+130
-29
@@ -31,6 +31,8 @@ import contextlib
|
|||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
@@ -51,6 +53,30 @@ 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_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)
|
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:
|
def _input_layer_index(name: str) -> int:
|
||||||
m = re.search(r"input_layer(\d+)$", name)
|
m = re.search(r"input_layer(\d+)$", name)
|
||||||
@@ -223,6 +249,11 @@ class HailoPredictor:
|
|||||||
for f in self._frame_infos):
|
for f in self._frame_infos):
|
||||||
self._frame_luts = self._build_frame_luts()
|
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)
|
h_shape = tuple(self._history_info.shape)
|
||||||
self.history_len, self.action_dim = h_shape[1], h_shape[2]
|
self.history_len, self.action_dim = h_shape[1], h_shape[2]
|
||||||
|
|
||||||
@@ -258,16 +289,8 @@ class HailoPredictor:
|
|||||||
|
|
||||||
# ---------------- 推理 ----------------
|
# ---------------- 推理 ----------------
|
||||||
|
|
||||||
def predict(self, images, action_history: np.ndarray,
|
def _prepare_inputs(self, images, action_history):
|
||||||
pred_token: np.ndarray = None,
|
"""输入校验 + 动作历史归一化(predict/predict_async 共用)"""
|
||||||
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 学习值。
|
|
||||||
"""
|
|
||||||
if not isinstance(images, (list, tuple)):
|
if not isinstance(images, (list, tuple)):
|
||||||
images = [images]
|
images = [images]
|
||||||
if len(images) != self.frame_len:
|
if len(images) != self.frame_len:
|
||||||
@@ -285,9 +308,12 @@ class HailoPredictor:
|
|||||||
mode=self._range_mode)
|
mode=self._range_mode)
|
||||||
elif self._action_mean is not None:
|
elif self._action_mean is not None:
|
||||||
action_history = (action_history - self._action_mean) / self._action_std
|
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
|
from hailo_platform import FormatType
|
||||||
bindings = self._configured.create_bindings()
|
|
||||||
|
|
||||||
if self._frame_luts is not None:
|
if self._frame_luts is not None:
|
||||||
# cv2.LUT 快路径:逐帧独立输入,直接绑定量化缓冲
|
# cv2.LUT 快路径:逐帧独立输入,直接绑定量化缓冲
|
||||||
@@ -338,15 +364,16 @@ class HailoPredictor:
|
|||||||
out_buf = np.zeros(out_shape, dtype=np.float32)
|
out_buf = np.zeros(out_shape, dtype=np.float32)
|
||||||
bindings.output(self._out_info.name).set_buffer(out_buf)
|
bindings.output(self._out_info.name).set_buffer(out_buf)
|
||||||
|
|
||||||
self._configured.run([bindings], timeout=10000)
|
def _postprocess(self, raw_result, history_tail=None) -> np.ndarray:
|
||||||
|
"""反量化 + 解码/反归一化 → (action_dim,)(predict/predict_async 共用)
|
||||||
|
|
||||||
result = np.asarray(
|
history_tail: 归一化空间的历史末步(predict_delta 残差模式用)。
|
||||||
bindings.output(self._out_info.name).get_buffer(),
|
"""
|
||||||
dtype=np.float32)
|
from hailo_platform import FormatType
|
||||||
if self._out_info.format.type == FormatType.UINT8:
|
if self._out_info.format.type == FormatType.UINT8:
|
||||||
q = self._out_info.quant_info
|
q = self._out_info.quant_info
|
||||||
result = (result - q.qp_zp) * q.qp_scale
|
raw_result = (raw_result - q.qp_zp) * q.qp_scale
|
||||||
result = result.reshape(-1)
|
result = raw_result.reshape(-1)
|
||||||
if self._onehot:
|
if self._onehot:
|
||||||
# 多头 logits → 期望值解码 → 原始量纲(箱生活在原始空间,
|
# 多头 logits → 期望值解码 → 原始量纲(箱生活在原始空间,
|
||||||
# 无需反归一化);权重 0 列固定为常量(与 predict.py 对齐),
|
# 无需反归一化);权重 0 列固定为常量(与 predict.py 对齐),
|
||||||
@@ -360,23 +387,97 @@ class HailoPredictor:
|
|||||||
fixed = fixed_zero_columns(w, self.action_dim)
|
fixed = fixed_zero_columns(w, self.action_dim)
|
||||||
result[w[:self.action_dim] <= 0.0] = (
|
result[w[:self.action_dim] <= 0.0] = (
|
||||||
fixed[w[:self.action_dim] <= 0.0])
|
fixed[w[:self.action_dim] <= 0.0])
|
||||||
if not np.all(np.isfinite(result)):
|
else:
|
||||||
print("WARNING: HEF output contains NaN/Inf, replacing with zeros")
|
if self._predict_delta:
|
||||||
result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
|
# 残差模式:Δ̂ + a_{t-1}(归一化空间),与 predict.py 对齐
|
||||||
return result
|
result = result + history_tail
|
||||||
if self._predict_delta:
|
if self._range_mode:
|
||||||
# 残差模式:Δ̂ + a_{t-1}(归一化空间),与 predict.py 对齐
|
result = denormalize_action(result, self.action_dim,
|
||||||
result = result + action_history[-1]
|
mode=self._range_mode)
|
||||||
if self._range_mode:
|
elif self._action_mean is not None:
|
||||||
result = denormalize_action(result, self.action_dim,
|
result = result * self._action_std + self._action_mean
|
||||||
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)):
|
if not np.all(np.isfinite(result)):
|
||||||
print("WARNING: HEF output contains NaN/Inf, replacing with zeros")
|
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)
|
result = np.nan_to_num(result, nan=0.0, posinf=0.0, neginf=0.0)
|
||||||
return result
|
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:
|
def _frame_float(self, image) -> np.ndarray:
|
||||||
"""Resize → float32 HWC;归一化 HEF(负零偏)再走 ImageNet 归一化"""
|
"""Resize → float32 HWC;归一化 HEF(负零偏)再走 ImageNet 归一化"""
|
||||||
img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
|
img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
|
||||||
|
|||||||
@@ -2,8 +2,7 @@
|
|||||||
|
|
||||||
从 .mt 文件中提取 @Action 块,解析为 (n_rows, n_cols) 的 float32 矩阵。
|
从 .mt 文件中提取 @Action 块,解析为 (n_rows, n_cols) 的 float32 矩阵。
|
||||||
支持 `-` 占位符(沿用上一行同列的值);首行 `-` 填 0,
|
支持 `-` 占位符(沿用上一行同列的值);首行 `-` 填 0,
|
||||||
与 pybot/bot4cj 运行时语义一致(2026-08 统一,见
|
与 pybot/bot4cj 运行时语义一致(2026-08 统一)。
|
||||||
docs/reports/bot4cj_python_migration_dev_doc.md §2.3)。
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import re
|
import re
|
||||||
|
|||||||
@@ -1279,6 +1279,13 @@ class _FakeConfigured:
|
|||||||
for b in bindings_list:
|
for b in bindings_list:
|
||||||
self._transform(b)
|
self._transform(b)
|
||||||
|
|
||||||
|
def run_async(self, bindings_list, callback=None):
|
||||||
|
# 模拟异步:同步完成 transform 后触发回调(completion_info 无异常)
|
||||||
|
for b in bindings_list:
|
||||||
|
self._transform(b)
|
||||||
|
if callback is not None:
|
||||||
|
callback(completion_info=types.SimpleNamespace(exception=None))
|
||||||
|
|
||||||
|
|
||||||
class _FakeInferModel:
|
class _FakeInferModel:
|
||||||
def __init__(self, transform):
|
def __init__(self, transform):
|
||||||
@@ -1582,6 +1589,40 @@ class TestHailoPredictor(unittest.TestCase):
|
|||||||
result = p.predict(self._frames(2), history)
|
result = p.predict(self._frames(2), history)
|
||||||
np.testing.assert_allclose(result, [3, 3, 3], atol=1e-6)
|
np.testing.assert_allclose(result, [3, 3, 3], atol=1e-6)
|
||||||
|
|
||||||
|
def test_predict_async_matches_predict(self):
|
||||||
|
from hailo_infer import HailoPredictor
|
||||||
|
infos_in = [
|
||||||
|
_FakeVStreamInfo("net/images", (1, 2, 3, 4, 4), _FakeFormatType.FLOAT32),
|
||||||
|
_FakeVStreamInfo("net/history", (1, 2, 3), _FakeFormatType.FLOAT32),
|
||||||
|
]
|
||||||
|
infos_out = [_FakeVStreamInfo("net/action", (1, 3), _FakeFormatType.FLOAT32)]
|
||||||
|
_install_fake_hailo(self, infos_in, infos_out,
|
||||||
|
_first_row_transform("net/history", "net/action"))
|
||||||
|
history = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
|
||||||
|
with HailoPredictor("fake.hef") as p:
|
||||||
|
sync = p.predict(self._frames(2), history)
|
||||||
|
job = p.predict_async(self._frames(2), history)
|
||||||
|
async_result = job.result(timeout=5.0)
|
||||||
|
np.testing.assert_allclose(async_result, sync, atol=1e-6)
|
||||||
|
|
||||||
|
def test_predict_async_buffer_pool_roundtrip(self):
|
||||||
|
"""连续提交超过缓冲池数量的任务,输出互不串扰"""
|
||||||
|
from hailo_infer import HailoPredictor
|
||||||
|
infos_in = [
|
||||||
|
_FakeVStreamInfo("net/images", (1, 2, 3, 4, 4), _FakeFormatType.FLOAT32),
|
||||||
|
_FakeVStreamInfo("net/history", (1, 2, 3), _FakeFormatType.FLOAT32),
|
||||||
|
]
|
||||||
|
infos_out = [_FakeVStreamInfo("net/action", (1, 3), _FakeFormatType.FLOAT32)]
|
||||||
|
_install_fake_hailo(self, infos_in, infos_out,
|
||||||
|
_first_row_transform("net/history", "net/action"))
|
||||||
|
with HailoPredictor("fake.hef") as p:
|
||||||
|
jobs = []
|
||||||
|
for _ in range(6): # 超过 _ASYNC_BUFFERS=4
|
||||||
|
jobs.append(p.predict_async(self._frames(2), [[1, 2, 3]] * 2))
|
||||||
|
for job in jobs:
|
||||||
|
np.testing.assert_allclose(job.result(timeout=5.0), [1, 2, 3],
|
||||||
|
atol=1e-6)
|
||||||
|
|
||||||
|
|
||||||
class TestDfcCompile(unittest.TestCase):
|
class TestDfcCompile(unittest.TestCase):
|
||||||
"""dfc_compile.py 的 HAR 输入层连通性检查(编译前拦截输入断连)"""
|
"""dfc_compile.py 的 HAR 输入层连通性检查(编译前拦截输入断连)"""
|
||||||
|
|||||||
+238
-18
@@ -17,12 +17,14 @@
|
|||||||
- 🧬 **智能体自我进化**:通过 `update_soul` / `update_identity` 工具动态更新角色定位与用户画像
|
- 🧬 **智能体自我进化**:通过 `update_soul` / `update_identity` 工具动态更新角色定位与用户画像
|
||||||
- 🤖 **子代理**:通过 `create_subagent` 工具派生出拥有独立上下文/角色设定的轻量子代理,主代理可分解长任务、隔离上下文,子代理可选用不同模型
|
- 🤖 **子代理**:通过 `create_subagent` 工具派生出拥有独立上下文/角色设定的轻量子代理,主代理可分解长任务、隔离上下文,子代理可选用不同模型
|
||||||
- 🔄 **技能热重载**:无需重启即可添加、修改或删除扩展技能
|
- 🔄 **技能热重载**:无需重启即可添加、修改或删除扩展技能
|
||||||
- 🔐 **SOCKS5 代理**:通过 `set_proxy` / `get_proxy` 工具为网页访问与 API 调用配置代理,可持久化
|
- 🔐 **SOCKS5 代理**:通过 `set_proxy` / `get_proxy` 工具为网页访问与 API 调用配置代理,可持久化;密码明文仅存内存不落盘,跨重启经 `password_env` 环境变量注入(对齐 `api_key_env` 惯例)
|
||||||
- 💾 **对话管理**:历史对话归档、开始新对话(重置 AI 记忆并重新加载提示词)
|
- 💾 **对话管理**:历史对话归档、开始新对话(重置 AI 记忆并重新加载提示词)
|
||||||
- 📦 **对话压缩**:消息数或 Token 数达到阈值时自动压缩,保持上下文质量
|
- 📦 **对话压缩**:消息数或 Token 数达到阈值时自动压缩,保持上下文质量
|
||||||
- 🛡️ **内存安全防护**:文件读取按需加载不全载、命令输出 1MB 限制、单行超长截断、工具结果多层防护防止 OOM
|
- 🛡️ **内存安全防护**:文件读取按需加载不全载、命令输出 1MB 限制、单行超长截断、工具结果多层防护防止 OOM
|
||||||
- 🔒 **错误自动恢复**:API 请求失败自动回滚消息,工具结果安全清洗(控制字符过滤、非字符串转换)
|
- 🔒 **错误自动恢复**:API 请求失败自动回滚消息,工具结果安全清洗(控制字符过滤、非字符串转换)
|
||||||
- 🛑 **可靠中断机制**:支持 Ctrl+C 实时中断,通过 `threading.Event` + HTTP 连接关闭 + 子进程 kill 三层取消,流式响应和长时间运行命令均可立即终止
|
- 🛑 **可靠中断机制**:支持 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,13 +57,28 @@ python aibox.py
|
|||||||
|
|
||||||
启动后,WebSocket 服务器将运行在 `ws://localhost:8765`(地址与端口可在 `config.json` 的 `system.websocket_host` / `system.websocket_port` 中修改)。
|
启动后,WebSocket 服务器将运行在 `ws://localhost:8765`(地址与端口可在 `config.json` 的 `system.websocket_host` / `system.websocket_port` 中修改)。
|
||||||
|
|
||||||
> 注意:本模块入口文件为 `aibox.py`(当前约 3018 行,集成 WebSocket 服务器、AI 引擎、工具/技能系统、任务调度器于一体;自 v3.0 起工具类与子系统已抽离至 `tools/`、`subsystems/`、`common/`、`tests/` 包,aibox.py 仅保留核心类与编排逻辑)。`client.py` 为示例 WebSocket 客户端,`voice_client.py` 为语音客户端,`llm_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/
|
||||||
├── aibox.py # 服务主程序(WebSocket 服务器 + AI 引擎 + 任务调度器 + 核心类 DeepSeekChat/WebSocketHandler/EnhancedAIManager)
|
├── aibox.py # 服务主程序(WebSocket 服务器 + AI 引擎 + 任务调度器 + 核心类 DeepSeekChat/WebSocketHandler/EnhancedAIManager)
|
||||||
|
├── model_config.py # 大模型接入点交互式管理工具(列表/添加/删除/连通性验证,不用手写配置)
|
||||||
├── client.py # 示例 WebSocket 客户端
|
├── client.py # 示例 WebSocket 客户端
|
||||||
├── voice_client.py # 语音客户端
|
├── voice_client.py # 语音客户端
|
||||||
├── llm_client.py # 大模型客户端抽象层(多厂商切换,DeepSeekClient 兼容 OpenAI 协议)
|
├── llm_client.py # 大模型客户端抽象层(多厂商切换,DeepSeekClient 兼容 OpenAI 协议)
|
||||||
@@ -79,8 +96,9 @@ aibox/
|
|||||||
│ ├── db.py # MemoryDatabase + MemoDatabase(SQLite,FTS5 trigram 中文检索 + LIKE 兜底)
|
│ ├── db.py # MemoryDatabase + MemoDatabase(SQLite,FTS5 trigram 中文检索 + LIKE 兜底)
|
||||||
│ ├── skill_manager.py # SkillManager(技能发现/加载/热重载)
|
│ ├── skill_manager.py # SkillManager(技能发现/加载/热重载)
|
||||||
│ ├── conversation_history.py # ConversationHistory(对话历史管理)
|
│ ├── conversation_history.py # ConversationHistory(对话历史管理)
|
||||||
│ └── conversation_compressor.py # ConversationCompressor(对话压缩,注入 llm_client 解耦)
|
│ ├── conversation_compressor.py # ConversationCompressor(对话压缩,注入 llm_client 解耦)
|
||||||
├── tools/ # 工具系统包(A 批 + B 批共 17 个工具类)
|
│ └── process_manager.py # ProcessManager(后台进程管理,全局单例 + 环形输出缓冲 + atexit 收尾)
|
||||||
|
├── tools/ # 工具系统包(A 批 + B 批 + 后台进程共 22 个工具类)
|
||||||
│ ├── base.py # Tool 抽象基类
|
│ ├── base.py # Tool 抽象基类
|
||||||
│ ├── registry.py # ToolRegistry 工具注册器
|
│ ├── registry.py # ToolRegistry 工具注册器
|
||||||
│ ├── basic_tools.py # A 批:ReadFile/WriteFile/ListFiles/SimpleWebCrawler/GetCurrentTime/Calculator
|
│ ├── basic_tools.py # A 批:ReadFile/WriteFile/ListFiles/SimpleWebCrawler/GetCurrentTime/Calculator
|
||||||
@@ -89,8 +107,9 @@ aibox/
|
|||||||
│ ├── db_tools.py # B 批:MemoryTool/MemoTool
|
│ ├── db_tools.py # B 批:MemoryTool/MemoTool
|
||||||
│ ├── prompt_tools.py # B 批:UpdateIdentity/UpdateSoul/SaveMemoryAndEndConversation
|
│ ├── prompt_tools.py # B 批:UpdateIdentity/UpdateSoul/SaveMemoryAndEndConversation
|
||||||
│ ├── proxy_tools.py # B 批:SetProxy/GetProxy
|
│ ├── proxy_tools.py # B 批:SetProxy/GetProxy
|
||||||
│ └── misc_tools.py # B 批:AskUserTool
|
│ ├── misc_tools.py # B 批:AskUserTool
|
||||||
├── tests/ # 单元测试套件(232 项全过)
|
│ └── process_tools.py # 后台进程:BgRun/BgWrite/BgList/BgOutput/BgKill(bg_run/bg_write/bg_list/bg_output/bg_kill)
|
||||||
|
├── tests/ # 单元测试套件(598 项全过)
|
||||||
│ ├── conftest.py # sys.path 配置 + websockets 桩
|
│ ├── conftest.py # sys.path 配置 + websockets 桩
|
||||||
│ ├── test_basic_tools.py # A 批工具 43 项
|
│ ├── test_basic_tools.py # A 批工具 43 项
|
||||||
│ ├── test_b_tools.py # B 批工具 58 项
|
│ ├── test_b_tools.py # B 批工具 58 项
|
||||||
@@ -99,7 +118,18 @@ aibox/
|
|||||||
│ ├── test_db.py # 记忆/备忘录数据库(含中文检索/迁移)
|
│ ├── test_db.py # 记忆/备忘录数据库(含中文检索/迁移)
|
||||||
│ ├── test_conversation_history.py # 对话历史
|
│ ├── test_conversation_history.py # 对话历史
|
||||||
│ ├── test_conversation_compressor.py # 对话压缩器
|
│ ├── test_conversation_compressor.py # 对话压缩器
|
||||||
│ └── test_skill_manager.py # 技能管理器
|
│ ├── test_skill_manager.py # 技能管理器
|
||||||
|
│ ├── test_once_mode.py # CLI 一次性对话模式(Batch A/B/C 15 项)
|
||||||
|
│ ├── test_models_profile.py # 多模型档案(Batch D 17 项)
|
||||||
|
│ ├── test_config.py # 配置脱敏(F1 3 项)
|
||||||
|
│ ├── test_f2_fixes.py # 超时清理/线程收尾(F2 5 项)
|
||||||
|
│ ├── test_g_fixes.py # 第三轮审计 Batch G(20 项)
|
||||||
|
│ ├── test_h_fixes.py # 第四轮审计 Batch H(26 项)
|
||||||
|
│ ├── test_i_fixes.py # 遗留收尾 Batch I(16 项)
|
||||||
|
│ ├── test_subagent_max_turns.py # 子代理最大轮次行为回归(14 项)
|
||||||
|
│ ├── test_k_fixes.py # 第五轮审计 Batch K(35 项)
|
||||||
|
│ ├── test_model_config.py # 模型接入点管理工具(17 项)
|
||||||
|
│ └── test_process_manager.py # 后台进程管理器与 bg_* 工具(39 项:生命周期/交互 I/O/收尾)
|
||||||
├── skills/ # 内置扩展技能(*.py 可执行脚本 / *.md Markdown 技能)
|
├── skills/ # 内置扩展技能(*.py 可执行脚本 / *.md Markdown 技能)
|
||||||
│ ├── hello-world.py
|
│ ├── hello-world.py
|
||||||
│ ├── aes-crypto.py
|
│ ├── aes-crypto.py
|
||||||
@@ -134,7 +164,8 @@ common ← subsystems ← tools ← aibox.py
|
|||||||
```bash
|
```bash
|
||||||
# 从仓库根目录运行
|
# 从仓库根目录运行
|
||||||
py -m pytest aibox/tests -q
|
py -m pytest aibox/tests -q
|
||||||
# 当前:232 项全过,约 11s
|
# 当前:598 项全过(Linux;test_execute_timeout 为 Windows 向用例,本机无 ping 预存失败)
|
||||||
|
# Windows 开发机:532 passed / 4 failed(平台环境项:symlink 权限、python3 缺失、POSIX 可执行位语义 ×2)/ 1 skipped
|
||||||
```
|
```
|
||||||
|
|
||||||
运行期数据存放于 `~/.aibox/`:
|
运行期数据存放于 `~/.aibox/`:
|
||||||
@@ -186,7 +217,7 @@ py -m pytest aibox/tests -q
|
|||||||
| `memory` | 保存、搜索(支持分页)、删除记忆 |
|
| `memory` | 保存、搜索(支持分页)、删除记忆 |
|
||||||
| `memo` | 创建、查看、完成、搜索任务 |
|
| `memo` | 创建、查看、完成、搜索任务 |
|
||||||
| `express_emotion` | 表达情绪并记录当前情绪状态(anger/joy/sadness/surprise),可指定 `persona` 从 `~/.aibox/skills/` 加载人格技能(如 angry/coquettish/tsundere)派生携带该人格的情绪子代理,保证人格独立、不污染主对话上下文 |
|
| `express_emotion` | 表达情绪并记录当前情绪状态(anger/joy/sadness/surprise),可指定 `persona` 从 `~/.aibox/skills/` 加载人格技能(如 angry/coquettish/tsundere)派生携带该人格的情绪子代理,保证人格独立、不污染主对话上下文 |
|
||||||
| `create_subagent` | 派生子代理:独立上下文/角色设定,仅拥有执行面工具(无记忆/人格/派生/提问能力),可选独立模型,用于长任务分解、上下文隔离、独立角色专注处理 |
|
| `create_subagent` | 派生子代理:独立上下文/角色设定,仅拥有执行面工具(无记忆/人格/派生/提问能力),可选独立模型(工具描述动态列出 `models` 节可用档案名,防 LLM 编造名称),用于长任务分解、上下文隔离、独立角色专注处理 |
|
||||||
| `search_skills` | 搜索扩展技能(支持 Markdown 和可执行技能),查到后派生子代理执行 |
|
| `search_skills` | 搜索扩展技能(支持 Markdown 和可执行技能),查到后派生子代理执行 |
|
||||||
| `reload_skills` | 热重载扩展技能 |
|
| `reload_skills` | 热重载扩展技能 |
|
||||||
| `create_skill` | 创建新技能(AI 自主编写 Python 脚本,实际写入经子代理完成) |
|
| `create_skill` | 创建新技能(AI 自主编写 Python 脚本,实际写入经子代理完成) |
|
||||||
@@ -196,7 +227,7 @@ py -m pytest aibox/tests -q
|
|||||||
| `update_soul` | 更新 AI 自我定义(`99_SOUL.md`,智能体进化) |
|
| `update_soul` | 更新 AI 自我定义(`99_SOUL.md`,智能体进化) |
|
||||||
| `save_memory_and_end_conversation` | 保存记忆并结束对话 |
|
| `save_memory_and_end_conversation` | 保存记忆并结束对话 |
|
||||||
| `get_current_time` | 获取当前时间 |
|
| `get_current_time` | 获取当前时间 |
|
||||||
| `set_proxy` | 设置 SOCKS5 代理(host/port/可选账号密码,可持久化,即时生效) |
|
| `set_proxy` | 设置 SOCKS5 代理(host/port/可选账号密码,可持久化,即时生效);`password` 明文仅本次会话有效不落盘,跨重启用 `password_env` 指定存有密码的环境变量名 |
|
||||||
| `get_proxy` | 查看当前 SOCKS5 代理配置 |
|
| `get_proxy` | 查看当前 SOCKS5 代理配置 |
|
||||||
|
|
||||||
### 子代理(worker)工具集
|
### 子代理(worker)工具集
|
||||||
@@ -210,6 +241,13 @@ py -m pytest aibox/tests -q
|
|||||||
| `web` | 访问网页、获取标题/文本/链接(返回 markdown 指引,改用 `curl` 命令行访问) |
|
| `web` | 访问网页、获取标题/文本/链接(返回 markdown 指引,改用 `curl` 命令行访问) |
|
||||||
| `calculator` | 数学计算 |
|
| `calculator` | 数学计算 |
|
||||||
| `get_current_time` | 获取当前时间 |
|
| `get_current_time` | 获取当前时间 |
|
||||||
|
| `bg_run` | 后台启动长驻进程(dev server / 训练任务 / watch),立即返回句柄不阻塞;stdin 为管道可交互;默认 `detach=false`(aibox 退出时终止),`detach=true` 豁免但脱离句柄管理;`log_file=true` 同步镜像全程输出到 `workspace/bg_<句柄>.log`(无界全档,可被 run_command 的 tail/cat 原生读取);名称参数为 `proc_name`(≤64 字符,防与 `execute_tool` 首参撞名) |
|
||||||
|
| `bg_write` | 向后台进程 stdin 写入一行(行协议交互:REPL / 问答式 CLI);返回 **watermark 水位**——`bg_output(since=watermark)` 得到的必然且仅是本条输入引发的响应(写读对齐,不含 `[stdin]` 回显行);`wait_lines=N` 同步等待 N 行响应(expect 式一步问答);单次上限 64KB,超时写入不阻塞工具线程;无 TTY,全屏程序(vim/top)不可用,建议启动时强制行缓冲(`python -i -u` / `PYTHONUNBUFFERED=1` / `stdbuf -oL`) |
|
||||||
|
| `bg_list` | 列出全部后台进程(运行中在前,已退出保留尾部输出供验尸,LRU 淘汰) |
|
||||||
|
| `bg_output` | 拉取输出:`tail=N` 尾部 N 行 / `since=游标` 增量续读(游标为上次返回的 `next_cursor`);单次默认 64KB、上限 1MB;环形缓冲(默认 256KB/进程)淘汰的早期行如实提示 |
|
||||||
|
| `bg_kill` | 终止后台进程(整棵进程组):三级优雅终止——关 stdin(行协议进程见 EOF 常自行退出)→ SIGTERM 宽限 → SIGKILL 兜底;`force=true` 直接 SIGKILL;对已退出进程幂等 |
|
||||||
|
|
||||||
|
> **后台进程语义(v3.8,T22 属主隔离)**:注册表为全局单例(跨会话/跨子代理可见——list/output 观察面全局共享;写/杀按属主隔离,仅启动者所在会话可 `bg_write`/`bg_kill` 自己的进程,防提示注入跨会话破坏他人 dev server;未标识属主的调用不受限);用户中断(Ctrl+C/取消事件)**不**影响后台进程;aibox 进程退出经 atexit 统一收尾(非 detach 全杀,SIGKILL/SIGTERM 硬杀不覆盖);单行输出截断 16KB;运行中上限 16 / 总条目上限 32 / 单次 stdin 写入上限 64KB(`background_processes` 配置节可调);POSIX 按进程组终止(`start_new_session`,孙进程不悬挂),Windows `taskkill /T /F` 收割进程树。
|
||||||
|
|
||||||
## 💬 WebSocket API
|
## 💬 WebSocket API
|
||||||
|
|
||||||
@@ -502,15 +540,69 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
|||||||
"broadcast_enabled": true
|
"broadcast_enabled": true
|
||||||
},
|
},
|
||||||
"api": {
|
"api": {
|
||||||
"provider": "deepseek",
|
"provider": "openai_compatible",
|
||||||
"deepseek_api_url": "https://api.deepseek.com/chat/completions",
|
"deepseek_api_url": "https://api.deepseek.com/chat/completions",
|
||||||
"deepseek_model": "deepseek-v4-flash",
|
"deepseek_model": "deepseek-v4-flash",
|
||||||
"reasoning_enabled": false,
|
"reasoning_enabled": false,
|
||||||
"api_key_env": "DEEPSEEK_API_KEY",
|
"api_key_env": "DEEPSEEK_API_KEY",
|
||||||
"timeout": 120,
|
"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_history_messages": 500,
|
||||||
"max_context_messages": 120,
|
"max_context_messages": 120,
|
||||||
"max_tokens": 8192,
|
"max_tokens": 8192,
|
||||||
@@ -533,6 +625,14 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
|||||||
"commands": {
|
"commands": {
|
||||||
"enabled": true
|
"enabled": true
|
||||||
},
|
},
|
||||||
|
"background_processes": {
|
||||||
|
"enabled": true,
|
||||||
|
"max_processes": 16,
|
||||||
|
"max_entries": 32,
|
||||||
|
"buffer_per_process": 262144,
|
||||||
|
"max_lifetime": 0,
|
||||||
|
"max_write_bytes": 65536
|
||||||
|
},
|
||||||
"web": {
|
"web": {
|
||||||
"timeout": 10,
|
"timeout": 10,
|
||||||
"delay": 0,
|
"delay": 0,
|
||||||
@@ -562,11 +662,14 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
|||||||
"llm_max_tokens": 800
|
"llm_max_tokens": 800
|
||||||
},
|
},
|
||||||
"proxy": {
|
"proxy": {
|
||||||
"type": "",
|
"socks5": {
|
||||||
"host": "",
|
"enabled": false,
|
||||||
"port": 0,
|
"host": "127.0.0.1",
|
||||||
"username": "",
|
"port": 1080,
|
||||||
"password": ""
|
"username": "",
|
||||||
|
"password": null,
|
||||||
|
"password_env": "SOCKS5_PASSWORD"
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
@@ -574,6 +677,8 @@ AI 可以通过 `update_soul` 工具动态更新 `99_SOUL.md`,通过 `update_i
|
|||||||
> 默认模型为 `deepseek-v4-flash`(通过 `api.deepseek_model` 修改)。设置 `api.reasoning_enabled: true` 可启用思考模式,此时会在请求中添加 `thinking: {type: "enabled"}` 参数,模型会先进行思考再回答,实时展示思考过程(灰色斜体区域),回答质量更高且**支持工具调用**(文件操作、命令执行、记忆搜索等工具均可在思考模式下正常使用)。
|
> 默认模型为 `deepseek-v4-flash`(通过 `api.deepseek_model` 修改)。设置 `api.reasoning_enabled: true` 可启用思考模式,此时会在请求中添加 `thinking: {type: "enabled"}` 参数,模型会先进行思考再回答,实时展示思考过程(灰色斜体区域),回答质量更高且**支持工具调用**(文件操作、命令执行、记忆搜索等工具均可在思考模式下正常使用)。
|
||||||
>
|
>
|
||||||
> API Key 从环境变量 `DEEPSEEK_API_KEY` 读取(亦支持 `~/.cert/deepseek.key` 文件)。
|
> 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` 工厂中注册,无需改动业务代码,为后续子代理/多模型共发预留了扩展点。
|
> **多厂商切换**:大模型调用已通过 `aibox/llm_client.py` 抽象层解耦,业务层不再直接依赖具体厂商。配置 `api.provider`(默认 `deepseek`)即可切换厂商,当前内置 `DeepSeekClient`(兼容 OpenAI Chat Completions 协议,Kimi / Qwen / GLM / Doubao 等绝大多数厂商均可复用)。新增厂商只需在此文件新增一个 `LLMClient` 子类并在 `create_llm_client` 工厂中注册,无需改动业务代码,为后续子代理/多模型共发预留了扩展点。
|
||||||
|
|
||||||
@@ -595,6 +700,14 @@ python aibox.py [选项]
|
|||||||
--memo-add "标题|内容|时间" 添加任务
|
--memo-add "标题|内容|时间" 添加任务
|
||||||
--memo-list 列出任务
|
--memo-list 列出任务
|
||||||
--memo-complete ID 完成任务
|
--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 秒)
|
||||||
```
|
```
|
||||||
|
|
||||||
### 使用示例
|
### 使用示例
|
||||||
@@ -616,6 +729,30 @@ python aibox.py --memo-add "提醒开会|下午3点会议室A|+1h"
|
|||||||
python aibox.py --memo-list
|
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` 全局超时。
|
||||||
|
|
||||||
## 🔄 架构流程
|
## 🔄 架构流程
|
||||||
|
|
||||||
```
|
```
|
||||||
@@ -795,6 +932,89 @@ elif command == "/mycommand":
|
|||||||
|
|
||||||
## 📄 更新日志
|
## 📄 更新日志
|
||||||
|
|
||||||
|
### v3.8
|
||||||
|
- ✅ **后台进程池(智能体运行/管理长驻进程)**:新增 `subsystems/process_manager.py`(`ProcessManager` 全局单例)+ `tools/process_tools.py` 5 个 worker 执行面工具(`bg_run` / `bg_write` / `bg_list` / `bg_output` / `bg_kill`),弥补 `run_command` 同步阻塞、超时即杀的不足——dev server / 训练任务 / watch 等长驻进程可后台启动、跨会话管理:
|
||||||
|
- **生命周期语义(三决策)**:用户中断(Ctrl+C/取消事件)不影响后台进程;aibox 退出经 atexit 统一收尾(非 detach 全杀,`detach=true` 豁免但脱离句柄管理);注册表跨会话/跨子代理共享(观察面全局可见;写/杀按属主隔离——T22,仅启动者所在会话可操作,定时任务 AI 为精简 `task` profile 且无 update_identity/update_soul/set_proxy)
|
||||||
|
- **交互式 I/O(stdin 管道 + watermark 写读对齐)**:stdin 挂管道,`bg_write` 向行协议进程(REPL / 问答式 CLI)写入一行,返回 **watermark**(`[stdin]` 回显行序号)——`bg_output(since=watermark)` 得到的必然且仅是该输入引发的响应(写读配对);`wait_lines=N` expect 式同步等待;回显行入环形缓冲与 log 镜像(tail 可见完整交互时序);写入 watchdog 线程兜底(进程不读 stdin 管道写满时不阻塞工具线程);单次写入上限 64KB
|
||||||
|
- **输出缓冲**:每进程有界环形缓冲(默认 256KB,`buffer_per_process` 可调),单行截断 16KB;`bg_output` 支持 `tail` 尾读与 `since` 游标增量续读(淘汰行如实提示),单次默认 64KB、上限 1MB;`bg_run log_file=true` 输出同步 tee 到 `workspace/bg_<句柄>.log`(无界全档,`[stdin]` 回显同入档,可被 run_command 的 tail/cat 原生读取)
|
||||||
|
- **三级优雅终止**:`bg_kill` 关 stdin(行协议进程见 EOF 常自行退出)→ SIGTERM 宽限(grace 比例 4:6)→ SIGKILL 兜底;`force=true` 直接 SIGKILL;shutdown 收尾同样先关 stdin
|
||||||
|
- **安全边界**:运行中上限 16 / 总条目上限 32(LRU 淘汰已退出条目);`background_processes` 配置节(含 `enabled` 双闸:`commands.enabled` 关闭时一并拒绝);`proc_name` 标签 ≤64 字符且拒绝控制字符(M5 惯例);POSIX `start_new_session` + 进程组终止(对齐执行器 L9 语义),Windows `taskkill /T /F` 收割进程树;TTL `max_lifetime` 惰性清理(默认 0 不启用);凭据类交互(sudo -S 喂密码)会留对话记录,敏感凭据走环境变量注入(工具描述已指引)
|
||||||
|
- **实现要点**:H2 reader 线程模式复用(EOF 即 poll 收割防僵尸);`list_procs` 惰性刷新全表状态;`proc_name` 命名规避 `ToolRegistry.execute_tool(name=...)` 首参撞名;编排者误调拦截表 `WORKER_ONLY_TOOLS` 同步纳入;无 TTY 的适用边界(全屏程序不可用、行缓冲指引)写入工具描述供模型自检
|
||||||
|
- **测试**:新增 `tests/test_process_manager.py` 39 项(生命周期/游标续读/环形淘汰/单行截断/上限/TTL/收尾/工具分发/禁用闸/watermark 对齐/回显/sync 等待/超时/EOF 优雅终止/log tee/注册表分发),Windows 开发机 532 passed / 4 failed(平台环境项,与代码无关)
|
||||||
|
|
||||||
|
### v3.7
|
||||||
|
- ✅ **第六轮安全审计 + Batch L 修复(2026-08-20,交叉验证轮:Batch K 回归验证 / 并发与生命周期 / 降级路径 / 修复间相互作用,4 路并行审读 + 主审逐条复核)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md) §16):
|
||||||
|
- **Batch K 回归验证**:J1~J16 逐项在位、零回归,无修复引入的缺陷;历届 Batch E~I 修复抽查(executors reader 收尾、1MB 输出上限、registry 流式预算、broadcast 快照等)均在位
|
||||||
|
- **Batch L(R1~R3 低危全落地)**:R1 `/memos` WS 命令 strptime 容错(畸形 reminder_time 按"待执行"展示,不再整体"处理失败");R2 移除 `handle_interrupt` 对全局 `ai_manager.command_executor` 的取消(WS 中断不再误杀 CLI 主 AI 子进程);R3 `unregister` 补齐回收(关闭 websocket + 取消在途处理任务,broadcast 超时路径与连接关闭路径行为对齐,半开连接不再泄漏)
|
||||||
|
- **测试**:新增 `tests/test_l_fixes.py` 6 例(R1 容错与语义不回归 ×2 / R2 范围收窄与中断能力不回归 ×2 / R3 回收与幂等 ×2);Linux 维持 490 passed,Windows 开发机 493 passed / 4 failed(均为平台环境项,与代码无关)
|
||||||
|
|
||||||
|
### v3.6
|
||||||
|
- ✅ **交互式大模型接入点管理([model_config.py](model_config.py))**:菜单式 列表 / 添加 / 删除 / 连通性验证,不用手写 `config.json`;Key 仅记录环境变量名不落盘(空 = 免密本地端点);删除自动清理其他档案的 fallback 引用;默认档案经 `api.default_model` 生效(详见 [模型接入点管理](#管理大模型接入点不用手写配置))
|
||||||
|
- ✅ **第五轮安全审计(2026-08-19,4 路并行审读 + Linux PoC 2/2 实证,发现 J1~J16:高危 2 / 中危 5 / 低危 6 / 信息 3)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md) §13~§14):
|
||||||
|
- **K1(J1 高危·任务系统投毒)**:畸形时间/数字字段经 CLI/WS/memo 工具三入口零校验入库 → 调度器失败热循环 + 每个新连接握手即崩——入库端统一校验(`reminder_time`/`repeat_end_time` strptime 预检、`repeat_interval` 强转 int、`repeat_type` 白名单)、调度器预检失败按到期执行(Q5 语义)、`get_pending_tasks_for_user` 解析失败按到期展示、`register` 移入 try/finally(会话登记不再泄漏);附带修复 `add_user_task` 缺失 `repeat_end_time` 参数透传
|
||||||
|
- **K2(J2 高危·前端 XSS)**:三页面 7 处 `marked.parse` 直渲(AI 输出/命令结果全链路 innerHTML)统一改经 `renderMarkdown()` + DOMPurify(缺失时退化为转义文本防 XSS 兜底);`entry1.addSystemMessage` 改 DOM 构建 + textContent;任务标题改 createTextNode;CDN 脚本(marked/hljs/socket.io/DOMPurify)补 SRI integrity(J13)
|
||||||
|
- **K3(中危)**:FTS 外部内容表删除换序——先删 `memories_fts` 再删 `memories`,索引残留消除(J3,PoC 实证);config/prompt/对话写路径统一原子写(`common/atomic.py`,mkstemp + fsync + os.replace)+ 损坏备份 `config.json.corrupt.<时间戳>` 拒绝覆盖(J5);PromptManager/ConfigManager 单例初始化纳入锁(J6,并发首例不再双重加载系统提示词)
|
||||||
|
- **K4(中危)**:`model_config.py` 连通性验证异常消息复用 `_sanitize_log_msg` 脱敏(J4,Bearer 令牌不再进终端);`broadcast` 改每客户端并发发送 + `asyncio.wait_for` 超时(`broadcast_send_timeout` 默认 5s),超时即 unregister——僵死客户端(半开 TCP)不再卡死广播与调度器循环(J7)
|
||||||
|
- **K5(低危随批)**:J8 交互输入容错(非法温度/超时友好报错);J9 MemoDatabase WAL + busy_timeout=10000 覆盖全部写路径;J10 memo 批量删除补实现(`memo_ids`);J11 `/new` 取消-重置竞态(排队 reset 等待 is_processing 清零);J12 skills 注入面(nmap hosts `shlex.quote`、screenshot PowerShell 单引号转义、SKILLS_FORMAT.md 新增"安全边界"节);J14 voice_client 死代码/僵尸连接/录音残留清理;J15 死代码移除(`full_response_queue` 等)+ 未知 task action 明确报错 + 广播门控对齐 + 热路径日志降级;J16 `--help` 输出 2000 上限 + 同名技能键确定性解决
|
||||||
|
- ✅ **测试**:490 项全过(Linux 环境;第五轮新增 Batch K 35 项 + 子代理 max_turns 回归 14 项;`test_execute_timeout`/`test_flat_api_still_reads_key_file` 为 Windows 向用例,Linux 预存失败与本轮无关)
|
||||||
|
|
||||||
|
### v3.5
|
||||||
|
- ✅ **安全审计与修复(第三轮 Batch G + 第四轮 Batch H + 遗留收尾 Batch I)**(完整记录见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md)):
|
||||||
|
- **Batch G(llm 加固 + 线程生命周期 + 取消传播 + DB 事务)**:
|
||||||
|
- llm_client:非瞬时 HTTP 错误体截断 `[:200]`(密钥/错误页不进异常栈)、Retry-After 上限 cap 3600s、fallback 档案缺 model 跳过并告警、fallback 去重、删除 `LLMClient._cancelled` 死状态、FailoverClient 并发契约文档化
|
||||||
|
- 子代理取消传播:`run_full_subagent` 共享父 `_cancel_event` + 0.2s 观察者线程转发深度中断(强关在途 HTTP、kill 子进程);`SubAgentTool` 构造注入父事件
|
||||||
|
- drain 线程重构:批次级单线程 + stop 事件 + try/finally 统一回收(消除每次工具调用泄漏一个常驻线程);`/new` 的 `reset_conversation` 移入线程池(不再同步获取 `_msg_lock` 冻结事件循环);broadcast 收窄 except + unregister 改 discard;`_processing_tasks` 改按连接的任务集合登记(关闭时取消全部在途任务)
|
||||||
|
- DB:FTS5 关键词双引号翻倍转义(`"` → `""`)、两处 `LIMIT` f-string 参数化、`mark_as_triggered` 读写收进单连接 `BEGIN IMMEDIATE` 事务(消灭并发触发竞态);历史索引裸 except 收窄 + 归档文件名同秒加序号
|
||||||
|
- **Batch H(第四轮审计,3 路并行审读 + Linux PoC 6/6 实证)**:
|
||||||
|
- **H1(高危)**:修复 N5 回归——子代理共享父取消事件后首轮 think 无条件 `clear()` 会吞掉父级已置位的取消信号(用户中断 / --once 超时被绕过),共享事件时跳过 clear + 父已取消直接早退
|
||||||
|
- `/reload_prompts` 移入线程池(N3 同类遗漏);技能扫描(串行 `--help` 子进程)启动前预热,不再在首条消息/首个任务时冻结事件循环
|
||||||
|
- 调度器任务级超时(`scheduler.task_timeout`,默认 600s)——AI think 卡死不再停摆整个调度循环
|
||||||
|
- `get_recent_memories(limit=None)` 默认上限(不再全表加载);`memo delete` 改物理删除(此前只标记完成);`PromptManager.reload` 本地构建后原子替换(并发读者不再读到空/半截系统提示词)
|
||||||
|
- read_file 的 sed 参数强转 int(防引号注入);`set_proxy` 落盘失败如实报告;SSE 非 UTF-8 行容错;persona 路径 realpath 校验(软链越界拒绝);1MB 截断无半个字符;compressor 除零防护;InstallSkillTool 指引命令路径引用;ask_user 缺参明确报错
|
||||||
|
- **POSIX 进程组终止**:executor 的 `Popen` 加 `start_new_session` + 组信号 kill(组 ID 在 Popen 后立即捕获)——`(sleep 300 &)` 形态的孙进程持管道不再悬挂 reader 线程/泄漏 fd(实证修复前孙进程残留)
|
||||||
|
- web 工具 scheme 白名单(http/https,`file://` 等拒绝);registry 日志打码改包含匹配(apiKey/passwd 变体);FailoverClient 日志 Bearer 令牌正则脱敏(网关回显 Authorization 头不再进日志)
|
||||||
|
- **Batch I(遗留收尾)**:
|
||||||
|
- **N9**:prompt 写入内容长度上限(`MAX_CONTENT_LENGTH=10000`,identity/soul/status 超长拒绝)——约束持久化注入面
|
||||||
|
- **F5**:情绪子代理取消传播——`run_light_subagent` 改走 stream_chat + 观察者线程,用户中断可即时终止(此前阻塞至请求超时)
|
||||||
|
- 提醒时间格式异常任务按"到期"返回(此前被 SQL 层字符串比较静默丢弃,任务永不执行);`mark_as_triggered` 加 `busy_timeout=10000`(多进程同库写竞争等待而非立即报 locked)
|
||||||
|
- 对话文件原子写(tmp + os.replace,崩溃不留半截 JSON);`load_conversation` realpath 越界防护;WS 命令专用线程池(`/new`、`/reload_prompts` 不再占用默认池);--once 硬超时兜底(软超时后宽限一个周期强杀进程)
|
||||||
|
- PromptManager 单例加锁;`limit=0` 语义修正(None 才不加 LIMIT)
|
||||||
|
- ✅ **测试**:424 项全过(Linux 环境;Batch G 20 + Batch H 26 + Batch I 16 项新增;`test_execute_timeout`/`test_flat_api_still_reads_key_file` 为 Windows 向用例,Linux 预存失败与本轮无关)
|
||||||
|
|
||||||
|
### v3.4
|
||||||
|
- ✅ **多模型档案加固(H1/H2/H3/M1/M2/M3 六项修复)**:
|
||||||
|
- **H1(凭据隔离)**:档案端点环境变量取不到时直接报错,**绝不回退读 `~/.cert/deepseek.key`**(该文件为 DeepSeek 专属凭据,跨厂商复用属于泄漏);key 文件回退仅对无档案的旧 `api` 扁平节生效
|
||||||
|
- **H2(schema 校验前移)**:档案缺 `api_url`/`model` 或类型错误时,解析期即抛出指明档案名与缺失字段的 `ValueError`——不再把档案名当模型名发出、不再延迟到 `requests InvalidURL` 或 API 400
|
||||||
|
- **H3(--model 服务模式生效)**:`--model` 升级为进程级覆盖,WebSocket 服务模式(主对话/每个会话实例/定时任务实例)与 `--once` 均生效,此前服务模式静默忽略
|
||||||
|
- **M2(防模型名幻觉)**:`create_subagent` 工具描述与 JSON Schema 动态携带 `models` 节可用档案名列表(默认档案标注),LLM 无从编造名称;`models` 为空时描述保持原样(零破坏)
|
||||||
|
- **M1(换模型同步派生引用)**:新增 `DeepSeekChat.set_client()` 统一替换 LLM 客户端——同步刷新 `compressor.llm_client` 与 `api_key`/`name` 快照;`create_subagent` 子代理换模型经此入口,压缩链不再残留默认模型
|
||||||
|
- **M3(reasoning 契约显式化)**:档案级 `reasoning_enabled` 从"构造后挂实例属性"改为 `LLMClient` 构造参数(三态:`True`/`False` 档案级显式、`None` 未配置回退全局),外部误设属性无从静默覆盖全局开关
|
||||||
|
- ✅ **L 批体验与兼容修复(L1/L2/L3/L5 四项)**:
|
||||||
|
- **L1(展示名按档案)**:`client.name` = 档案 `display_name` > 档案名,用 Kimi/Qwen 时对话上下文不再误标 "DeepSeek";旧扁平节保持 "DeepSeek"(零破坏)
|
||||||
|
- **L2(前缀逃逸)**:`--model model:NAME` 强制裸模型名语义、`--model profile:NAME` 强制档案语义(不存在时明确报错并列出可用档案),档案名与模型名歧义可逃逸
|
||||||
|
- **L3(档案级温度)**:档案 `temperature` 字段作为默认采样温度,优先级:CLI/调用方显式传入 > 档案字段 > 全局默认(0.7/0.3),主对话/轻量子代理/compressor 全链路生效
|
||||||
|
- **L5(思考流字段别名)**:`_parse_delta` 识别 `reasoning_content`/`reasoning`/`thinking`/`thinking_content` 四种常见思考字段(首个命中生效),其他厂商思考流不再静默丢失
|
||||||
|
- ✅ **L4a(chat_completion 跨档案失败转移)**:档案新字段 `fallback`(档案名列表,一层不递归);`create_llm_client` 构造 `FailoverClient` 包装链(主档案 + fallback)。瞬时故障(HTTP 429/5xx 抛 `TransientLLMError`、连接失败/超时)按链转移,认证/参数类 4xx 不转移(配置问题响亮暴露);失败档案 60s 冷却;fallback 档案缺 Key 等环境问题跳过 + warning 不阻断主档案;配置结构错误(类型/不存在/自引用)启动期硬性拦截。当前覆盖对话压缩与轻量子代理链路,流式(`stream_chat`)转移属 L4b
|
||||||
|
- ✅ **L4b(stream_chat 首事件前转移 + 中断联动)**:流式仅在**未产出任何事件**(reasoning/content/tool_call)时转移——重放对用户不可见;已产出事件后失败直接上抛(防前端重复输出/工具增量错乱)。`LLMClient.cancel()` 置 `_cancelled` 标志,`FailoverClient.cancel()` 广播整条链并置包装器标志——用户中断引发的连接异常与真实故障不可区分,凭标志跳过转移(否则用户越中断、系统越换厂商重试);`chat_completion` 同受标志保护。链耗尽抛最后异常(既有回滚/报错处理接住),全链冷却抛 RuntimeError
|
||||||
|
- ✅ **L4c(转移可观测性 + 预算控制 + Retry-After)**:`FailoverClient.get_stats()` 暴露 requests/failovers/cooldown_skips/exhausted 聚合计数;`on_failover(from, to, reason)` 回调(to=None 表示链耗尽,回调异常不影响主流程)供管理面/前端通知接入;档案字段 `failover_timeout_budget` 限制整链尝试总秒数(缺省不限,<=0 视为未配置),耗尽停止尝试剩余档案防超时叠加;429 响应 Retry-After(数字秒)经 `TransientLLMError.retry_after` 延长该档案冷却期(非数字格式忽略)
|
||||||
|
- ✅ **测试**:355 项全过(H1 隔离 4 + H2 校验 6 + H3 传导 6 + M2 提示 5 + M1 替换 2 + M3 契约 4 + L 批 16 + L4a 转移 15 + L4b 流式 12 + L4c 可观测 13 项新增)
|
||||||
|
|
||||||
|
### v3.3
|
||||||
|
- ✅ **安全审计与修复**(两轮审计报告见 [docs/reports/aibox_audit_20260817.md](../docs/reports/aibox_audit_20260817.md)):
|
||||||
|
- **Batch E(命令执行保护链修复)**:`CommandLineTool.execute_streaming` 纯委托注入的 executor(删除绕过 1MB/超时/中断保护的 150 行重复 Popen 实现);`SubprocessCommandExecutor` 读取循环线程化(后台 reader 线程 + 0.1s 超时轮询,静默命令超时不再失效);两处 `Popen` 改 `stdin=DEVNULL`(等待输入的命令不再永久挂起工具线程);write_file 指引 heredoc 定界符唯一化;registry 流式转发循环加 3600s 总超时兜底。PoC 回归 5/5 + 全量单测通过
|
||||||
|
- **Batch F(凭据保护 + 线程收尾)**:
|
||||||
|
- **F1**:SOCKS5 代理密码不再明文落盘——`_save_config` 落盘前强制脱敏(纵深防御),新增 `password_env` 环境变量注入(对齐 `api_key_env` 惯例),旧版明文配置加载时自动迁移清除;`registry` debug 日志对 `password`/`api_key`/`token`/`secret`/`authorization` 参数打码
|
||||||
|
- **F2**:registry 流式转发超时路径统一清理 `tool.streaming_output`/`tool.output_queue`(消除僵尸 daemon 线程经实例属性把输出混入后续会话队列的串流路径),`stream_tool_timeout` 提为实例属性(默认 3600s,测试可调);`SubprocessCommandExecutor` reader 线程收尾——正常 EOF 后显式 close 释放管道 fd,仍阻塞时有限 join 后放弃(daemon 兜底);实证修正:CPython 流锁下阻塞中的 readline 无法被 `close()` 中断(首版实现致超时被拖满命令时长,已按 PoC 回归修正)
|
||||||
|
- ✅ **测试**:272 项全过(Batch E PoC 回归 + F1 配置脱敏 3 项 + F2 超时清理/线程收尾 5 项新增)
|
||||||
|
|
||||||
|
### v3.2
|
||||||
|
- ✅ **CLI 一次性对话模式(`--once`)**(分三批落地):
|
||||||
|
- **Batch A**:`--once QUERY` / `--quiet` / `--once-timeout`;不启动 WebSocket,执行后退出;回答走 stdout、过程走 stderr,可管道组合;SIGINT 与超时复用三层取消机制
|
||||||
|
- **Batch B**:`--no-save-memory`(`enable_memory` 开关干净禁用 MemoryTool)、`--show-reasoning`(思考流默认静默);压缩提示改走 output_callback 不再污染 stdout;退出码 1(工具执行失败)
|
||||||
|
- **Batch C**:SQLite WAL(多 CLI 进程并发写锁缓解)、`--worker-only`(纯执行面 + 精简 system prompt 防人格污染)、`--model` / `--temperature` 覆盖、子代理工具失败标记透传(`_SUBAGENT_TOOL_FAILED_MARKER` → orchestrator stats)
|
||||||
|
- ✅ **多模型档案(`models` 配置节)**:配置多个 OpenAI Chat Completions 兼容端点(DeepSeek/Kimi/Qwen/本地 Ollama 等)按档案名切换;`--model` 支持档案名与裸模型名(同端点换模型)双语义;档案级 `api_key_env` 三级语义(回退 / 显式 / 空串免密);档案级 `reasoning_enabled` 优先于全局;`models: {}` 时零破坏回退旧 `api` 节;`llm_client.py` 增加 `openai_compatible` / `openai` provider 别名
|
||||||
|
- ✅ **测试**:264 项全过(once 模式 15 项 + 多模型档案 17 项新增)
|
||||||
|
|
||||||
### v3.1
|
### v3.1
|
||||||
- ✅ **修复 FTS5 中文搜索失效**([subsystems/db.py](subsystems/db.py)):
|
- ✅ **修复 FTS5 中文搜索失效**([subsystems/db.py](subsystems/db.py)):
|
||||||
- FTS5 虚拟表改用 `tokenize='trigram'` 分词器(需 SQLite >= 3.34,环境不支持时自动退回 unicode61)
|
- FTS5 虚拟表改用 `tokenize='trigram'` 分词器(需 SQLite >= 3.34,环境不支持时自动退回 unicode61)
|
||||||
|
|||||||
@@ -227,3 +227,14 @@ except ImportError:
|
|||||||
```
|
```
|
||||||
|
|
||||||
**`--help` 要完备**。包含参数说明和使用示例两部分。AI 通过阅读 `--help` 输出学会如何调用技能,信息越完整,AI 调用越准确。
|
**`--help` 要完备**。包含参数说明和使用示例两部分。AI 通过阅读 `--help` 输出学会如何调用技能,信息越完整,AI 调用越准确。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 安全边界(J12,审计 2026-08-19)
|
||||||
|
|
||||||
|
**技能即可信代码**。可执行技能由 AI(可被网页内容/用户消息诱导)以子进程方式运行,参数为不可信输入:
|
||||||
|
|
||||||
|
- **参数必须经 shell 引号转义**:凡技能内部把参数拼进 shell 命令/脚本串(subprocess shell、PowerShell、python-nmap 等内部拼串库),一律先 `shlex.quote`(POSIX)或对应语言转义(PowerShell 单引号翻倍)
|
||||||
|
- **禁止把凭据嵌入命令行参数**:命令行参数对同机进程可见(`ps`);密码/令牌应经环境变量或配置文件传递
|
||||||
|
- **ssh/sftp 类技能**:`AutoAddPolicy` 信任任意主机密钥存在 MITM 风险,仅限可信网络使用;生产环境改用 `known_hosts` 校验
|
||||||
|
- 内置技能示例(`skills/`)已按此规范处理;新增技能请遵守同一条例
|
||||||
+726
-231
File diff suppressed because it is too large.
Load diff
@@ -815,6 +815,14 @@ async def main():
|
|||||||
await cli.run()
|
await cli.run()
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
print(f"\n{Colors.SYSTEM}程序已退出{Colors.RESET}")
|
print(f"\n{Colors.SYSTEM}程序已退出{Colors.RESET}")
|
||||||
|
finally:
|
||||||
|
# M8(审计 2026-08-20):正常退出(EOF 等)此前不关 websocket,
|
||||||
|
# 服务端见异常断连——统一 finally 通知
|
||||||
|
if cli.client:
|
||||||
|
try:
|
||||||
|
await cli.client.disconnect()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -0,0 +1,49 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
common.atomic - 原子写文件工具(J5,审计 2026-08-19)
|
||||||
|
|
||||||
|
同目录 mkstemp 临时文件 + fsync + os.replace:进程中断/断电不留半截文件;
|
||||||
|
mkstemp 随机文件名避免多进程并发写同路径互踩(conversation_history 旧实现
|
||||||
|
的固定 .tmp 名在并发场景会互相覆盖)。仅依赖标准库,供 common/subsystems
|
||||||
|
各写路径复用。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
|
||||||
|
|
||||||
|
def atomic_write_text(path: str, content: str) -> None:
|
||||||
|
"""原子写文本文件。失败时清理临时文件并抛出原异常。"""
|
||||||
|
path = os.path.expanduser(path)
|
||||||
|
directory = os.path.dirname(path) or "."
|
||||||
|
os.makedirs(directory, exist_ok=True)
|
||||||
|
fd, tmp_path = tempfile.mkstemp(prefix=".aibox_tmp_", dir=directory)
|
||||||
|
try:
|
||||||
|
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||||
|
f.write(content)
|
||||||
|
f.flush()
|
||||||
|
os.fsync(f.fileno())
|
||||||
|
# T27(审计 2026-08-21):Windows 下目标被其他进程打开(杀毒扫描/
|
||||||
|
# 编辑器/并发读者)时 os.replace 抛 PermissionError。POSIX 无条件
|
||||||
|
# 成功不受影响;Windows 做短退避重试,高频写路径(每条消息都原子
|
||||||
|
# 写 current_conversation.json)一次冲突不再直接丢该条持久化。
|
||||||
|
last_exc = None
|
||||||
|
for attempt in range(3):
|
||||||
|
try:
|
||||||
|
os.replace(tmp_path, path)
|
||||||
|
last_exc = None
|
||||||
|
break
|
||||||
|
except PermissionError as e:
|
||||||
|
last_exc = e
|
||||||
|
if attempt < 2:
|
||||||
|
time.sleep(0.05 * (attempt + 1))
|
||||||
|
if last_exc is not None:
|
||||||
|
raise last_exc
|
||||||
|
except BaseException:
|
||||||
|
try:
|
||||||
|
os.unlink(tmp_path)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
raise
|
||||||
|
|
||||||
+113
-25
@@ -8,8 +8,13 @@ common.config - aibox 配置管理器
|
|||||||
|
|
||||||
import os
|
import os
|
||||||
import json
|
import json
|
||||||
|
import copy
|
||||||
|
import shutil
|
||||||
|
import threading
|
||||||
|
from datetime import datetime
|
||||||
from typing import Dict, Optional
|
from typing import Dict, Optional
|
||||||
|
|
||||||
|
from .atomic import atomic_write_text
|
||||||
from .i18n import i18n
|
from .i18n import i18n
|
||||||
class ConfigManager:
|
class ConfigManager:
|
||||||
"""配置管理器 - 单例模式,负责加载和保存配置"""
|
"""配置管理器 - 单例模式,负责加载和保存配置"""
|
||||||
@@ -17,15 +22,20 @@ class ConfigManager:
|
|||||||
_instance = None
|
_instance = None
|
||||||
_config = None
|
_config = None
|
||||||
_config_path = None
|
_config_path = None
|
||||||
|
_instance_lock = threading.Lock()
|
||||||
|
|
||||||
def __new__(cls):
|
def __new__(cls):
|
||||||
if cls._instance is None:
|
# J6(审计 2026-08-19):此前单例创建与初始化均无锁——多线程首例
|
||||||
cls._instance = super().__new__(cls)
|
# 并发可双实例、双重加载配置
|
||||||
|
with cls._instance_lock:
|
||||||
|
if cls._instance is None:
|
||||||
|
cls._instance = super().__new__(cls)
|
||||||
return cls._instance
|
return cls._instance
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
if self._config is None:
|
with self._instance_lock:
|
||||||
self._load_config()
|
if self._config is None:
|
||||||
|
self._load_config()
|
||||||
|
|
||||||
def _get_default_config_path(self) -> str:
|
def _get_default_config_path(self) -> str:
|
||||||
env_config = os.environ.get("AI_CONFIG_PATH")
|
env_config = os.environ.get("AI_CONFIG_PATH")
|
||||||
@@ -64,14 +74,36 @@ class ConfigManager:
|
|||||||
|
|
||||||
# 大模型调用接口
|
# 大模型调用接口
|
||||||
"api": {
|
"api": {
|
||||||
|
"provider": "openai_compatible", # 提供商:deepseek / openai_compatible(OpenAI Chat Completions 协议兼容)
|
||||||
"deepseek_api_url": "https://api.deepseek.com/chat/completions", # deepseek API URL
|
"deepseek_api_url": "https://api.deepseek.com/chat/completions", # deepseek API URL
|
||||||
"deepseek_model": "deepseek-v4-flash", # 模型名字
|
"deepseek_model": "deepseek-v4-flash", # 模型名字
|
||||||
"reasoning_enabled": False, # 开启思维链
|
"reasoning_enabled": False, # 开启思维链
|
||||||
"api_key_env": "DEEPSEEK_API_KEY", # API-KEY所在的环境变量
|
"api_key_env": "DEEPSEEK_API_KEY", # API-KEY所在的环境变量
|
||||||
"timeout": 120, # 调用大模型请求超时时间
|
"timeout": 120, # 调用大模型请求超时时间
|
||||||
"temperature": 0.7 # 大模型对话温度(类似思维发散程度)
|
"temperature": 0.7, # 大模型对话温度(类似思维发散程度)
|
||||||
|
"default_model": "deepseek" # 默认模型档案名(对应 models 节的 key)
|
||||||
},
|
},
|
||||||
|
|
||||||
|
# 多模型档案:可配置多个 OpenAI Chat Completions 兼容端点,按名称选择
|
||||||
|
# 示例(在 ~/.aibox/config.json 中添加):
|
||||||
|
# "models": {
|
||||||
|
# "kimi": {
|
||||||
|
# "api_url": "https://api.moonshot.cn/v1/chat/completions",
|
||||||
|
# "model": "kimi-k2",
|
||||||
|
# "api_key_env": "MOONSHOT_API_KEY", # 未配置回退 api.api_key_env;空串 "" 为免密端点
|
||||||
|
# "reasoning_enabled": false,
|
||||||
|
# "timeout": 120
|
||||||
|
# },
|
||||||
|
# "local": {
|
||||||
|
# "api_url": "http://localhost:11434/v1/chat/completions",
|
||||||
|
# "model": "qwen3:32b",
|
||||||
|
# "api_key_env": ""
|
||||||
|
# }
|
||||||
|
# }
|
||||||
|
# 选择优先级:--model <档案名> > api.default_model > 回退旧 api 扁平节。
|
||||||
|
# 默认留空:用户未配置档案时走旧 api 节(零破坏)。
|
||||||
|
"models": {},
|
||||||
|
|
||||||
# 大模型上下文
|
# 大模型上下文
|
||||||
"context": {
|
"context": {
|
||||||
"max_history_messages": 500, # 最大历史消息数量
|
"max_history_messages": 500, # 最大历史消息数量
|
||||||
@@ -160,12 +192,26 @@ class ConfigManager:
|
|||||||
print(i18n.get('config_loaded', config_path))
|
print(i18n.get('config_loaded', config_path))
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"⚠️ {i18n.get('error_prefix')}{e},使用默认配置")
|
print(f"⚠️ {i18n.get('error_prefix')}{e},使用默认配置")
|
||||||
|
# J5(审计 2026-08-19):损坏文件先备份再继续——此前直接改用
|
||||||
|
# 默认配置,任一后续保存会用"默认+本次改动"整体覆盖原文件,
|
||||||
|
# proxy/api/context 等全部手工配置永久丢失
|
||||||
|
self._backup_corrupted_config(config_path)
|
||||||
self._config = self._get_default_config()
|
self._config = self._get_default_config()
|
||||||
else:
|
else:
|
||||||
print(i18n.get('config_not_found', config_path))
|
print(i18n.get('config_not_found', config_path))
|
||||||
self._config = self._get_default_config()
|
self._config = self._get_default_config()
|
||||||
self._save_default_config(config_path)
|
self._save_default_config(config_path)
|
||||||
|
self._migrate_legacy_proxy_password()
|
||||||
|
|
||||||
|
def _backup_corrupted_config(self, config_path: str):
|
||||||
|
"""J5:损坏配置先复制备份(带时间戳),防止后续保存整体覆盖丢失"""
|
||||||
|
try:
|
||||||
|
backup = f"{config_path}.corrupt.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||||||
|
shutil.copy2(config_path, backup)
|
||||||
|
print(f"⚠️ 已备份损坏的配置文件到 {backup}")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"⚠️ 备份损坏配置文件失败: {e}")
|
||||||
|
|
||||||
def _merge_configs(self, default: Dict, user: Dict) -> Dict:
|
def _merge_configs(self, default: Dict, user: Dict) -> Dict:
|
||||||
result = default.copy()
|
result = default.copy()
|
||||||
for key, value in user.items():
|
for key, value in user.items():
|
||||||
@@ -177,9 +223,9 @@ class ConfigManager:
|
|||||||
|
|
||||||
def _save_default_config(self, config_path: str):
|
def _save_default_config(self, config_path: str):
|
||||||
try:
|
try:
|
||||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
atomic_write_text(
|
||||||
with open(config_path, 'w', encoding='utf-8') as f:
|
config_path,
|
||||||
json.dump(self._get_default_config(), f, ensure_ascii=False, indent=2)
|
json.dumps(self._get_default_config(), ensure_ascii=False, indent=2))
|
||||||
print(i18n.get('config_created', config_path))
|
print(i18n.get('config_created', config_path))
|
||||||
except:
|
except:
|
||||||
pass
|
pass
|
||||||
@@ -198,6 +244,28 @@ class ConfigManager:
|
|||||||
self._load_config()
|
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]]:
|
def get_proxy(self) -> Optional[Dict[str, str]]:
|
||||||
"""获取当前代理配置(返回 requests 可用的 proxies 字典)"""
|
"""获取当前代理配置(返回 requests 可用的 proxies 字典)"""
|
||||||
proxy_conf = self.get("proxy.socks5")
|
proxy_conf = self.get("proxy.socks5")
|
||||||
@@ -206,15 +274,20 @@ class ConfigManager:
|
|||||||
host = proxy_conf.get("host", "127.0.0.1")
|
host = proxy_conf.get("host", "127.0.0.1")
|
||||||
port = proxy_conf.get("port", 1080)
|
port = proxy_conf.get("port", 1080)
|
||||||
username = proxy_conf.get("username")
|
username = proxy_conf.get("username")
|
||||||
password = proxy_conf.get("password")
|
password = self._resolve_proxy_password(proxy_conf)
|
||||||
if username and password:
|
if username and password:
|
||||||
proxy_url = f"socks5://{username}:{password}@{host}:{port}"
|
proxy_url = f"socks5://{username}:{password}@{host}:{port}"
|
||||||
else:
|
else:
|
||||||
proxy_url = f"socks5://{host}:{port}"
|
proxy_url = f"socks5://{host}:{port}"
|
||||||
return {"http": proxy_url, "https": proxy_url}
|
return {"http": proxy_url, "https": proxy_url}
|
||||||
|
|
||||||
def set_proxy(self, host: str, port: int, username: str = None, password: str = None, enabled: bool = True, save: bool = True):
|
def set_proxy(self, host: str, port: int, username: str = None, password: str = None,
|
||||||
"""动态设置代理配置(内存中更新,并可选择写入配置文件)"""
|
password_env: str = None, enabled: bool = True, save: bool = True):
|
||||||
|
"""动态设置代理配置(内存中更新,并可选择写入配置文件)。
|
||||||
|
|
||||||
|
password 明文仅存内存(_save_config 落盘时强制脱敏);
|
||||||
|
跨重启持久化请用 password_env 指向存有密码的环境变量名。
|
||||||
|
"""
|
||||||
# 确保 proxy 节点存在
|
# 确保 proxy 节点存在
|
||||||
if "proxy" not in self._config:
|
if "proxy" not in self._config:
|
||||||
self._config["proxy"] = {}
|
self._config["proxy"] = {}
|
||||||
@@ -223,31 +296,46 @@ class ConfigManager:
|
|||||||
"host": host,
|
"host": host,
|
||||||
"port": port,
|
"port": port,
|
||||||
"username": username,
|
"username": username,
|
||||||
"password": password
|
"password": password,
|
||||||
|
"password_env": password_env
|
||||||
}
|
}
|
||||||
if save:
|
if save:
|
||||||
self._save_config()
|
return self._save_config()
|
||||||
return True
|
return True
|
||||||
|
|
||||||
def _save_config(self):
|
def _save_config(self) -> bool:
|
||||||
"""将当前配置保存到配置文件(覆盖)"""
|
"""将当前配置保存到配置文件(覆盖)。返回是否成功落盘。
|
||||||
|
|
||||||
|
落盘前对代理密码做纵深防御脱敏:无论调用路径如何,明文密码一律不写入文件。
|
||||||
|
L2(审计 2026-08-18):此前异常仅 print,set_proxy 无条件报"已保存",
|
||||||
|
磁盘只读等场景下用户以为已持久化,重启后配置丢失。
|
||||||
|
"""
|
||||||
config_path = self._config_path
|
config_path = self._config_path
|
||||||
if config_path:
|
if not config_path:
|
||||||
try:
|
return False
|
||||||
os.makedirs(os.path.dirname(config_path), exist_ok=True)
|
try:
|
||||||
with open(config_path, 'w', encoding='utf-8') as f:
|
safe_config = copy.deepcopy(self._config)
|
||||||
json.dump(self._config, f, ensure_ascii=False, indent=2)
|
socks = safe_config.get("proxy", {}).get("socks5")
|
||||||
print(f"✅ 配置已保存到 {config_path}")
|
if socks and socks.get("password"):
|
||||||
except Exception as e:
|
socks["password"] = None
|
||||||
print(f"⚠️ 保存配置失败: {e}")
|
# J5:原子写(mkstemp + fsync + os.replace),崩溃不留半截 JSON
|
||||||
|
atomic_write_text(
|
||||||
|
config_path,
|
||||||
|
json.dumps(safe_config, ensure_ascii=False, indent=2))
|
||||||
|
print(f"✅ 配置已保存到 {config_path}")
|
||||||
|
return True
|
||||||
|
except Exception as e:
|
||||||
|
print(f"⚠️ 保存配置失败: {e}")
|
||||||
|
return False
|
||||||
|
|
||||||
def get_proxy_info(self) -> Dict:
|
def get_proxy_info(self) -> Dict:
|
||||||
"""获取代理配置的详细信息(用于显示)"""
|
"""获取代理配置的详细信息(用于显示,不回传任何密码值)"""
|
||||||
proxy_conf = self.get("proxy.socks5", {})
|
proxy_conf = self.get("proxy.socks5", {})
|
||||||
return {
|
return {
|
||||||
"enabled": proxy_conf.get("enabled", False),
|
"enabled": proxy_conf.get("enabled", False),
|
||||||
"host": proxy_conf.get("host", ""),
|
"host": proxy_conf.get("host", ""),
|
||||||
"port": proxy_conf.get("port", 0),
|
"port": proxy_conf.get("port", 0),
|
||||||
"username": proxy_conf.get("username"),
|
"username": proxy_conf.get("username"),
|
||||||
"password": proxy_conf.get("password") is not None # 只显示是否存在
|
"password": self._resolve_proxy_password(proxy_conf) is not None, # 只显示是否存在
|
||||||
|
"password_env": proxy_conf.get("password_env") or ""
|
||||||
}
|
}
|
||||||
@@ -236,6 +236,31 @@ class Constants:
|
|||||||
def COMPRESSOR_LLM_MAX_TOKENS(cls):
|
def COMPRESSOR_LLM_MAX_TOKENS(cls):
|
||||||
return cls._get_config().get("compressor.llm_max_tokens", 1000)
|
return cls._get_config().get("compressor.llm_max_tokens", 1000)
|
||||||
|
|
||||||
|
# 后台进程管理(bg_run/bg_list/bg_output/bg_kill)
|
||||||
|
@property
|
||||||
|
def BG_PROCESSES_ENABLED(cls):
|
||||||
|
return cls._get_config().get("background_processes.enabled", True)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def BG_MAX_PROCESSES(cls):
|
||||||
|
return cls._get_config().get("background_processes.max_processes", 16)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def BG_MAX_ENTRIES(cls):
|
||||||
|
return cls._get_config().get("background_processes.max_entries", 32)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def BG_BUFFER_PER_PROCESS(cls):
|
||||||
|
return cls._get_config().get("background_processes.buffer_per_process", 262144)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def BG_MAX_LIFETIME(cls):
|
||||||
|
return cls._get_config().get("background_processes.max_lifetime", 0)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def BG_MAX_WRITE_BYTES(cls):
|
||||||
|
return cls._get_config().get("background_processes.max_write_bytes", 65536)
|
||||||
|
|
||||||
|
|
||||||
# 创建常量实例
|
# 创建常量实例
|
||||||
constants = Constants()
|
constants = Constants()
|
||||||
@@ -98,6 +98,7 @@ class I18n:
|
|||||||
'memo_overdue': "⚠️",
|
'memo_overdue': "⚠️",
|
||||||
'memo_pending': "⏰",
|
'memo_pending': "⏰",
|
||||||
'search_results': "🔍 搜索 '{}' 找到 {} 条记忆:",
|
'search_results': "🔍 搜索 '{}' 找到 {} 条记忆:",
|
||||||
|
'search_usage': "用法: /search 关键词",
|
||||||
'current_config': "⚙️ 当前配置:",
|
'current_config': "⚙️ 当前配置:",
|
||||||
'config_debug': " 调试级别: {}",
|
'config_debug': " 调试级别: {}",
|
||||||
'config_data_dir': " 数据目录: {}",
|
'config_data_dir': " 数据目录: {}",
|
||||||
@@ -110,6 +111,7 @@ class I18n:
|
|||||||
'enabled': "启用",
|
'enabled': "启用",
|
||||||
'disabled': "禁用",
|
'disabled': "禁用",
|
||||||
'config_reloaded': "✅ 配置已重新加载",
|
'config_reloaded': "✅ 配置已重新加载",
|
||||||
|
'config_reloaded_cli_kept': "(命令行进程级覆盖已保留)",
|
||||||
'prompts_reloaded': "✅ 提示词已重新加载 ({} 个)",
|
'prompts_reloaded': "✅ 提示词已重新加载 ({} 个)",
|
||||||
'checking_memos': "🔍 检查任务...",
|
'checking_memos': "🔍 检查任务...",
|
||||||
'overdue_found': "发现 {} 个到期任务:",
|
'overdue_found': "发现 {} 个到期任务:",
|
||||||
@@ -157,6 +159,7 @@ class I18n:
|
|||||||
'memo_add_failed': "❌ 添加失败: {}",
|
'memo_add_failed': "❌ 添加失败: {}",
|
||||||
'memo_completed': "✅ 任务 #{} 已完成",
|
'memo_completed': "✅ 任务 #{} 已完成",
|
||||||
'memo_complete_failed': "❌ 无法完成任务 #{}",
|
'memo_complete_failed': "❌ 无法完成任务 #{}",
|
||||||
|
'memo_complete_missing_id': "❌ 缺少 task_id,无法完成任务",
|
||||||
|
|
||||||
# 即时任务相关
|
# 即时任务相关
|
||||||
'immediate_memo_added': "⚡ 已添加为即时任务并立即检查",
|
'immediate_memo_added': "⚡ 已添加为即时任务并立即检查",
|
||||||
@@ -258,6 +261,7 @@ class I18n:
|
|||||||
'memo_overdue': "⚠️",
|
'memo_overdue': "⚠️",
|
||||||
'memo_pending': "⏰",
|
'memo_pending': "⏰",
|
||||||
'search_results': "🔍 Searching '{}' found {} memories:",
|
'search_results': "🔍 Searching '{}' found {} memories:",
|
||||||
|
'search_usage': "Usage: /search <keyword>",
|
||||||
'current_config': "⚙️ Current configuration:",
|
'current_config': "⚙️ Current configuration:",
|
||||||
'config_debug': " Debug level: {}",
|
'config_debug': " Debug level: {}",
|
||||||
'config_data_dir': " Data directory: {}",
|
'config_data_dir': " Data directory: {}",
|
||||||
@@ -270,6 +274,7 @@ class I18n:
|
|||||||
'enabled': "enabled",
|
'enabled': "enabled",
|
||||||
'disabled': "disabled",
|
'disabled': "disabled",
|
||||||
'config_reloaded': "✅ Configuration reloaded",
|
'config_reloaded': "✅ Configuration reloaded",
|
||||||
|
'config_reloaded_cli_kept': " (CLI process-level overrides preserved)",
|
||||||
'prompts_reloaded': "✅ Prompts reloaded ({})",
|
'prompts_reloaded': "✅ Prompts reloaded ({})",
|
||||||
'checking_memos': "🔍 Checking tasks...",
|
'checking_memos': "🔍 Checking tasks...",
|
||||||
'overdue_found': "Found {} overdue tasks:",
|
'overdue_found': "Found {} overdue tasks:",
|
||||||
@@ -317,6 +322,7 @@ class I18n:
|
|||||||
'memo_add_failed': "❌ Failed to add: {}",
|
'memo_add_failed': "❌ Failed to add: {}",
|
||||||
'memo_completed': "✅ Task #{} completed",
|
'memo_completed': "✅ Task #{} completed",
|
||||||
'memo_complete_failed': "❌ Cannot complete task #{}",
|
'memo_complete_failed': "❌ Cannot complete task #{}",
|
||||||
|
'memo_complete_missing_id': "❌ Missing task_id, cannot complete task",
|
||||||
|
|
||||||
# Immediate task related
|
# Immediate task related
|
||||||
'immediate_memo_added': "⚡ Added as immediate task and checked now",
|
'immediate_memo_added': "⚡ Added as immediate task and checked now",
|
||||||
|
|||||||
+31
-13
@@ -4,12 +4,14 @@
|
|||||||
<meta charset="UTF-8">
|
<meta charset="UTF-8">
|
||||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
<title>AI智能助手 - Socket.IO客户端</title>
|
<title>AI智能助手 - Socket.IO客户端</title>
|
||||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||||
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||||
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||||
<!-- Socket.IO 客户端库 -->
|
<!-- Socket.IO 客户端库 -->
|
||||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js"></script>
|
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js" integrity="sha384-7EyYLQZgWBi67fBtVxw60/OWl1kjsfrPFcaU0pp0nAh+i8FD068QogUvg85Ewy1k" crossorigin="anonymous"></script>
|
||||||
<style>
|
<style>
|
||||||
/* 样式保持不变,完全使用原有的样式 */
|
/* 样式保持不变,完全使用原有的样式 */
|
||||||
* {
|
* {
|
||||||
@@ -982,6 +984,15 @@
|
|||||||
</template>
|
</template>
|
||||||
|
|
||||||
<script>
|
<script>
|
||||||
|
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||||
|
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||||
|
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||||
|
function renderMarkdown(content) {
|
||||||
|
const html = marked.parse(content);
|
||||||
|
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||||
|
return html.replace(/</g, '<');
|
||||||
|
}
|
||||||
|
|
||||||
document.addEventListener('DOMContentLoaded', function() {
|
document.addEventListener('DOMContentLoaded', function() {
|
||||||
// 配置marked
|
// 配置marked
|
||||||
marked.setOptions({
|
marked.setOptions({
|
||||||
@@ -1124,7 +1135,7 @@
|
|||||||
textDiv.setAttribute('data-raw-content', streamData.content);
|
textDiv.setAttribute('data-raw-content', streamData.content);
|
||||||
|
|
||||||
if (markdownEnabled) {
|
if (markdownEnabled) {
|
||||||
textDiv.innerHTML = marked.parse(streamData.content);
|
textDiv.innerHTML = renderMarkdown(streamData.content);
|
||||||
textDiv.querySelectorAll('pre code').forEach((block) => {
|
textDiv.querySelectorAll('pre code').forEach((block) => {
|
||||||
hljs.highlightElement(block);
|
hljs.highlightElement(block);
|
||||||
});
|
});
|
||||||
@@ -1419,7 +1430,7 @@
|
|||||||
const rawContent = msgDiv.getAttribute('data-raw-content');
|
const rawContent = msgDiv.getAttribute('data-raw-content');
|
||||||
if (rawContent) {
|
if (rawContent) {
|
||||||
if (markdownEnabled) {
|
if (markdownEnabled) {
|
||||||
msgDiv.innerHTML = marked.parse(rawContent);
|
msgDiv.innerHTML = renderMarkdown(rawContent);
|
||||||
msgDiv.classList.add('markdown-body');
|
msgDiv.classList.add('markdown-body');
|
||||||
addCopyButtonsToCodeBlocks(msgDiv);
|
addCopyButtonsToCodeBlocks(msgDiv);
|
||||||
} else {
|
} else {
|
||||||
@@ -1438,6 +1449,7 @@
|
|||||||
let reconnectAttempts = 0;
|
let reconnectAttempts = 0;
|
||||||
const MAX_RECONNECT_ATTEMPTS = 5;
|
const MAX_RECONNECT_ATTEMPTS = 5;
|
||||||
let targetWsUrl = null;
|
let targetWsUrl = null;
|
||||||
|
let reconnectTimer = null; // F2(审计 2026-08-20):重连定时器句柄,手动断开时清除
|
||||||
|
|
||||||
function createSocketConnection() {
|
function createSocketConnection() {
|
||||||
targetWsUrl = document.getElementById('ws-url').value.trim();
|
targetWsUrl = document.getElementById('ws-url').value.trim();
|
||||||
@@ -1503,7 +1515,8 @@
|
|||||||
if (reason !== 'io client disconnect' && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
if (reason !== 'io client disconnect' && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
||||||
reconnectAttempts++;
|
reconnectAttempts++;
|
||||||
showSystemMessage(`⚠️ 连接断开,${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} 次重试...`);
|
showSystemMessage(`⚠️ 连接断开,${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} 次重试...`);
|
||||||
setTimeout(() => {
|
reconnectTimer = setTimeout(() => {
|
||||||
|
reconnectTimer = null;
|
||||||
if (!socket || !socket.connected) {
|
if (!socket || !socket.connected) {
|
||||||
console.log('尝试重连...');
|
console.log('尝试重连...');
|
||||||
socket = createSocketConnection();
|
socket = createSocketConnection();
|
||||||
@@ -1676,7 +1689,8 @@
|
|||||||
|
|
||||||
// 事件绑定
|
// 事件绑定
|
||||||
document.getElementById('connect-btn').addEventListener('click', function() {
|
document.getElementById('connect-btn').addEventListener('click', function() {
|
||||||
if (socket && socket.connected) {
|
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要断开,否则重复点击产生孤儿连接/双后端会话
|
||||||
|
if (socket) {
|
||||||
socket.disconnect();
|
socket.disconnect();
|
||||||
}
|
}
|
||||||
reconnectAttempts = 0;
|
reconnectAttempts = 0;
|
||||||
@@ -1684,6 +1698,8 @@
|
|||||||
});
|
});
|
||||||
|
|
||||||
document.getElementById('disconnect-btn').addEventListener('click', function() {
|
document.getElementById('disconnect-btn').addEventListener('click', function() {
|
||||||
|
// F2:清除已排程的重连定时器,防止倒计时到期把连接"复活"
|
||||||
|
if (reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
|
||||||
if (socket) {
|
if (socket) {
|
||||||
socket.disconnect();
|
socket.disconnect();
|
||||||
socket = null;
|
socket = null;
|
||||||
@@ -1717,11 +1733,13 @@
|
|||||||
|
|
||||||
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
||||||
const container = document.getElementById('messages-container');
|
const container = document.getElementById('messages-container');
|
||||||
|
// F1(审计 2026-08-20):welcome-message 在 container 内,
|
||||||
|
// 必须先 cloneNode 缓存再清空,否则 querySelector 返回 null 抛 TypeError
|
||||||
|
const welcomeEl = document.querySelector('.welcome-message');
|
||||||
|
const welcomeHtml = welcomeEl ? welcomeEl.cloneNode(true) : null;
|
||||||
container.innerHTML = '';
|
container.innerHTML = '';
|
||||||
|
if (welcomeHtml) container.appendChild(welcomeHtml);
|
||||||
const welcomeHtml = document.querySelector('.welcome-message').cloneNode(true);
|
|
||||||
container.appendChild(welcomeHtml);
|
|
||||||
|
|
||||||
clearAllStreamingMessages();
|
clearAllStreamingMessages();
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
+35
-11
@@ -5,13 +5,15 @@
|
|||||||
<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=no">
|
<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=no">
|
||||||
<title>小猫AI · 智能喵语伙伴</title>
|
<title>小猫AI · 智能喵语伙伴</title>
|
||||||
<!-- Font Awesome 6 (免费图标库) -->
|
<!-- Font Awesome 6 (免费图标库) -->
|
||||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||||
<!-- highlight.js 代码高亮 + marked markdown解析 -->
|
<!-- highlight.js 代码高亮 + marked markdown解析 -->
|
||||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||||
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||||
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||||
<!-- Socket.IO 客户端库 -->
|
<!-- Socket.IO 客户端库 -->
|
||||||
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js"></script>
|
<script src="https://cdn.socket.io/4.5.0/socket.io.min.js" integrity="sha384-7EyYLQZgWBi67fBtVxw60/OWl1kjsfrPFcaU0pp0nAh+i8FD068QogUvg85Ewy1k" crossorigin="anonymous"></script>
|
||||||
<style>
|
<style>
|
||||||
/* ========== 融合小猫风格 + WebSocket增强样式 ========== */
|
/* ========== 融合小猫风格 + WebSocket增强样式 ========== */
|
||||||
* {
|
* {
|
||||||
@@ -498,6 +500,15 @@
|
|||||||
|
|
||||||
<script>
|
<script>
|
||||||
(function(){
|
(function(){
|
||||||
|
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||||
|
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||||
|
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||||
|
function renderMarkdown(content) {
|
||||||
|
const html = marked.parse(content);
|
||||||
|
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||||
|
return html.replace(/</g, '<');
|
||||||
|
}
|
||||||
|
|
||||||
// ---------- DOM 元素 ----------
|
// ---------- DOM 元素 ----------
|
||||||
const bgFile = document.getElementById('bgFile');
|
const bgFile = document.getElementById('bgFile');
|
||||||
const roleImg = document.getElementById('roleImg');
|
const roleImg = document.getElementById('roleImg');
|
||||||
@@ -533,7 +544,12 @@
|
|||||||
const div = document.createElement('div');
|
const div = document.createElement('div');
|
||||||
div.className = isError ? 'error-msg' : 'system-msg';
|
div.className = isError ? 'error-msg' : 'system-msg';
|
||||||
div.style.cssText = 'background:rgba(0,0,0,0.6); border-radius:16px; padding:6px 12px; margin:6px 0; font-size:12px;';
|
div.style.cssText = 'background:rgba(0,0,0,0.6); border-radius:16px; padding:6px 12px; margin:6px 0; font-size:12px;';
|
||||||
div.innerHTML = `<i class="fas ${isError ? 'fa-exclamation-triangle' : 'fa-info-circle'}"></i> ${text}`;
|
// J2:text 为服务端内容(命令执行结果/AI 输出),禁止拼 innerHTML——
|
||||||
|
// 原实现 `${text}` 模板拼接使任意 HTML 注入执行;图标为静态字面量
|
||||||
|
const icon = document.createElement('i');
|
||||||
|
icon.className = `fas ${isError ? 'fa-exclamation-triangle' : 'fa-info-circle'}`;
|
||||||
|
div.appendChild(icon);
|
||||||
|
div.appendChild(document.createTextNode(' ' + text));
|
||||||
chatBox.appendChild(div);
|
chatBox.appendChild(div);
|
||||||
chatBox.scrollTop = chatBox.scrollHeight;
|
chatBox.scrollTop = chatBox.scrollHeight;
|
||||||
}
|
}
|
||||||
@@ -545,7 +561,7 @@
|
|||||||
if(textDiv) {
|
if(textDiv) {
|
||||||
textDiv.setAttribute('data-raw', stream.content);
|
textDiv.setAttribute('data-raw', stream.content);
|
||||||
if(markdownEnabled) {
|
if(markdownEnabled) {
|
||||||
textDiv.innerHTML = marked.parse(stream.content);
|
textDiv.innerHTML = renderMarkdown(stream.content);
|
||||||
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
||||||
} else textDiv.innerText = stream.content;
|
} else textDiv.innerText = stream.content;
|
||||||
}
|
}
|
||||||
@@ -572,7 +588,12 @@
|
|||||||
if(!streamingTasks[eid]) {
|
if(!streamingTasks[eid]) {
|
||||||
const sep = document.createElement('div');
|
const sep = document.createElement('div');
|
||||||
sep.className = 'task-separator';
|
sep.className = 'task-separator';
|
||||||
sep.innerHTML = `<i class="fas fa-bolt"></i> 定时任务 #${taskMeta.taskId} (${taskMeta.taskTitle || '执行'}) 第${taskMeta.triggerCount||1}次执行`;
|
// J2:taskMeta.taskTitle 为服务端可控内容,禁止拼 innerHTML
|
||||||
|
const bolt = document.createElement('i');
|
||||||
|
bolt.className = 'fas fa-bolt';
|
||||||
|
sep.appendChild(bolt);
|
||||||
|
sep.appendChild(document.createTextNode(
|
||||||
|
` 定时任务 #${taskMeta.taskId} (${taskMeta.taskTitle || '执行'}) 第${taskMeta.triggerCount||1}次执行`));
|
||||||
chatBox.appendChild(sep);
|
chatBox.appendChild(sep);
|
||||||
const taskDiv = document.createElement('div');
|
const taskDiv = document.createElement('div');
|
||||||
taskDiv.className = 'task-msg streaming';
|
taskDiv.className = 'task-msg streaming';
|
||||||
@@ -590,7 +611,7 @@
|
|||||||
if(st && st.element) {
|
if(st && st.element) {
|
||||||
const textDiv = st.element.querySelector('.message-text');
|
const textDiv = st.element.querySelector('.message-text');
|
||||||
if(textDiv) {
|
if(textDiv) {
|
||||||
if(markdownEnabled) textDiv.innerHTML = marked.parse(st.content);
|
if(markdownEnabled) textDiv.innerHTML = renderMarkdown(st.content);
|
||||||
else textDiv.innerText = st.content;
|
else textDiv.innerText = st.content;
|
||||||
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
textDiv.querySelectorAll('pre code').forEach(block => hljs.highlightElement(block));
|
||||||
}
|
}
|
||||||
@@ -613,7 +634,7 @@
|
|||||||
if(task.timer) clearTimeout(task.timer);
|
if(task.timer) clearTimeout(task.timer);
|
||||||
if(task.element) {
|
if(task.element) {
|
||||||
const txtDiv = task.element.querySelector('.message-text');
|
const txtDiv = task.element.querySelector('.message-text');
|
||||||
if(txtDiv && markdownEnabled) txtDiv.innerHTML = marked.parse(task.content);
|
if(txtDiv && markdownEnabled) txtDiv.innerHTML = renderMarkdown(task.content);
|
||||||
task.element.classList.remove('streaming');
|
task.element.classList.remove('streaming');
|
||||||
}
|
}
|
||||||
globalMsgCount++;
|
globalMsgCount++;
|
||||||
@@ -652,7 +673,8 @@
|
|||||||
}
|
}
|
||||||
|
|
||||||
function connectWebSocket() {
|
function connectWebSocket() {
|
||||||
if(socket && socket.connected) {
|
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要断开,否则重复点击产生孤儿连接/双后端会话
|
||||||
|
if(socket) {
|
||||||
socket.disconnect();
|
socket.disconnect();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -731,6 +753,8 @@
|
|||||||
|
|
||||||
function disconnectWebSocket() {
|
function disconnectWebSocket() {
|
||||||
manualDisconnect = true;
|
manualDisconnect = true;
|
||||||
|
// F2(审计 2026-08-20):清除已排程的重连定时器,防止倒计时到期把连接"复活"
|
||||||
|
if(reconnectTimer) { clearTimeout(reconnectTimer); reconnectTimer = null; }
|
||||||
if(socket) {
|
if(socket) {
|
||||||
socket.disconnect();
|
socket.disconnect();
|
||||||
socket = null;
|
socket = null;
|
||||||
|
|||||||
+31
-13
@@ -4,10 +4,12 @@
|
|||||||
<meta charset="UTF-8">
|
<meta charset="UTF-8">
|
||||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||||
<title>AI智能助手 - WebSocket直连</title>
|
<title>AI智能助手 - WebSocket直连</title>
|
||||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css" integrity="sha384-3B6NwesSXE7YJlcLI9RpRqGf2p/EgVH8BgoKTaUrmKNDkHPStTQ3EyoYjCGXaOTS" crossorigin="anonymous">
|
||||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css">
|
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/styles/github-dark.min.css" integrity="sha384-wH75j6z1lH97ZOpMOInqhgKzFkAInZPPSPlZpYKYTOqsaizPvhQZmAtLcPKXpLyH" crossorigin="anonymous">
|
||||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js"></script>
|
<!-- J13(审计 2026-08-19):CDN 脚本补 SRI integrity,防供应链劫持 -->
|
||||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js"></script>
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/9.1.6/marked.min.js" integrity="sha384-odPBjvtXVM/5hOYIr3A1dB+flh0c3wAT3bSesIOqEGmyUA4JoKf/YTWy0XKOYAY7" crossorigin="anonymous"></script>
|
||||||
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.8.0/highlight.min.js" integrity="sha384-g4mRvs7AO0/Ol5LxcGyz4Doe21pVhGNnC3EQw5shw+z+aXDN86HqUdwXWO+Gz2zI" crossorigin="anonymous"></script>
|
||||||
|
<script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.2.4/purify.min.js" integrity="sha384-eEu5CTj3qGvu9PdJuS+YlkNi7d2XxQROAFYOr59zgObtlcux1ae1Il3u7jvdCSWu" crossorigin="anonymous"></script>
|
||||||
<style>
|
<style>
|
||||||
/* 样式保持不变,完全使用原有的样式 */
|
/* 样式保持不变,完全使用原有的样式 */
|
||||||
* {
|
* {
|
||||||
@@ -951,6 +953,15 @@
|
|||||||
</template>
|
</template>
|
||||||
|
|
||||||
<script>
|
<script>
|
||||||
|
// J2(审计 2026-08-19):marked 输出经 DOMPurify 净化——AI 输出/命令
|
||||||
|
// 执行结果是不可信输入,含原始 HTML 即可注入脚本;DOMPurify 缺失时
|
||||||
|
// 退化为转义文本(只显示不执行,防 XSS 兜底)
|
||||||
|
function renderMarkdown(content) {
|
||||||
|
const html = marked.parse(content);
|
||||||
|
if (window.DOMPurify) return DOMPurify.sanitize(html);
|
||||||
|
return html.replace(/</g, '<');
|
||||||
|
}
|
||||||
|
|
||||||
document.addEventListener('DOMContentLoaded', function() {
|
document.addEventListener('DOMContentLoaded', function() {
|
||||||
// 配置marked
|
// 配置marked
|
||||||
marked.setOptions({
|
marked.setOptions({
|
||||||
@@ -1048,7 +1059,7 @@
|
|||||||
textDiv.setAttribute('data-raw-content', streamData.content);
|
textDiv.setAttribute('data-raw-content', streamData.content);
|
||||||
|
|
||||||
if (markdownEnabled) {
|
if (markdownEnabled) {
|
||||||
textDiv.innerHTML = marked.parse(streamData.content);
|
textDiv.innerHTML = renderMarkdown(streamData.content);
|
||||||
textDiv.querySelectorAll('pre code').forEach((block) => {
|
textDiv.querySelectorAll('pre code').forEach((block) => {
|
||||||
hljs.highlightElement(block);
|
hljs.highlightElement(block);
|
||||||
});
|
});
|
||||||
@@ -1336,7 +1347,7 @@
|
|||||||
const rawContent = msgDiv.getAttribute('data-raw-content');
|
const rawContent = msgDiv.getAttribute('data-raw-content');
|
||||||
if (rawContent) {
|
if (rawContent) {
|
||||||
if (markdownEnabled) {
|
if (markdownEnabled) {
|
||||||
msgDiv.innerHTML = marked.parse(rawContent);
|
msgDiv.innerHTML = renderMarkdown(rawContent);
|
||||||
msgDiv.classList.add('markdown-body');
|
msgDiv.classList.add('markdown-body');
|
||||||
addCopyButtonsToCodeBlocks(msgDiv);
|
addCopyButtonsToCodeBlocks(msgDiv);
|
||||||
} else {
|
} else {
|
||||||
@@ -1413,9 +1424,12 @@
|
|||||||
|
|
||||||
if (!manualDisconnect && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
if (!manualDisconnect && reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
|
||||||
reconnectAttempts++;
|
reconnectAttempts++;
|
||||||
showSystemMessage(`⚠️ 连接断开,${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS} 秒后重试...`);
|
// F6(审计 2026-08-20):实际延迟为 3000*N ms,文案与此对齐
|
||||||
|
showSystemMessage(`⚠️ 连接断开,${3 * reconnectAttempts} 秒后重试(${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS})...`);
|
||||||
setTimeout(() => {
|
setTimeout(() => {
|
||||||
if (!ws || ws.readyState !== WebSocket.OPEN) {
|
// F2(审计 2026-08-20):闭包到期时再判 manualDisconnect,
|
||||||
|
// 防止手动断开后已排程定时器把连接"复活"
|
||||||
|
if (!manualDisconnect && (!ws || ws.readyState !== WebSocket.OPEN)) {
|
||||||
console.log('尝试重连...');
|
console.log('尝试重连...');
|
||||||
ws = connectWebSocket();
|
ws = connectWebSocket();
|
||||||
}
|
}
|
||||||
@@ -1585,7 +1599,8 @@
|
|||||||
|
|
||||||
// 事件绑定
|
// 事件绑定
|
||||||
document.getElementById('connect-btn').addEventListener('click', function() {
|
document.getElementById('connect-btn').addEventListener('click', function() {
|
||||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
// F3(审计 2026-08-20):CONNECTING 中的旧连接也要关闭,否则重复点击产生孤儿连接/双后端会话
|
||||||
|
if (ws && ws.readyState !== WebSocket.CLOSED) {
|
||||||
ws.close();
|
ws.close();
|
||||||
}
|
}
|
||||||
reconnectAttempts = 0;
|
reconnectAttempts = 0;
|
||||||
@@ -1628,11 +1643,14 @@
|
|||||||
|
|
||||||
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
document.getElementById('clear-chat-btn').addEventListener('click', function() {
|
||||||
const container = document.getElementById('messages-container');
|
const container = document.getElementById('messages-container');
|
||||||
|
// F1(审计 2026-08-20):welcome-message 在 container 内,
|
||||||
|
// 必须先 cloneNode 缓存再清空,否则 querySelector 返回 null 抛 TypeError,
|
||||||
|
// 且后续 clearAllStreamingMessages/计数清零全部中断不执行
|
||||||
|
const welcomeEl = document.querySelector('.welcome-message');
|
||||||
|
const welcomeHtml = welcomeEl ? welcomeEl.cloneNode(true) : null;
|
||||||
container.innerHTML = '';
|
container.innerHTML = '';
|
||||||
|
if (welcomeHtml) container.appendChild(welcomeHtml);
|
||||||
const welcomeHtml = document.querySelector('.welcome-message').cloneNode(true);
|
|
||||||
container.appendChild(welcomeHtml);
|
|
||||||
|
|
||||||
clearAllStreamingMessages();
|
clearAllStreamingMessages();
|
||||||
messageCount = 0;
|
messageCount = 0;
|
||||||
toolCallCount = 0;
|
toolCallCount = 0;
|
||||||
|
|||||||
+560
-38
@@ -22,27 +22,63 @@ from typing import Dict, Iterator, List, Optional
|
|||||||
|
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
|
import re
|
||||||
|
import time
|
||||||
import requests
|
import requests
|
||||||
|
|
||||||
|
# L12(审计 2026-08-18):瞬时错误体可能回显网关/WAF 错误页中的
|
||||||
|
# Authorization 头,抹掉 Bearer 令牌后再进日志,防密钥落盘
|
||||||
|
_AUTH_HEADER_RE = re.compile(r"Bearer\s+[A-Za-z0-9._~+/=-]+")
|
||||||
|
# T13:错误响应体可能回显裸 API Key(如 sk-xxx),一并脱敏
|
||||||
|
_API_KEY_RE = re.compile(r"sk-[A-Za-z0-9_-]{16,}")
|
||||||
|
|
||||||
|
# T19(审计 2026-08-21):瞬时网络异常白名单。此前 FailoverClient 捕获
|
||||||
|
# requests.exceptions.RequestException 基类——档案 api_url 手写错(漏
|
||||||
|
# scheme 抛 MissingSchema)等永久性配置错误也被当瞬时:标记冷却 60s、
|
||||||
|
# 触发转移、计入 failovers,每次冷却到期再撞一次坏端点,统计与冷却表
|
||||||
|
# 被永久错误污染。显式枚举瞬时异常,其余 requests 异常直接上抛响亮暴露
|
||||||
|
_TRANSIENT_NET_ERRORS = (
|
||||||
|
requests.exceptions.ConnectionError,
|
||||||
|
requests.exceptions.Timeout,
|
||||||
|
requests.exceptions.ChunkedEncodingError,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _sanitize_log_msg(msg: str) -> str:
|
||||||
|
"""日志脱敏:替换错误消息中的 Bearer 令牌与裸 API Key。"""
|
||||||
|
if not msg:
|
||||||
|
return msg
|
||||||
|
msg = _AUTH_HEADER_RE.sub("Bearer ***", msg)
|
||||||
|
msg = _API_KEY_RE.sub("sk-***", msg)
|
||||||
|
return msg
|
||||||
|
|
||||||
|
|
||||||
class LLMClient(ABC):
|
class LLMClient(ABC):
|
||||||
"""大模型客户端抽象基类。子类需实现流式/非流式两种调用,并归一化为统一事件格式。"""
|
"""大模型客户端抽象基类。子类需实现流式/非流式两种调用,并归一化为统一事件格式。"""
|
||||||
|
|
||||||
def __init__(self, api_key: str, model: str, base_url: str,
|
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.api_key = api_key
|
||||||
self.model = model
|
self.model = model
|
||||||
self.base_url = base_url
|
self.base_url = base_url
|
||||||
self.timeout = timeout
|
self.timeout = timeout
|
||||||
self.config = config # 可选,用于读取代理等动态配置
|
self.config = config # 可选,用于读取代理等动态配置
|
||||||
self.logger = logger # 可选,用于日志输出
|
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() 强制关闭底层连接以中断阻塞读取
|
# 当前活动的 HTTP 响应引用,供 cancel() 强制关闭底层连接以中断阻塞读取
|
||||||
self._current_response = None
|
self._current_response = None
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def name(self) -> str:
|
def name(self) -> str:
|
||||||
"""厂商/模型展示名,业务层可用于上下文(如 history 里的 assistant 名)。"""
|
"""厂商/模型展示名,业务层可用于上下文(如 history 里的 assistant 名)。"""
|
||||||
return "LLM"
|
return self._display_name or "LLM"
|
||||||
|
|
||||||
# ---------- 工具方法 ----------
|
# ---------- 工具方法 ----------
|
||||||
|
|
||||||
@@ -55,12 +91,16 @@ class LLMClient(ABC):
|
|||||||
def _log(self, level: str, msg: str):
|
def _log(self, level: str, msg: str):
|
||||||
if self.logger is not None and hasattr(self.logger, level):
|
if self.logger is not None and hasattr(self.logger, level):
|
||||||
try:
|
try:
|
||||||
getattr(self.logger, level)(msg)
|
getattr(self.logger, level)(_sanitize_log_msg(msg))
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def cancel(self):
|
def cancel(self):
|
||||||
"""中断当前流式请求。线程安全,可从任意线程调用。"""
|
"""中断当前流式请求。线程安全,可从任意线程调用。
|
||||||
|
|
||||||
|
注:用户中断标志由 FailoverClient 维护(其 cancel() 广播整条链并置
|
||||||
|
自身 _cancelled,据此跳过转移)——底层 client 只负责强关连接。
|
||||||
|
"""
|
||||||
resp = self._current_response
|
resp = self._current_response
|
||||||
self._current_response = None
|
self._current_response = None
|
||||||
if resp is not None:
|
if resp is not None:
|
||||||
@@ -78,16 +118,18 @@ class LLMClient(ABC):
|
|||||||
# ---------- 抽象接口 ----------
|
# ---------- 抽象接口 ----------
|
||||||
|
|
||||||
@abstractmethod
|
@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:
|
max_tokens: int = 1024, **kwargs) -> str:
|
||||||
"""非流式对话,返回完整文本。用于对话压缩等离线摘要场景。"""
|
"""非流式对话,返回完整文本。用于对话压缩等离线摘要场景。
|
||||||
|
temperature=None 时按 default_temperature(档案级)→ 0.3 回退。"""
|
||||||
|
|
||||||
@abstractmethod
|
@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,
|
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||||
tool_choice: Optional[str] = None,
|
tool_choice: Optional[str] = None,
|
||||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||||
"""流式对话,逐条 yield 统一事件格式的 dict。业务层据此驱动输出/工具调用。"""
|
"""流式对话,逐条 yield 统一事件格式的 dict。业务层据此驱动输出/工具调用。
|
||||||
|
temperature=None 时按 default_temperature(档案级)→ 0.7 回退。"""
|
||||||
|
|
||||||
|
|
||||||
class DeepSeekClient(LLMClient):
|
class DeepSeekClient(LLMClient):
|
||||||
@@ -99,12 +141,18 @@ class DeepSeekClient(LLMClient):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, api_key: str, model: str, base_url: str,
|
def __init__(self, api_key: str, model: str, base_url: str,
|
||||||
timeout: int = 30, config=None, logger=None):
|
timeout: int = 30, config=None, logger=None,
|
||||||
super().__init__(api_key, model, base_url, timeout, config, logger)
|
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
|
@property
|
||||||
def name(self) -> str:
|
def name(self) -> str:
|
||||||
return "DeepSeek"
|
# L1:档案命中时展示档案名(或档案 display_name),旧扁平节保持 "DeepSeek"
|
||||||
|
return self._display_name or "DeepSeek"
|
||||||
|
|
||||||
def _headers(self) -> Dict[str, str]:
|
def _headers(self) -> Dict[str, str]:
|
||||||
return {
|
return {
|
||||||
@@ -112,8 +160,11 @@ class DeepSeekClient(LLMClient):
|
|||||||
"Content-Type": "application/json",
|
"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:
|
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 = {
|
payload = {
|
||||||
"model": self.model,
|
"model": self.model,
|
||||||
"messages": messages,
|
"messages": messages,
|
||||||
@@ -130,18 +181,40 @@ class DeepSeekClient(LLMClient):
|
|||||||
proxies=self._resolve_proxies(),
|
proxies=self._resolve_proxies(),
|
||||||
)
|
)
|
||||||
if response.status_code == 200:
|
if response.status_code == 200:
|
||||||
result = response.json()
|
try:
|
||||||
return result["choices"][0]["message"]["content"].strip()
|
result = response.json()
|
||||||
self._log("error", f"LLM 压缩请求失败: HTTP {response.status_code}")
|
return result["choices"][0]["message"]["content"].strip()
|
||||||
|
finally:
|
||||||
|
response.close()
|
||||||
|
# T9+T13:非 200 立即关连接(stream=True 才需要,但统一防御),
|
||||||
|
# 错误体脱敏后再入异常,防 API Key 经错误页回显进历史/日志
|
||||||
|
status = response.status_code
|
||||||
|
body = _sanitize_log_msg(response.text[:200])
|
||||||
|
response.close()
|
||||||
|
if status == 429 or status >= 500:
|
||||||
|
# L4a:瞬时故障显式抛出——FailoverClient 据此转移下一档案;
|
||||||
|
# 无包装器时由调用方 catch-all 接住,下游仍走失败分支(语义等价)
|
||||||
|
raise TransientLLMError(
|
||||||
|
f"LLM 请求瞬时失败: HTTP {status}: {body}",
|
||||||
|
retry_after=_retry_after_of(response))
|
||||||
|
self._log("error", f"LLM 压缩请求失败: HTTP {status}")
|
||||||
return None
|
return None
|
||||||
|
except TransientLLMError:
|
||||||
|
raise # L4a:瞬时故障上抛供 FailoverClient 转移,不吞
|
||||||
|
except _TRANSIENT_NET_ERRORS:
|
||||||
|
raise # T19:连接类瞬时故障同样上抛(原被吞成 None,无法转移;
|
||||||
|
# ChunkedEncodingError 一并对齐——响应体中途损坏属瞬时网络故障)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self._log("error", f"LLM 压缩请求异常: {e}")
|
self._log("error", f"LLM 压缩请求异常: {e}")
|
||||||
return None
|
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,
|
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||||
tool_choice: Optional[str] = None,
|
tool_choice: Optional[str] = None,
|
||||||
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
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 = {
|
payload = {
|
||||||
"model": self.model,
|
"model": self.model,
|
||||||
"messages": messages,
|
"messages": messages,
|
||||||
@@ -170,18 +243,32 @@ class DeepSeekClient(LLMClient):
|
|||||||
self._current_response = response
|
self._current_response = response
|
||||||
|
|
||||||
if response.status_code != 200:
|
if response.status_code != 200:
|
||||||
|
# T9:非 200 立即释放连接(stream=True 下不关闭会阻塞同域名
|
||||||
|
# in-flight 请求直到 GC);T13:错误体脱敏后再入异常
|
||||||
|
status = response.status_code
|
||||||
|
retry_after = _retry_after_of(response)
|
||||||
|
body = _sanitize_log_msg(response.text[:200])
|
||||||
|
response.close()
|
||||||
|
self._current_response = None
|
||||||
|
if status == 429 or status >= 500:
|
||||||
|
# L4b:瞬时故障显式分类(此时尚未产出任何事件,可安全转移重放)
|
||||||
|
raise TransientLLMError(
|
||||||
|
f"LLM 流式请求瞬时失败: HTTP {status}: {body}",
|
||||||
|
retry_after=retry_after)
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"HTTP {response.status_code}: {response.text}")
|
f"HTTP {status}: {body}")
|
||||||
|
|
||||||
|
saw_done = False
|
||||||
try:
|
try:
|
||||||
for line in response.iter_lines():
|
for line in response.iter_lines():
|
||||||
if not line:
|
if not line:
|
||||||
continue
|
continue
|
||||||
line_text = line.decode("utf-8")
|
line_text = line.decode("utf-8", errors="replace") # L3:非 UTF-8 行不中断流
|
||||||
if not line_text.startswith("data: "):
|
if not line_text.startswith("data: "):
|
||||||
continue
|
continue
|
||||||
data = line_text[6:]
|
data = line_text[6:]
|
||||||
if data == "[DONE]":
|
if data == "[DONE]":
|
||||||
|
saw_done = True
|
||||||
yield {"type": "done"}
|
yield {"type": "done"}
|
||||||
break
|
break
|
||||||
try:
|
try:
|
||||||
@@ -197,12 +284,28 @@ class DeepSeekClient(LLMClient):
|
|||||||
yield event
|
yield event
|
||||||
finally:
|
finally:
|
||||||
self._current_response = None
|
self._current_response = None
|
||||||
|
# T20(审计 2026-08-21):iter_lines 耗尽但从未收到 [DONE](代理缓冲
|
||||||
|
# 后提前 FIN / 服务端中途断开,末行半截 JSON 解析失败被 continue)
|
||||||
|
# 此前被静默视为成功——回答无声截断无任何信号。抛瞬时异常:首事件
|
||||||
|
# 前 FailoverClient 可转移重放;已产出事件则上抛(既有语义:已输出
|
||||||
|
# 内容绝不重放),由上层错误处理回滚/报错。用户中断(cancel 关连接
|
||||||
|
# 抛异常)不会走到这里。
|
||||||
|
if not saw_done:
|
||||||
|
raise TransientLLMError("LLM 流式响应不完整:未收到 [DONE] 结束标记(连接中途断开)")
|
||||||
|
|
||||||
|
# L5:思考内容的常见厂商字段别名(首个命中的非空字段生效)
|
||||||
|
REASONING_FIELDS = ("reasoning_content", "reasoning", "thinking", "thinking_content")
|
||||||
|
|
||||||
def _parse_delta(self, delta: Dict, reasoning_enabled: bool) -> Iterator[Dict]:
|
def _parse_delta(self, delta: Dict, reasoning_enabled: bool) -> Iterator[Dict]:
|
||||||
"""将单条 delta 归一化为统一事件。子类覆写以适配不同厂商字段。"""
|
"""将单条 delta 归一化为统一事件。子类覆写以适配不同厂商字段。"""
|
||||||
# 思考内容
|
# 思考内容:DeepSeek reasoning_content / OpenRouter reasoning /
|
||||||
if reasoning_enabled and delta.get("reasoning_content"):
|
# 部分 GLM thinking(_content) —— 此前仅识别 DeepSeek 字段,其他厂商思考流静默丢失
|
||||||
yield {"type": "reasoning", "text": delta["reasoning_content"]}
|
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"):
|
if delta.get("content"):
|
||||||
yield {"type": "content", "text": delta["content"]}
|
yield {"type": "content", "text": delta["content"]}
|
||||||
@@ -219,19 +322,279 @@ class DeepSeekClient(LLMClient):
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
def resolve_api_key(api_key: Optional[str], config=None) -> str:
|
class TransientLLMError(RuntimeError):
|
||||||
"""解析 API Key:优先显式传入,其次环境变量,最后 ~/.cert/deepseek.key。"""
|
"""瞬时 LLM 故障(HTTP 429/5xx)——FailoverClient 据此转移下一档案(L4a)。
|
||||||
|
|
||||||
|
L4c:retry_after 携带服务端 Retry-After 响应头(数字秒)建议冷却时长,
|
||||||
|
FailoverClient 以其替代默认冷却期(429 限流时尊重服务端节奏)。
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, message: str, retry_after: Optional[float] = None):
|
||||||
|
super().__init__(message)
|
||||||
|
self.retry_after = retry_after
|
||||||
|
|
||||||
|
|
||||||
|
def _retry_after_of(response) -> Optional[float]:
|
||||||
|
"""从 Retry-After 响应头解析冷却秒数(仅支持数字秒格式,其余忽略返回 None)。
|
||||||
|
|
||||||
|
P2:上限 cap 3600s——异常/恶意服务端返回超大值(如 999999)时,
|
||||||
|
避免档案被单条响应头冷却数天。
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
v = response.headers.get("Retry-After")
|
||||||
|
if v is None:
|
||||||
|
return None
|
||||||
|
return min(max(1.0, float(v)), 3600.0)
|
||||||
|
except (TypeError, ValueError, AttributeError):
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class FailoverClient:
|
||||||
|
"""跨档案失败转移包装器(L4a)。
|
||||||
|
|
||||||
|
持有序 client 链(主档案 + fallback 档案),对外实现与 LLMClient 相同的
|
||||||
|
消费面(未显式定义的属性经 __getattr__ 委托主 client)。
|
||||||
|
|
||||||
|
L4a 范围——chat_completion(非流式幂等)按链转移:
|
||||||
|
- 转移:HTTP 429/5xx(TransientLLMError)、连接失败/超时(requests 连接类异常)
|
||||||
|
- 不转移:认证/参数类 4xx(底层返回 None,直接透传——配置问题应响亮暴露)
|
||||||
|
失败档案进入冷却期(cooldown_seconds,默认 60s),期间直接跳过,
|
||||||
|
避免每次请求都撞刚失败的端点。链耗尽返回 None(与 chat_completion 失败契约一致)。
|
||||||
|
|
||||||
|
L4b 范围——stream_chat 仅**首事件前**故障转移(连接失败/HTTP 429/5xx,
|
||||||
|
尚未产出任何 reasoning/content/tool_call 事件,重放对用户不可见);
|
||||||
|
已产出事件后失败直接上抛(重放会导致前端重复输出/工具增量错乱);
|
||||||
|
用户中断(cancel() 置位)引发的连接异常绝不转移。
|
||||||
|
|
||||||
|
L4c 范围——可观测性与预算控制:
|
||||||
|
- get_stats() 暴露 requests/failovers/cooldown_skips/exhausted 聚合计数
|
||||||
|
- on_failover(from_name, to_name, reason) 回调(to_name=None 表示链耗尽),
|
||||||
|
供管理面/前端通知接入,回调异常不影响主流程
|
||||||
|
- timeout_budget 限制整条链尝试的总秒数(缺省不限),耗尽即停止尝试剩余档案,
|
||||||
|
防止长链故障时超时叠加(3 档案最坏 3×timeout)
|
||||||
|
- 429 响应带 Retry-After(数字秒)时按其延长该档案冷却期
|
||||||
|
|
||||||
|
并发契约(P4):实例非线程安全(_cancelled/_call_started/_stats 无锁)。
|
||||||
|
当前调用面满足"单线程串行调用 + cancel() 可从其他线程"的模式(compressor
|
||||||
|
在 _msg_lock 内同步调用、子代理各持独立实例);如需多线程共享调用,
|
||||||
|
需由调用方加锁或为实例补锁。
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, clients: List, logger=None, cooldown_seconds: float = 60.0,
|
||||||
|
timeout_budget: Optional[float] = None, on_failover=None):
|
||||||
|
self._clients = list(clients)
|
||||||
|
self.logger = logger
|
||||||
|
self.cooldown_seconds = cooldown_seconds
|
||||||
|
self._fail_at = {} # client.name → (失败时刻, 冷却秒数)
|
||||||
|
self._cancelled = False
|
||||||
|
self._stats = {"requests": 0, "failovers": 0, "cooldown_skips": 0, "exhausted": 0}
|
||||||
|
if timeout_budget is not None and timeout_budget <= 0:
|
||||||
|
self._log("warning", f"timeout_budget 非法({timeout_budget} <= 0),按未配置处理")
|
||||||
|
timeout_budget = None
|
||||||
|
self.timeout_budget = timeout_budget
|
||||||
|
self.on_failover = on_failover
|
||||||
|
self._call_started = None
|
||||||
|
|
||||||
|
def __getattr__(self, item):
|
||||||
|
# 未显式定义的属性/方法委托主 client(name/model/api_key/cancel 等);
|
||||||
|
# 下划线属性直接报错,防 __init__ 前访问 self._clients 的无限递归
|
||||||
|
if item.startswith("_"):
|
||||||
|
raise AttributeError(item)
|
||||||
|
return getattr(self._clients[0], item)
|
||||||
|
|
||||||
|
def _log(self, level: str, msg: str):
|
||||||
|
if self.logger is not None and hasattr(self.logger, level):
|
||||||
|
try:
|
||||||
|
getattr(self.logger, level)(_sanitize_log_msg(msg))
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _cool_key(client) -> int:
|
||||||
|
"""T6:冷却键按客户端实例(id)——client.name 取 display_name 优先且
|
||||||
|
无唯一性校验,两个 profile_name 不同但 display_name 相同的档案
|
||||||
|
(如主备都叫 "Kimi")冷却状态互相污染:A 失败冷却 → B 被误判
|
||||||
|
跳过(该转移时不转移)。P6 去重保证每档案恰好一个实例,实例被
|
||||||
|
_clients 持有生命周期内 id 稳定,与去重键语义一致。"""
|
||||||
|
return id(client)
|
||||||
|
|
||||||
|
def _cooling(self, client) -> bool:
|
||||||
|
entry = self._fail_at.get(self._cool_key(client))
|
||||||
|
if entry is None:
|
||||||
|
return False
|
||||||
|
t, cd = entry
|
||||||
|
return (time.monotonic() - t) < cd
|
||||||
|
|
||||||
|
def _mark_fail(self, client, cooldown: Optional[float] = None):
|
||||||
|
# L4c:429 Retry-After(经 TransientLLMError.retry_after)覆盖默认冷却时长
|
||||||
|
self._fail_at[self._cool_key(client)] = (time.monotonic(),
|
||||||
|
cooldown if cooldown and cooldown > 0
|
||||||
|
else self.cooldown_seconds)
|
||||||
|
|
||||||
|
def _budget_exhausted(self) -> bool:
|
||||||
|
if self.timeout_budget is None or self._call_started is None:
|
||||||
|
return False
|
||||||
|
return (time.monotonic() - self._call_started) >= self.timeout_budget
|
||||||
|
|
||||||
|
def get_stats(self) -> Dict:
|
||||||
|
"""L4c:转移统计快照(requests/failovers/cooldown_skips/exhausted)。"""
|
||||||
|
return dict(self._stats)
|
||||||
|
|
||||||
|
def _notify_failover(self, from_name: str, to_name, reason: str):
|
||||||
|
self._stats["failovers"] += 1
|
||||||
|
try:
|
||||||
|
if self.on_failover is not None:
|
||||||
|
self.on_failover(from_name, to_name, _sanitize_log_msg(reason))
|
||||||
|
except Exception as e:
|
||||||
|
self._log("warning", f"on_failover 回调异常(忽略): {e}")
|
||||||
|
|
||||||
|
def cancel(self):
|
||||||
|
"""L4b:中断链上全部 client(任一可能正在流式中)并置自身中断标志。
|
||||||
|
|
||||||
|
显式定义以覆盖 __getattr__ 委托——包装器必须同时通知整条链。
|
||||||
|
"""
|
||||||
|
self._cancelled = True
|
||||||
|
for client in self._clients:
|
||||||
|
try:
|
||||||
|
client.cancel()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def chat_completion(self, messages: List[Dict], temperature: float = None,
|
||||||
|
max_tokens: int = 1024, **kwargs):
|
||||||
|
self._cancelled = False # 新一轮请求,清除上一轮中断标志
|
||||||
|
self._call_started = time.monotonic()
|
||||||
|
self._stats["requests"] += 1
|
||||||
|
last_exc = None
|
||||||
|
tried = 0
|
||||||
|
for idx, client in enumerate(self._clients):
|
||||||
|
if self._cooling(client):
|
||||||
|
self._stats["cooldown_skips"] += 1
|
||||||
|
self._log("info", f"模型档案 {client.name} 处于失败冷却期,跳过")
|
||||||
|
continue
|
||||||
|
if self._budget_exhausted():
|
||||||
|
self._log("warning",
|
||||||
|
f"失败转移时间预算({self.timeout_budget}s)已耗尽,停止尝试剩余档案")
|
||||||
|
break
|
||||||
|
tried += 1
|
||||||
|
try:
|
||||||
|
result = client.chat_completion(messages, temperature=temperature,
|
||||||
|
max_tokens=max_tokens, **kwargs)
|
||||||
|
if result is not None:
|
||||||
|
return result
|
||||||
|
# None:认证/参数类 4xx 或响应异常——不可转移,直接透传
|
||||||
|
return None
|
||||||
|
except (TransientLLMError,) + _TRANSIENT_NET_ERRORS as e:
|
||||||
|
if self._cancelled:
|
||||||
|
return None # 用户中断引发的异常:不转移
|
||||||
|
last_exc = e
|
||||||
|
self._mark_fail(client, cooldown=getattr(e, "retry_after", None))
|
||||||
|
to_name = (self._clients[idx + 1].name
|
||||||
|
if idx + 1 < len(self._clients) else None)
|
||||||
|
self._notify_failover(client.name, to_name, str(e))
|
||||||
|
self._log("warning", f"模型档案 {client.name} 瞬时故障({e}),转移下一档案")
|
||||||
|
self._stats["exhausted"] += 1 # 末尾统一计数(链耗尽/预算耗尽/全冷却)
|
||||||
|
self._log("error", f"全部 {len(self._clients)} 个模型档案均不可用(尝试 {tried} 个):{last_exc}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
def stream_chat(self, messages: List[Dict], temperature: float = None,
|
||||||
|
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||||
|
tool_choice: Optional[str] = None,
|
||||||
|
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||||
|
"""L4b:流式转移——仅首事件前故障可安全重放(对用户不可见)。
|
||||||
|
|
||||||
|
T7:本方法是普通函数(非生成器)——生成器函数体延迟到首次 next()
|
||||||
|
才执行,若 cancel() 在创建生成器与首迭代之间被其他线程调用,
|
||||||
|
延迟执行的 `_cancelled = False` 会把用户的取消清掉(请求照常发出)。
|
||||||
|
标志重置与计数在返回生成器前同步完成,取消契约(类文档)恢复成立。
|
||||||
|
"""
|
||||||
|
self._cancelled = False # 新一轮请求,清除上一轮中断标志
|
||||||
|
self._call_started = time.monotonic()
|
||||||
|
self._stats["requests"] += 1
|
||||||
|
return self._stream_chat_impl(messages, temperature=temperature,
|
||||||
|
max_tokens=max_tokens, tools=tools,
|
||||||
|
tool_choice=tool_choice,
|
||||||
|
reasoning_enabled=reasoning_enabled, **extra)
|
||||||
|
|
||||||
|
def _stream_chat_impl(self, messages: List[Dict], temperature: float = None,
|
||||||
|
max_tokens: int = 8192, tools: Optional[List[Dict]] = None,
|
||||||
|
tool_choice: Optional[str] = None,
|
||||||
|
reasoning_enabled: bool = False, **extra) -> Iterator[Dict]:
|
||||||
|
last_exc = None
|
||||||
|
for idx, client in enumerate(self._clients):
|
||||||
|
if self._cooling(client):
|
||||||
|
self._stats["cooldown_skips"] += 1
|
||||||
|
self._log("info", f"模型档案 {client.name} 处于失败冷却期,跳过")
|
||||||
|
continue
|
||||||
|
if self._budget_exhausted():
|
||||||
|
self._log("warning",
|
||||||
|
f"失败转移时间预算({self.timeout_budget}s)已耗尽,停止尝试剩余档案")
|
||||||
|
break
|
||||||
|
produced_any = False
|
||||||
|
try:
|
||||||
|
for event in client.stream_chat(
|
||||||
|
messages, temperature=temperature, max_tokens=max_tokens,
|
||||||
|
tools=tools, tool_choice=tool_choice,
|
||||||
|
reasoning_enabled=reasoning_enabled, **extra):
|
||||||
|
produced_any = True
|
||||||
|
yield event
|
||||||
|
return # 自然结束(含 done 事件)视为成功
|
||||||
|
except (TransientLLMError,) + _TRANSIENT_NET_ERRORS as e:
|
||||||
|
if self._cancelled:
|
||||||
|
raise # 用户中断(cancel 强关连接):绝不转移
|
||||||
|
if produced_any:
|
||||||
|
raise # 已向消费方产出事件,重放会重复输出/工具增量错乱
|
||||||
|
last_exc = e
|
||||||
|
self._mark_fail(client, cooldown=getattr(e, "retry_after", None))
|
||||||
|
to_name = (self._clients[idx + 1].name
|
||||||
|
if idx + 1 < len(self._clients) else None)
|
||||||
|
self._notify_failover(client.name, to_name, str(e))
|
||||||
|
self._log("warning",
|
||||||
|
f"模型档案 {client.name} 流式瞬时故障(首事件前,{e}),转移下一档案")
|
||||||
|
self._stats["exhausted"] += 1
|
||||||
|
if last_exc is not None:
|
||||||
|
raise last_exc # 链耗尽:抛最后异常,由既有错误处理(回滚/报错)接住
|
||||||
|
raise RuntimeError(
|
||||||
|
f"全部 {len(self._clients)} 个模型档案不可用(冷却期或转移时间预算耗尽)")
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_api_key(api_key: Optional[str], config=None,
|
||||||
|
env_name: Optional[str] = None,
|
||||||
|
allow_missing: bool = False,
|
||||||
|
allow_key_file: bool = True) -> str:
|
||||||
|
"""解析 API Key:优先显式传入,其次环境变量,最后 ~/.cert/deepseek.key。
|
||||||
|
|
||||||
|
env_name: 档案级环境变量名(来自 models.<name>.api_key_env);
|
||||||
|
None 时回退 config 的 api.api_key_env。
|
||||||
|
allow_missing: 本地免密端点(models.<name>.api_key_env 为空串)时,
|
||||||
|
允许无 Key,返回空串(Bearer 空头,Ollama/vLLM 忽略)。
|
||||||
|
allow_key_file: 是否允许回退读取 ~/.cert/deepseek.key。该文件是 DeepSeek
|
||||||
|
专属凭据,模型档案端点(其他厂商)必须传 False——环境变量取不到
|
||||||
|
时直接报错,绝不把 A 厂商 Key 发给 B 厂商(跨厂商凭据泄漏)。
|
||||||
|
"""
|
||||||
if api_key:
|
if api_key:
|
||||||
return api_key
|
return api_key
|
||||||
env_name = None
|
if env_name is None:
|
||||||
if config is not None and hasattr(config, "get"):
|
if config is not None and hasattr(config, "get"):
|
||||||
env_name = config.get("api.api_key_env")
|
env_name = config.get("api.api_key_env")
|
||||||
env_name = env_name or "DEEPSEEK_API_KEY"
|
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)
|
env_key = os.environ.get(env_name)
|
||||||
if env_key:
|
if env_key:
|
||||||
return 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:
|
try:
|
||||||
home = os.environ.get("HOME", os.environ["USERPROFILE"])
|
# dict.get 的 default 参数会无条件先求值,直接写 os.environ["USERPROFILE"]
|
||||||
|
# 在无该变量的 Linux 上必抛 KeyError,导致 key 文件回退路径永远不可达
|
||||||
|
home = os.environ.get("HOME") or os.environ.get("USERPROFILE")
|
||||||
key_path = os.path.join(home, ".cert", "deepseek.key")
|
key_path = os.path.join(home, ".cert", "deepseek.key")
|
||||||
with open(key_path, "r") as f:
|
with open(key_path, "r") as f:
|
||||||
return f.read().strip()
|
return f.read().strip()
|
||||||
@@ -240,26 +603,184 @@ def resolve_api_key(api_key: Optional[str], config=None) -> str:
|
|||||||
f"请设置环境变量 {env_name} 或在配置文件中设置 API Key")
|
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,
|
def create_llm_client(config, api_key: Optional[str] = None,
|
||||||
logger=None, model: Optional[str] = None) -> LLMClient:
|
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:
|
Args:
|
||||||
config: 配置对象(需提供 get 方法)。
|
config: 配置对象(需提供 get 方法)。
|
||||||
api_key: 显式传入的 API Key;为 None 时按 resolve_api_key 解析。
|
api_key: 显式传入的 API Key;为 None 时按 resolve_api_key 解析。
|
||||||
logger: 可选日志对象。
|
logger: 可选日志对象。
|
||||||
model: 可选模型名覆盖;为 None 时使用配置中的 `api.deepseek_model`。
|
model: 档案名或裸模型名(支持 model:/profile: 前缀逃逸)。
|
||||||
子代理等场景可通过此参数显式指定与主代理不同的模型。
|
|
||||||
"""
|
"""
|
||||||
provider = (config.get("api.provider", "deepseek") or "deepseek").lower()
|
profile, model, profile_name = _resolve_model_profile(config, model)
|
||||||
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)
|
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"):
|
if provider in ("deepseek", "deepseek-chat", "openai_compatible", "openai"):
|
||||||
return DeepSeekClient(
|
client = DeepSeekClient(
|
||||||
api_key=resolved_key,
|
api_key=resolved_key,
|
||||||
model=model,
|
model=model,
|
||||||
base_url=base_url,
|
base_url=base_url,
|
||||||
@@ -267,6 +788,7 @@ def create_llm_client(config, api_key: Optional[str] = None,
|
|||||||
config=config,
|
config=config,
|
||||||
logger=logger,
|
logger=logger,
|
||||||
)
|
)
|
||||||
|
return client
|
||||||
|
|
||||||
# 其他 OpenAI 兼容厂商:可直接复用 DeepSeekClient(继承其协议解析),
|
# 其他 OpenAI 兼容厂商:可直接复用 DeepSeekClient(继承其协议解析),
|
||||||
# 或在此接入专有实现。新增厂商时在此扩展即可。
|
# 或在此接入专有实现。新增厂商时在此扩展即可。
|
||||||
|
|||||||
@@ -0,0 +1,492 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
大模型接入点(模型档案)交互式管理工具
|
||||||
|
|
||||||
|
不用手写 config.json:菜单式列出 / 添加 / 删除 / 连通性验证
|
||||||
|
`models` 节的模型接入点(OpenAI Chat Completions 兼容端点:
|
||||||
|
DeepSeek / Kimi / Qwen / GLM / 本地 Ollama 等)。
|
||||||
|
|
||||||
|
用法:
|
||||||
|
python model_config.py # 交互式菜单
|
||||||
|
python model_config.py --list # 非交互:列出当前接入点
|
||||||
|
|
||||||
|
安全约定(与 aibox 审计一致):API Key 一律不写入配置文件——
|
||||||
|
只保存环境变量名(api_key_env),或留空表示免密本地端点。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import sys
|
||||||
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
|
# 允许从仓库根目录或 aibox/ 目录启动
|
||||||
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||||
|
|
||||||
|
from common import ConfigManager # noqa: E402
|
||||||
|
from llm_client import _sanitize_log_msg # noqa: E402
|
||||||
|
|
||||||
|
PROFILE_NAME_RE = re.compile(r"^[A-Za-z0-9_-]+$")
|
||||||
|
|
||||||
|
# 常见厂商端点提示(首行是默认建议值)
|
||||||
|
PROVIDER_HINTS = {
|
||||||
|
"DeepSeek": "https://api.deepseek.com/v1/chat/completions",
|
||||||
|
"Kimi (Moonshot)": "https://api.moonshot.cn/v1/chat/completions",
|
||||||
|
"Qwen (阿里云百炼)": "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
|
||||||
|
"GLM (智谱)": "https://open.bigmodel.cn/api/paas/v4/chat/completions",
|
||||||
|
"Ollama 本地": "http://localhost:11434/v1/chat/completions",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 核心逻辑(可单测,不依赖 input) ====================
|
||||||
|
|
||||||
|
def get_models(config: ConfigManager) -> Dict:
|
||||||
|
"""取 models 节(不存在则初始化为空 dict)。"""
|
||||||
|
models = config.get("models") or {}
|
||||||
|
if not isinstance(models, dict):
|
||||||
|
models = {}
|
||||||
|
return models
|
||||||
|
|
||||||
|
|
||||||
|
def _api_section(config: ConfigManager) -> Dict:
|
||||||
|
"""M3(审计 2026-08-20):api 节被改坏为非 dict 时重置为空节——
|
||||||
|
否则 setdefault 后索引赋值抛 TypeError/AttributeError"""
|
||||||
|
api = config._config.get("api")
|
||||||
|
if not isinstance(api, dict):
|
||||||
|
api = config._config["api"] = {}
|
||||||
|
return api
|
||||||
|
|
||||||
|
|
||||||
|
def list_profiles(config: ConfigManager) -> List[Dict]:
|
||||||
|
"""列出所有档案的归一化展示数据。"""
|
||||||
|
models = get_models(config)
|
||||||
|
default = config.get("api.default_model") or ""
|
||||||
|
profiles = []
|
||||||
|
for name, profile in models.items():
|
||||||
|
if not isinstance(profile, dict):
|
||||||
|
continue
|
||||||
|
# M4(审计 2026-08-20):fallback 被改坏为字符串时原样下传,
|
||||||
|
# 展示层 join 成单字符误导——统一规整为列表
|
||||||
|
fallback = profile.get("fallback")
|
||||||
|
profiles.append({
|
||||||
|
"name": name,
|
||||||
|
"is_default": name == default,
|
||||||
|
"display_name": profile.get("display_name") or name,
|
||||||
|
"api_url": profile.get("api_url", ""),
|
||||||
|
"model": profile.get("model", ""),
|
||||||
|
"api_key_env": profile.get("api_key_env", ""),
|
||||||
|
"reasoning_enabled": bool(profile.get("reasoning_enabled", False)),
|
||||||
|
"temperature": profile.get("temperature"),
|
||||||
|
"timeout": profile.get("timeout"),
|
||||||
|
"fallback": fallback if isinstance(fallback, list) else [],
|
||||||
|
})
|
||||||
|
return profiles
|
||||||
|
|
||||||
|
|
||||||
|
def validate_new_profile(config: ConfigManager, name: str, api_url: str,
|
||||||
|
model: str) -> Optional[str]:
|
||||||
|
"""校验新档案字段,返回错误信息(合法返回 None)。"""
|
||||||
|
if not name:
|
||||||
|
return "档案名不能为空"
|
||||||
|
if not PROFILE_NAME_RE.match(name):
|
||||||
|
return f"档案名只能包含字母/数字/下划线/短横线(当前: {name})"
|
||||||
|
models = get_models(config)
|
||||||
|
if name in models:
|
||||||
|
return f"档案名 '{name}' 已存在"
|
||||||
|
if not api_url or not str(api_url).strip().lower().startswith(("http://", "https://")):
|
||||||
|
return "API 地址必须是 http:// 或 https:// 开头的完整 URL"
|
||||||
|
if not model or not str(model).strip():
|
||||||
|
return "模型名不能为空(如 deepseek-chat / gpt-4o / qwen-plus)"
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def add_profile(config: ConfigManager, name: str, api_url: str, model: str,
|
||||||
|
api_key_env: str = "", display_name: str = None,
|
||||||
|
temperature: float = None, timeout: int = None,
|
||||||
|
reasoning_enabled: bool = False,
|
||||||
|
fallback: Optional[List[str]] = None,
|
||||||
|
set_default: bool = False) -> List[str]:
|
||||||
|
"""添加档案。返回提示信息列表;字段非法时返回带 ❌ 的错误信息。"""
|
||||||
|
err = validate_new_profile(config, name, api_url, model)
|
||||||
|
if err:
|
||||||
|
return [f"❌ {err}"]
|
||||||
|
|
||||||
|
# M5(审计 2026-08-20):api_key_env 需为合法环境变量名;展示/地址字段
|
||||||
|
# 拒控制字符并限长——防终端注入回显(\x1b] 等序列写入配置后被 print)
|
||||||
|
api_key_env = (api_key_env or "").strip()
|
||||||
|
if api_key_env and not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", api_key_env):
|
||||||
|
return [f"❌ 环境变量名不合法: {api_key_env!r}(须为字母/数字/下划线且不以数字开头)"]
|
||||||
|
for label, val, limit in (("显示名", display_name, 64),
|
||||||
|
("模型名", model, 128),
|
||||||
|
("API 地址", api_url, 256)):
|
||||||
|
if val and (len(str(val)) > limit or not all(c.isprintable() for c in str(val))):
|
||||||
|
return [f"❌ {label}含控制字符或超过 {limit} 字符"]
|
||||||
|
|
||||||
|
models = get_models(config)
|
||||||
|
profile = {
|
||||||
|
"api_url": str(api_url).strip(),
|
||||||
|
"model": str(model).strip(),
|
||||||
|
}
|
||||||
|
if api_key_env:
|
||||||
|
profile["api_key_env"] = api_key_env.strip()
|
||||||
|
else:
|
||||||
|
# 空串 = 免密本地端点(Ollama/vLLM),显式声明
|
||||||
|
profile["api_key_env"] = ""
|
||||||
|
if display_name:
|
||||||
|
profile["display_name"] = display_name.strip()
|
||||||
|
if temperature is not None:
|
||||||
|
try:
|
||||||
|
profile["temperature"] = max(0.0, min(2.0, float(temperature)))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return ["❌ 温度必须是 0.0~2.0 的数字"]
|
||||||
|
if timeout is not None:
|
||||||
|
try:
|
||||||
|
profile["timeout"] = max(5, int(timeout))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return ["❌ 超时必须是整数秒"]
|
||||||
|
if reasoning_enabled:
|
||||||
|
profile["reasoning_enabled"] = True
|
||||||
|
if fallback:
|
||||||
|
# 只保留已存在的档案,避免悬空引用(L4a 启动期硬性报错)
|
||||||
|
known = set(models.keys())
|
||||||
|
invalid = [f for f in fallback if f not in known or f == name]
|
||||||
|
if invalid:
|
||||||
|
return [f"❌ 失败转移包含不存在的档案或自身: {invalid},请重新选择"]
|
||||||
|
profile["fallback"] = list(fallback)
|
||||||
|
|
||||||
|
models[name] = profile
|
||||||
|
config._config["models"] = models
|
||||||
|
if set_default:
|
||||||
|
_api_section(config)["default_model"] = name
|
||||||
|
|
||||||
|
saved = config._save_config()
|
||||||
|
msgs = [f"✅ 接入点 '{name}' 已添加({profile['api_url']},模型 {profile['model']})"]
|
||||||
|
if not saved:
|
||||||
|
msgs.append("⚠️ 保存失败(仅本次运行生效,重启后丢失)——请检查配置文件目录权限")
|
||||||
|
if set_default:
|
||||||
|
msgs.append(f"⭐ 已设为默认接入点(--model 不指定时使用)")
|
||||||
|
if not api_key_env:
|
||||||
|
msgs.append("ℹ️ 免密端点:请求不携带 API Key(适用于 Ollama/vLLM 等本地服务)")
|
||||||
|
else:
|
||||||
|
msgs.append(f"🔑 Key 从环境变量 {api_key_env} 读取(未设置时该档案不可用)")
|
||||||
|
return msgs
|
||||||
|
|
||||||
|
|
||||||
|
def delete_profile(config: ConfigManager, name: str) -> List[str]:
|
||||||
|
"""删除档案,并自动清理其他档案 fallback 中的悬空引用。"""
|
||||||
|
models = get_models(config)
|
||||||
|
if name not in models:
|
||||||
|
return [f"❌ 档案 '{name}' 不存在"]
|
||||||
|
|
||||||
|
del models[name]
|
||||||
|
|
||||||
|
# 清理其他档案对该档案的 fallback 引用(悬空引用会导致启动期报错)
|
||||||
|
cleaned = []
|
||||||
|
for other, profile in models.items():
|
||||||
|
if not isinstance(profile, dict):
|
||||||
|
continue
|
||||||
|
fb = profile.get("fallback")
|
||||||
|
# M4(审计 2026-08-20):fallback 为字符串时 `name in fb` 退化为子串
|
||||||
|
# 误判,且下方列表推导会把字符串拆成单字符数组写回配置
|
||||||
|
if not isinstance(fb, list):
|
||||||
|
continue
|
||||||
|
if name in fb:
|
||||||
|
profile["fallback"] = [f for f in fb if f != name]
|
||||||
|
cleaned.append(other)
|
||||||
|
config._config["models"] = models
|
||||||
|
|
||||||
|
# 若删除的是默认档案,清除默认标记(回退旧 api 节语义)
|
||||||
|
if config.get("api.default_model") == name:
|
||||||
|
_api_section(config).pop("default_model", None)
|
||||||
|
|
||||||
|
saved = config._save_config()
|
||||||
|
msgs = [f"✅ 接入点 '{name}' 已删除"]
|
||||||
|
if cleaned:
|
||||||
|
msgs.append(f"ℹ️ 已同时清理 {len(cleaned)} 个档案的失败转移引用: {', '.join(cleaned)}")
|
||||||
|
if not saved:
|
||||||
|
msgs.append("⚠️ 保存失败——请检查配置文件目录权限")
|
||||||
|
return msgs
|
||||||
|
|
||||||
|
|
||||||
|
def test_profile(config: ConfigManager, name: str, logger=None) -> str:
|
||||||
|
"""连通性验证:向该档案端点发一次最小请求。"""
|
||||||
|
models = get_models(config)
|
||||||
|
if name not in models:
|
||||||
|
return f"❌ 档案 '{name}' 不存在"
|
||||||
|
try:
|
||||||
|
from llm_client import create_llm_client
|
||||||
|
client = create_llm_client(config, api_key=None, logger=logger, model=name)
|
||||||
|
result = client.chat_completion(
|
||||||
|
[{"role": "user", "content": "ping"}], max_tokens=5)
|
||||||
|
if result is None or not str(result).strip():
|
||||||
|
return f"❌ 档案 '{name}' 验证失败:端点返回空响应(多为 4xx 认证/参数错误)"
|
||||||
|
return f"✅ 档案 '{name}' 连通正常:{str(result).strip()[:60]}"
|
||||||
|
except Exception as e:
|
||||||
|
# J4(审计 2026-08-19):异常消息可能含 response.text 片段,网关/WAF
|
||||||
|
# 错误页会回显 Authorization 头——脱敏 Bearer 令牌后再回显终端
|
||||||
|
return f"❌ 档案 '{name}' 验证失败:{_sanitize_log_msg(str(e))}"
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 交互式菜单 ====================
|
||||||
|
|
||||||
|
def _ask(prompt: str, default: str = "") -> str:
|
||||||
|
if default:
|
||||||
|
val = input(f"{prompt}(回车默认: {default}): ").strip()
|
||||||
|
return val if val else default
|
||||||
|
return input(f"{prompt}: ").strip()
|
||||||
|
|
||||||
|
|
||||||
|
def _ask_yes_no(prompt: str, default: bool = False) -> bool:
|
||||||
|
hint = "y/n" if not default else "Y/n"
|
||||||
|
while True:
|
||||||
|
val = input(f"{prompt} [{hint}] ").strip().lower()
|
||||||
|
if not val:
|
||||||
|
return default
|
||||||
|
if val in ("y", "yes", "是"):
|
||||||
|
return True
|
||||||
|
if val in ("n", "no", "否"):
|
||||||
|
return False
|
||||||
|
print(" 请输入 y 或 n")
|
||||||
|
|
||||||
|
|
||||||
|
def _choose(prompt: str, options: List[str], multi: bool = False) -> List[str]:
|
||||||
|
"""编号选择(单选或多选)。多选用逗号分隔,直接回车 = 不选。"""
|
||||||
|
for i, opt in enumerate(options, 1):
|
||||||
|
print(f" {i}. {opt}")
|
||||||
|
if multi:
|
||||||
|
while True:
|
||||||
|
raw = input(f"{prompt}(多选,逗号分隔,直接回车不选): ").strip()
|
||||||
|
if not raw:
|
||||||
|
return []
|
||||||
|
try:
|
||||||
|
idxs = [int(x) for x in raw.replace(",", ",").split(",") if x.strip()]
|
||||||
|
except ValueError:
|
||||||
|
print(" 请输入数字编号,逗号分隔")
|
||||||
|
continue
|
||||||
|
if all(1 <= i <= len(options) for i in idxs):
|
||||||
|
return [options[i - 1] for i in idxs]
|
||||||
|
print(f" 编号超出范围(1~{len(options)})")
|
||||||
|
else:
|
||||||
|
while True:
|
||||||
|
raw = input(f"{prompt}(1~{len(options)},直接回车不选): ").strip()
|
||||||
|
if not raw:
|
||||||
|
return []
|
||||||
|
try:
|
||||||
|
idx = int(raw)
|
||||||
|
except ValueError:
|
||||||
|
print(" 请输入数字编号")
|
||||||
|
continue
|
||||||
|
if 1 <= idx <= len(options):
|
||||||
|
return [options[idx - 1]]
|
||||||
|
print(f" 编号超出范围(1~{len(options)})")
|
||||||
|
|
||||||
|
|
||||||
|
# T-L17(审计 2026-08-21):按 api_url 域名推荐 Key 环境变量名——此前默认
|
||||||
|
# DEEPSEEK_API_KEY,一路回车会把 DeepSeek 凭据发给其他厂商(moonshot/
|
||||||
|
# aliyuncs 等),正是 resolve_api_key 文档声称要杜绝的跨厂商凭据外发。
|
||||||
|
# 未知域名留空强制输入
|
||||||
|
_ENV_NAME_BY_DOMAIN = {
|
||||||
|
"deepseek": "DEEPSEEK_API_KEY",
|
||||||
|
"moonshot": "MOONSHOT_API_KEY",
|
||||||
|
"aliyuncs": "DASHSCOPE_API_KEY",
|
||||||
|
"bigmodel": "ZHIPU_API_KEY",
|
||||||
|
"ark.cn": "ARK_API_KEY",
|
||||||
|
"openrouter": "OPENROUTER_API_KEY",
|
||||||
|
"openai.com": "OPENAI_API_KEY",
|
||||||
|
"siliconflow": "SILICONFLOW_API_KEY",
|
||||||
|
"volces": "ARK_API_KEY",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _suggest_env_name(api_url: str) -> str:
|
||||||
|
url = (api_url or "").lower()
|
||||||
|
for kw, env in _ENV_NAME_BY_DOMAIN.items():
|
||||||
|
if kw in url:
|
||||||
|
return env
|
||||||
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
def _sanitize_display(s) -> str:
|
||||||
|
"""T-L18(审计 2026-08-21):手写 config.json 可嵌 ANSI/控制转义
|
||||||
|
(\x1b]0;pwned\x07 等),展示层原样 print 每次列表即触发;写入侧有
|
||||||
|
isprintable 校验,读出侧同样过滤控制字符。"""
|
||||||
|
return "".join(ch for ch in str(s or "") if ch.isprintable())
|
||||||
|
|
||||||
|
|
||||||
|
def _print_profiles(profiles: List[Dict]):
|
||||||
|
if not profiles:
|
||||||
|
print(" (当前没有任何接入点。旧版 api 节仍可用,或选 2 添加)")
|
||||||
|
return
|
||||||
|
for p in profiles:
|
||||||
|
mark = " ⭐默认" if p["is_default"] else ""
|
||||||
|
key_desc = f"环境变量 {_sanitize_display(p['api_key_env'])}" if p["api_key_env"] else "免密"
|
||||||
|
fb = (f",失败转移 → "
|
||||||
|
f"{', '.join(_sanitize_display(f) for f in p['fallback'])}") if p["fallback"] else ""
|
||||||
|
print(f" {_sanitize_display(p['name'])}{mark} [{_sanitize_display(p['display_name'])}]")
|
||||||
|
print(f" 地址: {_sanitize_display(p['api_url'])} 模型: {_sanitize_display(p['model'])}")
|
||||||
|
print(f" Key: {key_desc} 温度: {p['temperature'] or '默认'} 超时: {p['timeout'] or '默认'}{fb}")
|
||||||
|
|
||||||
|
|
||||||
|
def _add_flow(config: ConfigManager):
|
||||||
|
print("\n── 添加接入点 ──")
|
||||||
|
print("常见厂商端点参考:")
|
||||||
|
for k, v in PROVIDER_HINTS.items():
|
||||||
|
print(f" {k}: {v}")
|
||||||
|
name = _ask("档案名(英文,如 deepseek/kimi/local)")
|
||||||
|
if not name:
|
||||||
|
print("❌ 已取消(档案名不能为空)")
|
||||||
|
return
|
||||||
|
if not PROFILE_NAME_RE.match(name):
|
||||||
|
print(f"❌ 档案名只能包含字母/数字/下划线/短横线")
|
||||||
|
return
|
||||||
|
if name in get_models(config):
|
||||||
|
print(f"❌ 档案名 '{name}' 已存在")
|
||||||
|
return
|
||||||
|
api_url = _ask("API 地址")
|
||||||
|
model = _ask("模型名(如 deepseek-chat / qwen-plus)")
|
||||||
|
|
||||||
|
print("Key 方式:")
|
||||||
|
print(" 1. 环境变量名(推荐,如 DEEPSEEK_API_KEY——Key 不落盘,重启不丢)")
|
||||||
|
print(" 2. 免密本地端点(如 Ollama,不需要 Key)")
|
||||||
|
key_mode = _ask("选择 [1/2]", "1")
|
||||||
|
api_key_env = ""
|
||||||
|
if key_mode == "1":
|
||||||
|
suggested = _suggest_env_name(api_url)
|
||||||
|
prompt = "环境变量名" + (f"(回车用 {suggested})" if suggested else "")
|
||||||
|
api_key_env = _ask(prompt, suggested)
|
||||||
|
if not api_key_env:
|
||||||
|
print("❌ 已取消(环境变量名不能为空;免密端点请选方式 2)")
|
||||||
|
return
|
||||||
|
elif key_mode != "2":
|
||||||
|
print("❌ 已取消(Key 方式无效)")
|
||||||
|
return
|
||||||
|
|
||||||
|
display_name = _ask("显示名(可选,回车用档案名)")
|
||||||
|
temperature_raw = _ask("温度 0.0~2.0(可选,越低越确定)")
|
||||||
|
timeout_raw = _ask("超时秒数(可选)")
|
||||||
|
reasoning = _ask_yes_no("启用思考模式?", default=False)
|
||||||
|
|
||||||
|
profiles = list_profiles(config)
|
||||||
|
fallback = []
|
||||||
|
if profiles:
|
||||||
|
print("失败转移:主接入点故障时按顺序尝试的备用接入点(可选)")
|
||||||
|
fallback = _choose("选择备用接入点", [p["name"] for p in profiles], multi=True)
|
||||||
|
|
||||||
|
set_default = _ask_yes_no("设为默认接入点?", default=(len(profiles) == 0))
|
||||||
|
|
||||||
|
print("\n── 确认 ──")
|
||||||
|
print(f" 档案名: {name}")
|
||||||
|
print(f" 地址: {api_url} 模型: {model}")
|
||||||
|
print(f" Key: {'环境变量 ' + api_key_env if api_key_env else '免密(不携带)'}")
|
||||||
|
if display_name:
|
||||||
|
print(f" 显示名: {display_name}")
|
||||||
|
if temperature_raw:
|
||||||
|
print(f" 温度: {temperature_raw}")
|
||||||
|
if timeout_raw:
|
||||||
|
print(f" 超时: {timeout_raw}s")
|
||||||
|
if reasoning:
|
||||||
|
print(" 思考模式: 开启")
|
||||||
|
if fallback:
|
||||||
|
print(f" 失败转移: {', '.join(fallback)}")
|
||||||
|
if not _ask_yes_no("确认保存?", default=True):
|
||||||
|
print("已取消")
|
||||||
|
return
|
||||||
|
|
||||||
|
# J8(审计 2026-08-19):原始字符串直传——此前交互层先 float()/int(),
|
||||||
|
# 非数字输入直接 traceback 退出,add_profile 内部的容错成为死代码
|
||||||
|
msgs = add_profile(config, name, api_url, model,
|
||||||
|
api_key_env=api_key_env, display_name=display_name,
|
||||||
|
temperature=temperature_raw or None,
|
||||||
|
timeout=timeout_raw or None,
|
||||||
|
reasoning_enabled=reasoning, fallback=fallback,
|
||||||
|
set_default=set_default)
|
||||||
|
for m in msgs:
|
||||||
|
print(m)
|
||||||
|
|
||||||
|
|
||||||
|
def _delete_flow(config: ConfigManager):
|
||||||
|
profiles = list_profiles(config)
|
||||||
|
if not profiles:
|
||||||
|
print("\n(当前没有任何接入点)")
|
||||||
|
return
|
||||||
|
print("\n── 删除接入点 ──")
|
||||||
|
picked = _choose("选择要删除的接入点", [p["name"] for p in profiles])
|
||||||
|
if not picked:
|
||||||
|
print("已取消")
|
||||||
|
return
|
||||||
|
name = picked[0]
|
||||||
|
if not _ask_yes_no(f"确认删除 '{name}'?(其他档案的失败转移引用会自动清理)", default=False):
|
||||||
|
print("已取消")
|
||||||
|
return
|
||||||
|
for m in delete_profile(config, name):
|
||||||
|
print(m)
|
||||||
|
|
||||||
|
|
||||||
|
def _test_flow(config: ConfigManager):
|
||||||
|
profiles = list_profiles(config)
|
||||||
|
if not profiles:
|
||||||
|
print("\n(当前没有任何接入点)")
|
||||||
|
return
|
||||||
|
print("\n── 连通性验证 ──")
|
||||||
|
picked = _choose("选择要验证的接入点", [p["name"] for p in profiles])
|
||||||
|
if not picked:
|
||||||
|
print("已取消")
|
||||||
|
return
|
||||||
|
print(" 正在请求端点(最长 20 秒)...")
|
||||||
|
print(test_profile(config, picked[0]))
|
||||||
|
|
||||||
|
|
||||||
|
def run_interactive(config: ConfigManager = None):
|
||||||
|
"""交互式主循环。config 可注入(测试用)。"""
|
||||||
|
config = config or ConfigManager()
|
||||||
|
config_path = getattr(config, "_config_path", None) or "(默认位置)"
|
||||||
|
print("=" * 56)
|
||||||
|
print(" aibox 大模型接入点管理工具")
|
||||||
|
print(f" 配置文件: {config_path}")
|
||||||
|
print(" 说明:API Key 只存环境变量名,不落盘;免密端点适合 Ollama 等本地服务")
|
||||||
|
print("=" * 56)
|
||||||
|
while True:
|
||||||
|
# M2(审计 2026-08-20):input 未捕获 EOFError/KeyboardInterrupt——
|
||||||
|
# 管道耗尽/Ctrl+C 此前直接 traceback;整循环包 try 同时覆盖子流程
|
||||||
|
# (_add_flow/_delete_flow/_test_flow)内的 input 调用
|
||||||
|
try:
|
||||||
|
print("\n请选择操作:")
|
||||||
|
print(" 1. 📋 列出接入点")
|
||||||
|
print(" 2. ➕ 添加接入点")
|
||||||
|
print(" 3. 🗑️ 删除接入点")
|
||||||
|
print(" 4. 🔍 验证接入点连通性")
|
||||||
|
print(" 0. 退出")
|
||||||
|
choice = input("> ").strip()
|
||||||
|
if choice == "1":
|
||||||
|
_print_profiles(list_profiles(config))
|
||||||
|
elif choice == "2":
|
||||||
|
_add_flow(config)
|
||||||
|
elif choice == "3":
|
||||||
|
_delete_flow(config)
|
||||||
|
elif choice == "4":
|
||||||
|
_test_flow(config)
|
||||||
|
elif choice in ("0", "q", "exit", "退出"):
|
||||||
|
print("再见")
|
||||||
|
return
|
||||||
|
else:
|
||||||
|
print("无效选择,请输入 0~4")
|
||||||
|
except (EOFError, KeyboardInterrupt):
|
||||||
|
print("\n再见")
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="aibox 大模型接入点管理工具")
|
||||||
|
parser.add_argument("--list", action="store_true", help="非交互:列出当前接入点")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
config = ConfigManager()
|
||||||
|
if args.list:
|
||||||
|
_print_profiles(list_profiles(config))
|
||||||
|
return
|
||||||
|
run_interactive(config)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
+25
-26
@@ -59,14 +59,14 @@ class AESCli:
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
def generate_random_key(self, key_size=16):
|
def generate_random_key(self, key_size=16):
|
||||||
"""生成随机密钥"""
|
"""生成随机密钥(S8:密钥内容不打印,避免进入会话日志;经 --save-key 持久化)"""
|
||||||
if key_size not in [16, 24, 32]:
|
if key_size not in [16, 24, 32]:
|
||||||
print(f"错误: 密钥大小必须为16、24或32,当前为{key_size}")
|
print(f"错误: 密钥大小必须为16、24或32,当前为{key_size}")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
self.key = get_random_bytes(key_size)
|
self.key = get_random_bytes(key_size)
|
||||||
key_hex = self.key.hex()
|
key_hex = self.key.hex()
|
||||||
print(f"已生成 {key_size} 字节随机密钥: {key_hex}")
|
print(f"已生成 {key_size} 字节随机密钥(内容不显示,请用 --save-key 保存到文件)")
|
||||||
return key_hex
|
return key_hex
|
||||||
|
|
||||||
def generate_random_iv(self):
|
def generate_random_iv(self):
|
||||||
@@ -330,6 +330,10 @@ class AESCli:
|
|||||||
|
|
||||||
with open(filepath, 'w', encoding='utf-8') as f:
|
with open(filepath, 'w', encoding='utf-8') as f:
|
||||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||||
|
try:
|
||||||
|
os.chmod(filepath, 0o600) # S8:密钥文件仅属主可读写
|
||||||
|
except OSError:
|
||||||
|
pass # 不支持 POSIX 权限的平台(如 Windows)跳过
|
||||||
|
|
||||||
print(f"密钥和IV已保存到: {filepath}")
|
print(f"密钥和IV已保存到: {filepath}")
|
||||||
return True
|
return True
|
||||||
@@ -395,38 +399,33 @@ def main():
|
|||||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||||
epilog="""
|
epilog="""
|
||||||
使用示例:
|
使用示例:
|
||||||
# 生成随机密钥和IV
|
# 生成随机密钥和IV(密钥不打印,经 --save-key 保存到文件)
|
||||||
aes-crypto --gen-key
|
aes-crypto --gen-key --gen-iv --save-key keys.json
|
||||||
aes-crypto --gen-iv
|
aes-crypto --gen-iv
|
||||||
|
|
||||||
# 文本加解密
|
# 文本加解密(密钥经文件提供,不走命令行)
|
||||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --encrypt-text "Hello World"
|
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --encrypt-text "Hello World"
|
||||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --decrypt-text "<密文base64>"
|
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --decrypt-text "<密文base64>"
|
||||||
|
|
||||||
# 文件加解密
|
# 文件加解密
|
||||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --encrypt-file input.txt
|
aes-crypto --load-key keys.json --encrypt-file input.txt
|
||||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --decrypt-file input.txt.encrypted
|
aes-crypto --load-key keys.json --decrypt-file input.txt.encrypted
|
||||||
|
|
||||||
# 批量加解密
|
# 批量加解密
|
||||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --batch-encrypt file1.txt file2.txt
|
aes-crypto --load-key keys.json --batch-encrypt file1.txt file2.txt
|
||||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --batch-decrypt *.encrypted
|
aes-crypto --load-key keys.json --batch-decrypt *.encrypted
|
||||||
|
|
||||||
# 保存和加载密钥
|
|
||||||
aes-crypto --key <密钥hex> --iv <IV/Nonce hex> --save-key keys.json
|
|
||||||
aes-crypto --load-key keys.json --encrypt-file input.txt
|
|
||||||
|
|
||||||
# 使用GCM模式
|
# 使用GCM模式
|
||||||
aes-crypto --mode GCM --gen-key --key-size 32
|
aes-crypto --mode GCM --gen-key --key-size 32 --save-key keys.json
|
||||||
aes-crypto --mode GCM --key <密钥hex> --iv <Nonce hex> --encrypt-file input.txt
|
aes-crypto --mode GCM --load-key keys.json --encrypt-file input.txt
|
||||||
|
|
||||||
# 从文件读取密钥
|
# 从文件读取密钥
|
||||||
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --encrypt-file input.txt
|
aes-crypto --key-file key.txt --iv <IV/Nonce hex> --encrypt-file input.txt
|
||||||
|
|
||||||
参数说明:
|
参数说明:
|
||||||
--mode CBC|GCM 加密模式 (默认: CBC)
|
--mode CBC|GCM 加密模式 (默认: CBC)
|
||||||
--key KEY 十六进制密钥 (16/24/32字节)
|
|
||||||
--iv IV 十六进制IV/Nonce (CBC:16字节, GCM:12字节)
|
--iv IV 十六进制IV/Nonce (CBC:16字节, GCM:12字节)
|
||||||
--key-file FILE 从文件读取密钥
|
--key-file FILE 从文件读取密钥(密钥仅经文件提供,避免进入进程列表与会话日志)
|
||||||
--load-key FILE 从JSON文件加载密钥和IV
|
--load-key FILE 从JSON文件加载密钥和IV
|
||||||
--save-key FILE 保存密钥和IV到JSON文件
|
--save-key FILE 保存密钥和IV到JSON文件
|
||||||
--gen-key 生成随机密钥
|
--gen-key 生成随机密钥
|
||||||
@@ -448,8 +447,7 @@ def main():
|
|||||||
parser.add_argument('--mode', choices=['CBC', 'GCM'], default='CBC',
|
parser.add_argument('--mode', choices=['CBC', 'GCM'], default='CBC',
|
||||||
help='加密模式 (默认: CBC)')
|
help='加密模式 (默认: CBC)')
|
||||||
|
|
||||||
# 密钥和IV设置
|
# 密钥和IV设置(S7:密钥仅经 --key-file/--load-key 文件途径提供,不进命令行)
|
||||||
parser.add_argument('--key', help='十六进制密钥 (16/24/32字节)')
|
|
||||||
parser.add_argument('--iv', help='十六进制IV/Nonce (CBC:16字节, GCM:12字节)')
|
parser.add_argument('--iv', help='十六进制IV/Nonce (CBC:16字节, GCM:12字节)')
|
||||||
parser.add_argument('--key-file', help='从文件读取密钥')
|
parser.add_argument('--key-file', help='从文件读取密钥')
|
||||||
parser.add_argument('--load-key', help='从JSON文件加载密钥和IV')
|
parser.add_argument('--load-key', help='从JSON文件加载密钥和IV')
|
||||||
@@ -491,27 +489,28 @@ def main():
|
|||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
# 处理密钥文件
|
# 处理密钥文件
|
||||||
|
key_input = None
|
||||||
if args.key_file:
|
if args.key_file:
|
||||||
try:
|
try:
|
||||||
with open(args.key_file, 'r') as f:
|
with open(args.key_file, 'r') as f:
|
||||||
args.key = f.read().strip()
|
key_input = f.read().strip()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"错误: 读取密钥文件失败: {e}")
|
print(f"错误: 读取密钥文件失败: {e}")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
# 设置密钥和IV
|
# 设置密钥和IV
|
||||||
if args.key:
|
if key_input:
|
||||||
if args.iv is None and not aes.iv:
|
if args.iv is None and not aes.iv:
|
||||||
print("错误: 需要同时提供密钥和IV/Nonce")
|
print("错误: 需要同时提供密钥和IV/Nonce")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
if args.iv:
|
if args.iv:
|
||||||
if not aes.set_key_iv(args.key, args.iv):
|
if not aes.set_key_iv(key_input, args.iv):
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
elif aes.iv:
|
elif aes.iv:
|
||||||
# 从加载的密钥文件中获取IV
|
# 从加载的密钥文件中获取IV
|
||||||
try:
|
try:
|
||||||
aes.key = bytes.fromhex(args.key)
|
aes.key = bytes.fromhex(key_input)
|
||||||
if len(aes.key) not in [16, 24, 32]:
|
if len(aes.key) not in [16, 24, 32]:
|
||||||
print(f"错误: 密钥长度必须为16、24或32字节,当前为{len(aes.key)}字节")
|
print(f"错误: 密钥长度必须为16、24或32字节,当前为{len(aes.key)}字节")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
@@ -548,7 +547,7 @@ def main():
|
|||||||
if args.encrypt_text or args.decrypt_text or \
|
if args.encrypt_text or args.decrypt_text or \
|
||||||
args.encrypt_file or args.decrypt_file or \
|
args.encrypt_file or args.decrypt_file or \
|
||||||
args.batch_encrypt or args.batch_decrypt:
|
args.batch_encrypt or args.batch_decrypt:
|
||||||
print("错误: 请先设置密钥和IV/Nonce (使用 --key 和 --iv 或 --load-key)")
|
print("错误: 请先设置密钥和IV/Nonce (使用 --key-file 和 --iv 或 --load-key)")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -18,6 +18,16 @@ def ensure_camera_dir():
|
|||||||
return CAMERA_DIR
|
return CAMERA_DIR
|
||||||
|
|
||||||
|
|
||||||
|
def validate_filename(filename):
|
||||||
|
"""S9(审计 2026-08-20):--filename 仅允许纯文件名——
|
||||||
|
拒绝路径分隔符/父目录引用/绝对路径,防止经 os.path.join 逃逸 CAMERA_DIR"""
|
||||||
|
if not filename:
|
||||||
|
return "文件名不能为空"
|
||||||
|
if filename in (".", "..") or "/" in filename or "\\" in filename or os.path.isabs(filename):
|
||||||
|
return f"非法文件名(不允许路径分隔符或父目录引用): {filename}"
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def list_cameras(max_test=10):
|
def list_cameras(max_test=10):
|
||||||
"""列出可用的摄像头设备"""
|
"""列出可用的摄像头设备"""
|
||||||
available_cameras = []
|
available_cameras = []
|
||||||
@@ -341,6 +351,10 @@ def main():
|
|||||||
ensure_camera_dir()
|
ensure_camera_dir()
|
||||||
|
|
||||||
if args.filename:
|
if args.filename:
|
||||||
|
filename_error = validate_filename(args.filename)
|
||||||
|
if filename_error:
|
||||||
|
print(f"❌ {filename_error}")
|
||||||
|
sys.exit(1)
|
||||||
save_path = os.path.join(CAMERA_DIR, args.filename)
|
save_path = os.path.join(CAMERA_DIR, args.filename)
|
||||||
else:
|
else:
|
||||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||||
|
|||||||
@@ -2,6 +2,7 @@
|
|||||||
"""人脸检测和识别工具 - 检测图像中的人脸并保存裁剪的人脸图像。"""
|
"""人脸检测和识别工具 - 检测图像中的人脸并保存裁剪的人脸图像。"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
import re
|
||||||
import sys
|
import sys
|
||||||
import argparse
|
import argparse
|
||||||
import cv2
|
import cv2
|
||||||
@@ -15,6 +16,16 @@ except ImportError:
|
|||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
# S11(审计 2026-08-20):person_name 白名单——仅允许文字/数字/._-,
|
||||||
|
# 拒绝路径分隔符与纯点名("."/".."),防止经 f"{name}.jpg" 拼接逃逸 known_faces 目录。
|
||||||
|
# \w 在 py3 默认 UNICODE,兼容中文人名(如 "张三")。
|
||||||
|
_PERSON_NAME_RE = re.compile(r"^(?!\.+$)[\w.-]{1,64}$")
|
||||||
|
|
||||||
|
|
||||||
|
def _valid_person_name(name):
|
||||||
|
return bool(name) and bool(_PERSON_NAME_RE.match(name))
|
||||||
|
|
||||||
|
|
||||||
class FaceDetector:
|
class FaceDetector:
|
||||||
"""人脸检测和识别工具"""
|
"""人脸检测和识别工具"""
|
||||||
|
|
||||||
@@ -198,6 +209,8 @@ class FaceDetector:
|
|||||||
def add_known_face(self, image_path: str, person_name: str) -> str:
|
def add_known_face(self, image_path: str, person_name: str) -> str:
|
||||||
"""添加已知人脸到数据库"""
|
"""添加已知人脸到数据库"""
|
||||||
try:
|
try:
|
||||||
|
if not _valid_person_name(person_name):
|
||||||
|
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||||
if not os.path.exists(image_path):
|
if not os.path.exists(image_path):
|
||||||
return f"❌ 文件不存在: {image_path}"
|
return f"❌ 文件不存在: {image_path}"
|
||||||
|
|
||||||
@@ -237,6 +250,8 @@ class FaceDetector:
|
|||||||
def remove_known_face(self, person_name: str) -> str:
|
def remove_known_face(self, person_name: str) -> str:
|
||||||
"""从已知人脸库中删除人脸"""
|
"""从已知人脸库中删除人脸"""
|
||||||
try:
|
try:
|
||||||
|
if not _valid_person_name(person_name):
|
||||||
|
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||||
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
||||||
|
|
||||||
if not os.path.exists(filepath):
|
if not os.path.exists(filepath):
|
||||||
@@ -260,6 +275,8 @@ class FaceDetector:
|
|||||||
def get_face_info(self, person_name: str) -> str:
|
def get_face_info(self, person_name: str) -> str:
|
||||||
"""获取已知人脸信息"""
|
"""获取已知人脸信息"""
|
||||||
try:
|
try:
|
||||||
|
if not _valid_person_name(person_name):
|
||||||
|
return f"❌ 非法人名标识: {person_name}(仅允许文字/数字/._-,≤64字符)"
|
||||||
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
filepath = os.path.join(self.known_faces_dir, f"{person_name}.jpg")
|
||||||
|
|
||||||
if not os.path.exists(filepath):
|
if not os.path.exists(filepath):
|
||||||
|
|||||||
+41
-10
@@ -3,6 +3,7 @@
|
|||||||
|
|
||||||
import sys
|
import sys
|
||||||
import os
|
import os
|
||||||
|
import re
|
||||||
import argparse
|
import argparse
|
||||||
import requests
|
import requests
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
@@ -35,6 +36,13 @@ MAX_CONTENT_PREVIEW = 1000
|
|||||||
MAX_LIST_PREVIEW = 20
|
MAX_LIST_PREVIEW = 20
|
||||||
MAX_STRING_LENGTH = 200
|
MAX_STRING_LENGTH = 200
|
||||||
|
|
||||||
|
# S3:分支/ref 白名单——首字符必须为字母或数字,防止以 '-' 开头的值被 git 解析为选项(选项注入)
|
||||||
|
_REF_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/@+-]{0,127}$")
|
||||||
|
|
||||||
|
|
||||||
|
def _valid_ref(name: str) -> bool:
|
||||||
|
return bool(name) and bool(_REF_RE.match(name))
|
||||||
|
|
||||||
|
|
||||||
# ==================== GitHub工具类 ====================
|
# ==================== GitHub工具类 ====================
|
||||||
|
|
||||||
@@ -79,19 +87,30 @@ class GitHubTool:
|
|||||||
return date_str
|
return date_str
|
||||||
|
|
||||||
def _run_git(self, repo_path: str, cmd: List[str]) -> Dict:
|
def _run_git(self, repo_path: str, cmd: List[str]) -> Dict:
|
||||||
|
env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
|
||||||
|
if GITHUB_TOKEN:
|
||||||
|
# 凭据仅经环境变量传递(git>=2.31 的 GIT_CONFIG_* 机制),
|
||||||
|
# 不进命令行 argv / 远端 URL / .git/config(S1,对齐 J12 条例)
|
||||||
|
env["GIT_CONFIG_COUNT"] = "1"
|
||||||
|
env["GIT_CONFIG_KEY_0"] = "http.extraheader"
|
||||||
|
env["GIT_CONFIG_VALUE_0"] = f"Authorization: token {GITHUB_TOKEN}"
|
||||||
try:
|
try:
|
||||||
result = subprocess.run(
|
result = subprocess.run(
|
||||||
["git"] + cmd, cwd=repo_path, capture_output=True,
|
["git"] + cmd, cwd=repo_path, capture_output=True,
|
||||||
text=True, timeout=60, env={**os.environ, "GIT_TERMINAL_PROMPT": "0"}
|
text=True, timeout=60, env=env
|
||||||
)
|
)
|
||||||
return {"success": result.returncode == 0, "stdout": result.stdout.strip(), "stderr": result.stderr.strip()}
|
return {"success": result.returncode == 0,
|
||||||
|
"stdout": self._scrub_token(result.stdout.strip()),
|
||||||
|
"stderr": self._scrub_token(result.stderr.strip())}
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return {"success": False, "error": str(e)}
|
return {"success": False, "error": self._scrub_token(str(e))}
|
||||||
|
|
||||||
def _auth_url(self, repo_url: str) -> str:
|
@staticmethod
|
||||||
if GITHUB_TOKEN and repo_url.startswith("https://"):
|
def _scrub_token(text: str) -> str:
|
||||||
return repo_url.replace("https://", f"https://{GITHUB_TOKEN}@")
|
"""输出脱敏:任何路径回显的 token 一律替换(S1)"""
|
||||||
return repo_url
|
if GITHUB_TOKEN and text:
|
||||||
|
return text.replace(GITHUB_TOKEN, "***")
|
||||||
|
return text
|
||||||
|
|
||||||
# ==================== 搜索功能 ====================
|
# ==================== 搜索功能 ====================
|
||||||
|
|
||||||
@@ -154,7 +173,9 @@ class GitHubTool:
|
|||||||
if os.path.exists(target_dir):
|
if os.path.exists(target_dir):
|
||||||
return f"⚠️ 仓库已存在: {target_dir}\n💡 使用 pull 命令拉取更新"
|
return f"⚠️ 仓库已存在: {target_dir}\n💡 使用 pull 命令拉取更新"
|
||||||
|
|
||||||
clone_url = self._auth_url(f"https://github.com/{repo}.git")
|
if branch and not _valid_ref(branch):
|
||||||
|
return f"❌ 非法分支名: {branch}"
|
||||||
|
clone_url = f"https://github.com/{repo}.git"
|
||||||
cmd = ["clone"]
|
cmd = ["clone"]
|
||||||
if depth > 0:
|
if depth > 0:
|
||||||
cmd.extend(["--depth", str(depth)])
|
cmd.extend(["--depth", str(depth)])
|
||||||
@@ -175,6 +196,8 @@ class GitHubTool:
|
|||||||
br = self._run_git(repo_path, ["branch", "--show-current"])
|
br = self._run_git(repo_path, ["branch", "--show-current"])
|
||||||
branch = br["stdout"] if br["success"] else None
|
branch = br["stdout"] if br["success"] else None
|
||||||
|
|
||||||
|
if branch and not _valid_ref(branch):
|
||||||
|
return f"❌ 非法分支名: {branch}"
|
||||||
result = self._run_git(repo_path, ["pull", "origin", branch] if branch else ["pull"])
|
result = self._run_git(repo_path, ["pull", "origin", branch] if branch else ["pull"])
|
||||||
if result["success"]:
|
if result["success"]:
|
||||||
return f"✅ 拉取成功\n {result['stdout'] or '已是最新'}"
|
return f"✅ 拉取成功\n {result['stdout'] or '已是最新'}"
|
||||||
@@ -188,6 +211,8 @@ class GitHubTool:
|
|||||||
br = self._run_git(repo_path, ["branch", "--show-current"])
|
br = self._run_git(repo_path, ["branch", "--show-current"])
|
||||||
branch = br["stdout"] if br["success"] else None
|
branch = br["stdout"] if br["success"] else None
|
||||||
|
|
||||||
|
if branch and not _valid_ref(branch):
|
||||||
|
return f"❌ 非法分支名: {branch}"
|
||||||
cmd = ["push"]
|
cmd = ["push"]
|
||||||
if force:
|
if force:
|
||||||
cmd.append("--force")
|
cmd.append("--force")
|
||||||
@@ -273,9 +298,15 @@ class GitHubTool:
|
|||||||
return "⚠️ 需要确认删除,请添加 --confirm 参数"
|
return "⚠️ 需要确认删除,请添加 --confirm 参数"
|
||||||
if not os.path.exists(repo_path):
|
if not os.path.exists(repo_path):
|
||||||
return f"❌ 路径不存在: {repo_path}"
|
return f"❌ 路径不存在: {repo_path}"
|
||||||
|
# S2:仅允许删除 WORKSPACE 内的 git 仓库,防止任意目录递归删除
|
||||||
|
workspace = os.path.realpath(WORKSPACE_DIR)
|
||||||
|
real_path = os.path.realpath(repo_path)
|
||||||
|
if not real_path.startswith(workspace + os.sep) \
|
||||||
|
or not os.path.isdir(os.path.join(real_path, ".git")):
|
||||||
|
return f"❌ 仅允许删除工作区({workspace})内的 git 仓库目录"
|
||||||
try:
|
try:
|
||||||
shutil.rmtree(repo_path)
|
shutil.rmtree(real_path)
|
||||||
return f"✅ 已删除: {repo_path}"
|
return f"✅ 已删除: {real_path}"
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return f"❌ 删除失败: {e}"
|
return f"❌ 删除失败: {e}"
|
||||||
|
|
||||||
|
|||||||
@@ -3,6 +3,7 @@
|
|||||||
一个基于python-nmap的局域网Ping扫描工具
|
一个基于python-nmap的局域网Ping扫描工具
|
||||||
"""
|
"""
|
||||||
import argparse
|
import argparse
|
||||||
|
import shlex
|
||||||
import sys
|
import sys
|
||||||
import socket
|
import socket
|
||||||
|
|
||||||
@@ -25,7 +26,9 @@ def scan_network(network="192.168.1.0/24", timeout=3, verbose=False):
|
|||||||
print("-" * 50)
|
print("-" * 50)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
result = nm.scan(hosts=network, arguments=arguments)
|
# J12(审计 2026-08-19):python-nmap 内部把 hosts 拼进 shell 命令
|
||||||
|
# 执行——LLM/AI 可控的 network 参数需 shlex.quote,防参数/命令注入
|
||||||
|
result = nm.scan(hosts=shlex.quote(network), arguments=arguments)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"❌ 扫描失败: {e}")
|
print(f"❌ 扫描失败: {e}")
|
||||||
print("💡 提示:请确保已安装 nmap (sudo apt install nmap) 并有足够权限")
|
print("💡 提示:请确保已安装 nmap (sudo apt install nmap) 并有足够权限")
|
||||||
|
|||||||
@@ -29,6 +29,14 @@ def get_timestamp_filename(prefix="screenshot", extension="png"):
|
|||||||
return f"{prefix}_{timestamp}.{extension}"
|
return f"{prefix}_{timestamp}.{extension}"
|
||||||
|
|
||||||
|
|
||||||
|
def validate_filename(filename):
|
||||||
|
"""S10(审计 2026-08-20):--filename 仅允许纯文件名——
|
||||||
|
拒绝路径分隔符/父目录引用/绝对路径,防止经 os.path.join 逃逸保存目录"""
|
||||||
|
if filename in (".", "..") or "/" in filename or "\\" in filename or os.path.isabs(filename):
|
||||||
|
return f"非法文件名(不允许路径分隔符或父目录引用): {filename}"
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def get_system():
|
def get_system():
|
||||||
"""获取当前操作系统"""
|
"""获取当前操作系统"""
|
||||||
return platform.system()
|
return platform.system()
|
||||||
@@ -41,6 +49,10 @@ def take_screenshot_windows(output_path, delay=0):
|
|||||||
if delay > 0:
|
if delay > 0:
|
||||||
time.sleep(delay)
|
time.sleep(delay)
|
||||||
|
|
||||||
|
# J12(审计 2026-08-19):output_path 经 PowerShell 单引号转义——
|
||||||
|
# 路径中的单引号翻倍(PS 字符串转义),AI/用户可控文件名无法闭合
|
||||||
|
# 字符串注入脚本(此前 `${output_path}` 直拼,`");calc;#` 即注入)
|
||||||
|
ps_path = str(output_path).replace("'", "''")
|
||||||
try:
|
try:
|
||||||
ps_script = f'''
|
ps_script = f'''
|
||||||
Add-Type -AssemblyName System.Windows.Forms
|
Add-Type -AssemblyName System.Windows.Forms
|
||||||
@@ -53,11 +65,11 @@ $bitmap = New-Object System.Drawing.Bitmap $bounds.Width, $bounds.Height
|
|||||||
$graphics = [System.Drawing.Graphics]::FromImage($bitmap)
|
$graphics = [System.Drawing.Graphics]::FromImage($bitmap)
|
||||||
$graphics.CopyFromScreen($bounds.X, $bounds.Y, 0, 0, $bounds.Size)
|
$graphics.CopyFromScreen($bounds.X, $bounds.Y, 0, 0, $bounds.Size)
|
||||||
|
|
||||||
$bitmap.Save("{output_path}", [System.Drawing.Imaging.ImageFormat]::Png)
|
$bitmap.Save('{ps_path}', [System.Drawing.Imaging.ImageFormat]::Png)
|
||||||
$graphics.Dispose()
|
$graphics.Dispose()
|
||||||
$bitmap.Dispose()
|
$bitmap.Dispose()
|
||||||
|
|
||||||
Write-Output "截图已保存: {output_path}"
|
Write-Output "截图已保存: {ps_path}"
|
||||||
'''
|
'''
|
||||||
|
|
||||||
result = subprocess.run(
|
result = subprocess.run(
|
||||||
@@ -150,6 +162,9 @@ def take_screenshot(output_dir=None, filename=None, delay=0):
|
|||||||
os.makedirs(save_dir, exist_ok=True)
|
os.makedirs(save_dir, exist_ok=True)
|
||||||
|
|
||||||
if filename:
|
if filename:
|
||||||
|
filename_error = validate_filename(filename)
|
||||||
|
if filename_error:
|
||||||
|
return False, filename_error, None
|
||||||
filepath = os.path.join(save_dir, filename)
|
filepath = os.path.join(save_dir, filename)
|
||||||
else:
|
else:
|
||||||
filepath = os.path.join(save_dir, get_timestamp_filename())
|
filepath = os.path.join(save_dir, get_timestamp_filename())
|
||||||
|
|||||||
@@ -27,9 +27,12 @@ def build_parser():
|
|||||||
--host SSH服务器主机名或IP地址 (必填)
|
--host SSH服务器主机名或IP地址 (必填)
|
||||||
--port SSH服务器端口号,默认 22
|
--port SSH服务器端口号,默认 22
|
||||||
-u, --username SSH登录用户名 (必填)
|
-u, --username SSH登录用户名 (必填)
|
||||||
-p, --password SSH登录密码 (如不提供则交互式输入)
|
|
||||||
--key SSH私钥文件路径 (用于密钥认证)
|
--key SSH私钥文件路径 (用于密钥认证)
|
||||||
|
|
||||||
|
密码认证:
|
||||||
|
密码经环境变量 AIBOX_SFTP_PASS 提供;未设置且未用 --key 时
|
||||||
|
交互式输入。不走命令行参数,避免凭据进入进程列表与会话日志(J12 条例)。
|
||||||
|
|
||||||
操作参数:
|
操作参数:
|
||||||
--upload LOCAL REMOTE 上传:本地文件/目录 -> 远程服务器路径
|
--upload LOCAL REMOTE 上传:本地文件/目录 -> 远程服务器路径
|
||||||
--download REMOTE LOCAL 下载:远程服务器文件/目录 -> 本地路径
|
--download REMOTE LOCAL 下载:远程服务器文件/目录 -> 本地路径
|
||||||
@@ -37,19 +40,19 @@ def build_parser():
|
|||||||
-v, --verbose 显示详细信息
|
-v, --verbose 显示详细信息
|
||||||
|
|
||||||
使用示例:
|
使用示例:
|
||||||
1. 上传文件(密码认证):
|
1. 上传文件(密码认证,环境变量提供密码):
|
||||||
sftp-transfer --host 192.168.1.100 -u root -p mypass --upload /tmp/test.txt /root/
|
AIBOX_SFTP_PASS=mypass sftp-transfer --host 192.168.1.100 -u root --upload /tmp/test.txt /root/
|
||||||
|
|
||||||
2. 下载文件(密钥认证):
|
2. 下载文件(密钥认证):
|
||||||
sftp-transfer --host 192.168.1.100 -u root --key ~/.ssh/id_rsa \\
|
sftp-transfer --host 192.168.1.100 -u root --key ~/.ssh/id_rsa \\
|
||||||
--download /root/test.txt /tmp/
|
--download /root/test.txt /tmp/
|
||||||
|
|
||||||
3. 递归上传目录:
|
3. 递归上传目录:
|
||||||
sftp-transfer --host 192.168.1.100 -u user -p pass -r \\
|
AIBOX_SFTP_PASS=pass sftp-transfer --host 192.168.1.100 -u user -r \\
|
||||||
--upload ./mydir /home/user/mydir
|
--upload ./mydir /home/user/mydir
|
||||||
|
|
||||||
4. 下载目录:
|
4. 下载目录:
|
||||||
sftp-transfer --host 192.168.1.100 -u user -p pass -r \\
|
AIBOX_SFTP_PASS=pass sftp-transfer --host 192.168.1.100 -u user -r \\
|
||||||
--download /home/user/backup ./backup
|
--download /home/user/backup ./backup
|
||||||
|
|
||||||
5. 交互式输入密码:
|
5. 交互式输入密码:
|
||||||
@@ -61,7 +64,6 @@ def build_parser():
|
|||||||
parser.add_argument('--host', required=True, help='SSH服务器主机名或IP地址')
|
parser.add_argument('--host', required=True, help='SSH服务器主机名或IP地址')
|
||||||
parser.add_argument('--port', type=int, default=22, help='SSH服务器端口号,默认22')
|
parser.add_argument('--port', type=int, default=22, help='SSH服务器端口号,默认22')
|
||||||
parser.add_argument('-u', '--username', required=True, help='SSH登录用户名')
|
parser.add_argument('-u', '--username', required=True, help='SSH登录用户名')
|
||||||
parser.add_argument('-p', '--password', help='SSH登录密码(如不提供则交互式输入)')
|
|
||||||
parser.add_argument('--key', help='SSH私钥文件路径(用于密钥认证)')
|
parser.add_argument('--key', help='SSH私钥文件路径(用于密钥认证)')
|
||||||
|
|
||||||
# 操作参数
|
# 操作参数
|
||||||
@@ -81,10 +83,20 @@ def log(verbose, msg):
|
|||||||
print(f"[SFTP] {msg}")
|
print(f"[SFTP] {msg}")
|
||||||
|
|
||||||
|
|
||||||
|
def safe_member_name(name):
|
||||||
|
"""S6(审计 2026-08-20):目录成员名校验——拒绝路径分隔符与 . / ..,
|
||||||
|
防止恶意 SFTP 服务端构造 listdir 条目(如 ../../etc/cron.d/x)在递归下载时逃逸本地目录"""
|
||||||
|
return name not in ('', '.', '..') and '/' not in name and '\\' not in name
|
||||||
|
|
||||||
|
|
||||||
def get_password(args):
|
def get_password(args):
|
||||||
"""获取密码:优先从命令行参数获取,否则交互式输入"""
|
"""获取密码:环境变量 AIBOX_SFTP_PASS 优先,否则交互式输入(S5:凭据不走命令行)"""
|
||||||
if args.password:
|
env_pass = os.environ.get("AIBOX_SFTP_PASS")
|
||||||
return args.password
|
if env_pass:
|
||||||
|
return env_pass
|
||||||
|
if args.key:
|
||||||
|
# 密钥认证时不强制索要密码(仅在私钥有口令时经交互输入)
|
||||||
|
return None
|
||||||
return getpass.getpass(f"请输入 {args.username}@{args.host} 的密码: ")
|
return getpass.getpass(f"请输入 {args.username}@{args.host} 的密码: ")
|
||||||
|
|
||||||
|
|
||||||
@@ -143,6 +155,9 @@ def upload_dir(sftp, local_dir, remote_dir, verbose):
|
|||||||
"""递归上传目录"""
|
"""递归上传目录"""
|
||||||
ensure_remote_dir(sftp, remote_dir, verbose)
|
ensure_remote_dir(sftp, remote_dir, verbose)
|
||||||
for item in os.listdir(local_dir):
|
for item in os.listdir(local_dir):
|
||||||
|
if not safe_member_name(item):
|
||||||
|
print(f" 警告: 跳过非法本地条目名: {item!r}", file=sys.stderr)
|
||||||
|
continue
|
||||||
local_item = os.path.join(local_dir, item)
|
local_item = os.path.join(local_dir, item)
|
||||||
remote_item = os.path.join(remote_dir, item)
|
remote_item = os.path.join(remote_dir, item)
|
||||||
if os.path.isdir(local_item):
|
if os.path.isdir(local_item):
|
||||||
@@ -168,6 +183,9 @@ def download_dir(sftp, remote_dir, local_dir, verbose):
|
|||||||
log(verbose, f" 创建本地目录: {local_dir}")
|
log(verbose, f" 创建本地目录: {local_dir}")
|
||||||
|
|
||||||
for item in sftp.listdir(remote_dir):
|
for item in sftp.listdir(remote_dir):
|
||||||
|
if not safe_member_name(item):
|
||||||
|
print(f" 警告: 跳过非法远端条目名: {item!r}", file=sys.stderr)
|
||||||
|
continue
|
||||||
remote_item = os.path.join(remote_dir, item).replace('\\', '/')
|
remote_item = os.path.join(remote_dir, item).replace('\\', '/')
|
||||||
local_item = os.path.join(local_dir, item)
|
local_item = os.path.join(local_dir, item)
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -2,6 +2,7 @@
|
|||||||
"""ssh-exec - 通过SSH远程执行单条命令的工具(支持密码认证)"""
|
"""ssh-exec - 通过SSH远程执行单条命令的工具(支持密码认证)"""
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
|
import os
|
||||||
import sys
|
import sys
|
||||||
import paramiko
|
import paramiko
|
||||||
|
|
||||||
@@ -14,24 +15,32 @@ def main():
|
|||||||
-H, --host SSH主机地址(默认: 127.0.0.1)
|
-H, --host SSH主机地址(默认: 127.0.0.1)
|
||||||
-P, --port SSH端口号(默认: 22)
|
-P, --port SSH端口号(默认: 22)
|
||||||
-u, --user SSH用户名(必填)
|
-u, --user SSH用户名(必填)
|
||||||
-p, --password SSH密码(必填)
|
|
||||||
-c, --command 要执行的远程命令(必填)
|
-c, --command 要执行的远程命令(必填)
|
||||||
|
|
||||||
|
密码认证:
|
||||||
|
密码经环境变量 AIBOX_SSH_PASS 提供(必填),不走命令行参数,
|
||||||
|
避免凭据进入进程列表与会话日志(J12 条例)。
|
||||||
|
|
||||||
使用示例:
|
使用示例:
|
||||||
%(prog)s -u qing -p 12345678 -c "ls -la"
|
AIBOX_SSH_PASS=12345678 %(prog)s -u qing -c "ls -la"
|
||||||
%(prog)s -u root -H 192.168.1.100 -P 22 -p mypass -c "df -h"
|
AIBOX_SSH_PASS=mypass %(prog)s -u root -H 192.168.1.100 -P 22 -c "df -h"
|
||||||
%(prog)s -u qing -H 127.0.0.1 -p 12345678 -c "uname -a"
|
AIBOX_SSH_PASS=12345678 %(prog)s -u qing -H 127.0.0.1 -c "uname -a"
|
||||||
""",
|
""",
|
||||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||||
)
|
)
|
||||||
parser.add_argument("-H", "--host", default="127.0.0.1", help="SSH主机地址(默认: 127.0.0.1)")
|
parser.add_argument("-H", "--host", default="127.0.0.1", help="SSH主机地址(默认: 127.0.0.1)")
|
||||||
parser.add_argument("-P", "--port", type=int, default=22, help="SSH端口号(默认: 22)")
|
parser.add_argument("-P", "--port", type=int, default=22, help="SSH端口号(默认: 22)")
|
||||||
parser.add_argument("-u", "--user", required=True, help="SSH用户名(必填)")
|
parser.add_argument("-u", "--user", required=True, help="SSH用户名(必填)")
|
||||||
parser.add_argument("-p", "--password", required=True, help="SSH密码(必填)")
|
|
||||||
parser.add_argument("-c", "--command", required=True, help="要执行的远程命令(必填)")
|
parser.add_argument("-c", "--command", required=True, help="要执行的远程命令(必填)")
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# 密码仅从环境变量读取(S4:凭据不进命令行)
|
||||||
|
password = os.environ.get("AIBOX_SSH_PASS")
|
||||||
|
if not password:
|
||||||
|
print("错误: 请通过环境变量 AIBOX_SSH_PASS 提供SSH密码", file=sys.stderr)
|
||||||
|
sys.exit(2)
|
||||||
|
|
||||||
client = paramiko.SSHClient()
|
client = paramiko.SSHClient()
|
||||||
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
||||||
|
|
||||||
@@ -40,7 +49,7 @@ def main():
|
|||||||
hostname=args.host,
|
hostname=args.host,
|
||||||
port=args.port,
|
port=args.port,
|
||||||
username=args.user,
|
username=args.user,
|
||||||
password=args.password,
|
password=password,
|
||||||
timeout=10,
|
timeout=10,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -51,7 +51,14 @@ create_apprun() {
|
|||||||
local app_dir="$1"
|
local app_dir="$1"
|
||||||
local exec_name="$2"
|
local exec_name="$2"
|
||||||
local terminal="$3"
|
local terminal="$3"
|
||||||
|
|
||||||
|
# S13(审计 2026-08-20):exec_name 进入 sed 替换串,含 / & 等字符会注入或破坏
|
||||||
|
# 替换表达式(如 "a/x" 使 s///g 错位、"&" 展开为匹配串)。白名单拦截,拒绝即退出。
|
||||||
|
if [[ ! "$exec_name" =~ ^[A-Za-z0-9._+-]+$ ]]; then
|
||||||
|
echo -e "${RED}错误: 可执行文件名含非法字符(仅允许 A-Za-z0-9._+-): $exec_name${NC}"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
|
||||||
cat > "$app_dir/AppRun" << 'EOF'
|
cat > "$app_dir/AppRun" << 'EOF'
|
||||||
#!/bin/bash
|
#!/bin/bash
|
||||||
HERE="$(dirname "$(readlink -f "$0")")"
|
HERE="$(dirname "$(readlink -f "$0")")"
|
||||||
|
|||||||
+27
-2
@@ -23,6 +23,27 @@ def get_api_key():
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def detect_image_mime(image_path):
|
||||||
|
"""S12(审计 2026-08-20):按魔数识别真实图片类型——
|
||||||
|
拒绝非图片文件(防任意文件读取外泄),并返回正确 MIME(此前硬编码 image/jpeg)"""
|
||||||
|
try:
|
||||||
|
with open(image_path, 'rb') as f:
|
||||||
|
head = f.read(16)
|
||||||
|
except Exception:
|
||||||
|
return None
|
||||||
|
if head.startswith(b'\xff\xd8\xff'):
|
||||||
|
return 'image/jpeg'
|
||||||
|
if head.startswith(b'\x89PNG\r\n\x1a\n'):
|
||||||
|
return 'image/png'
|
||||||
|
if head[:6] in (b'GIF87a', b'GIF89a'):
|
||||||
|
return 'image/gif'
|
||||||
|
if head.startswith(b'RIFF') and head[8:12] == b'WEBP':
|
||||||
|
return 'image/webp'
|
||||||
|
if head.startswith(b'BM'):
|
||||||
|
return 'image/bmp'
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def encode_image_to_base64(image_path):
|
def encode_image_to_base64(image_path):
|
||||||
"""将图片编码为base64"""
|
"""将图片编码为base64"""
|
||||||
with open(image_path, 'rb') as f:
|
with open(image_path, 'rb') as f:
|
||||||
@@ -37,7 +58,11 @@ def analyze_image(image_path, prompt, stream=False):
|
|||||||
|
|
||||||
if not os.path.exists(image_path):
|
if not os.path.exists(image_path):
|
||||||
return f"❌ 错误:图片文件不存在 - {image_path}"
|
return f"❌ 错误:图片文件不存在 - {image_path}"
|
||||||
|
|
||||||
|
mime = detect_image_mime(image_path)
|
||||||
|
if not mime:
|
||||||
|
return f"❌ 错误:文件不是有效图片(仅支持 JPEG/PNG/GIF/WebP/BMP) - {image_path}"
|
||||||
|
|
||||||
try:
|
try:
|
||||||
base64_image = encode_image_to_base64(image_path)
|
base64_image = encode_image_to_base64(image_path)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -55,7 +80,7 @@ def analyze_image(image_path, prompt, stream=False):
|
|||||||
{
|
{
|
||||||
"type": "image_url",
|
"type": "image_url",
|
||||||
"image_url": {
|
"image_url": {
|
||||||
"url": f"data:image/jpeg;base64,{base64_image}"
|
"url": f"data:{mime};base64,{base64_image}"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -24,6 +24,11 @@ from .db import MemoryDatabase, MemoDatabase
|
|||||||
from .skill_manager import SkillManager
|
from .skill_manager import SkillManager
|
||||||
from .conversation_history import ConversationHistory
|
from .conversation_history import ConversationHistory
|
||||||
from .conversation_compressor import ConversationCompressor
|
from .conversation_compressor import ConversationCompressor
|
||||||
|
from .process_manager import (
|
||||||
|
ProcessManager,
|
||||||
|
get_process_manager,
|
||||||
|
reset_process_manager,
|
||||||
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"CommandExecutor",
|
"CommandExecutor",
|
||||||
@@ -34,4 +39,7 @@ __all__ = [
|
|||||||
"SkillManager",
|
"SkillManager",
|
||||||
"ConversationHistory",
|
"ConversationHistory",
|
||||||
"ConversationCompressor",
|
"ConversationCompressor",
|
||||||
|
"ProcessManager",
|
||||||
|
"get_process_manager",
|
||||||
|
"reset_process_manager",
|
||||||
]
|
]
|
||||||
@@ -223,49 +223,6 @@ class ConversationCompressor:
|
|||||||
|
|
||||||
return conversation_turns
|
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:
|
def _build_compression_prompt(self, messages: List[Dict]) -> str:
|
||||||
"""构建压缩提示词 - 包含工具调用信息"""
|
"""构建压缩提示词 - 包含工具调用信息"""
|
||||||
conversation_lines = []
|
conversation_lines = []
|
||||||
@@ -368,25 +325,57 @@ class ConversationCompressor:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def _validate_message_structure(self, messages: List[Dict]) -> List[Dict]:
|
def _validate_message_structure(self, messages: List[Dict]) -> List[Dict]:
|
||||||
"""验证并修复消息结构,移除孤立的 tool 消息"""
|
"""验证并修复消息结构(tool 配对双向完整性)。
|
||||||
|
|
||||||
|
OpenAI 协议要求 assistant.tool_calls 与对应的 tool 结果消息成对出现,
|
||||||
|
缺任一侧后续 API 调用直接 400:
|
||||||
|
1. 孤立 tool 消息(tool_call_id 无对应 assistant)→ 移除(既有行为)
|
||||||
|
2. T23(审计 2026-08-21):压缩块边界按条数切,可能落在带 tool_calls
|
||||||
|
的 assistant 消息与其 tool 结果之间——保留侧 assistant 的
|
||||||
|
tool_calls 失去配对 → 剥离失配的 tool_calls 条目;全部失配时删除
|
||||||
|
字段(content 为空则置空串,防产生非法空消息)
|
||||||
|
"""
|
||||||
if not messages:
|
if not messages:
|
||||||
return messages
|
return messages
|
||||||
|
|
||||||
# 找出所有 assistant 消息的 tool_call_ids
|
# tool 侧实际存在的 id
|
||||||
assistant_tool_ids = set()
|
present_tool_ids = {
|
||||||
|
m.get("tool_call_id")
|
||||||
|
for m in messages
|
||||||
|
if m.get("role") == "tool" and m.get("tool_call_id")
|
||||||
|
}
|
||||||
|
|
||||||
|
# 第一遍:assistant 剥离失去配对的 tool_calls
|
||||||
|
first_pass = []
|
||||||
|
kept_assistant_ids = set()
|
||||||
for msg in messages:
|
for msg in messages:
|
||||||
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
||||||
for tc in msg.get("tool_calls", []):
|
calls = msg.get("tool_calls", [])
|
||||||
tc_id = tc.get("id")
|
kept_calls = [tc for tc in calls
|
||||||
if tc_id:
|
if tc.get("id") and tc.get("id") in present_tool_ids]
|
||||||
assistant_tool_ids.add(tc_id)
|
kept_assistant_ids.update(tc.get("id") for tc in kept_calls)
|
||||||
|
if len(kept_calls) == len(calls):
|
||||||
|
first_pass.append(msg) # 全部有配对,原样保留
|
||||||
|
continue
|
||||||
|
new_msg = dict(msg)
|
||||||
|
if kept_calls:
|
||||||
|
new_msg["tool_calls"] = kept_calls
|
||||||
|
print(f" ⚠️ 剥离部分失去配对 tool 结果的 tool_calls(T23)")
|
||||||
|
else:
|
||||||
|
new_msg.pop("tool_calls", None)
|
||||||
|
if not new_msg.get("content"):
|
||||||
|
new_msg["content"] = ""
|
||||||
|
print(f" ⚠️ 剥离失去配对 tool 结果的 assistant.tool_calls(T23)")
|
||||||
|
first_pass.append(new_msg)
|
||||||
|
else:
|
||||||
|
first_pass.append(msg)
|
||||||
|
|
||||||
# 过滤掉没有对应 assistant 的 tool 消息
|
# 第二遍:移除孤立 tool 消息(其 id 不在任何保留的 assistant tool_calls 中)
|
||||||
filtered = []
|
filtered = []
|
||||||
for msg in messages:
|
for msg in first_pass:
|
||||||
if msg.get("role") == "tool":
|
if msg.get("role") == "tool":
|
||||||
tool_call_id = msg.get("tool_call_id")
|
tool_call_id = msg.get("tool_call_id")
|
||||||
if tool_call_id and tool_call_id not in assistant_tool_ids:
|
if tool_call_id and tool_call_id not in kept_assistant_ids:
|
||||||
print(f" ⚠️ 移除孤立的 tool 消息: {tool_call_id}")
|
print(f" ⚠️ 移除孤立的 tool 消息: {tool_call_id}")
|
||||||
continue
|
continue
|
||||||
filtered.append(msg)
|
filtered.append(msg)
|
||||||
@@ -514,7 +503,8 @@ class ConversationCompressor:
|
|||||||
final_tokens = self.estimate_tokens([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)}")
|
print(f"\n✅ {i18n.get('conversation_compressed', original_count, final_count)}")
|
||||||
print(f" Token: {original_tokens} -> {final_tokens} (压缩率: {final_tokens/original_tokens:.1%})")
|
# L6:压缩块全空时 original_tokens 可能为 0,除零防护
|
||||||
|
print(f" Token: {original_tokens} -> {final_tokens} (压缩率: {final_tokens / max(original_tokens, 1):.1%})")
|
||||||
|
|
||||||
# 检查是否压缩到阈值以下
|
# 检查是否压缩到阈值以下
|
||||||
if final_count >= threshold:
|
if final_count >= threshold:
|
||||||
|
|||||||
@@ -15,12 +15,14 @@ subsystems.conversation_history - 对话历史管理器
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
import time
|
import time
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Dict, List, Optional
|
from typing import TYPE_CHECKING, Dict, List, Optional
|
||||||
|
|
||||||
from common import ConfigManager, Logger, constants
|
from common import ConfigManager, Logger, constants
|
||||||
|
from common.atomic import atomic_write_text
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
# 仅用于类型注解,避免运行时循环导入;MemoryDatabase 来自同包 subsystems.db
|
# 仅用于类型注解,避免运行时循环导入;MemoryDatabase 来自同包 subsystems.db
|
||||||
@@ -61,14 +63,24 @@ class ConversationHistory:
|
|||||||
try:
|
try:
|
||||||
with open(self.history_file, 'r', encoding='utf-8') as f:
|
with open(self.history_file, 'r', encoding='utf-8') as f:
|
||||||
return json.load(f)
|
return json.load(f)
|
||||||
except:
|
except Exception as e:
|
||||||
|
# N16:收窄为 Exception——裸 except 连 KeyboardInterrupt/SystemExit
|
||||||
|
# 也吞掉;损坏时重置为空索引并告警
|
||||||
|
self.logger.warning(f"历史索引文件损坏,已重置: {e}")
|
||||||
return {"conversations": [], "last_id": 0}
|
return {"conversations": [], "last_id": 0}
|
||||||
return {"conversations": [], "last_id": 0}
|
return {"conversations": [], "last_id": 0}
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _atomic_write_json(path: Path, data) -> None:
|
||||||
|
"""原子写 JSON(I 批,审计 2026-08-18):先写同目录临时文件再
|
||||||
|
os.replace——此前 'w' 直写,进程中断/崩溃留半截 JSON(current 损坏
|
||||||
|
后仅静默置 None,历史内容丢失)。J5(2026-08-19):委托 common.atomic
|
||||||
|
(mkstemp 随机名),修复旧实现固定 .tmp 名在多进程并发写同路径互踩。"""
|
||||||
|
atomic_write_text(str(path), json.dumps(data, ensure_ascii=False, indent=2))
|
||||||
|
|
||||||
def _save_history_index(self):
|
def _save_history_index(self):
|
||||||
"""保存历史索引"""
|
"""保存历史索引"""
|
||||||
with open(self.history_file, 'w', encoding='utf-8') as f:
|
self._atomic_write_json(self.history_file, self.history_index)
|
||||||
json.dump(self.history_index, f, ensure_ascii=False, indent=2)
|
|
||||||
|
|
||||||
def _load_current_conversation(self):
|
def _load_current_conversation(self):
|
||||||
"""加载当前的对话(如果有)"""
|
"""加载当前的对话(如果有)"""
|
||||||
@@ -83,8 +95,7 @@ class ConversationHistory:
|
|||||||
def save_current_conversation(self):
|
def save_current_conversation(self):
|
||||||
"""保存当前对话"""
|
"""保存当前对话"""
|
||||||
if self.current_conversation:
|
if self.current_conversation:
|
||||||
with open(self.current_file, 'w', encoding='utf-8') as f:
|
self._atomic_write_json(self.current_file, self.current_conversation)
|
||||||
json.dump(self.current_conversation, f, ensure_ascii=False, indent=2)
|
|
||||||
|
|
||||||
def start_new_conversation(self, topic: str = None) -> Dict:
|
def start_new_conversation(self, topic: str = None) -> Dict:
|
||||||
"""开始新对话"""
|
"""开始新对话"""
|
||||||
@@ -132,12 +143,17 @@ class ConversationHistory:
|
|||||||
importance=2
|
importance=2
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# N16:归档文件名加序号防同秒覆盖——DATE_FORMAT 精度为秒,
|
||||||
|
# 同一秒内两次归档(如测试/快速重置)会静默覆盖前一份
|
||||||
timestamp = datetime.now().strftime(constants.DATE_FORMAT)
|
timestamp = datetime.now().strftime(constants.DATE_FORMAT)
|
||||||
archive_filename = f"conversation_{timestamp}.json"
|
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
|
archive_path = self.archive_dir / archive_filename
|
||||||
|
|
||||||
with open(archive_path, 'w', encoding='utf-8') as f:
|
self._atomic_write_json(archive_path, self.current_conversation)
|
||||||
json.dump(self.current_conversation, f, ensure_ascii=False, indent=2)
|
|
||||||
|
|
||||||
self.history_index["conversations"].append({
|
self.history_index["conversations"].append({
|
||||||
"id": self.current_conversation["id"],
|
"id": self.current_conversation["id"],
|
||||||
@@ -341,7 +357,13 @@ class ConversationHistory:
|
|||||||
if file_name == constants.CURRENT_CONVERSATION_FILE and self.current_conversation:
|
if file_name == constants.CURRENT_CONVERSATION_FILE and self.current_conversation:
|
||||||
return self.current_conversation
|
return self.current_conversation
|
||||||
|
|
||||||
file_path = self.archive_dir / file_name
|
# 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():
|
if file_path.exists():
|
||||||
with open(file_path, 'r', encoding='utf-8') as f:
|
with open(file_path, 'r', encoding='utf-8') as f:
|
||||||
return json.load(f)
|
return json.load(f)
|
||||||
|
|||||||
+212
-90
@@ -33,9 +33,18 @@ class MemoryDatabase:
|
|||||||
self.logger = Logger("MemoryDB")
|
self.logger = Logger("MemoryDB")
|
||||||
self._init_database()
|
self._init_database()
|
||||||
|
|
||||||
|
def _connect(self) -> sqlite3.Connection:
|
||||||
|
"""统一连接工厂(T25,审计 2026-08-21):对齐 MemoDatabase J9 修复——
|
||||||
|
busy_timeout 覆盖全部路径。两类默认共用同一 memories.db,此前
|
||||||
|
MemoryDatabase 直连只靠 sqlite3 默认 5s 等待,CLI 与服务进程并发
|
||||||
|
写时锁冲突即抛 database is locked(记忆保存/检索失败)。"""
|
||||||
|
conn = sqlite3.connect(self.db_path)
|
||||||
|
conn.execute("PRAGMA busy_timeout = 10000")
|
||||||
|
return conn
|
||||||
|
|
||||||
def _init_database(self):
|
def _init_database(self):
|
||||||
"""初始化数据库表"""
|
"""初始化数据库表"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
# WAL 模式:读写并发不互斥,缓解多 CLI 进程并发写锁等待
|
# WAL 模式:读写并发不互斥,缓解多 CLI 进程并发写锁等待
|
||||||
conn.execute("PRAGMA journal_mode=WAL")
|
conn.execute("PRAGMA journal_mode=WAL")
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
@@ -118,7 +127,7 @@ class MemoryDatabase:
|
|||||||
tags_str = ",".join(tags) if tags else ""
|
tags_str = ",".join(tags) if tags else ""
|
||||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||||
|
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
INSERT INTO memories (topic, summary, content, start_time, end_time, created_at, tags, importance)
|
INSERT INTO memories (topic, summary, content, start_time, end_time, created_at, tags, importance)
|
||||||
@@ -139,13 +148,16 @@ class MemoryDatabase:
|
|||||||
|
|
||||||
def delete_memory(self, memory_id: int) -> bool:
|
def delete_memory(self, memory_id: int) -> bool:
|
||||||
"""删除记忆"""
|
"""删除记忆"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
# J3(审计 2026-08-19):必须先删 fts 再删内容表——外部内容表
|
||||||
|
# (content=memories)删除索引条目需从内容表回读旧值,内容行
|
||||||
|
# 先删会导致索引条目永久残留(MATCH 仍可命中已删 rowid)
|
||||||
|
cursor.execute('DELETE FROM memories_fts WHERE rowid = ?', (memory_id,))
|
||||||
cursor.execute('DELETE FROM memories WHERE id = ?', (memory_id,))
|
cursor.execute('DELETE FROM memories WHERE id = ?', (memory_id,))
|
||||||
deleted = cursor.rowcount > 0
|
deleted = cursor.rowcount > 0
|
||||||
|
|
||||||
if deleted:
|
if deleted:
|
||||||
cursor.execute('DELETE FROM memories_fts WHERE rowid = ?', (memory_id,))
|
|
||||||
conn.commit()
|
conn.commit()
|
||||||
self.logger.info(f"已删除记忆 ID: {memory_id}")
|
self.logger.info(f"已删除记忆 ID: {memory_id}")
|
||||||
|
|
||||||
@@ -163,7 +175,7 @@ class MemoryDatabase:
|
|||||||
if limit is None:
|
if limit is None:
|
||||||
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
||||||
|
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
@@ -176,7 +188,10 @@ class MemoryDatabase:
|
|||||||
def _search_via_fts(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]:
|
def _search_via_fts(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]:
|
||||||
"""FTS5 全文检索(关键词含特殊字符导致 MATCH 语法错误时返回空,走 LIKE 兜底)"""
|
"""FTS5 全文检索(关键词含特殊字符导致 MATCH 语法错误时返回空,走 LIKE 兜底)"""
|
||||||
joiner = " AND " if match_all else " OR "
|
joiner = " AND " if match_all else " OR "
|
||||||
query = joiner.join([f'"{kw}"' for kw in keywords])
|
# N7:短语内双引号翻倍转义——关键词含 " 时若不翻倍会提前闭合短语
|
||||||
|
# 边界,破坏 MATCH 语法或改变匹配语义(FTS5 短语引号转义规则)
|
||||||
|
terms = [f'"{kw.replace(chr(34), chr(34) * 2)}"' for kw in keywords]
|
||||||
|
query = joiner.join(terms)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
@@ -229,14 +244,15 @@ class MemoryDatabase:
|
|||||||
return results
|
return results
|
||||||
|
|
||||||
def get_recent_memories(self, limit: int = 10) -> List[Dict]:
|
def get_recent_memories(self, limit: int = 10) -> List[Dict]:
|
||||||
"""获取最近的记忆。limit=None 时不限制返回条数(供分页使用)。"""
|
"""获取最近的记忆。limit=None 时使用默认上限(M4:此前 None 无 LIMIT
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
全表加载,记忆库增长后单次调用 fetchall 全量入内存;与 search_memories
|
||||||
|
的 None→默认上限语义对齐)。"""
|
||||||
|
if limit is None:
|
||||||
|
limit = constants.MEMORY_MAX_SEARCH_RESULTS
|
||||||
|
with self._connect() as conn:
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
if limit is None:
|
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC LIMIT ?', (limit,))
|
||||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC')
|
|
||||||
else:
|
|
||||||
cursor.execute('SELECT * FROM memories ORDER BY end_time DESC LIMIT ?', (limit,))
|
|
||||||
|
|
||||||
results = []
|
results = []
|
||||||
for row in cursor.fetchall():
|
for row in cursor.fetchall():
|
||||||
@@ -249,7 +265,7 @@ class MemoryDatabase:
|
|||||||
|
|
||||||
def get_memory_by_id(self, memory_id: int) -> Optional[Dict]:
|
def get_memory_by_id(self, memory_id: int) -> Optional[Dict]:
|
||||||
"""根据ID获取记忆"""
|
"""根据ID获取记忆"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute('SELECT * FROM memories WHERE id = ?', (memory_id,))
|
cursor.execute('SELECT * FROM memories WHERE id = ?', (memory_id,))
|
||||||
@@ -267,14 +283,15 @@ class MemoryDatabase:
|
|||||||
if not memory_ids:
|
if not memory_ids:
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
placeholders = ','.join(['?'] * len(memory_ids))
|
placeholders = ','.join(['?'] * len(memory_ids))
|
||||||
|
# J3:同 delete_memory——先删 fts 再删内容表,防索引条目残留
|
||||||
|
cursor.execute(f'DELETE FROM memories_fts WHERE rowid IN ({placeholders})', memory_ids)
|
||||||
cursor.execute(f'DELETE FROM memories WHERE id IN ({placeholders})', memory_ids)
|
cursor.execute(f'DELETE FROM memories WHERE id IN ({placeholders})', memory_ids)
|
||||||
deleted_count = cursor.rowcount
|
deleted_count = cursor.rowcount
|
||||||
|
|
||||||
if deleted_count > 0:
|
if deleted_count > 0:
|
||||||
cursor.execute(f'DELETE FROM memories_fts WHERE rowid IN ({placeholders})', memory_ids)
|
|
||||||
conn.commit()
|
conn.commit()
|
||||||
self.logger.info(f"已批量删除 {deleted_count} 条记忆")
|
self.logger.info(f"已批量删除 {deleted_count} 条记忆")
|
||||||
|
|
||||||
@@ -296,10 +313,19 @@ class MemoDatabase:
|
|||||||
self.logger = Logger("MemoDB")
|
self.logger = Logger("MemoDB")
|
||||||
self._init_database()
|
self._init_database()
|
||||||
|
|
||||||
|
def _connect(self) -> sqlite3.Connection:
|
||||||
|
"""统一连接工厂(J9,审计 2026-08-19):busy_timeout 覆盖全部写路径,
|
||||||
|
不再只有 mark_as_triggered 一条路径等待锁释放"""
|
||||||
|
conn = sqlite3.connect(self.db_path)
|
||||||
|
conn.execute("PRAGMA busy_timeout = 10000")
|
||||||
|
return conn
|
||||||
|
|
||||||
def _init_database(self):
|
def _init_database(self):
|
||||||
"""初始化备忘录表"""
|
"""初始化备忘录表"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
# J9:WAL 显式设置——与 MemoryDatabase 同库文件时不再依赖创建顺序
|
||||||
|
cursor.execute("PRAGMA journal_mode=WAL")
|
||||||
|
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
CREATE TABLE IF NOT EXISTS memos (
|
CREATE TABLE IF NOT EXISTS memos (
|
||||||
@@ -345,10 +371,37 @@ class MemoDatabase:
|
|||||||
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
||||||
is_immediate: bool = False) -> int:
|
is_immediate: bool = False) -> int:
|
||||||
"""添加任务 - 所有任务都是对话任务,不再需要metadata和任务类型"""
|
"""添加任务 - 所有任务都是对话任务,不再需要metadata和任务类型"""
|
||||||
|
# J1(审计 2026-08-19):入库校验——畸形时间/数字字段此前原样入库,
|
||||||
|
# 下游无防护解析点引发调度器失败热循环与新连接握手崩溃(审计报告
|
||||||
|
# §13.1 J1)。CLI/WS/memo 工具三入口均经本方法,校验一处覆盖。
|
||||||
|
# W12(审计 2026-08-20):空标题任务此前可入库(CLI "|x" / 调用方疏忽)
|
||||||
|
if not isinstance(title, str) or not title.strip():
|
||||||
|
raise ValueError("任务标题不能为空")
|
||||||
|
if reminder_time:
|
||||||
|
try:
|
||||||
|
datetime.strptime(reminder_time, constants.DATETIME_FORMAT)
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
raise ValueError(
|
||||||
|
f"提醒时间格式非法: {reminder_time!r}(应为 {constants.DATETIME_FORMAT})")
|
||||||
|
if repeat_end_time:
|
||||||
|
try:
|
||||||
|
datetime.strptime(repeat_end_time, constants.DATETIME_FORMAT)
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
raise ValueError(
|
||||||
|
f"重复结束时间格式非法: {repeat_end_time!r}(应为 {constants.DATETIME_FORMAT})")
|
||||||
|
if repeat_type not in ("none", "daily", "weekly", "monthly", "custom"):
|
||||||
|
raise ValueError(f"无效的重复类型: {repeat_type!r}")
|
||||||
|
try:
|
||||||
|
repeat_interval = int(repeat_interval)
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
raise ValueError(f"重复间隔必须是整数: {repeat_interval!r}")
|
||||||
|
if repeat_type == "custom" and repeat_interval <= 0:
|
||||||
|
raise ValueError("自定义重复需要设置有效的 repeat_interval (>0)")
|
||||||
|
|
||||||
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
created_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||||
tags_str = ",".join(tags) if tags else ""
|
tags_str = ",".join(tags) if tags else ""
|
||||||
|
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
INSERT INTO memos
|
INSERT INTO memos
|
||||||
@@ -373,7 +426,7 @@ class MemoDatabase:
|
|||||||
|
|
||||||
def get_memo(self, memo_id: int) -> Optional[Dict]:
|
def get_memo(self, memo_id: int) -> Optional[Dict]:
|
||||||
"""获取单个任务"""
|
"""获取单个任务"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||||
@@ -388,7 +441,7 @@ class MemoDatabase:
|
|||||||
|
|
||||||
def get_pending_memos(self, limit: int = None, before_time: str = None, include_immediate: bool = False) -> List[Dict]:
|
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:
|
with self._connect() as conn:
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
@@ -398,11 +451,14 @@ class MemoDatabase:
|
|||||||
# 如果传入了 before_time,使用它,否则使用当前时间
|
# 如果传入了 before_time,使用它,否则使用当前时间
|
||||||
compare_time = before_time if before_time else now
|
compare_time = before_time if before_time else now
|
||||||
|
|
||||||
|
# I 批(审计 2026-08-18):去掉 SQL 层 reminder_time <= ? 预过滤——
|
||||||
|
# 字符串比较会把格式异常的时间(如 "不是时间")静默排除在 SQL 层,
|
||||||
|
# Python 二次验证永远看不到它们(Q5:格式异常任务按到期返回)。
|
||||||
|
# 时间过滤改由下方 Python 解析做权威判断(本地小库全扫描可忽略)。
|
||||||
query = '''
|
query = '''
|
||||||
SELECT * FROM memos
|
SELECT * FROM memos
|
||||||
WHERE is_completed = 0
|
WHERE is_completed = 0
|
||||||
AND reminder_time IS NOT NULL
|
AND reminder_time IS NOT NULL
|
||||||
AND reminder_time <= ? -- 严格小于等于比较时间
|
|
||||||
AND (repeat_end_time IS NULL OR repeat_end_time >= ?)
|
AND (repeat_end_time IS NULL OR repeat_end_time >= ?)
|
||||||
'''
|
'''
|
||||||
|
|
||||||
@@ -411,10 +467,13 @@ class MemoDatabase:
|
|||||||
|
|
||||||
query += ' ORDER BY reminder_time ASC, priority DESC'
|
query += ' ORDER BY reminder_time ASC, priority DESC'
|
||||||
|
|
||||||
params = [compare_time, compare_time]
|
params = [compare_time]
|
||||||
|
|
||||||
if limit:
|
# I 批:limit=0 语义修正——此前 `if limit:` 使 0 等同"不限",
|
||||||
query += f' LIMIT {limit}'
|
# 与 None=不限的惯例冲突;改为 None 才不加 LIMIT
|
||||||
|
if limit is not None:
|
||||||
|
query += ' LIMIT ?'
|
||||||
|
params.append(max(0, int(limit)))
|
||||||
|
|
||||||
cursor.execute(query, params)
|
cursor.execute(query, params)
|
||||||
|
|
||||||
@@ -435,7 +494,11 @@ class MemoDatabase:
|
|||||||
# 调试日志
|
# 调试日志
|
||||||
self.logger.debug(f"过滤掉未到期任务 #{memo['id']}: {memo['reminder_time']} > {compare_time}")
|
self.logger.debug(f"过滤掉未到期任务 #{memo['id']}: {memo['reminder_time']} > {compare_time}")
|
||||||
except (ValueError, TypeError) as e:
|
except (ValueError, TypeError) as e:
|
||||||
self.logger.warning(f"解析任务 #{memo['id']} 的提醒时间失败: {e}")
|
# I 批(审计 2026-08-18):格式异常此前被静默丢弃——
|
||||||
|
# 任务永不执行且用户无感知。改为按"到期"返回并告警,
|
||||||
|
# 至少保证任务被处理而非永久丢失
|
||||||
|
self.logger.warning(f"解析任务 #{memo['id']} 的提醒时间失败,按到期处理: {e}")
|
||||||
|
results.append(memo)
|
||||||
else:
|
else:
|
||||||
# 没有提醒时间的任务不应该在待执行列表中
|
# 没有提醒时间的任务不应该在待执行列表中
|
||||||
self.logger.warning(f"任务 #{memo['id']} 没有提醒时间,但被返回为待执行")
|
self.logger.warning(f"任务 #{memo['id']} 没有提醒时间,但被返回为待执行")
|
||||||
@@ -444,7 +507,7 @@ class MemoDatabase:
|
|||||||
|
|
||||||
def get_immediate_memos(self, limit: int = None) -> List[Dict]:
|
def get_immediate_memos(self, limit: int = None) -> List[Dict]:
|
||||||
"""获取待执行的即时任务"""
|
"""获取待执行的即时任务"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
@@ -455,10 +518,14 @@ class MemoDatabase:
|
|||||||
ORDER BY created_at ASC
|
ORDER BY created_at ASC
|
||||||
'''
|
'''
|
||||||
|
|
||||||
if limit:
|
params = []
|
||||||
query += f' LIMIT {limit}'
|
# 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)
|
cursor.execute(query, params)
|
||||||
|
|
||||||
results = []
|
results = []
|
||||||
for row in cursor.fetchall():
|
for row in cursor.fetchall():
|
||||||
@@ -476,7 +543,7 @@ class MemoDatabase:
|
|||||||
|
|
||||||
def get_all_memos(self, include_completed: bool = False, limit: int = 50) -> List[Dict]:
|
def get_all_memos(self, include_completed: bool = False, limit: int = 50) -> List[Dict]:
|
||||||
"""获取所有任务"""
|
"""获取所有任务"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
@@ -508,71 +575,91 @@ class MemoDatabase:
|
|||||||
|
|
||||||
def mark_as_triggered(self, memo_id: int) -> Optional[str]:
|
def mark_as_triggered(self, memo_id: int) -> Optional[str]:
|
||||||
"""标记任务已被触发,并根据重复规则计算下一次提醒时间"""
|
"""标记任务已被触发,并根据重复规则计算下一次提醒时间"""
|
||||||
memo = self.get_memo(memo_id)
|
# N10:读取与更新收进同一连接的 BEGIN IMMEDIATE 事务。此前先
|
||||||
if not memo:
|
# get_memo(独立连接)后 UPDATE,check-then-act 之间存在竞态:
|
||||||
return None
|
# 调度器并发触发同一任务会双计数/双推进 repeat。IMMEDIATE 在事务
|
||||||
|
# 开始即取写锁,串行化并发触发。
|
||||||
now = datetime.now()
|
now = datetime.now()
|
||||||
next_reminder = None
|
next_reminder = None
|
||||||
|
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
# J9:busy_timeout 已由 _connect() 工厂统一设置(全部写路径一致),
|
||||||
|
# 多进程同库(CLI --memo-add + 服务调度)写竞争等待锁释放而非立即失败
|
||||||
|
cursor.execute("BEGIN IMMEDIATE")
|
||||||
|
try:
|
||||||
|
cursor.execute('SELECT * FROM memos WHERE id = ?', (memo_id,))
|
||||||
|
row = cursor.fetchone()
|
||||||
|
if row is None:
|
||||||
|
conn.rollback()
|
||||||
|
return None
|
||||||
|
memo = dict(row)
|
||||||
|
if memo.get('tags'):
|
||||||
|
memo['tags'] = memo['tags'].split(',')
|
||||||
|
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
UPDATE memos
|
UPDATE memos
|
||||||
SET last_triggered_at = ?, trigger_count = trigger_count + 1
|
SET last_triggered_at = ?, trigger_count = trigger_count + 1
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||||
|
|
||||||
# 处理重复任务
|
# 处理重复任务
|
||||||
if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now):
|
if memo['repeat_type'] != 'none' and not self._should_end_repeat(memo, now):
|
||||||
# 使用当前时间作为基准来计算下一次执行时间
|
# 使用当前时间作为基准来计算下一次执行时间
|
||||||
next_reminder = self._calculate_next_reminder(memo, now)
|
next_reminder = self._calculate_next_reminder(memo, now)
|
||||||
|
|
||||||
if next_reminder:
|
if next_reminder:
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
UPDATE memos
|
UPDATE memos
|
||||||
SET reminder_time = ?
|
SET reminder_time = ?
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id))
|
''', (next_reminder.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||||
|
|
||||||
self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}")
|
self.logger.info(f"任务 #{memo_id} 已更新下一次执行时间: {next_reminder}")
|
||||||
else:
|
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('''
|
cursor.execute('''
|
||||||
UPDATE memos
|
UPDATE memos
|
||||||
SET is_completed = 1, completed_at = ?
|
SET is_completed = 1, completed_at = ?
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||||
|
|
||||||
self.logger.info(f"任务 #{memo_id} 重复结束,已标记为完成")
|
self.logger.info(f"即时任务 #{memo_id} 已完成")
|
||||||
elif memo['is_immediate']:
|
else:
|
||||||
# 即时任务执行后标记为完成
|
# 一次性任务,标记为完成
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
UPDATE memos
|
UPDATE memos
|
||||||
SET is_completed = 1, completed_at = ?
|
SET is_completed = 1, completed_at = ?
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
''', (now.strftime(constants.DATETIME_FORMAT), memo_id))
|
||||||
|
|
||||||
self.logger.info(f"即时任务 #{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.commit()
|
conn.rollback()
|
||||||
|
raise
|
||||||
|
|
||||||
return next_reminder.strftime(constants.DATETIME_FORMAT) if next_reminder else None
|
return next_reminder.strftime(constants.DATETIME_FORMAT) if next_reminder else None
|
||||||
|
|
||||||
def _should_end_repeat(self, memo: Dict, current_time: datetime) -> bool:
|
def _should_end_repeat(self, memo: Dict, current_time: datetime) -> bool:
|
||||||
"""检查是否应该结束重复"""
|
"""检查是否应该结束重复"""
|
||||||
if memo['repeat_end_time']:
|
if memo['repeat_end_time']:
|
||||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
try:
|
||||||
|
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
# J1:存量脏数据防御——解析失败视为无结束时间,不打断触发事务
|
||||||
|
return False
|
||||||
if current_time > end_time:
|
if current_time > end_time:
|
||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
@@ -602,24 +689,37 @@ class MemoDatabase:
|
|||||||
next_time = current_time.replace(year=current_time.year + 1, month=1, day=1) - timedelta(days=1)
|
next_time = current_time.replace(year=current_time.year + 1, month=1, day=1) - timedelta(days=1)
|
||||||
else:
|
else:
|
||||||
next_time = current_time.replace(month=current_time.month + 1, day=1) - timedelta(days=1)
|
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:
|
elif memo['repeat_type'] == 'custom':
|
||||||
unit = memo.get('repeat_interval_unit', 'days')
|
# J1:repeat_interval 类型异常(存量脏数据)按 0 处理,
|
||||||
if unit == 'minutes':
|
# 不再抛 TypeError 打崩触发事务引发热循环
|
||||||
# 每N分钟重复:当前时间 + N分钟
|
try:
|
||||||
next_time = current_time + timedelta(minutes=memo['repeat_interval'])
|
repeat_interval = int(memo.get('repeat_interval') or 0)
|
||||||
elif unit == 'hours':
|
except (ValueError, TypeError):
|
||||||
# 每N小时重复:当前时间 + N小时
|
repeat_interval = 0
|
||||||
next_time = current_time + timedelta(hours=memo['repeat_interval'])
|
if repeat_interval > 0:
|
||||||
else: # days
|
unit = memo.get('repeat_interval_unit', 'days')
|
||||||
# 每N天重复:当前时间 + N天
|
if unit == 'minutes':
|
||||||
next_time = current_time + timedelta(days=memo['repeat_interval'])
|
# 每N分钟重复:当前时间 + N分钟
|
||||||
|
next_time = current_time + timedelta(minutes=repeat_interval)
|
||||||
|
elif unit == 'hours':
|
||||||
|
# 每N小时重复:当前时间 + N小时
|
||||||
|
next_time = current_time + timedelta(hours=repeat_interval)
|
||||||
|
else: # days
|
||||||
|
# 每N天重复:当前时间 + N天
|
||||||
|
next_time = current_time + timedelta(days=repeat_interval)
|
||||||
|
else:
|
||||||
|
return None
|
||||||
else:
|
else:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# 检查是否超过重复结束时间
|
# 检查是否超过重复结束时间
|
||||||
if memo['repeat_end_time']:
|
if memo['repeat_end_time']:
|
||||||
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
try:
|
||||||
if next_time > end_time:
|
end_time = datetime.strptime(memo['repeat_end_time'], constants.DATETIME_FORMAT)
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
# J1:存量脏数据防御——解析失败视为无结束时间
|
||||||
|
end_time = None
|
||||||
|
if end_time and next_time > end_time:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return next_time
|
return next_time
|
||||||
@@ -636,7 +736,7 @@ class MemoDatabase:
|
|||||||
|
|
||||||
completed_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
completed_at = datetime.now().strftime(constants.DATETIME_FORMAT)
|
||||||
|
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
UPDATE memos
|
UPDATE memos
|
||||||
@@ -652,9 +752,18 @@ class MemoDatabase:
|
|||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
def delete_memo(self, memo_id: int) -> bool:
|
||||||
|
"""物理删除任务(M5:delete 语义此前被 complete_memo 顶替——只标记
|
||||||
|
完成不删行,重复任务还会被推进而非移除,stats 计数失真)"""
|
||||||
|
with self._connect() as conn:
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute("DELETE FROM memos WHERE id = ?", (memo_id,))
|
||||||
|
conn.commit()
|
||||||
|
return cursor.rowcount > 0
|
||||||
|
|
||||||
def get_memo_stats(self) -> Dict:
|
def get_memo_stats(self) -> Dict:
|
||||||
"""获取任务统计信息"""
|
"""获取任务统计信息"""
|
||||||
with sqlite3.connect(self.db_path) as conn:
|
with self._connect() as conn:
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
cursor.execute('SELECT COUNT(*) FROM memos')
|
cursor.execute('SELECT COUNT(*) FROM memos')
|
||||||
@@ -687,8 +796,9 @@ class MemoDatabase:
|
|||||||
}
|
}
|
||||||
|
|
||||||
def add_user_task(self, title: str, content: str = None, reminder_time: str = None,
|
def add_user_task(self, title: str, content: str = None, reminder_time: str = None,
|
||||||
repeat_type: str = 'none', repeat_interval: int = 0,
|
repeat_type: str = 'none', repeat_interval: int = 0,
|
||||||
repeat_interval_unit: str = 'days', is_immediate: bool = False) -> int:
|
repeat_interval_unit: str = 'days', repeat_end_time: str = None,
|
||||||
|
is_immediate: bool = False) -> int:
|
||||||
"""添加用户任务(前端直接调用的任务)- 所有任务都是对话任务"""
|
"""添加用户任务(前端直接调用的任务)- 所有任务都是对话任务"""
|
||||||
if not constants.MEMO_ALLOW_USER_TASKS:
|
if not constants.MEMO_ALLOW_USER_TASKS:
|
||||||
raise Exception("用户任务功能已禁用")
|
raise Exception("用户任务功能已禁用")
|
||||||
@@ -700,11 +810,23 @@ class MemoDatabase:
|
|||||||
repeat_type=repeat_type,
|
repeat_type=repeat_type,
|
||||||
repeat_interval=repeat_interval,
|
repeat_interval=repeat_interval,
|
||||||
repeat_interval_unit=repeat_interval_unit,
|
repeat_interval_unit=repeat_interval_unit,
|
||||||
|
repeat_end_time=repeat_end_time,
|
||||||
created_by="user",
|
created_by="user",
|
||||||
source="websocket",
|
source="websocket",
|
||||||
is_immediate=is_immediate
|
is_immediate=is_immediate
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _is_overdue_time(reminder_time: str, now: datetime) -> bool:
|
||||||
|
"""J1:展示路径时间解析防御——格式异常按到期处理(对齐 Q5 语义),
|
||||||
|
不再抛 ValueError 打崩新连接握手流程"""
|
||||||
|
if not reminder_time:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
return datetime.strptime(reminder_time, constants.DATETIME_FORMAT) < now
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
return True
|
||||||
|
|
||||||
def get_pending_tasks_for_user(self) -> List[Dict]:
|
def get_pending_tasks_for_user(self) -> List[Dict]:
|
||||||
"""获取待执行任务列表(用于前端显示)"""
|
"""获取待执行任务列表(用于前端显示)"""
|
||||||
pending = self.get_pending_memos(include_immediate=False)
|
pending = self.get_pending_memos(include_immediate=False)
|
||||||
@@ -719,7 +841,7 @@ class MemoDatabase:
|
|||||||
"content": memo['content'],
|
"content": memo['content'],
|
||||||
"reminder_time": memo['reminder_time'],
|
"reminder_time": memo['reminder_time'],
|
||||||
"priority": memo['priority'],
|
"priority": memo['priority'],
|
||||||
"is_overdue": datetime.strptime(memo['reminder_time'], constants.DATETIME_FORMAT) < now if memo['reminder_time'] else False,
|
"is_overdue": self._is_overdue_time(memo['reminder_time'], now),
|
||||||
"repeat_type": memo['repeat_type'],
|
"repeat_type": memo['repeat_type'],
|
||||||
"repeat_interval": memo.get('repeat_interval', 0),
|
"repeat_interval": memo.get('repeat_interval', 0),
|
||||||
"repeat_interval_unit": memo.get('repeat_interval_unit', 'days'),
|
"repeat_interval_unit": memo.get('repeat_interval_unit', 'days'),
|
||||||
|
|||||||
+267
-106
@@ -13,6 +13,7 @@ subsystems.executors - 命令执行器
|
|||||||
|
|
||||||
import os
|
import os
|
||||||
import queue
|
import queue
|
||||||
|
import signal
|
||||||
import subprocess
|
import subprocess
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
@@ -21,6 +22,36 @@ from typing import Dict, Optional, Tuple
|
|||||||
|
|
||||||
from common import ConfigManager, constants, Logger
|
from common import ConfigManager, constants, Logger
|
||||||
|
|
||||||
|
# T21(审计 2026-08-21):env 覆盖黑名单——工具 schema 中 env 是开放
|
||||||
|
# object,LLM(可被网页内容提示注入诱导)可传 LD_PRELOAD/PYTHONPATH/
|
||||||
|
# PATH 等在 shell=True 之上进一步劫持进程加载(POSIX LD_PRELOAD 可加载
|
||||||
|
# 任意共享库)。命中黑名单的覆盖键拒绝并告警
|
||||||
|
ENV_DENY_EXACT = frozenset({
|
||||||
|
"PATH", "PYTHONPATH", "PYTHONHOME", "SHELL", "IFS",
|
||||||
|
"LD_LIBRARY_PATH", "LD_PRELOAD", "LD_AUDIT",
|
||||||
|
"DYLD_LIBRARY_PATH", "DYLD_INSERT_LIBRARIES",
|
||||||
|
"PSMODULEPATH", "COMSPEC",
|
||||||
|
})
|
||||||
|
ENV_DENY_PREFIXES = ("LD_", "DYLD_")
|
||||||
|
|
||||||
|
|
||||||
|
def filter_env_overrides(env) -> Tuple[Dict, list]:
|
||||||
|
"""T21:过滤 env 覆盖项,返回 (安全覆盖, 被拒绝键列表)。
|
||||||
|
|
||||||
|
键大小写不敏感匹配(Windows 环境变量不区分大小写);
|
||||||
|
LD_*/DYLD_* 前缀整体拒绝(加载器劫持面)。
|
||||||
|
"""
|
||||||
|
if not env:
|
||||||
|
return {}, []
|
||||||
|
safe, rejected = {}, []
|
||||||
|
for key, value in dict(env).items():
|
||||||
|
ku = str(key).upper()
|
||||||
|
if ku in ENV_DENY_EXACT or ku.startswith(ENV_DENY_PREFIXES):
|
||||||
|
rejected.append(key)
|
||||||
|
else:
|
||||||
|
safe[key] = value
|
||||||
|
return safe, rejected
|
||||||
|
|
||||||
|
|
||||||
class CommandExecutor(ABC):
|
class CommandExecutor(ABC):
|
||||||
"""命令执行器抽象基类 - 负责执行系统命令"""
|
"""命令执行器抽象基类 - 负责执行系统命令"""
|
||||||
@@ -70,63 +101,206 @@ class SubprocessCommandExecutor(CommandExecutor):
|
|||||||
|
|
||||||
def __init__(self, logger=None):
|
def __init__(self, logger=None):
|
||||||
super().__init__(logger)
|
super().__init__(logger)
|
||||||
self._current_process = None
|
self._current_process = None # 兼容保留:最近一次注册的进程(诊断用)
|
||||||
|
self._current_pgid = None # L9:Popen 后立即捕获,shell 死亡后 getpgid 会失败
|
||||||
self._process_lock = threading.Lock()
|
self._process_lock = threading.Lock()
|
||||||
|
# T5:单槽 _current_process 在并发执行时互相覆盖——先启动的执行
|
||||||
|
# 超时/取消会误杀后来者的进程组(killpg 连坐孙进程),先结束的
|
||||||
|
# finally 清引用导致后来者永远无法取消。改为令牌多槽注册表,
|
||||||
|
# 终止只作用于自己槽位内的 (proc, pgid)。
|
||||||
|
self._proc_slots: Dict[int, Tuple] = {}
|
||||||
|
self._token_seq = 0
|
||||||
|
|
||||||
def cancel(self):
|
def _register_proc(self, process):
|
||||||
"""取消当前正在运行的命令(kill子进程)。线程安全。"""
|
"""T5:注册本次执行的进程槽,返回 (token, pgid)"""
|
||||||
with self._process_lock:
|
with self._process_lock:
|
||||||
proc = self._current_process
|
self._token_seq += 1
|
||||||
if proc is not None:
|
token = self._token_seq
|
||||||
|
pgid = None
|
||||||
|
if os.name == "posix":
|
||||||
|
try:
|
||||||
|
pgid = os.getpgid(process.pid)
|
||||||
|
except Exception:
|
||||||
|
pgid = None
|
||||||
|
self._proc_slots[token] = (process, pgid)
|
||||||
|
# 兼容属性指向最新注册(诊断/旧测试观察用,不参与终止逻辑)
|
||||||
|
self._current_process = process
|
||||||
|
self._current_pgid = pgid
|
||||||
|
return token, pgid
|
||||||
|
|
||||||
|
def _unregister_proc(self, token):
|
||||||
|
"""T5:仅移除自己的槽位;全部槽位清空时同步清兼容属性"""
|
||||||
|
with self._process_lock:
|
||||||
|
self._proc_slots.pop(token, None)
|
||||||
|
if not self._proc_slots:
|
||||||
|
self._current_process = None
|
||||||
|
self._current_pgid = None
|
||||||
|
|
||||||
|
def _terminate_tree(self, proc, force: bool = False, pgid=None):
|
||||||
|
"""终止整个进程组(L9:POSIX 下 shell=True 只杀 shell 时,孙进程
|
||||||
|
(`(x &)` 等)仍持管道写端——reader 阻塞无 EOF、fd/线程悬挂至孙
|
||||||
|
进程退出,写满 64KB 管道缓冲后孙进程也卡死。start_new_session 使
|
||||||
|
shell 成为会话首进程,按组信号可一次覆盖整棵进程树)。
|
||||||
|
组 ID 在 Popen 后立即保存:shell 死亡后 os.getpgid(pid) 会抛
|
||||||
|
ProcessLookupError,届时无法再从进程反查组。
|
||||||
|
T5:pgid 由调用方显式传入(注册槽位时捕获),不再读共享
|
||||||
|
_current_pgid——并发执行时那是别人的进程组。
|
||||||
|
T11(审计 2026-08-21):Windows 无进程组概念,单进程 terminate/
|
||||||
|
kill 收割不到 shell=True 生成的 cmd.exe /c 之下 spawn 的孙进程
|
||||||
|
(超时/取消后孙进程成孤儿,持续占用 CPU/端口/文件锁)。用
|
||||||
|
taskkill /T 递归终止整棵进程树;Windows 无 SIGTERM 等价物,
|
||||||
|
taskkill 不带 /F 仅对有消息窗口的 GUI 进程有效,控制台进程
|
||||||
|
杀不掉,故统一 /F 强制,force 参数在 Windows 上不区分。"""
|
||||||
|
try:
|
||||||
|
if os.name == "posix":
|
||||||
|
sig = signal.SIGKILL if force else signal.SIGTERM
|
||||||
|
if pgid is None:
|
||||||
|
try:
|
||||||
|
pgid = os.getpgid(proc.pid)
|
||||||
|
except Exception:
|
||||||
|
pgid = None
|
||||||
|
if pgid is not None:
|
||||||
|
os.killpg(pgid, sig)
|
||||||
|
else:
|
||||||
|
proc.kill() if force else proc.terminate()
|
||||||
|
else:
|
||||||
|
# T11:Windows 树收割——taskkill /T /F 递归终止子孙进程,
|
||||||
|
# 失败(进程已退出/无权限)时兜底单进程 kill
|
||||||
|
try:
|
||||||
|
subprocess.run(
|
||||||
|
["taskkill", "/PID", str(proc.pid), "/T", "/F"],
|
||||||
|
capture_output=True, timeout=10)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
if proc.poll() is None:
|
||||||
|
proc.kill()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
except Exception:
|
||||||
try:
|
try:
|
||||||
if proc.poll() is None:
|
proc.kill()
|
||||||
proc.kill()
|
|
||||||
self.logger.info("已终止正在运行的子进程")
|
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None) -> Tuple[str, bool, bool]:
|
def cancel(self):
|
||||||
"""从进程读取输出,限制总大小和超时。
|
"""取消本执行器当前注册的全部运行中命令(kill 子进程组)。线程安全。
|
||||||
返回 (output_str, timeout_occurred, truncated)
|
T5:多槽令牌注册表——并发执行时定向终止各自槽位,互不误杀;
|
||||||
|
无参调用语义为"终止本执行器全部在途命令"(中断/关闭路径)。"""
|
||||||
|
with self._process_lock:
|
||||||
|
slots = list(self._proc_slots.values())
|
||||||
|
for proc, pgid in slots:
|
||||||
|
try:
|
||||||
|
if proc.poll() is None:
|
||||||
|
self._terminate_tree(proc, force=True, pgid=pgid)
|
||||||
|
self.logger.info("已终止正在运行的子进程组")
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _read_with_limit(self, process, timeout: int, start_time: float, output_queue=None,
|
||||||
|
pgid=None) -> Tuple[str, bool, bool]:
|
||||||
|
"""从进程读取输出,限制总大小和超时。返回 (output_str, timeout_occurred, truncated)。
|
||||||
|
|
||||||
|
H2 修复:阻塞 readline 在独立 reader 线程执行,主循环带超时轮询队列。
|
||||||
|
此前超时检查只在两次 readline 之间执行,无输出命令会让 readline
|
||||||
|
阻塞到进程自然退出,timeout 形同虚设(审计 PoC 实证)。
|
||||||
"""
|
"""
|
||||||
output_chunks = []
|
output_chunks = []
|
||||||
total_bytes = 0
|
total_bytes = 0
|
||||||
timeout_occurred = False
|
timeout_occurred = False
|
||||||
truncated = 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:
|
while True:
|
||||||
# 检查超时
|
# 检查超时(readline 阻塞不再影响本判断)
|
||||||
|
# J16(审计 2026-08-19):timeout<=0 语义为"不设超时"(有意为之,
|
||||||
|
# 调用方显式传 0 禁用);registry 层 3600s 流式兜底仍生效
|
||||||
if timeout > 0 and time.time() - start_time > timeout:
|
if timeout > 0 and time.time() - start_time > timeout:
|
||||||
process.terminate()
|
self._terminate_tree(process, pgid=pgid)
|
||||||
timeout_occurred = True
|
timeout_occurred = True
|
||||||
break
|
break
|
||||||
|
|
||||||
# 检查输出大小
|
# 检查输出大小
|
||||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
||||||
truncated = True
|
truncated = True
|
||||||
process.terminate()
|
self._terminate_tree(process, pgid=pgid)
|
||||||
break
|
break
|
||||||
|
|
||||||
# 读取一行
|
try:
|
||||||
line = process.stdout.readline()
|
line = line_queue.get(timeout=0.1)
|
||||||
if not line and process.poll() is not None:
|
shell_dead_since = None
|
||||||
break
|
except queue.Empty:
|
||||||
if line:
|
if reader_done.is_set():
|
||||||
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:
|
|
||||||
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace'))
|
|
||||||
total_bytes = self.MAX_OUTPUT_BYTES
|
|
||||||
truncated = True
|
|
||||||
process.terminate()
|
|
||||||
break
|
break
|
||||||
|
# shell 已被 kill/退出但 reader 仍阻塞(如 Windows shell=True
|
||||||
|
# 下 cancel 只杀死包装进程,孙进程仍持有管道写端导致无 EOF):
|
||||||
|
# 宽限 1 秒排空残留输出后放弃等待
|
||||||
|
if process.poll() is not None:
|
||||||
|
if shell_dead_since is None:
|
||||||
|
shell_dead_since = time.time()
|
||||||
|
elif time.time() - shell_dead_since > 1.0:
|
||||||
|
# L9:shell 已死但孙进程仍持管道写端(`(sleep 300 &) ;
|
||||||
|
# echo done` 形态)。只靠 daemon reader 兜底会悬挂线程
|
||||||
|
# 与管道 fd 至孙进程退出,写满 64KB 缓冲后孙进程自身
|
||||||
|
# 也卡死。此处强杀进程组收尾(nohup 重定向输出的
|
||||||
|
# 真正脱离场景不持管道,不会走到这里)。
|
||||||
|
self._terminate_tree(process, force=True, pgid=pgid)
|
||||||
|
break
|
||||||
|
continue
|
||||||
|
|
||||||
output_chunks.append(line)
|
line_bytes = len(line.encode('utf-8', errors='replace'))
|
||||||
total_bytes += line_bytes
|
if total_bytes + line_bytes > self.MAX_OUTPUT_BYTES and not truncated:
|
||||||
|
allowed = self.MAX_OUTPUT_BYTES - total_bytes
|
||||||
|
if allowed > 0:
|
||||||
|
# L5:decode 用 errors='ignore'——按字节截断可能落在多字节
|
||||||
|
# 字符中间,errors='replace' 会在行尾留下半个字符(�)
|
||||||
|
output_chunks.append(line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='ignore'))
|
||||||
|
total_bytes = self.MAX_OUTPUT_BYTES
|
||||||
|
truncated = True
|
||||||
|
self._terminate_tree(process, pgid=pgid)
|
||||||
|
break
|
||||||
|
|
||||||
if output_queue:
|
output_chunks.append(line)
|
||||||
output_queue.put(("line", line.rstrip('\n') + "\n"))
|
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
|
return ''.join(output_chunks), timeout_occurred, truncated
|
||||||
|
|
||||||
@@ -135,48 +309,63 @@ class SubprocessCommandExecutor(CommandExecutor):
|
|||||||
if not self.is_enabled():
|
if not self.is_enabled():
|
||||||
return "❌ 命令执行功能已禁用"
|
return "❌ 命令执行功能已禁用"
|
||||||
|
|
||||||
|
# T24:非流式路径没有 registry 3600s 流式兜底,timeout<=0("不设
|
||||||
|
# 超时")会让挂起命令永久阻塞调用线程——下限夹取到 1s 快速失败。
|
||||||
|
# 流式路径保留 0=不限语义(registry 兜底在位)
|
||||||
|
try:
|
||||||
|
timeout = max(1, int(timeout))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
timeout = 30
|
||||||
|
|
||||||
# 默认工作目录为workspace,避免在启动目录散落文件
|
# 默认工作目录为workspace,避免在启动目录散落文件
|
||||||
if work_dir is None or work_dir == ".":
|
if work_dir is None or work_dir == ".":
|
||||||
work_dir = constants.WORKSPACE_DIR
|
work_dir = constants.WORKSPACE_DIR
|
||||||
|
|
||||||
self.logger.info(f"执行命令: {command} (work_dir={work_dir})")
|
self.logger.info(f"执行命令: {command} (work_dir={work_dir})")
|
||||||
|
|
||||||
# 准备环境变量
|
# 准备环境变量(T21:黑名单键拒绝,防加载器/解释器劫持)
|
||||||
exec_env = os.environ.copy()
|
exec_env = os.environ.copy()
|
||||||
if env:
|
if env:
|
||||||
exec_env.update(env)
|
safe_env, rejected_keys = filter_env_overrides(env)
|
||||||
|
if rejected_keys:
|
||||||
|
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||||
|
exec_env.update(safe_env)
|
||||||
|
|
||||||
process = None
|
process = None
|
||||||
|
token = None
|
||||||
|
pgid = None
|
||||||
try:
|
try:
|
||||||
process = subprocess.Popen(
|
process = subprocess.Popen(
|
||||||
command,
|
command,
|
||||||
shell=True,
|
shell=True,
|
||||||
stdout=subprocess.PIPE,
|
stdout=subprocess.PIPE,
|
||||||
stderr=subprocess.STDOUT,
|
stderr=subprocess.STDOUT,
|
||||||
stdin=subprocess.PIPE,
|
stdin=subprocess.DEVNULL, # H3 修复:不挂 stdin,防等待输入命令永久挂起
|
||||||
text=True,
|
text=True,
|
||||||
bufsize=1,
|
bufsize=1,
|
||||||
universal_newlines=True,
|
universal_newlines=True,
|
||||||
|
encoding="utf-8", # T16:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出
|
||||||
|
errors="replace", # T16:非法字节不抛异常中断 reader 线程
|
||||||
env=exec_env,
|
env=exec_env,
|
||||||
cwd=os.path.expanduser(work_dir)
|
cwd=os.path.expanduser(work_dir),
|
||||||
|
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||||
|
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||||
|
start_new_session=(os.name == "posix")
|
||||||
)
|
)
|
||||||
# 保存进程引用以便cancel()可以kill它
|
# T5:注册本执行专属槽位,终止/取消只作用于自己的 (proc, pgid)
|
||||||
with self._process_lock:
|
token, pgid = self._register_proc(process)
|
||||||
self._current_process = process
|
|
||||||
|
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
output_str, timeout_occurred, truncated = self._read_with_limit(process, timeout, start_time)
|
output_str, timeout_occurred, truncated = self._read_with_limit(
|
||||||
|
process, timeout, start_time, pgid=pgid)
|
||||||
|
|
||||||
# 确保进程终止
|
# 确保进程终止
|
||||||
if process.poll() is None:
|
if process.poll() is None:
|
||||||
if timeout_occurred or truncated:
|
self._terminate_tree(process, force=timeout_occurred or truncated, pgid=pgid)
|
||||||
process.kill()
|
|
||||||
else:
|
|
||||||
process.terminate()
|
|
||||||
try:
|
try:
|
||||||
process.wait(timeout=2)
|
process.wait(timeout=2)
|
||||||
except:
|
except:
|
||||||
process.kill()
|
self._terminate_tree(process, force=True, pgid=pgid)
|
||||||
|
|
||||||
return_code = process.returncode if process.returncode is not None else -1
|
return_code = process.returncode if process.returncode is not None else -1
|
||||||
|
|
||||||
@@ -204,10 +393,10 @@ class SubprocessCommandExecutor(CommandExecutor):
|
|||||||
return f"❌ 执行错误: {str(e)}"
|
return f"❌ 执行错误: {str(e)}"
|
||||||
finally:
|
finally:
|
||||||
if process and process.poll() is None:
|
if process and process.poll() is None:
|
||||||
process.kill()
|
self._terminate_tree(process, force=True, pgid=pgid)
|
||||||
# 清除当前进程引用
|
# T5:仅注销自己的槽位,不再无差别清空共享引用
|
||||||
with self._process_lock:
|
if token is not None:
|
||||||
self._current_process = None
|
self._unregister_proc(token)
|
||||||
|
|
||||||
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30,
|
def execute_streaming(self, command: str, output_queue: queue.Queue, timeout: int = 30,
|
||||||
work_dir: str = None, env: Dict = None, **kwargs):
|
work_dir: str = None, env: Dict = None, **kwargs):
|
||||||
@@ -228,79 +417,52 @@ class SubprocessCommandExecutor(CommandExecutor):
|
|||||||
output_queue.put(("line", f"📁 工作目录: {work_dir}\n"))
|
output_queue.put(("line", f"📁 工作目录: {work_dir}\n"))
|
||||||
output_queue.put(("line", "-" * 50 + "\n"))
|
output_queue.put(("line", "-" * 50 + "\n"))
|
||||||
|
|
||||||
# 准备环境变量
|
# 准备环境变量(T21:黑名单键拒绝,防加载器/解释器劫持)
|
||||||
exec_env = os.environ.copy()
|
exec_env = os.environ.copy()
|
||||||
if env:
|
if env:
|
||||||
exec_env.update(env)
|
safe_env, rejected_keys = filter_env_overrides(env)
|
||||||
|
if rejected_keys:
|
||||||
|
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||||
|
exec_env.update(safe_env)
|
||||||
|
|
||||||
process = None
|
process = None
|
||||||
|
token = None
|
||||||
|
pgid = None
|
||||||
try:
|
try:
|
||||||
process = subprocess.Popen(
|
process = subprocess.Popen(
|
||||||
command,
|
command,
|
||||||
shell=True,
|
shell=True,
|
||||||
stdout=subprocess.PIPE,
|
stdout=subprocess.PIPE,
|
||||||
stderr=subprocess.STDOUT,
|
stderr=subprocess.STDOUT,
|
||||||
stdin=subprocess.PIPE,
|
stdin=subprocess.DEVNULL, # H3 修复:不挂 stdin,防等待输入命令永久挂起
|
||||||
text=True,
|
text=True,
|
||||||
bufsize=1,
|
bufsize=1,
|
||||||
universal_newlines=True,
|
universal_newlines=True,
|
||||||
|
encoding="utf-8", # T16:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出
|
||||||
|
errors="replace", # T16:非法字节不抛异常中断 reader 线程
|
||||||
env=exec_env,
|
env=exec_env,
|
||||||
cwd=os.path.expanduser(work_dir)
|
cwd=os.path.expanduser(work_dir),
|
||||||
|
# L9:POSIX 下 shell 成为会话首进程,超时/取消可按进程组
|
||||||
|
# 一次终止整棵进程树(含 & 后台孙进程),消除管道悬挂
|
||||||
|
start_new_session=(os.name == "posix")
|
||||||
)
|
)
|
||||||
self.process = process
|
self.process = process
|
||||||
# 保存进程引用以便cancel()可以kill它
|
# T5:注册本执行专属槽位,终止/取消只作用于自己的 (proc, pgid)
|
||||||
with self._process_lock:
|
token, pgid = self._register_proc(process)
|
||||||
self._current_process = process
|
|
||||||
|
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
output_lines = []
|
# H2 修复:复用线程化读取(reader 线程 + 超时轮询),行经 output_queue 流式转发
|
||||||
total_bytes = 0
|
output_str, timeout_occurred, truncated = self._read_with_limit(
|
||||||
timeout_occurred = False
|
process, timeout, start_time, output_queue=output_queue, pgid=pgid)
|
||||||
truncated = False
|
|
||||||
|
|
||||||
# 实时读取输出
|
if timeout_occurred:
|
||||||
while True:
|
output_queue.put(("line", f"\n⏰ 命令执行超时({timeout}秒)\n"))
|
||||||
# 检查超时
|
elif truncated:
|
||||||
if timeout > 0 and time.time() - start_time > timeout:
|
output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断\n"))
|
||||||
process.terminate()
|
|
||||||
timeout_occurred = True
|
|
||||||
output_queue.put(("line", f"\n⏰ 命令执行超时({timeout}秒)\n"))
|
|
||||||
break
|
|
||||||
|
|
||||||
# 检查输出大小
|
|
||||||
if total_bytes >= self.MAX_OUTPUT_BYTES:
|
|
||||||
truncated = True
|
|
||||||
process.terminate()
|
|
||||||
output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断\n"))
|
|
||||||
break
|
|
||||||
|
|
||||||
# 读取一行输出
|
|
||||||
line = process.stdout.readline()
|
|
||||||
if not line and process.poll() is not None:
|
|
||||||
break
|
|
||||||
if line:
|
|
||||||
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:
|
|
||||||
truncated_line = line.encode('utf-8', errors='replace')[:allowed].decode('utf-8', errors='replace')
|
|
||||||
output_queue.put(("line", truncated_line))
|
|
||||||
output_lines.append(truncated_line)
|
|
||||||
total_bytes = self.MAX_OUTPUT_BYTES
|
|
||||||
truncated = True
|
|
||||||
process.terminate()
|
|
||||||
output_queue.put(("line", f"\n📦 输出超过 {self.MAX_OUTPUT_BYTES/1024/1024:.0f}MB 限制,结果已截断\n"))
|
|
||||||
break
|
|
||||||
|
|
||||||
stripped = line.rstrip('\n')
|
|
||||||
output_queue.put(("line", stripped + "\n"))
|
|
||||||
output_lines.append(line)
|
|
||||||
total_bytes += line_bytes
|
|
||||||
|
|
||||||
# 如果超时或截断,确保进程被终止
|
# 如果超时或截断,确保进程被终止
|
||||||
if (timeout_occurred or truncated) and process.poll() is None:
|
if (timeout_occurred or truncated) and process.poll() is None:
|
||||||
process.kill()
|
self._terminate_tree(process, force=True, pgid=pgid)
|
||||||
try:
|
try:
|
||||||
process.wait(timeout=2)
|
process.wait(timeout=2)
|
||||||
except:
|
except:
|
||||||
@@ -317,10 +479,8 @@ class SubprocessCommandExecutor(CommandExecutor):
|
|||||||
else:
|
else:
|
||||||
output_queue.put(("line", f"✅ 命令执行完成,返回码: {return_code}\n"))
|
output_queue.put(("line", f"✅ 命令执行完成,返回码: {return_code}\n"))
|
||||||
|
|
||||||
# 准备完整结果
|
|
||||||
full_output = "".join(output_lines)
|
|
||||||
data = {
|
data = {
|
||||||
"输出": full_output.strip() if full_output.strip() else "(无输出)",
|
"输出": output_str.strip() if output_str.strip() else "(无输出)",
|
||||||
"返回码": return_code
|
"返回码": return_code
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -339,11 +499,12 @@ class SubprocessCommandExecutor(CommandExecutor):
|
|||||||
output_queue.put(("complete", f"❌ 执行错误: {str(e)}"))
|
output_queue.put(("complete", f"❌ 执行错误: {str(e)}"))
|
||||||
finally:
|
finally:
|
||||||
if process and process.poll() is None:
|
if process and process.poll() is None:
|
||||||
process.terminate()
|
self._terminate_tree(process, pgid=pgid)
|
||||||
try:
|
try:
|
||||||
process.wait(timeout=2)
|
process.wait(timeout=2)
|
||||||
except:
|
except:
|
||||||
process.kill()
|
self._terminate_tree(process, force=True, pgid=pgid)
|
||||||
self.process = None
|
self.process = None
|
||||||
with self._process_lock:
|
# T5:仅注销自己的槽位,不再无差别清空共享引用
|
||||||
self._current_process = None
|
if token is not None:
|
||||||
|
self._unregister_proc(token)
|
||||||
@@ -0,0 +1,779 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
subsystems.process_manager - 后台进程管理器
|
||||||
|
|
||||||
|
让智能体能够启动并管理长驻后台进程(dev server / 训练任务 / watch 等),
|
||||||
|
弥补 run_command 同步阻塞、超时即杀的不足:
|
||||||
|
- bg 语义:start() 立即返回句柄,进程独立于对话轮次存活
|
||||||
|
- 有界环形输出缓冲:每进程上限 buffer_per_process(默认 256KB),
|
||||||
|
解决 nohup+日志文件方案输出无界增长的问题
|
||||||
|
- 注册表:自增句柄跨会话/跨子代理可见(全局单例)
|
||||||
|
- 生命周期:惰性状态收割(reader 线程 EOF 即 poll 收割防僵尸)+
|
||||||
|
TTL 惰性清理 + 退出统一收尾(非 detach 进程默认全杀)
|
||||||
|
|
||||||
|
安全边界(对齐 SubprocessCommandExecutor 既有惯例):
|
||||||
|
- 输出单行截断 16KB,防超长单行撑爆内存
|
||||||
|
- 运行中进程数上限(默认 16)+ 总条目上限(默认 32,已退出条目
|
||||||
|
保留尾部输出供验尸,LRU 淘汰)
|
||||||
|
- POSIX start_new_session + 进程组信号(L9 语义:孙进程不悬挂);
|
||||||
|
Windows 无进程组概念,回退单进程 terminate/kill
|
||||||
|
- 中断语义:用户 Ctrl+C / 取消事件不影响后台进程(后台的本意就是
|
||||||
|
脱离本轮对话生命周期);aibox 进程退出时经 atexit 统一收尾,
|
||||||
|
detach=true 的进程豁免
|
||||||
|
|
||||||
|
依赖:common (ConfigManager, constants, Logger) + 标准库
|
||||||
|
不反向依赖 aibox 任何符号。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import atexit
|
||||||
|
import collections
|
||||||
|
import os
|
||||||
|
import signal
|
||||||
|
import subprocess
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
|
from common import ConfigManager, constants, Logger
|
||||||
|
from subsystems.executors import filter_env_overrides
|
||||||
|
|
||||||
|
# 单行输出最大字节数,超出截断(防超长单行撑爆环形缓冲内存预算)
|
||||||
|
MAX_LINE_BYTES = 16 * 1024
|
||||||
|
|
||||||
|
# 单次 bg_output 拉取的硬上限(对齐执行器 1MB 输出限制)
|
||||||
|
MAX_FETCH_BYTES = 1 * 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
|
class _ProcEntry:
|
||||||
|
"""单个后台进程的注册表条目(内部类)"""
|
||||||
|
|
||||||
|
__slots__ = (
|
||||||
|
"handle", "name", "command", "popen", "pgid", "started_at",
|
||||||
|
"detach", "max_lifetime", "owner", "status", "exit_code", "end_time",
|
||||||
|
"lines", "buffer_bytes", "max_buffer", "next_seq", "evicted",
|
||||||
|
"lock", "reader_done", "reader", "stdin_lock", "log_fh",
|
||||||
|
)
|
||||||
|
|
||||||
|
# 状态取值:running / exited(自然退出)/ killed(被管理器终止)/ expired(TTL 到期)
|
||||||
|
def __init__(self, handle: int, name: str, command: str, popen, pgid,
|
||||||
|
detach: bool, max_lifetime: float, buffer_bytes: int,
|
||||||
|
owner: Optional[str] = None):
|
||||||
|
self.handle = handle
|
||||||
|
self.name = name
|
||||||
|
self.command = command
|
||||||
|
self.popen = popen
|
||||||
|
self.pgid = pgid
|
||||||
|
self.started_at = time.time()
|
||||||
|
self.detach = detach
|
||||||
|
self.max_lifetime = max_lifetime
|
||||||
|
self.owner = owner
|
||||||
|
self.status = "running"
|
||||||
|
self.exit_code = None
|
||||||
|
self.end_time = None
|
||||||
|
# 环形输出缓冲:deque of (seq, line, nbytes),seq 单调递增不回收,
|
||||||
|
# 淘汰最旧行仍保持 seq 连续性——bg_output 的 since 游标据此续读
|
||||||
|
self.lines = collections.deque()
|
||||||
|
self.buffer_bytes = 0
|
||||||
|
self.max_buffer = max(1024, int(buffer_bytes))
|
||||||
|
self.next_seq = 1
|
||||||
|
self.evicted = 0
|
||||||
|
self.lock = threading.RLock() # T8:可重入,write_input 持锁期间可调用 _record_line
|
||||||
|
self.reader_done = threading.Event()
|
||||||
|
self.reader = None
|
||||||
|
# 交互写入:单写锁保序(bg_write 串行化,防交错输入)
|
||||||
|
self.stdin_lock = threading.Lock()
|
||||||
|
# 可选全档镜像文件句柄(log_file=True 时由 reader 线程 tee 写入)
|
||||||
|
self.log_fh = None
|
||||||
|
|
||||||
|
|
||||||
|
class ProcessManager:
|
||||||
|
"""后台进程管理器:启动 / 列举 / 增量拉取输出 / 终止 / 退出收尾"""
|
||||||
|
|
||||||
|
def __init__(self, logger: Optional[Logger] = None,
|
||||||
|
max_processes: int = None, max_entries: int = None,
|
||||||
|
buffer_per_process: int = None, max_lifetime: float = None):
|
||||||
|
self.logger = logger or Logger("ProcessManager")
|
||||||
|
self.config = ConfigManager()
|
||||||
|
self.max_processes = int(max_processes if max_processes is not None
|
||||||
|
else constants.BG_MAX_PROCESSES)
|
||||||
|
self.max_entries = int(max_entries if max_entries is not None
|
||||||
|
else constants.BG_MAX_ENTRIES)
|
||||||
|
self.buffer_per_process = int(buffer_per_process if buffer_per_process is not None
|
||||||
|
else constants.BG_BUFFER_PER_PROCESS)
|
||||||
|
self.max_lifetime = float(max_lifetime if max_lifetime is not None
|
||||||
|
else constants.BG_MAX_LIFETIME)
|
||||||
|
self.max_write_bytes = int(constants.BG_MAX_WRITE_BYTES)
|
||||||
|
self._procs: Dict[int, _ProcEntry] = {}
|
||||||
|
self._next_id = 1
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
self._shutdown_done = False
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 启动
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def start(self, command: str, name: str = None, work_dir: str = None,
|
||||||
|
env: Dict = None, detach: bool = False,
|
||||||
|
log_file: bool = False, owner: str = None) -> Dict:
|
||||||
|
"""后台启动命令,立即返回。成功返回 {ok, handle, pid, name, log_file},
|
||||||
|
失败返回 {ok: False, error}。
|
||||||
|
|
||||||
|
stdin 为管道:后续可经 write() 交互写入(行协议 REPL 问答)。
|
||||||
|
log_file=True 时输出同步 tee 到 workspace/bg_<句柄>.log(无界
|
||||||
|
全档,可被原生工具 tail/cat 读取;环形缓冲仍为 bg_output 视图)。"""
|
||||||
|
if not command or not isinstance(command, str) or not command.strip():
|
||||||
|
return {"ok": False, "error": "缺少或空的 command"}
|
||||||
|
|
||||||
|
# 名称规整(M5 惯例):标签仅用于展示回显,拒绝控制字符防终端/前端注入
|
||||||
|
if name is not None:
|
||||||
|
name = str(name).strip()
|
||||||
|
if not name.isprintable() or len(name) > 64:
|
||||||
|
return {"ok": False, "error": "name 含控制字符或超过 64 字符"}
|
||||||
|
if not name:
|
||||||
|
name = None
|
||||||
|
if name is None:
|
||||||
|
name = command.strip()[:40]
|
||||||
|
|
||||||
|
if work_dir is None or work_dir == ".":
|
||||||
|
work_dir = constants.WORKSPACE_DIR
|
||||||
|
|
||||||
|
with self._lock:
|
||||||
|
if self._shutdown_done:
|
||||||
|
return {"ok": False, "error": "管理器已关闭"}
|
||||||
|
# TTL 惰性清理 + 状态收割
|
||||||
|
self._sweep_locked()
|
||||||
|
running = [e for e in self._procs.values() if e.status == "running"]
|
||||||
|
if len(running) >= self.max_processes:
|
||||||
|
return {"ok": False,
|
||||||
|
"error": f"运行中后台进程已达上限({self.max_processes}),"
|
||||||
|
f"请先用 bg_kill 清理"}
|
||||||
|
# 总条目上限:淘汰最老的已退出条目(保留其信息的价值最低)
|
||||||
|
if len(self._procs) >= self.max_entries:
|
||||||
|
exited = sorted(
|
||||||
|
(e for e in self._procs.values() if e.status != "running"),
|
||||||
|
key=lambda e: e.started_at)
|
||||||
|
while self._procs and len(self._procs) >= self.max_entries and exited:
|
||||||
|
victim = exited.pop(0)
|
||||||
|
self._procs.pop(victim.handle, None)
|
||||||
|
if len(self._procs) >= self.max_entries:
|
||||||
|
return {"ok": False,
|
||||||
|
"error": f"后台进程条目已达上限({self.max_entries})且全部运行中"}
|
||||||
|
handle = self._next_id
|
||||||
|
self._next_id += 1
|
||||||
|
|
||||||
|
exec_env = os.environ.copy()
|
||||||
|
if env:
|
||||||
|
# T21:黑名单键拒绝(对齐 executors,防 LD_PRELOAD/PYTHONPATH 劫持)
|
||||||
|
safe_env, rejected_keys = filter_env_overrides(env)
|
||||||
|
if rejected_keys:
|
||||||
|
self.logger.warning(f"env 覆盖被黑名单拒绝(T21): {rejected_keys}")
|
||||||
|
exec_env.update(safe_env)
|
||||||
|
|
||||||
|
try:
|
||||||
|
popen = subprocess.Popen(
|
||||||
|
command,
|
||||||
|
shell=True,
|
||||||
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=subprocess.STDOUT,
|
||||||
|
# 交互式 I/O:stdin 挂管道供 write() 写入。与 H3 的差异:
|
||||||
|
# 后台进程由管理器显式管理生命周期,不存在"等待输入的
|
||||||
|
# 命令永久挂起无人知晓"问题——不读 stdin 的进程不受影响,
|
||||||
|
# 读 stdin 的进程获得交互能力,kill 路径三级优雅终止兜底。
|
||||||
|
stdin=subprocess.PIPE,
|
||||||
|
text=True,
|
||||||
|
bufsize=1,
|
||||||
|
universal_newlines=True,
|
||||||
|
encoding="utf-8", # T17:显式 UTF-8,避免 Windows cp936 解码 UTF-8 输出致 reader 静默退出
|
||||||
|
errors="replace",
|
||||||
|
env=exec_env,
|
||||||
|
cwd=os.path.expanduser(work_dir),
|
||||||
|
# 对齐 L9:POSIX 下 shell 成为会话首进程,kill 按进程组
|
||||||
|
# 一次覆盖整棵进程树(含 & 后台孙进程)
|
||||||
|
start_new_session=(os.name == "posix"),
|
||||||
|
)
|
||||||
|
except FileNotFoundError as e:
|
||||||
|
return {"ok": False, "error": f"命令未找到: {e}"}
|
||||||
|
except PermissionError as e:
|
||||||
|
return {"ok": False, "error": f"权限不足: {e}"}
|
||||||
|
except Exception as e:
|
||||||
|
return {"ok": False, "error": f"启动失败: {e}"}
|
||||||
|
|
||||||
|
pgid = None
|
||||||
|
if os.name == "posix":
|
||||||
|
try:
|
||||||
|
pgid = os.getpgid(popen.pid)
|
||||||
|
except Exception:
|
||||||
|
pgid = None
|
||||||
|
|
||||||
|
entry = _ProcEntry(
|
||||||
|
handle, name, command.strip(), popen, pgid,
|
||||||
|
detach=bool(detach), max_lifetime=self.max_lifetime,
|
||||||
|
buffer_bytes=self.buffer_per_process, owner=owner,
|
||||||
|
)
|
||||||
|
# 可选全档镜像(失败降级为仅环形缓冲,不阻断启动)
|
||||||
|
log_path = None
|
||||||
|
if log_file:
|
||||||
|
try:
|
||||||
|
log_dir = constants.WORKSPACE_DIR
|
||||||
|
os.makedirs(log_dir, exist_ok=True)
|
||||||
|
log_path = os.path.join(log_dir, f"bg_{handle}.log")
|
||||||
|
entry.log_fh = open(log_path, "a", encoding="utf-8",
|
||||||
|
errors="replace", buffering=1)
|
||||||
|
except Exception as e:
|
||||||
|
entry.log_fh = None
|
||||||
|
log_path = None
|
||||||
|
self.logger.warning(f"输出镜像文件开启失败,仅保留环形缓冲: {e}")
|
||||||
|
with self._lock:
|
||||||
|
# start() 竞态兜底:期间并发 shutdown 则立即回收刚拉起的进程
|
||||||
|
if self._shutdown_done:
|
||||||
|
self._close_stdin(entry)
|
||||||
|
self._terminate_tree(entry, force=True)
|
||||||
|
try:
|
||||||
|
popen.wait(timeout=2)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return {"ok": False, "error": "管理器已关闭"}
|
||||||
|
self._procs[handle] = entry
|
||||||
|
|
||||||
|
entry.reader = threading.Thread(
|
||||||
|
target=self._reader, args=(entry,), daemon=True,
|
||||||
|
name=f"bg-proc-{handle}-reader")
|
||||||
|
entry.reader.start()
|
||||||
|
self.logger.info(f"后台进程已启动: #{handle} pid={popen.pid} "
|
||||||
|
f"detach={entry.detach} cmd={entry.command[:60]}")
|
||||||
|
return {"ok": True, "handle": handle, "pid": popen.pid, "name": name,
|
||||||
|
"detach": entry.detach, "log_file": log_path}
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 输出读取
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def get_output(self, handle, tail: int = None, since: int = None,
|
||||||
|
max_bytes: int = 65536) -> Dict:
|
||||||
|
"""拉取后台进程输出。
|
||||||
|
|
||||||
|
- tail=N:取最近 N 行(缺省 50)
|
||||||
|
- since=S:取行序号 > S 的全部缓冲行(增量续读,游标为上次
|
||||||
|
返回的 next_cursor;被环形淘汰的早期行以 evicted 提示)
|
||||||
|
- max_bytes:本次拉取字节上限(缺省 64KB,硬顶 1MB)
|
||||||
|
"""
|
||||||
|
entry = self._lookup(handle)
|
||||||
|
if entry is None:
|
||||||
|
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||||
|
|
||||||
|
try:
|
||||||
|
max_bytes = int(max_bytes or 65536)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
max_bytes = 65536
|
||||||
|
max_bytes = max(1024, min(max_bytes, MAX_FETCH_BYTES))
|
||||||
|
|
||||||
|
with entry.lock:
|
||||||
|
snapshot = list(entry.lines)
|
||||||
|
evicted = entry.evicted
|
||||||
|
next_seq = entry.next_seq
|
||||||
|
|
||||||
|
selected = []
|
||||||
|
evicted_note = 0
|
||||||
|
if since is not None:
|
||||||
|
try:
|
||||||
|
since = int(since)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return {"ok": False, "error": "since 游标必须是整数"}
|
||||||
|
selected = [(s, l) for (s, l, _) in snapshot if s > since]
|
||||||
|
# 请求范围早于环形缓冲起点:被淘汰部分如实告知
|
||||||
|
oldest = snapshot[0][0] if snapshot else next_seq
|
||||||
|
if since + 1 < oldest:
|
||||||
|
evicted_note = oldest - since - 1
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
tail = int(tail) if tail is not None else 50
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
tail = 50
|
||||||
|
tail = max(1, min(tail, 100000))
|
||||||
|
selected = [(s, l) for (s, l, _) in snapshot[-tail:]]
|
||||||
|
|
||||||
|
# 字节上限截断(保留最新行优先:从尾部回溯收集)
|
||||||
|
kept = []
|
||||||
|
total = 0
|
||||||
|
truncated = False
|
||||||
|
for s, line in reversed(selected):
|
||||||
|
nb = len(line.encode("utf-8", errors="replace"))
|
||||||
|
if total + nb > max_bytes:
|
||||||
|
truncated = True
|
||||||
|
break
|
||||||
|
kept.append((s, line))
|
||||||
|
total += nb
|
||||||
|
kept.reverse()
|
||||||
|
|
||||||
|
output = "".join(line for _, line in kept)
|
||||||
|
result = {
|
||||||
|
"ok": True,
|
||||||
|
"handle": entry.handle,
|
||||||
|
"name": entry.name,
|
||||||
|
"status": entry.status,
|
||||||
|
"exit_code": entry.exit_code,
|
||||||
|
"output": output,
|
||||||
|
"lines": len(kept),
|
||||||
|
"next_cursor": kept[-1][0] if kept else max(since or 0, next_seq - 1),
|
||||||
|
"truncated": truncated,
|
||||||
|
}
|
||||||
|
if evicted_note > 0 or (since is None and evicted > 0):
|
||||||
|
result["evicted"] = evicted_note if since is not None else evicted
|
||||||
|
return result
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 交互写入(stdin 管道 + watermark 对齐)
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def write(self, handle, text: str, append_newline: bool = True,
|
||||||
|
wait_lines: int = 0, wait_timeout: float = 5.0,
|
||||||
|
write_timeout: float = 5.0, owner: str = None) -> Dict:
|
||||||
|
"""向后台进程 stdin 写入一行(交互式问答)。
|
||||||
|
|
||||||
|
返回 {ok, handle, watermark, sent_bytes[, wait]}:
|
||||||
|
- watermark:本条输入的 [stdin] 回显行序号;bg_output(since=
|
||||||
|
watermark) 得到的必然且仅是本条输入引发的响应(写读对齐)
|
||||||
|
- wait_lines>0 时同步等待 N 行响应(wait 子 dict 含 output/
|
||||||
|
timed_out),expect 式一步问答
|
||||||
|
失败返回 {ok: False, error}(进程已退出/写入超时/超长等)。
|
||||||
|
|
||||||
|
T22 属主隔离:条目与调用方都携带 owner 且不一致时拒绝——他人
|
||||||
|
stdin 不容注入。任一侧未标识(None,如测试/管理面直调)不隔离。"""
|
||||||
|
entry = self._lookup(handle)
|
||||||
|
if entry is None:
|
||||||
|
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||||
|
denial = self._check_owner(entry, owner, "写入")
|
||||||
|
if denial:
|
||||||
|
return denial
|
||||||
|
if not isinstance(text, str) or text == "":
|
||||||
|
return {"ok": False, "error": "缺少或空的 text"}
|
||||||
|
data = text if (not append_newline or text.endswith("\n")) else text + "\n"
|
||||||
|
payload = data.encode("utf-8")
|
||||||
|
if len(payload) > self.max_write_bytes:
|
||||||
|
return {"ok": False,
|
||||||
|
"error": f"单次写入超过 {self.max_write_bytes} 字节上限(超长内容建议落文件后让进程从文件读)"}
|
||||||
|
|
||||||
|
with self._lock:
|
||||||
|
self._update_status_locked(entry)
|
||||||
|
if entry.status != "running":
|
||||||
|
return {"ok": False, "error": f"进程已结束({entry.status}),无法写入"}
|
||||||
|
|
||||||
|
# watermark:取号与回显行记账必须在同一临界区,保证 watermark 等于
|
||||||
|
# 回显行真实序号,bg_output(since=watermark) 正确对齐本条响应(T8)
|
||||||
|
echo = "[stdin] " + text.replace("\n", "⏎")
|
||||||
|
with entry.lock:
|
||||||
|
watermark = entry.next_seq
|
||||||
|
# 回显行进环形缓冲 + log 镜像(多行输入折叠为单行展示)
|
||||||
|
self._record_line(entry, echo)
|
||||||
|
|
||||||
|
# 写入放 watchdog 线程:进程不读 stdin 且管道缓冲(~64KB)写满时
|
||||||
|
# write() 会阻塞,不能拖死工具线程;超时后 daemon 线程由系统兜底
|
||||||
|
error_box = {}
|
||||||
|
done = threading.Event()
|
||||||
|
|
||||||
|
def _do_write():
|
||||||
|
try:
|
||||||
|
with entry.stdin_lock:
|
||||||
|
stdin = entry.popen.stdin
|
||||||
|
if stdin is None or stdin.closed:
|
||||||
|
raise BrokenPipeError("stdin 已关闭")
|
||||||
|
stdin.write(data)
|
||||||
|
stdin.flush()
|
||||||
|
except Exception as e:
|
||||||
|
error_box["e"] = e
|
||||||
|
finally:
|
||||||
|
done.set()
|
||||||
|
|
||||||
|
t = threading.Thread(target=_do_write, daemon=True,
|
||||||
|
name=f"bg-proc-{entry.handle}-writer")
|
||||||
|
t.start()
|
||||||
|
try:
|
||||||
|
write_timeout = max(0.1, min(60.0, float(write_timeout)))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
write_timeout = 5.0
|
||||||
|
if not done.wait(timeout=write_timeout):
|
||||||
|
return {"ok": False,
|
||||||
|
"error": f"写入超时(>{write_timeout}s):进程可能不读取 stdin,管道缓冲已满"}
|
||||||
|
if error_box:
|
||||||
|
with self._lock:
|
||||||
|
self._update_status_locked(entry)
|
||||||
|
return {"ok": False,
|
||||||
|
"error": f"写入失败: {error_box['e']}"}
|
||||||
|
|
||||||
|
result = {"ok": True, "handle": entry.handle, "watermark": watermark,
|
||||||
|
"sent_bytes": len(payload)}
|
||||||
|
if wait_lines and int(wait_lines) > 0:
|
||||||
|
result["wait"] = self._wait_lines(entry, watermark,
|
||||||
|
int(wait_lines), wait_timeout)
|
||||||
|
return result
|
||||||
|
|
||||||
|
def _wait_lines(self, entry: _ProcEntry, watermark: int, count: int,
|
||||||
|
timeout: float) -> Dict:
|
||||||
|
"""等待 watermark 之后出现 count 行新输出(进程退出且 reader 排空
|
||||||
|
即提前返回,如实标记 timed_out)。"""
|
||||||
|
try:
|
||||||
|
timeout = max(0.1, min(60.0, float(timeout)))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
timeout = 5.0
|
||||||
|
deadline = time.time() + timeout
|
||||||
|
|
||||||
|
def _snap():
|
||||||
|
with entry.lock:
|
||||||
|
return [(s, l) for (s, l, _) in entry.lines if s > watermark]
|
||||||
|
|
||||||
|
lines = _snap()
|
||||||
|
timed_out = False
|
||||||
|
while len(lines) < count:
|
||||||
|
if entry.popen.poll() is not None and entry.reader_done.is_set():
|
||||||
|
time.sleep(0.1) # 排空 reader 尾部在途行后取终值
|
||||||
|
lines = _snap()
|
||||||
|
break
|
||||||
|
if time.time() > deadline:
|
||||||
|
timed_out = True
|
||||||
|
break
|
||||||
|
time.sleep(0.05)
|
||||||
|
lines = _snap()
|
||||||
|
|
||||||
|
# 字节上限(复用单次拉取硬顶)
|
||||||
|
kept, total, truncated = [], 0, False
|
||||||
|
for s, line in lines:
|
||||||
|
nb = len(line.encode("utf-8", errors="replace"))
|
||||||
|
if total + nb > MAX_FETCH_BYTES:
|
||||||
|
truncated = True
|
||||||
|
break
|
||||||
|
kept.append((s, line))
|
||||||
|
total += nb
|
||||||
|
out = {
|
||||||
|
"lines": len(kept),
|
||||||
|
"output": "".join(l for _, l in kept),
|
||||||
|
"timed_out": timed_out,
|
||||||
|
"next_cursor": kept[-1][0] if kept else watermark,
|
||||||
|
}
|
||||||
|
if truncated:
|
||||||
|
out["truncated"] = True
|
||||||
|
return out
|
||||||
|
|
||||||
|
def _close_stdin(self, entry: _ProcEntry):
|
||||||
|
"""关闭 stdin 管道(行协议进程见 EOF 常自行退出;幂等)"""
|
||||||
|
try:
|
||||||
|
with entry.stdin_lock:
|
||||||
|
stdin = entry.popen.stdin
|
||||||
|
if stdin is not None and not stdin.closed:
|
||||||
|
stdin.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 列举 / 终止 / 收尾
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def list_procs(self) -> List[Dict]:
|
||||||
|
"""列出全部条目(运行中在前,其余按句柄序)"""
|
||||||
|
with self._lock:
|
||||||
|
self._sweep_locked()
|
||||||
|
# 惰性收割:列举即刷新全部条目状态(此前仅 kill/get_output 路径更新)
|
||||||
|
for entry in self._procs.values():
|
||||||
|
self._update_status_locked(entry)
|
||||||
|
entries = sorted(self._procs.values(),
|
||||||
|
key=lambda e: (e.status != "running", e.handle))
|
||||||
|
return [self._entry_info(e) for e in entries]
|
||||||
|
|
||||||
|
def kill(self, handle, force: bool = False, grace: float = 3.0,
|
||||||
|
owner: str = None) -> Dict:
|
||||||
|
"""终止后台进程。SIGTERM → 宽限轮询 → SIGKILL;force=True 直接 SIGKILL。
|
||||||
|
对已退出条目幂等,返回当前状态。
|
||||||
|
|
||||||
|
T22 属主隔离:条目与调用方都携带 owner 且不一致时拒绝——防止
|
||||||
|
提示注入的会话/任务跨属主杀掉他人的 dev server。shutdown/TTL
|
||||||
|
清理属管理器特权,不走本校验。"""
|
||||||
|
entry = self._lookup(handle)
|
||||||
|
if entry is None:
|
||||||
|
return {"ok": False, "error": f"句柄 #{handle} 不存在"}
|
||||||
|
denial = self._check_owner(entry, owner, "终止")
|
||||||
|
if denial:
|
||||||
|
return denial
|
||||||
|
|
||||||
|
with self._lock:
|
||||||
|
self._update_status_locked(entry)
|
||||||
|
if entry.status != "running":
|
||||||
|
return {"ok": True, "handle": entry.handle, "status": entry.status,
|
||||||
|
"exit_code": entry.exit_code, "already_finished": True}
|
||||||
|
|
||||||
|
try:
|
||||||
|
grace = min(30.0, max(0.1, float(grace)))
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
grace = 3.0
|
||||||
|
|
||||||
|
if force:
|
||||||
|
self._terminate_tree(entry, force=True)
|
||||||
|
else:
|
||||||
|
# 三级优雅终止:关 stdin(行协议进程见 EOF 常自行退出)→
|
||||||
|
# SIGTERM 宽限 → SIGKILL 兜底
|
||||||
|
self._close_stdin(entry)
|
||||||
|
deadline = time.time() + grace * 0.4
|
||||||
|
while time.time() < deadline and entry.popen.poll() is None:
|
||||||
|
time.sleep(0.05)
|
||||||
|
if entry.popen.poll() is None:
|
||||||
|
self._terminate_tree(entry)
|
||||||
|
deadline = time.time() + grace * 0.6
|
||||||
|
while time.time() < deadline and entry.popen.poll() is None:
|
||||||
|
time.sleep(0.05)
|
||||||
|
if entry.popen.poll() is None:
|
||||||
|
self._terminate_tree(entry, force=True)
|
||||||
|
try:
|
||||||
|
entry.popen.wait(timeout=2)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
with self._lock:
|
||||||
|
self._update_status_locked(entry)
|
||||||
|
if entry.status == "running":
|
||||||
|
# 进程仍未死(极端:不可杀的 D 状态),如实报告
|
||||||
|
return {"ok": False, "handle": entry.handle,
|
||||||
|
"error": "进程未能终止(可能处于不可中断状态)"}
|
||||||
|
entry.status = "killed"
|
||||||
|
return {"ok": True, "handle": entry.handle, "status": entry.status,
|
||||||
|
"exit_code": entry.exit_code}
|
||||||
|
|
||||||
|
def shutdown(self) -> Dict:
|
||||||
|
"""退出收尾:终止全部非 detach 运行中进程(SIGTERM→SIGKILL),
|
||||||
|
清空注册表;detach 进程保留运行、仅注销登记。幂等。"""
|
||||||
|
with self._lock:
|
||||||
|
if self._shutdown_done:
|
||||||
|
return {"killed": 0, "detached": 0}
|
||||||
|
self._shutdown_done = True
|
||||||
|
entries = list(self._procs.values())
|
||||||
|
self._procs.clear()
|
||||||
|
|
||||||
|
killed = detached = 0
|
||||||
|
for entry in entries:
|
||||||
|
self._update_status(entry)
|
||||||
|
if entry.status != "running":
|
||||||
|
continue
|
||||||
|
if entry.detach:
|
||||||
|
detached += 1
|
||||||
|
self.logger.info(f"后台进程 #{entry.handle}(detach)保留运行: "
|
||||||
|
f"pid={entry.popen.pid}")
|
||||||
|
continue
|
||||||
|
# 优雅收尾:先关 stdin(行协议进程见 EOF 常自行退出)再走信号
|
||||||
|
self._close_stdin(entry)
|
||||||
|
try:
|
||||||
|
entry.popen.wait(timeout=1)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
if entry.popen.poll() is None:
|
||||||
|
self._terminate_tree(entry)
|
||||||
|
try:
|
||||||
|
entry.popen.wait(timeout=3)
|
||||||
|
except Exception:
|
||||||
|
self._terminate_tree(entry, force=True)
|
||||||
|
try:
|
||||||
|
entry.popen.wait(timeout=2)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
entry.status = "killed"
|
||||||
|
killed += 1
|
||||||
|
if killed or detached:
|
||||||
|
self.logger.info(f"后台进程收尾: 终止 {killed} 个, detach 保留 {detached} 个")
|
||||||
|
return {"killed": killed, "detached": detached}
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
# 内部实现
|
||||||
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
|
def _check_owner(self, entry: _ProcEntry, caller_owner: Optional[str],
|
||||||
|
action: str) -> Optional[Dict]:
|
||||||
|
"""T22:属主校验——条目与调用方都携带 owner 且不一致时返回拒绝
|
||||||
|
dict(写/杀权限隔离);任一侧未标识(None,如测试/管理面直调)
|
||||||
|
不隔离,保持向后兼容。观察面(list/get_output)不经过本校验。"""
|
||||||
|
if (entry.owner is not None and caller_owner is not None
|
||||||
|
and entry.owner != caller_owner):
|
||||||
|
self.logger.warning(
|
||||||
|
f"属主校验拒绝:owner={caller_owner} 尝试{action} "
|
||||||
|
f"#{entry.handle}(属主 {entry.owner})")
|
||||||
|
return {"ok": False,
|
||||||
|
"error": f"进程 #{entry.handle} 属于其他会话(属主不匹配),无权{action}"}
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _lookup(self, handle) -> Optional[_ProcEntry]:
|
||||||
|
try:
|
||||||
|
handle = int(handle)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return None
|
||||||
|
with self._lock:
|
||||||
|
entry = self._procs.get(handle)
|
||||||
|
if entry is not None:
|
||||||
|
self._update_status(entry)
|
||||||
|
return entry
|
||||||
|
|
||||||
|
def _update_status(self, entry: _ProcEntry):
|
||||||
|
with self._lock:
|
||||||
|
self._update_status_locked(entry)
|
||||||
|
|
||||||
|
def _update_status_locked(self, entry: _ProcEntry):
|
||||||
|
"""惰性收割:poll() 即 waitpid(WNOHANG),自然退出时记录状态。"""
|
||||||
|
if entry.status != "running":
|
||||||
|
return
|
||||||
|
rc = entry.popen.poll()
|
||||||
|
if rc is not None:
|
||||||
|
entry.exit_code = rc
|
||||||
|
entry.end_time = time.time()
|
||||||
|
entry.status = "exited"
|
||||||
|
|
||||||
|
def _sweep_locked(self):
|
||||||
|
"""TTL 惰性清理:超期运行中进程标记 expired 并终止(max_lifetime<=0 不启用)。
|
||||||
|
调用方须持有 self._lock。"""
|
||||||
|
if self.max_lifetime <= 0:
|
||||||
|
return
|
||||||
|
now = time.time()
|
||||||
|
for entry in list(self._procs.values()):
|
||||||
|
if entry.status == "running" and now - entry.started_at > self.max_lifetime:
|
||||||
|
self._terminate_tree(entry)
|
||||||
|
entry.status = "expired"
|
||||||
|
entry.end_time = now
|
||||||
|
entry.exit_code = entry.popen.poll()
|
||||||
|
self.logger.info(f"后台进程 #{entry.handle} TTL 到期已终止")
|
||||||
|
|
||||||
|
def _terminate_tree(self, entry: _ProcEntry, force: bool = False):
|
||||||
|
"""按进程组终止整棵进程树(对齐 executors L9/T11 语义与 Windows 回退)"""
|
||||||
|
proc = entry.popen
|
||||||
|
try:
|
||||||
|
if os.name == "posix":
|
||||||
|
sig = signal.SIGKILL if force else signal.SIGTERM
|
||||||
|
pgid = entry.pgid
|
||||||
|
if pgid is None:
|
||||||
|
pgid = os.getpgid(proc.pid)
|
||||||
|
os.killpg(pgid, sig)
|
||||||
|
else:
|
||||||
|
# T11(审计 2026-08-21):Windows 无进程组,单进程 kill 收割
|
||||||
|
# 不到 spawn 的孙进程(dev server/训练任务成孤儿继续占用端口
|
||||||
|
# /文件锁)。taskkill /T /F 递归终止整棵树,失败兜底单进程 kill
|
||||||
|
try:
|
||||||
|
subprocess.run(
|
||||||
|
["taskkill", "/PID", str(proc.pid), "/T", "/F"],
|
||||||
|
capture_output=True, timeout=10)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
if proc.poll() is None:
|
||||||
|
proc.kill()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
except Exception:
|
||||||
|
try:
|
||||||
|
proc.kill()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _reader(self, entry: _ProcEntry):
|
||||||
|
"""后台 reader 线程(对齐 executors H2 模式):阻塞 readline 入环形
|
||||||
|
缓冲(+ 可选 log 镜像);EOF/异常即 poll() 收割防僵尸,并显式关闭
|
||||||
|
管道与镜像文件句柄。"""
|
||||||
|
try:
|
||||||
|
for line in iter(entry.popen.stdout.readline, ""):
|
||||||
|
self._record_line(entry, line)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
entry.reader_done.set()
|
||||||
|
try:
|
||||||
|
entry.popen.poll() # 立即收割退出子进程,防僵尸悬挂
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
entry.popen.stdout.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
if entry.log_fh is not None:
|
||||||
|
try:
|
||||||
|
entry.log_fh.flush()
|
||||||
|
entry.log_fh.close()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def _record_line(self, entry: _ProcEntry, line: str):
|
||||||
|
"""统一记录:环形缓冲 + 可选 log 全档镜像(stdin 回显同样入档)"""
|
||||||
|
if entry.log_fh is not None:
|
||||||
|
try:
|
||||||
|
entry.log_fh.write(line)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
self._append_line(entry, line)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _append_line(entry: _ProcEntry, line: str):
|
||||||
|
# 单行截断:环形缓冲内存预算不被超长单行击穿
|
||||||
|
nbytes = len(line.encode("utf-8", errors="replace"))
|
||||||
|
if nbytes > MAX_LINE_BYTES:
|
||||||
|
allowed = line.encode("utf-8", errors="replace")[:MAX_LINE_BYTES]
|
||||||
|
line = allowed.decode("utf-8", errors="ignore") + "…[单行截断]\n"
|
||||||
|
nbytes = len(line.encode("utf-8", errors="replace"))
|
||||||
|
with entry.lock:
|
||||||
|
entry.lines.append((entry.next_seq, line, nbytes))
|
||||||
|
entry.next_seq += 1
|
||||||
|
entry.buffer_bytes += nbytes
|
||||||
|
while entry.buffer_bytes > entry.max_buffer and len(entry.lines) > 1:
|
||||||
|
_, old, old_nb = entry.lines.popleft()
|
||||||
|
entry.buffer_bytes -= old_nb
|
||||||
|
entry.evicted += 1
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _entry_info(entry: _ProcEntry) -> Dict:
|
||||||
|
now = time.time()
|
||||||
|
return {
|
||||||
|
"handle": entry.handle,
|
||||||
|
"name": entry.name,
|
||||||
|
"pid": entry.popen.pid,
|
||||||
|
"status": entry.status,
|
||||||
|
"exit_code": entry.exit_code,
|
||||||
|
"detach": entry.detach,
|
||||||
|
"started_at": entry.started_at,
|
||||||
|
"runtime_sec": round((entry.end_time or now) - entry.started_at, 1),
|
||||||
|
"buffered_lines": len(entry.lines),
|
||||||
|
"evicted_lines": entry.evicted,
|
||||||
|
"command": entry.command,
|
||||||
|
"owner": entry.owner, # T22:属主标识(观察面可见,写/杀需匹配)
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# ----------------------------------------------------------------------
|
||||||
|
# 模块级单例(工具层共享;测试经 reset_process_manager() 重置)
|
||||||
|
# ----------------------------------------------------------------------
|
||||||
|
|
||||||
|
_manager: Optional[ProcessManager] = None
|
||||||
|
_manager_lock = threading.Lock()
|
||||||
|
_atexit_registered = False
|
||||||
|
|
||||||
|
|
||||||
|
def get_process_manager(logger: Optional[Logger] = None) -> ProcessManager:
|
||||||
|
"""获取全局 ProcessManager 单例(懒加载)。
|
||||||
|
|
||||||
|
后台进程注册表必须跨会话/跨子代理共享——任何 worker 子代理启动的
|
||||||
|
进程,后续任意子代理都能 bg_list/bg_output/bg_kill 管理。"""
|
||||||
|
global _manager, _atexit_registered
|
||||||
|
with _manager_lock:
|
||||||
|
if _manager is None:
|
||||||
|
_manager = ProcessManager(logger)
|
||||||
|
if not _atexit_registered:
|
||||||
|
# 退出语义:aibox 进程退出(含 --once sys.exit / Ctrl+C)时
|
||||||
|
# 非 detach 后台进程统一终止;detach 进程保留运行仅注销登记。
|
||||||
|
# SIGKILL/SIGTERM 硬杀无法覆盖(与既有线程池等资源同级)。
|
||||||
|
atexit.register(_shutdown_atexit)
|
||||||
|
_atexit_registered = True
|
||||||
|
return _manager
|
||||||
|
|
||||||
|
|
||||||
|
def reset_process_manager(manager: Optional[ProcessManager] = None) -> None:
|
||||||
|
"""替换/清空全局单例(测试隔离用)。传入 None 仅解除引用,不触发收尾。"""
|
||||||
|
global _manager
|
||||||
|
with _manager_lock:
|
||||||
|
_manager = manager
|
||||||
|
|
||||||
|
|
||||||
|
def _shutdown_atexit():
|
||||||
|
m = _manager
|
||||||
|
if m is not None:
|
||||||
|
try:
|
||||||
|
m.shutdown()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
@@ -18,10 +18,12 @@ subsystems.prompt_manager - 提示词管理器
|
|||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import shutil
|
import shutil
|
||||||
|
import threading
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Dict, List, Optional
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
from common import ConfigManager
|
from common import ConfigManager
|
||||||
|
from common.atomic import atomic_write_text
|
||||||
|
|
||||||
|
|
||||||
class PromptManager:
|
class PromptManager:
|
||||||
@@ -35,6 +37,7 @@ class PromptManager:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
_instance = None
|
_instance = None
|
||||||
|
_instance_lock = threading.Lock() # I 批:并发首例化竞态防护
|
||||||
_prompts = []
|
_prompts = []
|
||||||
_prompts_dir = None
|
_prompts_dir = None
|
||||||
|
|
||||||
@@ -43,27 +46,43 @@ class PromptManager:
|
|||||||
STATUS_FILE = "100_STATUS.md"
|
STATUS_FILE = "100_STATUS.md"
|
||||||
SOUL_FILE = "99_SOUL.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)
|
# 基础技能文件名模式:N_xxx.md(N 为数字,如 1_memory_search.md、10_task_handling.md)
|
||||||
_SKILL_FILE_RE = re.compile(r'^(\d+)_')
|
_SKILL_FILE_RE = re.compile(r'^(\d+)_')
|
||||||
|
|
||||||
def __new__(cls):
|
def __new__(cls):
|
||||||
if cls._instance is None:
|
with cls._instance_lock:
|
||||||
cls._instance = super().__new__(cls)
|
if cls._instance is None:
|
||||||
|
cls._instance = super().__new__(cls)
|
||||||
return cls._instance
|
return cls._instance
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
if not self._prompts:
|
# J6(审计 2026-08-19):__new__ 的锁只防双实例,多线程首例并发
|
||||||
self._load_prompts()
|
# 同时进入 __init__ 会使 _load_prompts 双重执行,_prompts(类属性)
|
||||||
|
# 整份翻倍——初始化检查与加载收进同一把锁
|
||||||
|
with self._instance_lock:
|
||||||
|
if not self._prompts:
|
||||||
|
self._load_prompts()
|
||||||
|
|
||||||
def _get_prompts_dir(self) -> str:
|
def _get_prompts_dir(self) -> str:
|
||||||
config = ConfigManager()
|
config = ConfigManager()
|
||||||
prompts_dir = config.get("system.prompts_dir", "~/.aibox/prompts")
|
prompts_dir = config.get("system.prompts_dir", "~/.aibox/prompts")
|
||||||
return os.path.expanduser(prompts_dir)
|
return os.path.expanduser(prompts_dir)
|
||||||
|
|
||||||
def _load_prompts(self):
|
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()
|
self._prompts_dir = self._get_prompts_dir()
|
||||||
os.makedirs(self._prompts_dir, exist_ok=True)
|
os.makedirs(self._prompts_dir, exist_ok=True)
|
||||||
self._prompts = []
|
|
||||||
|
|
||||||
# 确保核心文件存在
|
# 确保核心文件存在
|
||||||
self._ensure_identity_file()
|
self._ensure_identity_file()
|
||||||
@@ -86,17 +105,14 @@ class PromptManager:
|
|||||||
with open(filepath, 'r', encoding='utf-8') as f:
|
with open(filepath, 'r', encoding='utf-8') as f:
|
||||||
content = f.read().strip()
|
content = f.read().strip()
|
||||||
if content:
|
if content:
|
||||||
self._prompts.append({
|
prompts.append({
|
||||||
'file': filename,
|
'file': filename,
|
||||||
'content': content
|
'content': content
|
||||||
})
|
})
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"⚠️ 读取提示词文件 {filename} 失败: {e}")
|
print(f"⚠️ 读取提示词文件 {filename} 失败: {e}")
|
||||||
|
|
||||||
essential_files = [self.IDENTITY_FILE, self.STATUS_FILE, self.SOUL_FILE]
|
if not prompts:
|
||||||
other_files = [f for f in prompt_files if f not in essential_files] if prompt_files else []
|
|
||||||
|
|
||||||
if not self._prompts:
|
|
||||||
self._create_default_prompts()
|
self._create_default_prompts()
|
||||||
|
|
||||||
def _get_template_dir(self) -> str:
|
def _get_template_dir(self) -> str:
|
||||||
@@ -148,6 +164,9 @@ class PromptManager:
|
|||||||
|
|
||||||
def update_identity(self, content: str) -> bool:
|
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)
|
identity_path = os.path.join(self._prompts_dir, self.IDENTITY_FILE)
|
||||||
try:
|
try:
|
||||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||||
@@ -158,8 +177,8 @@ class PromptManager:
|
|||||||
【更新记录】
|
【更新记录】
|
||||||
- 上次更新:{timestamp}
|
- 上次更新:{timestamp}
|
||||||
"""
|
"""
|
||||||
with open(identity_path, 'w', encoding='utf-8') as f:
|
# J5:原子写,崩溃不留半截提示词文件
|
||||||
f.write(updated_content)
|
atomic_write_text(identity_path, updated_content)
|
||||||
self.reload()
|
self.reload()
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -168,6 +187,9 @@ class PromptManager:
|
|||||||
|
|
||||||
def update_status(self, content: str) -> bool:
|
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)
|
status_path = os.path.join(self._prompts_dir, self.STATUS_FILE)
|
||||||
try:
|
try:
|
||||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||||
@@ -178,8 +200,8 @@ class PromptManager:
|
|||||||
【更新记录】
|
【更新记录】
|
||||||
- 上次更新:{timestamp}
|
- 上次更新:{timestamp}
|
||||||
"""
|
"""
|
||||||
with open(status_path, 'w', encoding='utf-8') as f:
|
# J5:原子写,崩溃不留半截提示词文件
|
||||||
f.write(updated_content)
|
atomic_write_text(status_path, updated_content)
|
||||||
self.reload()
|
self.reload()
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -188,6 +210,9 @@ class PromptManager:
|
|||||||
|
|
||||||
def update_soul(self, content: str) -> bool:
|
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)
|
soul_path = os.path.join(self._prompts_dir, self.SOUL_FILE)
|
||||||
try:
|
try:
|
||||||
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||||
@@ -198,8 +223,8 @@ class PromptManager:
|
|||||||
【更新记录】
|
【更新记录】
|
||||||
- 上次更新:{timestamp}
|
- 上次更新:{timestamp}
|
||||||
"""
|
"""
|
||||||
with open(soul_path, 'w', encoding='utf-8') as f:
|
# J5:原子写,崩溃不留半截提示词文件
|
||||||
f.write(updated_content)
|
atomic_write_text(soul_path, updated_content)
|
||||||
self.reload()
|
self.reload()
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -240,10 +265,6 @@ class PromptManager:
|
|||||||
return """我是DeepSeek,一个智能AI助手。
|
return """我是DeepSeek,一个智能AI助手。
|
||||||
请创建 99_SOUL.md 文件来定义我的完整身份、角色、性格和能力边界。"""
|
请创建 99_SOUL.md 文件来定义我的完整身份、角色、性格和能力边界。"""
|
||||||
|
|
||||||
def get_skill_prompts(self) -> List[Dict]:
|
|
||||||
"""获取所有基础技能提示词(以1_开头的文件)"""
|
|
||||||
return [p for p in self._prompts if p['file'].startswith(self.SKILL_PREFIX)]
|
|
||||||
|
|
||||||
def get_soul_prompt(self) -> Optional[Dict]:
|
def get_soul_prompt(self) -> Optional[Dict]:
|
||||||
"""获取灵魂提示词(99_SOUL.md)"""
|
"""获取灵魂提示词(99_SOUL.md)"""
|
||||||
for p in self._prompts:
|
for p in self._prompts:
|
||||||
@@ -296,9 +317,10 @@ class PromptManager:
|
|||||||
return [p for p in self._prompts if self._skill_sort_key(p['file']) is not None]
|
return [p for p in self._prompts if self._skill_sort_key(p['file']) is not None]
|
||||||
|
|
||||||
def reload(self):
|
def reload(self):
|
||||||
"""重新加载所有提示词"""
|
"""重新加载所有提示词(M6:构建新列表后原子替换,并发读者始终读到完整列表)"""
|
||||||
self._prompts = []
|
new_prompts: List[Dict] = []
|
||||||
self._load_prompts()
|
self._load_prompts(new_prompts)
|
||||||
|
self._prompts = new_prompts
|
||||||
return len(self._prompts)
|
return len(self._prompts)
|
||||||
|
|
||||||
def get_prompt_count(self) -> int:
|
def get_prompt_count(self) -> int:
|
||||||
|
|||||||
@@ -27,6 +27,9 @@ class SkillManager:
|
|||||||
2. Markdown 文件 (.md):第一行为简介(# 标题或普通文本),其余为技能内容
|
2. Markdown 文件 (.md):第一行为简介(# 标题或普通文本),其余为技能内容
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
# J16(审计 2026-08-19):--help 输出上限——内容会进系统提示词
|
||||||
|
MAX_HELP_LENGTH = 2000
|
||||||
|
|
||||||
def __init__(self, logger: Optional[Logger] = None):
|
def __init__(self, logger: Optional[Logger] = None):
|
||||||
self.logger = logger or Logger("SkillManager")
|
self.logger = logger or Logger("SkillManager")
|
||||||
# name -> {path, file, summary, content, is_markdown}
|
# name -> {path, file, summary, content, is_markdown}
|
||||||
@@ -38,7 +41,11 @@ class SkillManager:
|
|||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
def _run_help(self, skill_path: str) -> str:
|
def _run_help(self, skill_path: str) -> str:
|
||||||
"""执行 skill_path --help 并返回完整输出(stdout 优先,否则 stderr)。"""
|
"""执行 skill_path --help 并返回完整输出(stdout 优先,否则 stderr)。
|
||||||
|
|
||||||
|
J16(审计 2026-08-19):输出截断到 MAX_HELP_LENGTH——技能 --help
|
||||||
|
输出会进系统提示词,异常/恶意脚本可撑爆上下文。
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
result = subprocess.run(
|
result = subprocess.run(
|
||||||
[skill_path, "--help"],
|
[skill_path, "--help"],
|
||||||
@@ -50,6 +57,8 @@ class SkillManager:
|
|||||||
output = result.stdout.strip()
|
output = result.stdout.strip()
|
||||||
if not output:
|
if not output:
|
||||||
output = result.stderr.strip()
|
output = result.stderr.strip()
|
||||||
|
if len(output) > self.MAX_HELP_LENGTH:
|
||||||
|
output = output[:self.MAX_HELP_LENGTH] + "\n…(已截断)"
|
||||||
return output
|
return output
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.logger.warning(f"运行 {skill_path} --help 失败: {e}")
|
self.logger.warning(f"运行 {skill_path} --help 失败: {e}")
|
||||||
@@ -114,12 +123,18 @@ class SkillManager:
|
|||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
def _load_skills(self):
|
def _load_skills(self):
|
||||||
"""扫描技能目录,加载所有技能文件(可执行文件或 Markdown 文件)。"""
|
"""扫描技能目录,加载所有技能文件(可执行文件或 Markdown 文件)。
|
||||||
|
|
||||||
|
T26(审计 2026-08-21):本地构建 new_skills 完成后一次性原子替换,
|
||||||
|
对齐 PromptManager M6——此前先 self.skills = {} 清空再逐个重建,
|
||||||
|
reload 期间(每个可执行技能 --help 最多 10s)并发读者会看到空/半截
|
||||||
|
技能表(技能"消失"再"逐渐出现")。"""
|
||||||
skills_dir = constants.SKILLS_DIR
|
skills_dir = constants.SKILLS_DIR
|
||||||
self.skills = {}
|
new_skills = {}
|
||||||
|
|
||||||
if not os.path.exists(skills_dir):
|
if not os.path.exists(skills_dir):
|
||||||
self.logger.info(f"技能目录不存在: {skills_dir}")
|
self.logger.info(f"技能目录不存在: {skills_dir}")
|
||||||
|
self.skills = new_skills
|
||||||
return
|
return
|
||||||
|
|
||||||
skill_files = [
|
skill_files = [
|
||||||
@@ -129,16 +144,30 @@ class SkillManager:
|
|||||||
and not f.startswith('.')
|
and not f.startswith('.')
|
||||||
]
|
]
|
||||||
|
|
||||||
for skill_file in skill_files:
|
# J16(审计 2026-08-19):同名键冲突(foo 可执行文件与 foo.md)此前
|
||||||
|
# 由 listdir 顺序决定谁覆盖谁,行为不确定。改为可执行优先 + 确定性:
|
||||||
|
# 先加载可执行技能,Markdown 遇到已占用的同名键时跳过并告警
|
||||||
|
executable_files = [
|
||||||
|
f for f in skill_files
|
||||||
|
if not self._is_markdown_file(f)
|
||||||
|
]
|
||||||
|
markdown_files = [f for f in skill_files if self._is_markdown_file(f)]
|
||||||
|
|
||||||
|
for skill_file in executable_files + markdown_files:
|
||||||
skill_path = os.path.join(skills_dir, skill_file)
|
skill_path = os.path.join(skills_dir, skill_file)
|
||||||
skill_name = os.path.splitext(skill_file)[0]
|
skill_name = os.path.splitext(skill_file)[0]
|
||||||
is_markdown = self._is_markdown_file(skill_file)
|
is_markdown = self._is_markdown_file(skill_file)
|
||||||
|
|
||||||
if is_markdown:
|
if is_markdown:
|
||||||
|
if skill_name in new_skills:
|
||||||
|
self.logger.warning(
|
||||||
|
f"Markdown 技能 '{skill_file}' 与同名可执行技能冲突,已跳过"
|
||||||
|
f"(可执行技能优先)")
|
||||||
|
continue
|
||||||
# Markdown 格式技能:直接读取内容
|
# Markdown 格式技能:直接读取内容
|
||||||
summary, content = self._read_markdown_skill(skill_path)
|
summary, content = self._read_markdown_skill(skill_path)
|
||||||
if content:
|
if content:
|
||||||
self.skills[skill_name] = {
|
new_skills[skill_name] = {
|
||||||
"path": skill_path,
|
"path": skill_path,
|
||||||
"file": skill_file,
|
"file": skill_file,
|
||||||
"summary": summary,
|
"summary": summary,
|
||||||
@@ -159,7 +188,7 @@ class SkillManager:
|
|||||||
help_output = self._run_help(skill_path)
|
help_output = self._run_help(skill_path)
|
||||||
summary = self._first_nonempty_line(help_output) or f"扩展技能: {skill_name}"
|
summary = self._first_nonempty_line(help_output) or f"扩展技能: {skill_name}"
|
||||||
|
|
||||||
self.skills[skill_name] = {
|
new_skills[skill_name] = {
|
||||||
"path": skill_path,
|
"path": skill_path,
|
||||||
"file": skill_file,
|
"file": skill_file,
|
||||||
"summary": summary,
|
"summary": summary,
|
||||||
@@ -169,7 +198,8 @@ class SkillManager:
|
|||||||
}
|
}
|
||||||
self.logger.info(f"已加载可执行技能: {skill_name} — {summary}")
|
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 技能)")
|
self.skills = new_skills
|
||||||
|
self.logger.info(f"共加载 {len(new_skills)} 个技能({len([s for s in new_skills.values() if s['is_markdown']])} 个 Markdown 技能)")
|
||||||
|
|
||||||
def reload(self) -> int:
|
def reload(self) -> int:
|
||||||
"""重新扫描技能目录,返回加载的技能数量。"""
|
"""重新扫描技能目录,返回加载的技能数量。"""
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -13,6 +13,8 @@ from pathlib import Path
|
|||||||
# 把 aibox 目录加入 sys.path(无需安装为包即可 import tools/common)
|
# 把 aibox 目录加入 sys.path(无需安装为包即可 import tools/common)
|
||||||
AIBOX_DIR = Path(__file__).resolve().parent.parent
|
AIBOX_DIR = Path(__file__).resolve().parent.parent
|
||||||
sys.path.insert(0, str(AIBOX_DIR))
|
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,测试环境可能未安装)
|
# 桩掉 websockets(aibox.py 顶部 import websockets,测试环境可能未安装)
|
||||||
if "websockets" not in sys.modules:
|
if "websockets" not in sys.modules:
|
||||||
|
|||||||
+155
-1
@@ -15,7 +15,7 @@ from datetime import datetime, timedelta
|
|||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from common import constants
|
from common import ConfigManager, constants
|
||||||
from subsystems import (
|
from subsystems import (
|
||||||
MemoryDatabase, MemoDatabase, PromptManager, SkillManager,
|
MemoryDatabase, MemoDatabase, PromptManager, SkillManager,
|
||||||
SubprocessCommandExecutor,
|
SubprocessCommandExecutor,
|
||||||
@@ -425,6 +425,160 @@ def test_get_proxy_when_disabled():
|
|||||||
assert "未启用" in r or "启用" in r
|
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 ====================
|
# ==================== AskUserTool ====================
|
||||||
|
|
||||||
def test_ask_user_interface():
|
def test_ask_user_interface():
|
||||||
|
|||||||
@@ -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,171 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""F2 修复回归测试:N6 registry 超时清理 + N4 reader 线程收尾。
|
||||||
|
|
||||||
|
N6:流式转发超时后必须清理 tool.streaming_output / tool.output_queue,
|
||||||
|
防止僵尸 daemon 线程经实例属性把输出混入后续会话队列。
|
||||||
|
N4:_read_with_limit 超时/宽限 break 后必须 close stdout + join reader,
|
||||||
|
防止 Windows shell=True 孙进程持管道场景下 reader 线程与 fd 泄漏。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import queue
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from subsystems.executors import SubprocessCommandExecutor
|
||||||
|
from tools.base import Tool
|
||||||
|
from tools.registry import ToolRegistry
|
||||||
|
|
||||||
|
|
||||||
|
# ---------- N6:registry 超时路径清理 ----------
|
||||||
|
|
||||||
|
class _StuckStreamTool(Tool):
|
||||||
|
"""execute_streaming 永不结束(模拟卡死的工具,如等待外部资源)"""
|
||||||
|
|
||||||
|
def get_name(self):
|
||||||
|
return "stuck_tool"
|
||||||
|
|
||||||
|
def get_description(self):
|
||||||
|
return "test only"
|
||||||
|
|
||||||
|
def get_parameters(self):
|
||||||
|
return {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
def execute(self, **kwargs):
|
||||||
|
time.sleep(30)
|
||||||
|
return "should not reach"
|
||||||
|
|
||||||
|
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||||
|
self.output_queue = output_queue
|
||||||
|
self.streaming_output = True
|
||||||
|
try:
|
||||||
|
time.sleep(30) # 卡死,永不产出 complete
|
||||||
|
finally:
|
||||||
|
self.streaming_output = False
|
||||||
|
|
||||||
|
|
||||||
|
def test_registry_timeout_cleans_tool_state():
|
||||||
|
"""转发总超时后:返回超时消息,且工具实例属性被清理"""
|
||||||
|
reg = ToolRegistry()
|
||||||
|
reg.stream_tool_timeout = 1 # 测试调小(默认 3600)
|
||||||
|
tool = _StuckStreamTool()
|
||||||
|
reg.register(tool)
|
||||||
|
reg.set_output_queue(queue.Queue())
|
||||||
|
|
||||||
|
t0 = time.time()
|
||||||
|
result = reg.execute_tool("stuck_tool")
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
|
||||||
|
assert "超时" in result, f"应返回超时消息,实际: {result}"
|
||||||
|
assert elapsed < 10, f"应 ~1s 返回,实际 {elapsed:.1f}s"
|
||||||
|
# N6 核心:超时路径不得跳过清理
|
||||||
|
assert tool.output_queue is None, "超时后 tool.output_queue 必须置 None(防僵尸输出串流)"
|
||||||
|
assert tool.streaming_output is False, "超时后 tool.streaming_output 必须复位"
|
||||||
|
|
||||||
|
|
||||||
|
def test_registry_normal_path_cleans_tool_state():
|
||||||
|
"""正常流式完成后:清理同样执行,结果透传"""
|
||||||
|
reg = ToolRegistry()
|
||||||
|
|
||||||
|
class _OkTool(_StuckStreamTool):
|
||||||
|
def get_name(self):
|
||||||
|
return "ok_tool"
|
||||||
|
|
||||||
|
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||||
|
self.output_queue = output_queue
|
||||||
|
self.streaming_output = True
|
||||||
|
try:
|
||||||
|
output_queue.put(("line", "hello"))
|
||||||
|
output_queue.put(("complete", "done-result"))
|
||||||
|
finally:
|
||||||
|
pass
|
||||||
|
|
||||||
|
tool = _OkTool()
|
||||||
|
reg.register(tool)
|
||||||
|
reg.set_output_queue(queue.Queue())
|
||||||
|
|
||||||
|
result = reg.execute_tool("ok_tool")
|
||||||
|
assert result == "done-result"
|
||||||
|
assert tool.output_queue is None
|
||||||
|
assert tool.streaming_output is False
|
||||||
|
|
||||||
|
|
||||||
|
# ---------- N4:reader 线程收尾 ----------
|
||||||
|
|
||||||
|
class _FakeStdout:
|
||||||
|
"""模拟管道 stdout,可配置阻塞时长"""
|
||||||
|
|
||||||
|
def __init__(self, block_seconds: float = 8):
|
||||||
|
self.closed = False
|
||||||
|
self._block = block_seconds
|
||||||
|
|
||||||
|
def readline(self):
|
||||||
|
time.sleep(self._block) # > join 宽限 2s 即模拟"无 EOF 阻塞"
|
||||||
|
return ""
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
self.closed = True
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeProcess:
|
||||||
|
def __init__(self, block_seconds: float = 8):
|
||||||
|
self.stdout = _FakeStdout(block_seconds)
|
||||||
|
|
||||||
|
def poll(self):
|
||||||
|
return None # 进程"活着"
|
||||||
|
|
||||||
|
def terminate(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def kill(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def test_read_with_limit_timeout_returns_promptly():
|
||||||
|
"""超时 break 后 reader 仍阻塞:只做有限 join,不 close、不挂死。
|
||||||
|
|
||||||
|
CPython 流锁决定了此时 close(stdout) 会阻塞到 readline 返回,
|
||||||
|
故收尾策略为放弃等待(daemon 兜底),主流程必须及时返回。
|
||||||
|
"""
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
proc = _FakeProcess(block_seconds=8)
|
||||||
|
|
||||||
|
t0 = time.time()
|
||||||
|
output, timeout_occurred, truncated = ex._read_with_limit(
|
||||||
|
proc, timeout=1, start_time=time.time())
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
|
||||||
|
assert timeout_occurred is True
|
||||||
|
assert truncated is False
|
||||||
|
# 超时 1s + join 宽限 2s,绝不能等到 block 8s
|
||||||
|
assert elapsed < 5, f"应收尾后及时返回,实际 {elapsed:.1f}s"
|
||||||
|
|
||||||
|
|
||||||
|
def test_read_with_limit_grace_path_returns_promptly():
|
||||||
|
"""shell 死亡宽限 break 后 reader 仍阻塞:同样及时返回"""
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
proc = _FakeProcess(block_seconds=8)
|
||||||
|
proc.poll = lambda: 0 # shell 已死,但 reader 无 EOF(孙进程持管道)
|
||||||
|
|
||||||
|
t0 = time.time()
|
||||||
|
output, timeout_occurred, truncated = ex._read_with_limit(
|
||||||
|
proc, timeout=0, start_time=time.time()) # timeout=0 不触发超时
|
||||||
|
elapsed = time.time() - t0
|
||||||
|
|
||||||
|
assert timeout_occurred is False
|
||||||
|
# 宽限 1s + join 2s 上界,不能等到 block 8s
|
||||||
|
assert elapsed < 5, f"宽限+收尾应 ~3s 内返回,实际 {elapsed:.1f}s"
|
||||||
|
|
||||||
|
|
||||||
|
def test_read_with_limit_closes_stdout_after_normal_eof():
|
||||||
|
"""reader 正常收到 EOF 退出后:显式 close stdout 释放管道 fd"""
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
proc = _FakeProcess(block_seconds=0) # readline 立即返回 ""(EOF)
|
||||||
|
|
||||||
|
output, timeout_occurred, truncated = ex._read_with_limit(
|
||||||
|
proc, timeout=5, start_time=time.time())
|
||||||
|
|
||||||
|
assert timeout_occurred is False
|
||||||
|
assert output == ""
|
||||||
|
assert proc.stdout.closed is True, "reader 正常退出后必须 close stdout 释放 fd"
|
||||||
@@ -0,0 +1,359 @@
|
|||||||
|
"""aibox 第三轮审计 Batch G 修复回归测试(2026-08-18,审计报告 §9.4)
|
||||||
|
|
||||||
|
覆盖:
|
||||||
|
- P1:非瞬时 HTTP 错误响应体截断(密钥/超大错误页不进异常栈)
|
||||||
|
- P2:Retry-After 解析上限 cap 3600s
|
||||||
|
- P3:fallback 档案缺 model → 跳过并告警(与缺 Key 行为对齐)
|
||||||
|
- P5:LLMClient 死状态 _cancelled 已删除
|
||||||
|
- P6:fallback 列表重复项去重
|
||||||
|
- N2:drain 线程批次结束即回收(不再泄漏常驻线程)
|
||||||
|
- N5:子代理共享父代理取消事件
|
||||||
|
- N7:FTS5 关键词内双引号翻倍转义
|
||||||
|
- N8:registry 流式转发受 TOOL_MAX_RESULT_LENGTH 约束
|
||||||
|
- N10:mark_as_triggered 单事务(缺失 id 路径回归)
|
||||||
|
- N12:persona 路径边界校验(拒绝越界路径)
|
||||||
|
- N16:同秒归档不再覆盖(文件名加序号)
|
||||||
|
|
||||||
|
N3(/new 移线程池)/N13(broadcast)/N14(任务登记集合)为 async WebSocket
|
||||||
|
路径,依赖事件循环与真实连接对象,由现有 WebSocket 测试面回归,不在此单测。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import queue
|
||||||
|
import sys
|
||||||
|
import threading
|
||||||
|
import types
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
import aibox # noqa: E402
|
||||||
|
from aibox import DeepSeekChat, run_full_subagent # noqa: E402
|
||||||
|
|
||||||
|
from llm_client import ( # noqa: E402
|
||||||
|
DeepSeekClient,
|
||||||
|
FailoverClient,
|
||||||
|
_retry_after_of,
|
||||||
|
create_llm_client,
|
||||||
|
)
|
||||||
|
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||||
|
from subsystems.conversation_history import ConversationHistory # noqa: E402
|
||||||
|
from tools.base import Tool # noqa: E402
|
||||||
|
from tools.registry import ToolRegistry # noqa: E402
|
||||||
|
from common import constants # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== P1/P2:错误响应体截断 + Retry-After cap ====================
|
||||||
|
|
||||||
|
class TestP1P2LlmClient:
|
||||||
|
def test_p1_non_transient_error_body_truncated(self):
|
||||||
|
"""400 响应体 1000 字符 → 异常消息仅含前 200 字符"""
|
||||||
|
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||||
|
resp = mock.MagicMock()
|
||||||
|
resp.status_code = 400
|
||||||
|
resp.text = "A" * 1000
|
||||||
|
resp.headers = {}
|
||||||
|
with mock.patch("llm_client.requests.post", return_value=resp):
|
||||||
|
with pytest.raises(RuntimeError) as ei:
|
||||||
|
list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
msg = str(ei.value)
|
||||||
|
assert "HTTP 400" in msg
|
||||||
|
assert len(msg) < 300 # 200 字符正文 + 状态行前缀
|
||||||
|
assert msg.count("A") <= 200
|
||||||
|
|
||||||
|
def test_p2_retry_after_capped(self):
|
||||||
|
"""异常/恶意 Retry-After=999999 → cap 到 3600s"""
|
||||||
|
resp = mock.MagicMock()
|
||||||
|
resp.headers = {"Retry-After": "999999"}
|
||||||
|
assert _retry_after_of(resp) == 3600.0
|
||||||
|
|
||||||
|
def test_p2_retry_after_normal_value_kept(self):
|
||||||
|
resp = mock.MagicMock()
|
||||||
|
resp.headers = {"Retry-After": "30"}
|
||||||
|
assert _retry_after_of(resp) == 30.0
|
||||||
|
|
||||||
|
def test_p2_retry_after_invalid_returns_none(self):
|
||||||
|
resp = mock.MagicMock()
|
||||||
|
resp.headers = {"Retry-After": "not-a-number"}
|
||||||
|
assert _retry_after_of(resp) is None
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== P3/P5/P6:fallback 链构造 ====================
|
||||||
|
|
||||||
|
def _g_cfg():
|
||||||
|
"""主档案 main 的 fallback 链:缺 model 的 broken + 重复 dupe + 正常 ok2"""
|
||||||
|
return {
|
||||||
|
"models": {
|
||||||
|
"main": {
|
||||||
|
"api_url": "https://a.example/v1/chat/completions",
|
||||||
|
"model": "m1", "api_key_env": "MAIN_KEY",
|
||||||
|
"fallback": ["nomodel", "dupe", "dupe", "ok2"],
|
||||||
|
},
|
||||||
|
"nomodel": { # P3:缺 model 字段
|
||||||
|
"api_url": "https://b.example/v1/chat/completions",
|
||||||
|
"api_key_env": "B_KEY",
|
||||||
|
},
|
||||||
|
"dupe": {
|
||||||
|
"api_url": "https://c.example/v1/chat/completions",
|
||||||
|
"model": "c1", "api_key_env": "C_KEY",
|
||||||
|
},
|
||||||
|
"ok2": {
|
||||||
|
"api_url": "https://d.example/v1/chat/completions",
|
||||||
|
"model": "d1", "api_key_env": "D_KEY",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class TestP3P5P6FallbackChain:
|
||||||
|
def test_p3_p6_fallback_missing_model_skipped_and_deduped(self):
|
||||||
|
"""缺 model 的 fallback 跳过;重复项去重;链 = [main, dupe, ok2]"""
|
||||||
|
env = {k: f"sk-{k.lower()}" for k in
|
||||||
|
("MAIN_KEY", "B_KEY", "C_KEY", "D_KEY")}
|
||||||
|
with mock.patch.dict(os.environ, env):
|
||||||
|
c = create_llm_client(_g_cfg(), model="main")
|
||||||
|
assert isinstance(c, FailoverClient)
|
||||||
|
names = [cl.name for cl in c._clients]
|
||||||
|
assert names == ["main", "dupe", "ok2"]
|
||||||
|
|
||||||
|
def test_p5_dead_cancel_flag_removed(self):
|
||||||
|
"""LLMClient 不再持有只写不读的 _cancelled 死状态"""
|
||||||
|
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||||
|
assert not hasattr(c, "_cancelled")
|
||||||
|
c.cancel() # 无在途响应时为安全 no-op
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N2:drain 线程生命周期 ====================
|
||||||
|
|
||||||
|
class _EchoTool(Tool):
|
||||||
|
def get_name(self):
|
||||||
|
return "g_echo"
|
||||||
|
|
||||||
|
def get_description(self):
|
||||||
|
return "测试用回显工具"
|
||||||
|
|
||||||
|
def get_parameters(self):
|
||||||
|
return {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
def execute(self, **kwargs):
|
||||||
|
return "ok"
|
||||||
|
|
||||||
|
|
||||||
|
class TestN2DrainThread:
|
||||||
|
def test_drain_thread_reclaimed_after_batch(self):
|
||||||
|
"""_handle_tool_calls 返回后 drain 线程必须已退出(finally join)"""
|
||||||
|
from test_once_mode import _make_ai
|
||||||
|
ai = _make_ai()
|
||||||
|
ai.tool_registry.register(_EchoTool(ai.logger))
|
||||||
|
tool_calls = [{"id": "c1", "function": {"name": "g_echo", "arguments": "{}"}}]
|
||||||
|
ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None)
|
||||||
|
drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"]
|
||||||
|
assert drains == [] or all(not t.is_alive() for t in drains)
|
||||||
|
|
||||||
|
@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning")
|
||||||
|
def test_drain_thread_reclaimed_on_tool_exception(self):
|
||||||
|
"""工具执行抛异常时 drain 线程同样被回收(finally 路径)"""
|
||||||
|
from test_once_mode import _make_ai
|
||||||
|
|
||||||
|
class _BoomTool(_EchoTool):
|
||||||
|
def get_name(self):
|
||||||
|
return "g_boom"
|
||||||
|
|
||||||
|
def execute(self, **kwargs):
|
||||||
|
raise RuntimeError("boom")
|
||||||
|
|
||||||
|
ai = _make_ai()
|
||||||
|
ai.tool_registry.register(_BoomTool(ai.logger))
|
||||||
|
tool_calls = [{"id": "c1", "function": {"name": "g_boom", "arguments": "{}"}}]
|
||||||
|
ai._handle_tool_calls(tool_calls, output_callback=lambda t, c: None)
|
||||||
|
drains = [t for t in threading.enumerate() if t.name == "tool-output-drain"]
|
||||||
|
assert drains == [] or all(not t.is_alive() for t in drains)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N5:子代理取消传播 ====================
|
||||||
|
|
||||||
|
def _fake_sub_chat():
|
||||||
|
fake = mock.MagicMock()
|
||||||
|
fake.think_and_respond = mock.MagicMock(return_value="任务完成报告")
|
||||||
|
fake.stats = {}
|
||||||
|
fake.messages = []
|
||||||
|
fake._cancel_event = threading.Event()
|
||||||
|
return fake
|
||||||
|
|
||||||
|
|
||||||
|
class TestN5SubagentCancel:
|
||||||
|
def test_subagent_shares_parent_cancel_event(self):
|
||||||
|
"""cancel_event 传入 → 子代理取消事件被替换为父事件(检查点联动)"""
|
||||||
|
fake = _fake_sub_chat()
|
||||||
|
parent_event = threading.Event()
|
||||||
|
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||||
|
mock.patch("aibox.create_llm_client"), \
|
||||||
|
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||||
|
run_full_subagent("r", "t", None, cancel_event=parent_event)
|
||||||
|
assert fake._cancel_event is parent_event
|
||||||
|
|
||||||
|
def test_subagent_keeps_own_event_when_not_provided(self):
|
||||||
|
"""cancel_event=None → 兼容旧行为,子代理保留自身事件"""
|
||||||
|
fake = _fake_sub_chat()
|
||||||
|
own = fake._cancel_event
|
||||||
|
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||||
|
mock.patch("aibox.create_llm_client"), \
|
||||||
|
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||||
|
run_full_subagent("r", "t", None)
|
||||||
|
assert fake._cancel_event is own
|
||||||
|
|
||||||
|
def test_subagent_tool_registered_with_parent_event(self):
|
||||||
|
"""DeepSeekChat 注册 SubAgentTool 时注入自身 _cancel_event(N5 链路)"""
|
||||||
|
from test_once_mode import _make_ai
|
||||||
|
ai = _make_ai()
|
||||||
|
tool = next((t for t in ai.tool_registry.get_all_tools()
|
||||||
|
if t.get_name() == "create_subagent"), None)
|
||||||
|
assert tool is not None
|
||||||
|
assert tool.cancel_event is ai._cancel_event
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N7:FTS5 引号转义 ====================
|
||||||
|
|
||||||
|
class TestN7FtsQuoteEscape:
|
||||||
|
def test_keyword_with_quote_does_not_break_search(self, tmp_path):
|
||||||
|
"""关键词含双引号 → MATCH 语法合法(翻倍转义),不抛异常"""
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||||
|
db.add_memory(topic="引号话题", summary="摘要",
|
||||||
|
content='含"引号"的记忆内容',
|
||||||
|
start_time="2026-08-18 00:00:00",
|
||||||
|
end_time="2026-08-18 00:00:01", importance=2)
|
||||||
|
res = db.search_memories(['带"引号"的关键词'])
|
||||||
|
assert isinstance(res, list) # 不抛 OperationalError 即通过
|
||||||
|
|
||||||
|
def test_normal_keyword_still_matches(self, tmp_path):
|
||||||
|
"""常规关键词检索行为不回归"""
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||||
|
db.add_memory(topic="普通话题", summary="天气摘要",
|
||||||
|
content="今天天气晴朗",
|
||||||
|
start_time="2026-08-18 00:00:00",
|
||||||
|
end_time="2026-08-18 00:00:01", importance=2)
|
||||||
|
res = db.search_memories(["天气"])
|
||||||
|
assert len(res) >= 1
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N8:流式转发长度上限 ====================
|
||||||
|
|
||||||
|
class _BigStreamTool(Tool):
|
||||||
|
"""流式工具:发出 3 行 × 40000 字符(总 120000 > 100000 上限)"""
|
||||||
|
|
||||||
|
def get_name(self):
|
||||||
|
return "g_bigstream"
|
||||||
|
|
||||||
|
def get_description(self):
|
||||||
|
return "测试用超长流式工具"
|
||||||
|
|
||||||
|
def get_parameters(self):
|
||||||
|
return {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
def execute(self, **kwargs):
|
||||||
|
return "done"
|
||||||
|
|
||||||
|
def execute_streaming(self, output_queue, **kwargs):
|
||||||
|
for _ in range(3):
|
||||||
|
output_queue.put(("line", "x" * 40000))
|
||||||
|
output_queue.put(("complete", "done"))
|
||||||
|
|
||||||
|
|
||||||
|
class TestN8StreamingLimit:
|
||||||
|
def test_streaming_forward_truncated_to_budget(self):
|
||||||
|
registry = ToolRegistry()
|
||||||
|
registry.register(_BigStreamTool())
|
||||||
|
captured = queue.Queue()
|
||||||
|
registry.set_output_queue(captured)
|
||||||
|
|
||||||
|
result = registry.execute_tool("g_bigstream")
|
||||||
|
|
||||||
|
lines = []
|
||||||
|
while not captured.empty():
|
||||||
|
msg_type, content = captured.get_nowait()
|
||||||
|
if msg_type == "line":
|
||||||
|
lines.append(content)
|
||||||
|
budget = constants.TOOL_MAX_RESULT_LENGTH
|
||||||
|
total = sum(len(ln) for ln in lines)
|
||||||
|
# 截断标记行本身占少量额外字符(含提示文案)
|
||||||
|
assert total <= budget + 200
|
||||||
|
assert any("已停止转发" in ln or "已截断" in ln for ln in lines)
|
||||||
|
assert result == "done"
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N10:mark_as_triggered 单事务 ====================
|
||||||
|
|
||||||
|
class TestN10MarkTriggered:
|
||||||
|
def test_missing_id_returns_none_via_transaction(self, tmp_path):
|
||||||
|
"""事务内读取不存在行 → rollback + None(不再走独立 get_memo)"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "memo.db"))
|
||||||
|
assert db.mark_as_triggered(99999) is None
|
||||||
|
|
||||||
|
def test_single_shot_memo_completes(self, tmp_path):
|
||||||
|
"""一次性任务触发 → completed + 返回 None(无下次提醒)"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "memo.db"))
|
||||||
|
mid = db.add_memo(title="一次性", content="c",
|
||||||
|
reminder_time="2026-08-18 00:00:00")
|
||||||
|
assert mid is not None
|
||||||
|
assert db.mark_as_triggered(mid) is None
|
||||||
|
memo = db.get_memo(mid)
|
||||||
|
assert memo["is_completed"] == 1
|
||||||
|
assert memo["trigger_count"] == 1
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N12:persona 路径边界 ====================
|
||||||
|
|
||||||
|
class TestN12PersonaBoundary:
|
||||||
|
def _tool(self):
|
||||||
|
return aibox.ExpressEmotionTool()
|
||||||
|
|
||||||
|
def test_traversal_persona_rejected(self, tmp_path):
|
||||||
|
"""persona 含 ../ 越界 → 拒绝(返回空串,不读文件)"""
|
||||||
|
tool = self._tool()
|
||||||
|
skills = tmp_path / "skills"
|
||||||
|
skills.mkdir()
|
||||||
|
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||||
|
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||||
|
assert tool._load_persona("../../etc/passwd", None) == ""
|
||||||
|
assert tool._load_persona("../outside", None) == ""
|
||||||
|
|
||||||
|
def test_absolute_persona_rejected(self, tmp_path):
|
||||||
|
"""persona 为绝对路径 → 解析结果不在技能目录内,拒绝"""
|
||||||
|
tool = self._tool()
|
||||||
|
skills = tmp_path / "skills"
|
||||||
|
skills.mkdir()
|
||||||
|
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||||
|
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||||
|
assert tool._load_persona(str(tmp_path / "evil"), None) == ""
|
||||||
|
|
||||||
|
def test_normal_persona_loads(self, tmp_path):
|
||||||
|
"""合法 persona(技能目录内)正常加载"""
|
||||||
|
tool = self._tool()
|
||||||
|
skills = tmp_path / "skills"
|
||||||
|
skills.mkdir()
|
||||||
|
(skills / "ok.md").write_text("人格内容", encoding="utf-8")
|
||||||
|
with mock.patch.object(type(aibox.constants), "SKILLS_DIR",
|
||||||
|
new_callable=mock.PropertyMock(return_value=str(skills))):
|
||||||
|
assert "人格内容" in tool._load_persona("ok", None)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N16:同秒归档防覆盖 ====================
|
||||||
|
|
||||||
|
class TestN16ArchiveCollision:
|
||||||
|
def test_same_second_archive_no_overwrite(self, tmp_path):
|
||||||
|
"""同一秒内两次归档 → 文件名加序号,互不覆盖"""
|
||||||
|
hm = ConversationHistory(save_dir=str(tmp_path))
|
||||||
|
hm.start_new_conversation(topic="t1")
|
||||||
|
hm.add_message("user", "hi", speaker="u")
|
||||||
|
p1 = hm.archive_current_conversation()
|
||||||
|
|
||||||
|
hm.start_new_conversation(topic="t2")
|
||||||
|
hm.add_message("user", "hi2", speaker="u")
|
||||||
|
p2 = hm.archive_current_conversation()
|
||||||
|
|
||||||
|
assert p1 != p2
|
||||||
|
assert os.path.exists(p1) and os.path.exists(p2)
|
||||||
@@ -0,0 +1,413 @@
|
|||||||
|
"""aibox 第四轮审计 Batch H 修复回归测试(2026-08-18,审计报告 §十一)
|
||||||
|
|
||||||
|
覆盖:
|
||||||
|
- H1:子代理共享父取消事件后首轮 think 不再清除父事件(N5 回归)+ 父已取消时跳过启动
|
||||||
|
- M3:调度器任务级超时兜底(AI think 卡死不再停摆调度循环)
|
||||||
|
- M4:get_recent_memories(limit=None) 使用默认上限(不再全表加载)
|
||||||
|
- M5:memo delete 物理删除(此前只标记完成)
|
||||||
|
- M6:PromptManager.reload 原子替换(并发读者不再读到空/半截列表)
|
||||||
|
- L1:read_file 的 sed 参数强转 int(防引号注入)
|
||||||
|
- L2:set_proxy 落盘失败如实报告(不再静默报"已保存")
|
||||||
|
- L3:SSE 非 UTF-8 行不中断流(errors=replace)
|
||||||
|
- L4:persona 路径校验解析符号链接(软链越界拒绝)
|
||||||
|
- L5:1MB 截断不再产生半个字符(�)
|
||||||
|
- L6:compressor 空内容压缩不再除零
|
||||||
|
- L7:InstallSkillTool 指引命令路径引用(防拆词/注入)
|
||||||
|
- L8:ask_user 缺参返回明确错误(不再 TypeError)
|
||||||
|
- L10:web 工具仅允许 http/https(file:// 等拒绝)
|
||||||
|
- L11:registry 日志打码改包含匹配(apiKey/access_token 等变体)
|
||||||
|
- L12:FailoverClient 日志 Bearer 令牌脱敏
|
||||||
|
|
||||||
|
M1(/reload_prompts 移线程池)/M2(技能扫描启动前预热)为 async WebSocket /
|
||||||
|
启动路径,依赖事件循环与真实连接对象,由既有 WebSocket 测试面与启动逻辑回归。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
import queue
|
||||||
|
import sys
|
||||||
|
import threading
|
||||||
|
import types
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
import aibox # noqa: E402
|
||||||
|
from aibox import DeepSeekChat, TaskScheduler, run_full_subagent # noqa: E402
|
||||||
|
|
||||||
|
from llm_client import DeepSeekClient, FailoverClient, _sanitize_log_msg # noqa: E402
|
||||||
|
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||||
|
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||||
|
from subsystems.conversation_compressor import ConversationCompressor # noqa: E402
|
||||||
|
from common import constants # noqa: E402
|
||||||
|
|
||||||
|
from test_conversation_compressor import FakeLLMClient # noqa: E402
|
||||||
|
from test_once_mode import _make_ai # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def _fake_sub_chat():
|
||||||
|
fake = mock.MagicMock()
|
||||||
|
fake.think_and_respond = mock.MagicMock(return_value="任务完成报告")
|
||||||
|
fake.stats = {}
|
||||||
|
fake.messages = []
|
||||||
|
fake._cancel_event = threading.Event()
|
||||||
|
return fake
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== H1:子代理取消事件清除回归 ====================
|
||||||
|
|
||||||
|
class TestH1SubagentCancelClear:
|
||||||
|
def test_shared_event_not_cleared_on_first_think(self):
|
||||||
|
"""共享父事件后首轮 think 不清除事件——父级已置位的取消信号不被吞掉"""
|
||||||
|
ai = _make_ai()
|
||||||
|
parent_event = threading.Event()
|
||||||
|
parent_event.set()
|
||||||
|
ai._cancel_event = parent_event
|
||||||
|
ai._shared_cancel_event = True
|
||||||
|
with pytest.raises(InterruptedError):
|
||||||
|
ai._think_and_respond_impl("hi")
|
||||||
|
assert parent_event.is_set()
|
||||||
|
|
||||||
|
def test_own_event_still_cleared_on_first_think(self):
|
||||||
|
"""未共享(普通会话)时首轮仍清除自身事件——原语义不回归"""
|
||||||
|
import contextlib
|
||||||
|
ai = _make_ai()
|
||||||
|
ai._cancel_event.set()
|
||||||
|
with contextlib.suppress(Exception):
|
||||||
|
ai._think_and_respond_impl("hi")
|
||||||
|
assert not ai._cancel_event.is_set()
|
||||||
|
|
||||||
|
def test_subagent_not_started_when_parent_already_cancelled(self):
|
||||||
|
"""父事件已置位 → 直接返回取消提示,不启动子代理"""
|
||||||
|
fake = _fake_sub_chat()
|
||||||
|
parent_event = threading.Event()
|
||||||
|
parent_event.set()
|
||||||
|
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||||
|
mock.patch("aibox.create_llm_client"), \
|
||||||
|
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||||
|
result = run_full_subagent("r", "t", None, cancel_event=parent_event)
|
||||||
|
assert "取消" in result
|
||||||
|
fake.think_and_respond.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== M3:调度器任务超时 ====================
|
||||||
|
|
||||||
|
class _BlockingAI:
|
||||||
|
"""think 卡死直到被 cancel 的伪 AI(模拟 LLM 挂起)"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.cancelled = threading.Event()
|
||||||
|
self.cancel_called = 0
|
||||||
|
|
||||||
|
def cancel(self):
|
||||||
|
self.cancel_called += 1
|
||||||
|
self.cancelled.set()
|
||||||
|
|
||||||
|
def think_and_respond(self, input_text, output_callback=None,
|
||||||
|
is_continuation=False):
|
||||||
|
self.cancelled.wait(timeout=15)
|
||||||
|
raise InterruptedError("用户中断")
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeMemoDB:
|
||||||
|
def get_memo(self, memo_id):
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
class TestM3SchedulerTimeout:
|
||||||
|
def test_stuck_task_times_out_and_scheduler_continues(self):
|
||||||
|
from aibox import ConfigManager
|
||||||
|
loop = asyncio.new_event_loop()
|
||||||
|
try:
|
||||||
|
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||||
|
fake_ai = _BlockingAI()
|
||||||
|
memo = {"id": 7, "title": "卡死任务", "content": "内容", "trigger_count": 0}
|
||||||
|
orig_get = ConfigManager.get
|
||||||
|
|
||||||
|
def fake_get(self, key, default=None):
|
||||||
|
# 只把任务超时压到 0.3s,其余配置读取走原实现(防止
|
||||||
|
# DATETIME_FORMAT 等属性读到错误类型)
|
||||||
|
if key == "scheduler.task_timeout":
|
||||||
|
return 0.3
|
||||||
|
return orig_get(self, key, default)
|
||||||
|
|
||||||
|
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||||
|
mock.patch.object(scheduler, "_create_task_ai", return_value=fake_ai):
|
||||||
|
result = loop.run_until_complete(
|
||||||
|
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||||
|
assert "超时" in result
|
||||||
|
assert fake_ai.cancel_called >= 1
|
||||||
|
finally:
|
||||||
|
loop.close()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== M4/M5:记忆上限与任务删除 ====================
|
||||||
|
|
||||||
|
class TestM4M5Db:
|
||||||
|
def test_recent_memories_none_capped(self, tmp_path, monkeypatch):
|
||||||
|
"""limit=None 不再全表加载,使用默认上限(对齐 search_memories)"""
|
||||||
|
from test_conversation_compressor import patch_const
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
for i in range(7):
|
||||||
|
db.add_memory(topic=f"t{i}", summary=f"s{i}", content="c",
|
||||||
|
start_time="2026-08-18 00:00:00",
|
||||||
|
end_time="2026-08-18 00:00:01", importance=1)
|
||||||
|
patch_const(monkeypatch, "MEMORY_MAX_SEARCH_RESULTS", 5)
|
||||||
|
res = db.get_recent_memories(limit=None)
|
||||||
|
assert len(res) == 5
|
||||||
|
|
||||||
|
def test_delete_memo_physically_removes(self, tmp_path):
|
||||||
|
"""delete_memo 真实删除行(此前 complete_memo 只标记完成)"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = db.add_memo(title="待删任务", content="内容")
|
||||||
|
assert db.get_memo(mid) is not None
|
||||||
|
assert db.delete_memo(mid) is True
|
||||||
|
assert db.get_memo(mid) is None
|
||||||
|
assert db.delete_memo(mid) is False # 已不存在
|
||||||
|
|
||||||
|
def test_memo_tool_delete_uses_physical_delete(self, tmp_path):
|
||||||
|
"""MemoTool delete 走 delete_memo:行消失而非仅标记完成"""
|
||||||
|
from tools import MemoTool
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = db.add_memo(title="重复任务", content="内容", repeat_type="daily")
|
||||||
|
tool = MemoTool(db, None)
|
||||||
|
result = tool.execute(action="delete", memo_id=mid, confirm=True)
|
||||||
|
assert "已删除" in result
|
||||||
|
assert db.get_memo(mid) is None
|
||||||
|
stats = db.get_memo_stats()
|
||||||
|
assert stats["total"] == 0
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== M6:PromptManager 原子 reload ====================
|
||||||
|
|
||||||
|
class TestM6PromptReloadAtomic:
|
||||||
|
def test_reload_builds_new_list_then_swaps(self):
|
||||||
|
"""reload 在本地新列表构建完成后才原子替换——读者始终看到完整列表"""
|
||||||
|
pm = PromptManager()
|
||||||
|
old_list = pm._prompts
|
||||||
|
seen = []
|
||||||
|
|
||||||
|
def fake_load(target):
|
||||||
|
seen.append(target)
|
||||||
|
target.append({"file": "X_test.md", "content": "y"})
|
||||||
|
|
||||||
|
with mock.patch.object(pm, "_load_prompts", side_effect=fake_load):
|
||||||
|
n = pm.reload()
|
||||||
|
assert len(seen) == 1
|
||||||
|
assert seen[0] is not old_list # 构建用的是新列表
|
||||||
|
assert pm._prompts is seen[0] # 完成后才替换
|
||||||
|
assert n == 1
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L1:read_file 参数强转 ====================
|
||||||
|
|
||||||
|
class TestL1ReadFileCoercion:
|
||||||
|
def test_non_integer_lines_rejected(self, tmp_path):
|
||||||
|
from tools import ReadFileTool
|
||||||
|
f = tmp_path / "f.txt"
|
||||||
|
f.write_text("line1\nline2\n", encoding="utf-8")
|
||||||
|
tool = ReadFileTool()
|
||||||
|
r = tool.execute(filepath=str(f), start_line="1'; rm -rf /; echo '",
|
||||||
|
end_line=5)
|
||||||
|
assert "参数错误" in r
|
||||||
|
|
||||||
|
def test_valid_integers_still_work(self, tmp_path):
|
||||||
|
from tools import ReadFileTool
|
||||||
|
f = tmp_path / "f.txt"
|
||||||
|
f.write_text("line1\nline2\n", encoding="utf-8")
|
||||||
|
tool = ReadFileTool()
|
||||||
|
r = tool.execute(filepath=str(f), start_line=1, end_line=2)
|
||||||
|
assert "sed -n" in r and "参数错误" not in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L2:set_proxy 落盘失败如实报告 ====================
|
||||||
|
|
||||||
|
class TestL2ProxySaveFailure:
|
||||||
|
def test_set_proxy_returns_false_on_save_failure(self, tmp_path):
|
||||||
|
from common import ConfigManager
|
||||||
|
cm = ConfigManager()
|
||||||
|
cm._config_path = str(tmp_path / "config.json")
|
||||||
|
# J5(审计 2026-08-19):写路径改为原子写(mkstemp+os.replace),
|
||||||
|
# 磁盘失败改由 patch 原子写函数模拟(builtins.open 已不在路径上)
|
||||||
|
with mock.patch("common.config.atomic_write_text",
|
||||||
|
side_effect=PermissionError("disk ro")):
|
||||||
|
ok = cm.set_proxy("127.0.0.1", 1080, save=True)
|
||||||
|
assert ok is False
|
||||||
|
|
||||||
|
def test_proxy_tool_reports_save_failure(self, tmp_path):
|
||||||
|
from tools import SetProxyTool
|
||||||
|
tool = SetProxyTool(None)
|
||||||
|
# J5:磁盘失败经原子写函数模拟
|
||||||
|
with mock.patch("common.config.atomic_write_text",
|
||||||
|
side_effect=PermissionError("disk ro")):
|
||||||
|
r = tool.execute(host="127.0.0.1", port=1080, save=True,
|
||||||
|
skip_check=True, enabled=True)
|
||||||
|
assert "保存失败" in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L3:SSE 非 UTF-8 行不中断流 ====================
|
||||||
|
|
||||||
|
class TestL3SseDecode:
|
||||||
|
def test_non_utf8_line_does_not_break_stream(self):
|
||||||
|
c = DeepSeekClient(api_key="k", model="m", base_url="https://x")
|
||||||
|
resp = mock.MagicMock()
|
||||||
|
resp.status_code = 200
|
||||||
|
resp.iter_lines.return_value = [
|
||||||
|
b"data: \xff\xfe",
|
||||||
|
b'data: {"choices":[{"delta":{"content":"ok"}}]}',
|
||||||
|
b"data: [DONE]",
|
||||||
|
]
|
||||||
|
with mock.patch("llm_client.requests.post", return_value=resp):
|
||||||
|
events = list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
assert any(e.get("type") == "content" and e.get("text") == "ok"
|
||||||
|
for e in events)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L4:persona 软链越界 ====================
|
||||||
|
|
||||||
|
class TestL4PersonaSymlink:
|
||||||
|
def test_symlink_outside_skills_dir_rejected(self, tmp_path, monkeypatch):
|
||||||
|
from test_conversation_compressor import patch_const
|
||||||
|
from aibox import ExpressEmotionTool
|
||||||
|
skills_dir = tmp_path / "skills"
|
||||||
|
skills_dir.mkdir()
|
||||||
|
outside = tmp_path / "outside.md"
|
||||||
|
outside.write_text("秘密内容", encoding="utf-8")
|
||||||
|
(skills_dir / "evil.md").symlink_to(outside)
|
||||||
|
tool = ExpressEmotionTool()
|
||||||
|
patch_const(monkeypatch, "SKILLS_DIR", str(skills_dir))
|
||||||
|
content = tool._load_persona("evil", "")
|
||||||
|
assert content == ""
|
||||||
|
|
||||||
|
def test_normal_persona_still_loads(self, tmp_path, monkeypatch):
|
||||||
|
from test_conversation_compressor import patch_const
|
||||||
|
from aibox import ExpressEmotionTool
|
||||||
|
skills_dir = tmp_path / "skills"
|
||||||
|
skills_dir.mkdir()
|
||||||
|
(skills_dir / "happy.md").write_text("快乐人格", encoding="utf-8")
|
||||||
|
tool = ExpressEmotionTool()
|
||||||
|
patch_const(monkeypatch, "SKILLS_DIR", str(skills_dir))
|
||||||
|
content = tool._load_persona("happy", "")
|
||||||
|
assert content == "快乐人格"
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L5:截断无半个字符 ====================
|
||||||
|
|
||||||
|
class TestL5NoPartialChar:
|
||||||
|
def test_truncated_output_has_no_replacement_char(self):
|
||||||
|
from subsystems.executors import SubprocessCommandExecutor
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
r = ex.execute("python3 -c \"print('中' * 600000)\"", timeout=60)
|
||||||
|
assert "截断" in r
|
||||||
|
assert "\ufffd" not in r
|
||||||
|
|
||||||
|
@pytest.mark.skipif(os.name != "posix", reason="Linux 进程组语义")
|
||||||
|
def test_grandchild_holding_pipe_is_reaped(self):
|
||||||
|
"""`(sleep N &) ; echo done` 形态:shell 退出后孙进程持管道写端,
|
||||||
|
宽限窗口到期强杀进程组——不悬挂 reader、孙进程被清理"""
|
||||||
|
import subprocess as sp
|
||||||
|
from subsystems.executors import SubprocessCommandExecutor
|
||||||
|
marker = f"sleep {12340}"
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
r = ex.execute(f"({marker} &) ; echo done", timeout=8)
|
||||||
|
assert "done" in r
|
||||||
|
leftover = sp.run(["pgrep", "-f", marker],
|
||||||
|
capture_output=True, text=True).stdout.strip()
|
||||||
|
assert leftover == ""
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L6:compressor 除零 ====================
|
||||||
|
|
||||||
|
class TestL6CompressorDivZero:
|
||||||
|
def test_empty_content_compress_no_div0(self, monkeypatch):
|
||||||
|
comp = ConversationCompressor(llm_client=FakeLLMClient(),
|
||||||
|
logger=None)
|
||||||
|
from test_conversation_compressor import patch_const
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||||
|
messages = [{"role": "system", "content": "SYS"}] + [
|
||||||
|
{"role": "user" if i % 2 == 0 else "assistant", "content": ""}
|
||||||
|
for i in range(15)
|
||||||
|
]
|
||||||
|
result = comp.compress_conversation(messages)
|
||||||
|
assert isinstance(result, list)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L7:InstallSkillTool 路径引用 ====================
|
||||||
|
|
||||||
|
class TestL7SkillPathQuoting:
|
||||||
|
def test_guide_commands_quote_paths(self):
|
||||||
|
from tools import InstallSkillTool
|
||||||
|
tool = InstallSkillTool()
|
||||||
|
r = tool.execute(skill_path="/tmp/a b; rm -rf /")
|
||||||
|
assert "ls -l '/tmp/a b; rm -rf /'" in r
|
||||||
|
assert "cp '/tmp/a b; rm -rf /' " in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L8:ask_user 缺参 ====================
|
||||||
|
|
||||||
|
class TestL8AskUser:
|
||||||
|
def test_missing_question_returns_error(self):
|
||||||
|
from tools import AskUserTool
|
||||||
|
tool = AskUserTool()
|
||||||
|
assert "需要提供问题" in tool.execute()
|
||||||
|
|
||||||
|
def test_question_passed_through(self):
|
||||||
|
from tools import AskUserTool
|
||||||
|
tool = AskUserTool()
|
||||||
|
assert "测试问题" in tool.execute(question="测试问题")
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L10:web 工具 scheme 白名单 ====================
|
||||||
|
|
||||||
|
class TestL10WebScheme:
|
||||||
|
def test_file_url_rejected(self):
|
||||||
|
from tools import SimpleWebCrawlerTool
|
||||||
|
tool = SimpleWebCrawlerTool()
|
||||||
|
r = tool.execute(action="get", url="file:///etc/shadow")
|
||||||
|
assert "仅支持 http/https" in r
|
||||||
|
|
||||||
|
def test_http_url_accepted(self):
|
||||||
|
from tools import SimpleWebCrawlerTool
|
||||||
|
tool = SimpleWebCrawlerTool()
|
||||||
|
r = tool.execute(action="get", url="https://example.com")
|
||||||
|
assert "curl" in r and "仅支持" not in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L11:registry 打码包含匹配 ====================
|
||||||
|
|
||||||
|
class TestL11SanitizeSubstring:
|
||||||
|
def test_variant_keys_redacted(self):
|
||||||
|
from tools.registry import _sanitize_kwargs
|
||||||
|
out = _sanitize_kwargs({
|
||||||
|
"apiKey": "x", "access_token": "y", "passwd": "z",
|
||||||
|
"url": "https://ok",
|
||||||
|
})
|
||||||
|
assert out == {"apiKey": "***", "access_token": "***",
|
||||||
|
"passwd": "***", "url": "https://ok"}
|
||||||
|
|
||||||
|
def test_normal_keys_untouched(self):
|
||||||
|
from tools.registry import _sanitize_kwargs
|
||||||
|
out = _sanitize_kwargs({"filepath": "/tmp/a", "question": "hi"})
|
||||||
|
assert out == {"filepath": "/tmp/a", "question": "hi"}
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== L12:failover 日志 Bearer 脱敏 ====================
|
||||||
|
|
||||||
|
class TestL12LogSanitize:
|
||||||
|
def test_sanitize_msg_redacts_bearer(self):
|
||||||
|
msg = "瞬时故障: HTTP 401: Bearer sk-abc123xyz"
|
||||||
|
out = _sanitize_log_msg(msg)
|
||||||
|
assert "sk-abc123xyz" not in out
|
||||||
|
assert "Bearer ***" in out
|
||||||
|
|
||||||
|
def test_failover_log_redacts_bearer(self):
|
||||||
|
logger = mock.MagicMock()
|
||||||
|
c = FailoverClient([], logger=logger)
|
||||||
|
c._log("warning", "瞬时故障(Bearer sk-secret456)")
|
||||||
|
logged = logger.warning.call_args[0][0]
|
||||||
|
assert "sk-secret456" not in logged
|
||||||
|
assert "Bearer ***" in logged
|
||||||
@@ -0,0 +1,237 @@
|
|||||||
|
"""aibox 第四轮审计 Batch I 修复回归测试(2026-08-18,审计报告 §十二)
|
||||||
|
|
||||||
|
覆盖:
|
||||||
|
- N9:prompt 写入内容长度上限(持久化注入面约束)
|
||||||
|
- Q7b:load_conversation 越界路径拒绝(realpath 前缀校验)
|
||||||
|
- F5:run_light_subagent 取消传播(流式可中断 + 父已取消早退)
|
||||||
|
- Q5:提醒时间格式异常的任务按到期返回(不再静默丢弃)
|
||||||
|
- Q8:对话文件原子写(无 .tmp 残留、JSON 完整)
|
||||||
|
- F6:WebSocketHandler 命令专用线程池存在且可释放
|
||||||
|
- limit=0 语义修正(None 才不加 LIMIT)
|
||||||
|
|
||||||
|
Q4(busy_timeout 10s)为时序性行为(默认 5s 等待已能通过短竞争),
|
||||||
|
不做慢速单测;F9(--once 硬超时)为进程级行为,由 CLI 集成面覆盖。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import queue
|
||||||
|
import sqlite3
|
||||||
|
import sys
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
import types
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
import aibox # noqa: E402
|
||||||
|
from aibox import run_light_subagent # noqa: E402
|
||||||
|
|
||||||
|
from subsystems.db import MemoryDatabase, MemoDatabase # noqa: E402
|
||||||
|
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||||
|
from subsystems.conversation_history import ConversationHistory # noqa: E402
|
||||||
|
from tools.prompt_tools import ( # noqa: E402
|
||||||
|
UpdateIdentityTool,
|
||||||
|
UpdateSoulTool,
|
||||||
|
SaveMemoryAndEndConversationTool,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== N9:prompt 写入长度上限 ====================
|
||||||
|
|
||||||
|
class TestN9PromptContentCap:
|
||||||
|
def test_update_identity_oversized_rejected(self, tmp_path, monkeypatch):
|
||||||
|
"""manager 层兜底:超长身份内容拒绝写入"""
|
||||||
|
pm = PromptManager()
|
||||||
|
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||||
|
assert pm.update_identity("x" * (PromptManager.MAX_CONTENT_LENGTH + 1)) is False
|
||||||
|
assert not (tmp_path / pm.IDENTITY_FILE).exists()
|
||||||
|
|
||||||
|
def test_update_identity_tool_friendly_error(self, tmp_path):
|
||||||
|
"""工具层:超长内容返回明确提示而非写入"""
|
||||||
|
pm = PromptManager()
|
||||||
|
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||||
|
tool = UpdateIdentityTool(pm, None)
|
||||||
|
r = tool.execute(content="x" * (PromptManager.MAX_CONTENT_LENGTH + 1))
|
||||||
|
assert "内容过长" in r
|
||||||
|
|
||||||
|
def test_update_soul_tool_friendly_error(self, tmp_path):
|
||||||
|
pm = PromptManager()
|
||||||
|
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||||
|
tool = UpdateSoulTool(pm, None)
|
||||||
|
r = tool.execute(content="y" * (PromptManager.MAX_CONTENT_LENGTH + 1))
|
||||||
|
assert "内容过长" in r
|
||||||
|
|
||||||
|
def test_save_memory_status_oversized_rejected(self, tmp_path):
|
||||||
|
"""save_memory_and_end_conversation:状态内容超长拒绝"""
|
||||||
|
pm = PromptManager()
|
||||||
|
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||||
|
tool = SaveMemoryAndEndConversationTool(
|
||||||
|
exit_callback=lambda *a: None,
|
||||||
|
memory_db=None, prompt_manager=pm, logger=None)
|
||||||
|
r = tool.execute(summary="这是一段足够长的对话总结" * 20,
|
||||||
|
current_status="x" * PromptManager.MAX_CONTENT_LENGTH)
|
||||||
|
assert "状态内容过长" in r
|
||||||
|
|
||||||
|
def test_normal_content_still_written(self, tmp_path):
|
||||||
|
pm = PromptManager()
|
||||||
|
with mock.patch.object(pm, "_prompts_dir", str(tmp_path)):
|
||||||
|
assert pm.update_identity("正常身份内容") is True
|
||||||
|
assert "正常身份内容" in (tmp_path / pm.IDENTITY_FILE).read_text(encoding="utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== Q7b:load_conversation 路径校验 ====================
|
||||||
|
|
||||||
|
class TestQ7bLoadConversationBoundary:
|
||||||
|
def test_traversal_rejected(self, tmp_path):
|
||||||
|
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||||
|
secret = tmp_path / "secret.json"
|
||||||
|
secret.write_text('{"leak": true}', encoding="utf-8")
|
||||||
|
assert h.load_conversation("../secret.json") is None
|
||||||
|
|
||||||
|
def test_normal_file_still_loads(self, tmp_path):
|
||||||
|
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||||
|
h.current_conversation = {
|
||||||
|
"id": "1", "start_time": "2026-08-18 00:00:00", "end_time": None,
|
||||||
|
"topic": "t", "messages": [{"role": "user", "content": "hi"}],
|
||||||
|
"tool_calls": [], "stats": {},
|
||||||
|
}
|
||||||
|
h.save_current_conversation()
|
||||||
|
h.archive_current_conversation()
|
||||||
|
archives = list(h.archive_dir.glob("conversation_*.json"))
|
||||||
|
assert archives
|
||||||
|
loaded = h.load_conversation(archives[0].name)
|
||||||
|
assert loaded is not None and loaded["id"] == "1"
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== F5:run_light_subagent 取消传播 ====================
|
||||||
|
|
||||||
|
class _BlockingLightClient:
|
||||||
|
"""流式产出首段后卡住,直到 cancel() 才中断的伪 LLM 客户端"""
|
||||||
|
|
||||||
|
model = "fake-model"
|
||||||
|
default_temperature = None
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.cancelled = threading.Event()
|
||||||
|
|
||||||
|
def cancel(self):
|
||||||
|
self.cancelled.set()
|
||||||
|
|
||||||
|
def stream_chat(self, messages, **kwargs):
|
||||||
|
yield {"type": "content", "text": "开头"}
|
||||||
|
while not self.cancelled.wait(timeout=0.1):
|
||||||
|
pass
|
||||||
|
raise RuntimeError("连接已中断(cancel)")
|
||||||
|
|
||||||
|
|
||||||
|
class TestF5LightSubagentCancel:
|
||||||
|
def test_cancelled_before_start_returns_early(self):
|
||||||
|
"""父事件已置位 → 直接返回取消提示,不创建客户端"""
|
||||||
|
parent = threading.Event()
|
||||||
|
parent.set()
|
||||||
|
with mock.patch("aibox.create_llm_client") as mk:
|
||||||
|
r = run_light_subagent("r", "t", None, cancel_event=parent)
|
||||||
|
assert "取消" in r
|
||||||
|
mk.assert_not_called()
|
||||||
|
|
||||||
|
def test_cancel_during_stream_aborts(self):
|
||||||
|
"""流式中父级取消 → 观察者线程 cancel 客户端 → 立即中断返回"""
|
||||||
|
parent = threading.Event()
|
||||||
|
fake = _BlockingLightClient()
|
||||||
|
with mock.patch("aibox.create_llm_client", return_value=fake):
|
||||||
|
result_holder = {}
|
||||||
|
|
||||||
|
def _run():
|
||||||
|
result_holder["r"] = run_light_subagent(
|
||||||
|
"r", "t", None, cancel_event=parent)
|
||||||
|
|
||||||
|
th = threading.Thread(target=_run, daemon=True)
|
||||||
|
th.start()
|
||||||
|
time.sleep(0.5) # 等待首段产出、观察者就绪
|
||||||
|
parent.set()
|
||||||
|
th.join(timeout=5)
|
||||||
|
assert not th.is_alive() # 未挂死
|
||||||
|
assert "失败" in result_holder.get("r", "")
|
||||||
|
|
||||||
|
def test_express_emotion_registered_with_parent_event(self):
|
||||||
|
"""DeepSeekChat 注册 ExpressEmotionTool 时注入自身取消事件"""
|
||||||
|
from test_once_mode import _make_ai
|
||||||
|
ai = _make_ai()
|
||||||
|
tool = next((t for t in ai.tool_registry.get_all_tools()
|
||||||
|
if t.get_name() == "express_emotion"), None)
|
||||||
|
assert tool is not None
|
||||||
|
assert tool.cancel_event is ai._cancel_event
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== Q5:格式异常任务不丢弃 ====================
|
||||||
|
|
||||||
|
class TestQ5MalformedReminderTime:
|
||||||
|
def test_malformed_time_task_returned_as_due(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
# J1(审计 2026-08-19):add_memo 已拒绝畸形时间(入库校验),
|
||||||
|
# 模拟校验落地前的存量脏数据需绕过 add_memo 直接 SQL 插入
|
||||||
|
import sqlite3
|
||||||
|
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||||
|
cur = conn.cursor()
|
||||||
|
cur.execute('''INSERT INTO memos
|
||||||
|
(title, content, created_at, reminder_time, is_completed, source,
|
||||||
|
repeat_type, repeat_interval, repeat_interval_unit, is_immediate)
|
||||||
|
VALUES (?, ?, ?, ?, 0, 'test', 'none', 0, 'days', 0)''',
|
||||||
|
("格式异常", "c", "2026-08-18 00:00:00", "不是时间"))
|
||||||
|
mid = cur.lastrowid
|
||||||
|
pending = db.get_pending_memos(before_time="2026-08-18 00:00:00")
|
||||||
|
assert any(m["id"] == mid for m in pending)
|
||||||
|
|
||||||
|
def test_normal_pending_still_filtered(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
db.add_memo(title="未来任务", content="c", reminder_time="2099-01-01 00:00:00")
|
||||||
|
pending = db.get_pending_memos(before_time="2026-08-18 00:00:00")
|
||||||
|
assert pending == []
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== Q8:对话文件原子写 ====================
|
||||||
|
|
||||||
|
class TestQ8AtomicWrites:
|
||||||
|
def test_no_tmp_leftover_and_valid_json(self, tmp_path):
|
||||||
|
h = ConversationHistory(save_dir=str(tmp_path / "save"), load_history=False)
|
||||||
|
h.current_conversation = {
|
||||||
|
"id": "1", "start_time": "2026-08-18 00:00:00", "end_time": None,
|
||||||
|
"topic": "t", "messages": [], "tool_calls": [], "stats": {},
|
||||||
|
}
|
||||||
|
h.save_current_conversation()
|
||||||
|
assert h.current_file.exists() # 文件名可能被用户配置覆盖,用实例属性
|
||||||
|
import json
|
||||||
|
json.loads(h.current_file.read_text(encoding="utf-8")) # 不抛即合法 JSON
|
||||||
|
h.archive_current_conversation() # 归档后 current 文件按设计移除
|
||||||
|
assert list(Path(tmp_path).rglob("*.tmp")) == []
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== F6:WS 命令专用线程池 ====================
|
||||||
|
|
||||||
|
class TestF6CmdExecutor:
|
||||||
|
def test_handler_has_dedicated_executor(self):
|
||||||
|
from aibox import WebSocketHandler
|
||||||
|
h = WebSocketHandler(ai_manager=None)
|
||||||
|
assert h._cmd_executor is not None
|
||||||
|
h.shutdown()
|
||||||
|
assert h._cmd_executor is None
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== limit=0 语义 ====================
|
||||||
|
|
||||||
|
class TestLimitZeroSemantics:
|
||||||
|
def test_pending_memos_limit_zero_returns_none(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
db.add_memo(title="t", content="c", reminder_time="2026-01-01 00:00:00")
|
||||||
|
assert db.get_pending_memos(limit=0, before_time="2026-08-18 00:00:00") == []
|
||||||
|
|
||||||
|
def test_immediate_memos_limit_zero_returns_none(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
db.add_memo(title="t", content="c", is_immediate=True)
|
||||||
|
assert db.get_immediate_memos(limit=0) == []
|
||||||
@@ -0,0 +1,540 @@
|
|||||||
|
"""aibox 第五轮审计(Batch K)修复回归测试(2026-08-19)
|
||||||
|
|
||||||
|
波次 1:
|
||||||
|
- J3:FTS 外部内容表删除换序后索引无残留(delete_memory / delete_memories_by_ids)
|
||||||
|
- J4:model_config 连通性验证异常消息 Bearer 令牌脱敏
|
||||||
|
- J6:PromptManager/ConfigManager 单例并发初始化不双重加载
|
||||||
|
- J8:add_profile 对非法温度/超时的容错(交互层不再抢先 float()/int())
|
||||||
|
|
||||||
|
波次 2(J1 任务系统投毒):
|
||||||
|
- add_memo 入库校验:畸形时间/重复间隔/重复类型拒绝,三入口一处覆盖
|
||||||
|
- 存量脏数据防御:调度器预检、mark_as_triggered、展示路径不再打崩
|
||||||
|
- WS task add 友好错误 + register 握手异常可回收
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sqlite3
|
||||||
|
import sys
|
||||||
|
import threading
|
||||||
|
import types
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
import aibox # noqa: E402
|
||||||
|
import model_config # noqa: E402
|
||||||
|
from common import ConfigManager # noqa: E402
|
||||||
|
from subsystems.db import MemoDatabase, MemoryDatabase # noqa: E402
|
||||||
|
from subsystems.prompt_manager import PromptManager # noqa: E402
|
||||||
|
|
||||||
|
from test_model_config import _cfg # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J3:FTS 删除顺序 ====================
|
||||||
|
|
||||||
|
def _fts_hits(db_path, keyword):
|
||||||
|
conn = sqlite3.connect(str(db_path))
|
||||||
|
try:
|
||||||
|
return conn.execute(
|
||||||
|
"SELECT rowid FROM memories_fts WHERE memories_fts MATCH ?",
|
||||||
|
(f'"{keyword}"',)).fetchall()
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
def _add_mem(db, content):
|
||||||
|
return db.add_memory(topic="话题", summary="摘要", content=content,
|
||||||
|
start_time="2026-08-19 00:00:00",
|
||||||
|
end_time="2026-08-19 00:00:01", importance=2)
|
||||||
|
|
||||||
|
|
||||||
|
class TestJ3FtsDeleteOrder:
|
||||||
|
def test_delete_memory_leaves_no_fts_orphan(self, tmp_path):
|
||||||
|
"""单条删除后 FTS 索引不得残留(修复前 MATCH 仍命中已删 rowid)"""
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||||
|
mid = _add_mem(db, "机密内容甲")
|
||||||
|
assert _fts_hits(tmp_path / "mem.db", "机密内容甲")
|
||||||
|
assert db.delete_memory(mid) is True
|
||||||
|
assert _fts_hits(tmp_path / "mem.db", "机密内容甲") == []
|
||||||
|
|
||||||
|
def test_batch_delete_leaves_no_fts_orphan(self, tmp_path):
|
||||||
|
"""批量删除同样不留索引残留"""
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||||
|
ids = [_add_mem(db, f"机密内容乙{i}") for i in range(3)]
|
||||||
|
assert db.delete_memories_by_ids(ids) == 3
|
||||||
|
for i in range(3):
|
||||||
|
assert _fts_hits(tmp_path / "mem.db", f"机密内容乙{i}") == []
|
||||||
|
|
||||||
|
def test_search_after_delete_no_hit(self, tmp_path):
|
||||||
|
"""删除后搜索(FTS 路径)不再命中已删内容"""
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||||
|
mid = _add_mem(db, "独特关键词丙")
|
||||||
|
assert db.search_memories(["独特关键词丙"])
|
||||||
|
db.delete_memory(mid)
|
||||||
|
assert db.search_memories(["独特关键词丙"]) == []
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J4:验证异常脱敏 ====================
|
||||||
|
|
||||||
|
class TestJ4TestProfileSanitize:
|
||||||
|
def _models_cfg(self, tmp_path):
|
||||||
|
return _cfg(tmp_path, {"models": {"bad": {
|
||||||
|
"api_url": "https://api.example.com/v1/chat/completions",
|
||||||
|
"model": "m", "api_key_env": "K"}}})
|
||||||
|
|
||||||
|
def test_bearer_token_masked_in_failure_message(self, tmp_path):
|
||||||
|
"""网关错误页回显 Authorization 头时,令牌不得出现在终端回显"""
|
||||||
|
cm = self._models_cfg(tmp_path)
|
||||||
|
err = RuntimeError("HTTP 502: Authorization: Bearer sk-secret12345")
|
||||||
|
with mock.patch("llm_client.create_llm_client", side_effect=err):
|
||||||
|
result = model_config.test_profile(cm, "bad")
|
||||||
|
assert "sk-secret12345" not in result
|
||||||
|
assert "Bearer ***" in result
|
||||||
|
|
||||||
|
def test_plain_error_message_preserved(self, tmp_path):
|
||||||
|
"""无敏感内容的错误消息原样保留(脱敏不误伤)"""
|
||||||
|
cm = self._models_cfg(tmp_path)
|
||||||
|
with mock.patch("llm_client.create_llm_client",
|
||||||
|
side_effect=RuntimeError("连接被拒绝")):
|
||||||
|
result = model_config.test_profile(cm, "bad")
|
||||||
|
assert "连接被拒绝" in result
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J6:单例并发初始化 ====================
|
||||||
|
|
||||||
|
class TestJ6SingletonInitRace:
|
||||||
|
def test_prompt_manager_concurrent_init_no_duplication(self, tmp_path):
|
||||||
|
"""多线程首例并发:单实例且提示词只加载一份(修复前整份翻倍)"""
|
||||||
|
PromptManager._instance = None
|
||||||
|
PromptManager._prompts = []
|
||||||
|
PromptManager._prompts_dir = None
|
||||||
|
prompts_dir = tmp_path / "prompts"
|
||||||
|
prompts_dir.mkdir()
|
||||||
|
try:
|
||||||
|
barrier = threading.Barrier(8)
|
||||||
|
instances = [None] * 8
|
||||||
|
|
||||||
|
def worker(k):
|
||||||
|
barrier.wait()
|
||||||
|
instances[k] = PromptManager()
|
||||||
|
|
||||||
|
with mock.patch.object(ConfigManager, "get",
|
||||||
|
return_value=str(prompts_dir)):
|
||||||
|
threads = [threading.Thread(target=worker, args=(k,))
|
||||||
|
for k in range(8)]
|
||||||
|
for t in threads:
|
||||||
|
t.start()
|
||||||
|
for t in threads:
|
||||||
|
t.join()
|
||||||
|
|
||||||
|
assert len({id(i) for i in instances}) == 1
|
||||||
|
files = [p["file"] for p in PromptManager._prompts]
|
||||||
|
assert files, "提示词应已加载"
|
||||||
|
assert len(files) == len(set(files)), f"提示词被重复加载: {files}"
|
||||||
|
finally:
|
||||||
|
PromptManager._instance = None
|
||||||
|
PromptManager._prompts = []
|
||||||
|
PromptManager._prompts_dir = None
|
||||||
|
|
||||||
|
def test_config_manager_concurrent_init_single_load(self, tmp_path):
|
||||||
|
"""多线程首例并发:单实例且 _load_config 只执行一次"""
|
||||||
|
cfg_file = tmp_path / "config.json"
|
||||||
|
calls = {"n": 0}
|
||||||
|
original_load = ConfigManager._load_config
|
||||||
|
|
||||||
|
def counting_load(self):
|
||||||
|
calls["n"] += 1
|
||||||
|
original_load(self)
|
||||||
|
|
||||||
|
ConfigManager._instance = None
|
||||||
|
ConfigManager._config = None
|
||||||
|
ConfigManager._config_path = None
|
||||||
|
try:
|
||||||
|
barrier = threading.Barrier(8)
|
||||||
|
instances = [None] * 8
|
||||||
|
|
||||||
|
def worker(k):
|
||||||
|
barrier.wait()
|
||||||
|
instances[k] = ConfigManager()
|
||||||
|
|
||||||
|
with mock.patch.dict(os.environ, {"AI_CONFIG_PATH": str(cfg_file)}), \
|
||||||
|
mock.patch.object(ConfigManager, "_load_config", counting_load):
|
||||||
|
threads = [threading.Thread(target=worker, args=(k,))
|
||||||
|
for k in range(8)]
|
||||||
|
for t in threads:
|
||||||
|
t.start()
|
||||||
|
for t in threads:
|
||||||
|
t.join()
|
||||||
|
|
||||||
|
assert len({id(i) for i in instances}) == 1
|
||||||
|
assert calls["n"] == 1
|
||||||
|
finally:
|
||||||
|
ConfigManager._instance = None
|
||||||
|
ConfigManager._config = None
|
||||||
|
ConfigManager._config_path = None
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J8:交互输入容错 ====================
|
||||||
|
|
||||||
|
class TestJ8NumericTolerance:
|
||||||
|
def test_invalid_temperature_returns_friendly_error(self, tmp_path):
|
||||||
|
"""非数字温度 → 友好错误而非 ValueError(交互层不再抢先 float())"""
|
||||||
|
cm = _cfg(tmp_path, {"models": {}})
|
||||||
|
msgs = model_config.add_profile(
|
||||||
|
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||||
|
api_key_env="K", temperature="not-a-number")
|
||||||
|
assert any(m.startswith("❌") and "温度" in m for m in msgs)
|
||||||
|
assert "p1" not in model_config.get_models(cm)
|
||||||
|
|
||||||
|
def test_invalid_timeout_returns_friendly_error(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {"models": {}})
|
||||||
|
msgs = model_config.add_profile(
|
||||||
|
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||||
|
api_key_env="K", timeout="abc")
|
||||||
|
assert any(m.startswith("❌") and "超时" in m for m in msgs)
|
||||||
|
assert "p1" not in model_config.get_models(cm)
|
||||||
|
|
||||||
|
def test_string_numbers_accepted_and_clamped(self, tmp_path):
|
||||||
|
"""数字字符串正常接受并钳制(交互层原始字符串直传语义)"""
|
||||||
|
cm = _cfg(tmp_path, {"models": {}})
|
||||||
|
msgs = model_config.add_profile(
|
||||||
|
cm, "p1", "https://api.example.com/v1/chat/completions", "m1",
|
||||||
|
api_key_env="K", temperature="0.7", timeout="120")
|
||||||
|
assert all(not m.startswith("❌") for m in msgs)
|
||||||
|
profile = model_config.get_models(cm)["p1"]
|
||||||
|
assert profile["temperature"] == 0.7
|
||||||
|
assert profile["timeout"] == 120
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J1 波次 2:任务系统投毒 ====================
|
||||||
|
|
||||||
|
def _insert_memo_raw(db_path, reminder_time="2026-08-18 00:00:00",
|
||||||
|
repeat_type="none", repeat_interval=0,
|
||||||
|
repeat_interval_unit="days", repeat_end_time=None):
|
||||||
|
"""直接 SQL 插入任务行——模拟校验落地前的存量脏数据(绕过 add_memo)"""
|
||||||
|
with sqlite3.connect(str(db_path)) as conn:
|
||||||
|
cur = conn.cursor()
|
||||||
|
cur.execute('''INSERT INTO memos
|
||||||
|
(title, content, created_at, reminder_time, tags, priority, created_by,
|
||||||
|
is_completed, source, repeat_type, repeat_interval, repeat_interval_unit,
|
||||||
|
repeat_end_time, is_immediate)
|
||||||
|
VALUES ('t', 'c', '2026-08-18 00:00:00', ?, '', 1, 'test',
|
||||||
|
0, 'test', ?, ?, ?, ?, 0)''',
|
||||||
|
(reminder_time, repeat_type, repeat_interval, repeat_interval_unit,
|
||||||
|
repeat_end_time))
|
||||||
|
return cur.lastrowid
|
||||||
|
|
||||||
|
|
||||||
|
class TestJ1AddMemoValidation:
|
||||||
|
def test_malformed_reminder_time_rejected(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
with pytest.raises(ValueError, match="提醒时间格式非法"):
|
||||||
|
db.add_memo(title="t", reminder_time="不是时间")
|
||||||
|
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||||
|
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||||
|
|
||||||
|
def test_malformed_repeat_end_time_rejected(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
with pytest.raises(ValueError, match="重复结束时间格式非法"):
|
||||||
|
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||||
|
repeat_type="daily", repeat_end_time="不是时间")
|
||||||
|
|
||||||
|
def test_invalid_repeat_type_rejected(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
with pytest.raises(ValueError, match="无效的重复类型"):
|
||||||
|
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||||
|
repeat_type="every-second")
|
||||||
|
|
||||||
|
def test_non_numeric_repeat_interval_rejected(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
with pytest.raises(ValueError, match="重复间隔必须是整数"):
|
||||||
|
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||||
|
repeat_type="custom", repeat_interval="垃圾")
|
||||||
|
|
||||||
|
def test_custom_repeat_requires_positive_interval(self, tmp_path):
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
with pytest.raises(ValueError, match="repeat_interval"):
|
||||||
|
db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||||
|
repeat_type="custom", repeat_interval=0)
|
||||||
|
|
||||||
|
def test_numeric_string_interval_coerced(self, tmp_path):
|
||||||
|
"""数字字符串被强转接受(WS 传 JSON 数字/字符串均安全)"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = db.add_memo(title="t", reminder_time="2026-08-18 00:00:00",
|
||||||
|
repeat_type="custom", repeat_interval="30",
|
||||||
|
repeat_interval_unit="minutes")
|
||||||
|
assert db.get_memo(mid)["repeat_interval"] == 30
|
||||||
|
|
||||||
|
def test_add_user_task_propagates_validation(self, tmp_path):
|
||||||
|
"""WS 路径(add_user_task 薄包装 add_memo)同样被校验拦截"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
with pytest.raises(ValueError, match="提醒时间格式非法"):
|
||||||
|
db.add_user_task(title="t", reminder_time="不是时间")
|
||||||
|
|
||||||
|
|
||||||
|
class TestJ1DirtyDataDefense:
|
||||||
|
def test_mark_as_triggered_completes_malformed_one_shot(self, tmp_path):
|
||||||
|
"""存量畸形 reminder_time 一次性任务:触发不再抛异常,正常完成"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||||
|
assert db.mark_as_triggered(mid) is None
|
||||||
|
assert db.get_memo(mid)["is_completed"] == 1
|
||||||
|
|
||||||
|
def test_mark_as_triggered_handles_garbage_repeat_interval(self, tmp_path):
|
||||||
|
"""存量 custom 任务 repeat_interval 为文本:按非重复处理,不打崩事务"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="2026-08-18 00:00:00",
|
||||||
|
repeat_type="custom", repeat_interval="垃圾")
|
||||||
|
assert db.mark_as_triggered(mid) is None
|
||||||
|
assert db.get_memo(mid)["is_completed"] == 1
|
||||||
|
|
||||||
|
def test_malformed_repeat_end_time_treated_as_no_end(self, tmp_path):
|
||||||
|
"""存量畸形 repeat_end_time:解析失败视为无结束时间,不抛异常"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="2026-08-19 00:00:00",
|
||||||
|
repeat_type="daily", repeat_end_time="不是时间")
|
||||||
|
memo = db.get_memo(mid)
|
||||||
|
from datetime import datetime
|
||||||
|
base = datetime.strptime("2026-08-19 00:00:00", "%Y-%m-%d %H:%M:%S")
|
||||||
|
assert db._should_end_repeat(memo, base) is False
|
||||||
|
assert db._calculate_next_reminder(memo, base) is not None
|
||||||
|
|
||||||
|
def test_pending_tasks_for_user_malformed_time_overdue(self, tmp_path):
|
||||||
|
"""展示路径:畸形 reminder_time 行按到期展示,不再打崩握手"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||||
|
tasks = db.get_pending_tasks_for_user()
|
||||||
|
row = next(t for t in tasks if t["id"] == mid)
|
||||||
|
assert row["is_overdue"] is True
|
||||||
|
|
||||||
|
|
||||||
|
class TestJ1WsAndScheduler:
|
||||||
|
def test_ws_task_add_rejects_malformed_repeat_end(self, tmp_path):
|
||||||
|
"""WS task add 传畸形 repeat_end_time:友好错误且不入库"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
ai_manager = mock.MagicMock()
|
||||||
|
ai_manager.memo_db = db
|
||||||
|
handler = aibox.WebSocketHandler(ai_manager)
|
||||||
|
try:
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
asyncio.run(handler.handle_task_message(ws, {
|
||||||
|
"type": "task", "action": "add", "title": "t", "delay": 5,
|
||||||
|
"repeat_type": "daily", "repeat_end_time": "不是时间"}))
|
||||||
|
sent = json.loads(ws.send.call_args[0][0])
|
||||||
|
assert sent["type"] == "error"
|
||||||
|
assert "任务添加失败" in sent["content"]
|
||||||
|
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||||
|
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
def test_scheduler_precheck_malformed_time_executes_not_loops(self, tmp_path):
|
||||||
|
"""调度器预检:畸形 reminder_time 按到期执行并标记触发(不再热循环)"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
mid = _insert_memo_raw(tmp_path / "m.db", reminder_time="不是时间")
|
||||||
|
loop = asyncio.new_event_loop()
|
||||||
|
executed = {"n": 0}
|
||||||
|
try:
|
||||||
|
scheduler = aibox.TaskScheduler(loop=loop, logger=None)
|
||||||
|
scheduler.broadcast_enabled = False
|
||||||
|
|
||||||
|
async def fake_execute(memo, memo_db):
|
||||||
|
executed["n"] += 1
|
||||||
|
return "ok"
|
||||||
|
|
||||||
|
scheduler.execute_task_as_conversation = fake_execute
|
||||||
|
loop.run_until_complete(scheduler.check_and_execute_tasks(db))
|
||||||
|
finally:
|
||||||
|
loop.close()
|
||||||
|
assert executed["n"] == 1
|
||||||
|
assert db.get_memo(mid)["is_completed"] == 1
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J5 波次 4:原子写 + 损坏备份 ====================
|
||||||
|
|
||||||
|
class TestJ5AtomicWrites:
|
||||||
|
def test_save_config_atomic_no_tmp_leftover(self, tmp_path):
|
||||||
|
"""_save_config 原子写:无 .aibox_tmp_ 残留且 JSON 合法"""
|
||||||
|
cm = ConfigManager()
|
||||||
|
cm._config_path = str(tmp_path / "config.json")
|
||||||
|
cm._config = {"api": {"deepseek_model": "m"}, "models": {}}
|
||||||
|
assert cm._save_config() is True
|
||||||
|
import json as _json
|
||||||
|
data = _json.loads((tmp_path / "config.json").read_text(encoding="utf-8"))
|
||||||
|
assert data["api"]["deepseek_model"] == "m"
|
||||||
|
assert list(tmp_path.glob(".aibox_tmp_*")) == []
|
||||||
|
|
||||||
|
def test_corrupted_config_backed_up_before_fallback(self, tmp_path):
|
||||||
|
"""损坏配置加载:备份 .corrupt.* 文件后回退默认,原内容不丢"""
|
||||||
|
cfg_file = tmp_path / "config.json"
|
||||||
|
cfg_file.write_text('{"api": {"broken": ', encoding="utf-8") # 半截 JSON
|
||||||
|
ConfigManager._instance = None
|
||||||
|
ConfigManager._config = None
|
||||||
|
ConfigManager._config_path = None
|
||||||
|
try:
|
||||||
|
with mock.patch.dict(os.environ, {"AI_CONFIG_PATH": str(cfg_file)}):
|
||||||
|
cm = ConfigManager()
|
||||||
|
assert cm._config is not None
|
||||||
|
backups = list(tmp_path.glob("config.json.corrupt.*"))
|
||||||
|
assert len(backups) == 1
|
||||||
|
assert backups[0].read_text(encoding="utf-8") == '{"api": {"broken": '
|
||||||
|
finally:
|
||||||
|
ConfigManager._instance = None
|
||||||
|
ConfigManager._config = None
|
||||||
|
ConfigManager._config_path = None
|
||||||
|
|
||||||
|
def test_prompt_write_atomic(self, tmp_path):
|
||||||
|
"""prompt 更新原子写:内容生效且无临时文件残留"""
|
||||||
|
from subsystems.prompt_manager import PromptManager
|
||||||
|
PromptManager._instance = None
|
||||||
|
PromptManager._prompts = []
|
||||||
|
PromptManager._prompts_dir = None
|
||||||
|
prompts_dir = tmp_path / "prompts"
|
||||||
|
prompts_dir.mkdir()
|
||||||
|
try:
|
||||||
|
with mock.patch.object(ConfigManager, "get",
|
||||||
|
return_value=str(prompts_dir)):
|
||||||
|
pm = PromptManager()
|
||||||
|
assert pm.update_identity("新身份内容") is True
|
||||||
|
identity = prompts_dir / "00_IDENTITY.md"
|
||||||
|
assert "新身份内容" in identity.read_text(encoding="utf-8")
|
||||||
|
assert list(tmp_path.glob(".aibox_tmp_*")) == []
|
||||||
|
finally:
|
||||||
|
PromptManager._instance = None
|
||||||
|
PromptManager._prompts = []
|
||||||
|
PromptManager._prompts_dir = None
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J7 波次 4:broadcast 超时回收 ====================
|
||||||
|
|
||||||
|
class TestJ7BroadcastTimeout:
|
||||||
|
def test_stuck_client_timed_out_and_unregistered(self):
|
||||||
|
"""半开 TCP 僵死客户端:broadcast 在超时内返回并回收该连接"""
|
||||||
|
from aibox import WebSocketHandler
|
||||||
|
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||||
|
handler.broadcast_send_timeout = 0.1
|
||||||
|
try:
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
|
||||||
|
async def stuck_send(*args, **kwargs):
|
||||||
|
await asyncio.sleep(30)
|
||||||
|
|
||||||
|
ws.send.side_effect = stuck_send
|
||||||
|
handler.connected_clients.add(ws)
|
||||||
|
asyncio.run(handler.broadcast({"type": "ping"}))
|
||||||
|
assert ws not in handler.connected_clients
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
def test_healthy_client_stays_connected(self):
|
||||||
|
"""正常客户端:发送成功且不被回收"""
|
||||||
|
from aibox import WebSocketHandler
|
||||||
|
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||||
|
handler.broadcast_send_timeout = 1.0
|
||||||
|
try:
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
handler.connected_clients.add(ws)
|
||||||
|
asyncio.run(handler.broadcast({"type": "ping"}))
|
||||||
|
assert ws in handler.connected_clients
|
||||||
|
assert ws.send.called
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J9 波次 5:MemoDatabase 连接层 ====================
|
||||||
|
|
||||||
|
class TestJ9MemoDbConnect:
|
||||||
|
def test_wal_enabled_after_init(self, tmp_path):
|
||||||
|
"""WAL 显式设置——不再依赖与 MemoryDatabase 的创建顺序"""
|
||||||
|
MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||||
|
mode = conn.execute("PRAGMA journal_mode").fetchone()[0]
|
||||||
|
assert mode == "wal"
|
||||||
|
|
||||||
|
def test_connection_factory_sets_busy_timeout(self, tmp_path):
|
||||||
|
"""统一连接工厂:busy_timeout 覆盖全部写路径"""
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
conn = db._connect()
|
||||||
|
try:
|
||||||
|
assert conn.execute("PRAGMA busy_timeout").fetchone()[0] == 10000
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J10 波次 5:memo 批量删除 ====================
|
||||||
|
|
||||||
|
class TestJ10BatchDelete:
|
||||||
|
def test_batch_delete_via_memo_tool(self, tmp_path):
|
||||||
|
"""schema 声明的 memo_ids 批量删除真正生效(此前静默忽略)"""
|
||||||
|
from tools.db_tools import MemoTool
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
ids = [db.add_memo(title=f"任务{i}",
|
||||||
|
reminder_time="2026-08-18 00:00:00") for i in range(3)]
|
||||||
|
tool = MemoTool(db)
|
||||||
|
r = tool.execute(action="delete", confirm=True, memo_ids=ids)
|
||||||
|
assert "3/3" in r
|
||||||
|
with sqlite3.connect(str(tmp_path / "m.db")) as conn:
|
||||||
|
assert conn.execute("SELECT COUNT(*) FROM memos").fetchone()[0] == 0
|
||||||
|
|
||||||
|
def test_batch_delete_requires_list(self, tmp_path):
|
||||||
|
from tools.db_tools import MemoTool
|
||||||
|
db = MemoDatabase(db_path=str(tmp_path / "m.db"))
|
||||||
|
tool = MemoTool(db)
|
||||||
|
r = tool.execute(action="delete", confirm=True, memo_ids="1")
|
||||||
|
assert "列表" in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J15 波次 5:未知 task action ====================
|
||||||
|
|
||||||
|
class TestJ15UnknownTaskAction:
|
||||||
|
def test_unknown_action_returns_error(self):
|
||||||
|
"""未知 task action 明确报错(此前静默吞掉)"""
|
||||||
|
from aibox import WebSocketHandler
|
||||||
|
handler = WebSocketHandler(ai_manager=mock.MagicMock())
|
||||||
|
try:
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
asyncio.run(handler.handle_task_message(ws, {"action": "explode"}))
|
||||||
|
sent = json.loads(ws.send.call_args[0][0])
|
||||||
|
assert sent["type"] == "error"
|
||||||
|
assert "未知的任务操作" in sent["content"]
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== J16 波次 5:skill_manager ====================
|
||||||
|
|
||||||
|
class TestJ16SkillManager:
|
||||||
|
def test_markdown_skipped_when_executable_same_name(self, tmp_path, monkeypatch):
|
||||||
|
"""同名键冲突:可执行技能确定性优先(此前由 listdir 顺序决定)"""
|
||||||
|
from common import constants
|
||||||
|
from subsystems.skill_manager import SkillManager
|
||||||
|
skills_dir = tmp_path / "skills"
|
||||||
|
skills_dir.mkdir()
|
||||||
|
exe = skills_dir / "foo"
|
||||||
|
exe.write_text("#!/bin/sh\necho '可执行技能简介'", encoding="utf-8")
|
||||||
|
exe.chmod(0o755)
|
||||||
|
(skills_dir / "foo.md").write_text("# Markdown 简介\n内容", encoding="utf-8")
|
||||||
|
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||||
|
property(lambda self: str(skills_dir)))
|
||||||
|
sm = SkillManager()
|
||||||
|
assert "foo" in sm.skills
|
||||||
|
assert sm.skills["foo"]["is_executable"] is True
|
||||||
|
|
||||||
|
def test_help_output_truncated(self, tmp_path, monkeypatch):
|
||||||
|
"""--help 输出截断(内容进系统提示词,防撑爆上下文)"""
|
||||||
|
from common import constants
|
||||||
|
from subsystems.skill_manager import SkillManager
|
||||||
|
skills_dir = tmp_path / "skills"
|
||||||
|
skills_dir.mkdir()
|
||||||
|
exe = skills_dir / "big"
|
||||||
|
exe.write_text("#!/bin/sh\necho '" + "x" * 5000 + "'", encoding="utf-8")
|
||||||
|
exe.chmod(0o755)
|
||||||
|
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||||
|
property(lambda self: str(skills_dir)))
|
||||||
|
sm = SkillManager()
|
||||||
|
assert len(sm.skills["big"]["content"]) <= SkillManager.MAX_HELP_LENGTH + 20
|
||||||
@@ -0,0 +1,157 @@
|
|||||||
|
"""aibox 第六轮审计(Batch L)修复回归测试(2026-08-20)
|
||||||
|
|
||||||
|
R1:/memos WS 命令对畸形 reminder_time 不再整体报错(J1 加固漏点)
|
||||||
|
R2:handle_interrupt 只取消本会话 AI,不再触碰 CLI 主 AI 的全局 executor
|
||||||
|
R3:unregister 补齐连接回收(关闭 websocket + 取消在途处理任务)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import sys
|
||||||
|
import types
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
import aibox # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def _make_handler():
|
||||||
|
handler = aibox.WebSocketHandler(ai_manager=mock.MagicMock())
|
||||||
|
return handler
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== R1:/memos 畸形时间容错 ====================
|
||||||
|
|
||||||
|
class TestR1MemosCommand:
|
||||||
|
def test_malformed_reminder_time_shows_pending(self):
|
||||||
|
"""存量脏数据(畸形 reminder_time):/memos 按"待执行"展示,不抛 ValueError"""
|
||||||
|
handler = _make_handler()
|
||||||
|
try:
|
||||||
|
handler.ai_manager.memo_db.get_all_memos.return_value = [
|
||||||
|
{"id": 1, "title": "脏数据任务", "reminder_time": "not-a-time",
|
||||||
|
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||||
|
{"id": 2, "title": "正常过期任务", "reminder_time": "2020-01-01 00:00:00",
|
||||||
|
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||||
|
]
|
||||||
|
sent = []
|
||||||
|
|
||||||
|
async def fake_send(ws, msg):
|
||||||
|
sent.append(msg)
|
||||||
|
|
||||||
|
handler.send_message = fake_send
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
asyncio.run(handler.handle_command(ws, {"command": "/memos"}))
|
||||||
|
|
||||||
|
results = [m for m in sent if m.get("type") == "command_result"]
|
||||||
|
errors = [m for m in sent if m.get("type") == "error"]
|
||||||
|
assert len(results) == 1
|
||||||
|
assert not errors # 不再落入顶层兜底"处理失败"
|
||||||
|
content = results[0]["content"]
|
||||||
|
assert "脏数据任务" in content and "正常过期任务" in content
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
def test_valid_reminder_time_still_marks_overdue(self):
|
||||||
|
"""合法过期时间仍标 overdue(语义不回归)"""
|
||||||
|
handler = _make_handler()
|
||||||
|
try:
|
||||||
|
handler.ai_manager.memo_db.get_all_memos.return_value = [
|
||||||
|
{"id": 1, "title": "过期任务", "reminder_time": "2020-01-01 00:00:00",
|
||||||
|
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||||
|
{"id": 2, "title": "未来任务", "reminder_time": "2099-01-01 00:00:00",
|
||||||
|
"repeat_type": "none", "repeat_interval": 0, "is_immediate": False},
|
||||||
|
]
|
||||||
|
sent = []
|
||||||
|
|
||||||
|
async def fake_send(ws, msg):
|
||||||
|
sent.append(msg)
|
||||||
|
|
||||||
|
handler.send_message = fake_send
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
asyncio.run(handler.handle_command(ws, {"command": "/memos"}))
|
||||||
|
|
||||||
|
content = next(m["content"] for m in sent
|
||||||
|
if m.get("type") == "command_result")
|
||||||
|
overdue = aibox.i18n.get('memo_overdue')
|
||||||
|
pending = aibox.i18n.get('memo_pending')
|
||||||
|
overdue_line = next(l for l in content.splitlines() if "过期任务" in l)
|
||||||
|
pending_line = next(l for l in content.splitlines() if "未来任务" in l)
|
||||||
|
assert overdue in overdue_line and pending not in overdue_line
|
||||||
|
assert pending in pending_line and overdue not in pending_line
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== R2:interrupt 不再误杀 CLI 主 AI 子进程 ====================
|
||||||
|
|
||||||
|
class TestR2InterruptScope:
|
||||||
|
def test_interrupt_does_not_touch_global_executor(self):
|
||||||
|
"""WS 中断不再取消 ai_manager.command_executor(CLI 主 AI 的执行器)"""
|
||||||
|
handler = _make_handler()
|
||||||
|
try:
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
asyncio.run(handler.handle_interrupt(ws, {}))
|
||||||
|
handler.ai_manager.command_executor.cancel.assert_not_called()
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
def test_interrupt_still_cancels_session_ai(self):
|
||||||
|
"""本会话 AI 仍被正常取消(中断能力不回归)"""
|
||||||
|
handler = _make_handler()
|
||||||
|
try:
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
ai = mock.MagicMock()
|
||||||
|
handler.client_sessions[ws] = {"ai": ai, "is_processing": True}
|
||||||
|
asyncio.run(handler.handle_interrupt(ws, {}))
|
||||||
|
ai.cancel.assert_called_once()
|
||||||
|
handler.ai_manager.command_executor.cancel.assert_not_called()
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== R3:unregister 补齐连接回收 ====================
|
||||||
|
|
||||||
|
class TestR3UnregisterCleanup:
|
||||||
|
def test_unregister_closes_websocket_and_cancels_tasks(self):
|
||||||
|
"""broadcast 超时路径走 unregister 后:websocket 被关闭、在途任务被取消、登记清空"""
|
||||||
|
handler = _make_handler()
|
||||||
|
try:
|
||||||
|
ws = mock.AsyncMock()
|
||||||
|
|
||||||
|
async def scenario():
|
||||||
|
task = asyncio.create_task(asyncio.sleep(30))
|
||||||
|
handler.connected_clients.add(ws)
|
||||||
|
handler.client_sessions[ws] = {"ai": None, "is_processing": True}
|
||||||
|
handler._processing_tasks[ws] = {task}
|
||||||
|
await handler.unregister(ws)
|
||||||
|
await asyncio.sleep(0) # 让 cancel 传播
|
||||||
|
assert task.done() # 取消即视为完成(cancelled 或已收尾)
|
||||||
|
|
||||||
|
asyncio.run(scenario())
|
||||||
|
assert ws.close.called
|
||||||
|
assert ws not in handler.connected_clients
|
||||||
|
assert ws not in handler.client_sessions
|
||||||
|
assert ws not in handler._processing_tasks
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
|
|
||||||
|
def test_unregister_idempotent_and_mock_safe(self):
|
||||||
|
"""重复注销幂等;close 抛异常的对象不炸(AttributeError 被兜底)"""
|
||||||
|
class _BadCloseWS:
|
||||||
|
remote_address = ("127.0.0.1", 1)
|
||||||
|
|
||||||
|
async def close(self):
|
||||||
|
raise AttributeError("no close")
|
||||||
|
|
||||||
|
handler = _make_handler()
|
||||||
|
try:
|
||||||
|
ws = _BadCloseWS()
|
||||||
|
handler.connected_clients.add(ws)
|
||||||
|
handler.client_sessions[ws] = {"ai": None}
|
||||||
|
asyncio.run(handler.unregister(ws))
|
||||||
|
asyncio.run(handler.unregister(ws)) # 二次注销不抛
|
||||||
|
assert ws not in handler.connected_clients
|
||||||
|
assert ws not in handler.client_sessions
|
||||||
|
finally:
|
||||||
|
handler.shutdown()
|
||||||
@@ -0,0 +1,200 @@
|
|||||||
|
"""aibox 模型接入点管理工具(model_config.py)单元测试(2026-08-18)
|
||||||
|
|
||||||
|
覆盖:
|
||||||
|
- 列表归一化(默认标记/免密/失败转移展示字段)
|
||||||
|
- 添加:字段校验(档案名/地址/模型名/重复)、Key 两种方式、温度/超时钳制、
|
||||||
|
失败转移引用校验、设为默认、落盘(含脱敏)
|
||||||
|
- 删除:物理移除、fallback 悬空引用清理、默认档案标记清除
|
||||||
|
- 连通性验证:成功 / 失败 / 档案不存在
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import types
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
import model_config # noqa: E402
|
||||||
|
from common import ConfigManager # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def _cfg(tmp_path, initial: dict = None) -> ConfigManager:
|
||||||
|
"""构造指向临时配置文件的 ConfigManager(避免碰真实 ~/.aibox)"""
|
||||||
|
cm = ConfigManager()
|
||||||
|
cm._config_path = str(tmp_path / "config.json")
|
||||||
|
cm._config = initial if initial is not None else {}
|
||||||
|
return cm
|
||||||
|
|
||||||
|
|
||||||
|
def _base_models():
|
||||||
|
return {
|
||||||
|
"models": {
|
||||||
|
"deepseek": {
|
||||||
|
"api_url": "https://api.deepseek.com/v1/chat/completions",
|
||||||
|
"model": "deepseek-chat",
|
||||||
|
"api_key_env": "DEEPSEEK_API_KEY",
|
||||||
|
},
|
||||||
|
"local": {
|
||||||
|
"api_url": "http://localhost:11434/v1/chat/completions",
|
||||||
|
"model": "qwen3",
|
||||||
|
"api_key_env": "",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"api": {"default_model": "deepseek"},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 列表 ====================
|
||||||
|
|
||||||
|
class TestListProfiles:
|
||||||
|
def test_list_normalized(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
profiles = model_config.list_profiles(cm)
|
||||||
|
assert len(profiles) == 2
|
||||||
|
ds = next(p for p in profiles if p["name"] == "deepseek")
|
||||||
|
assert ds["is_default"] is True
|
||||||
|
assert ds["api_key_env"] == "DEEPSEEK_API_KEY"
|
||||||
|
local = next(p for p in profiles if p["name"] == "local")
|
||||||
|
assert local["api_key_env"] == ""
|
||||||
|
assert local["is_default"] is False
|
||||||
|
|
||||||
|
def test_list_empty_models_section(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {})
|
||||||
|
assert model_config.list_profiles(cm) == []
|
||||||
|
|
||||||
|
def test_list_ignores_malformed_entries(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {"models": {"bad": "not-a-dict"}})
|
||||||
|
assert model_config.list_profiles(cm) == []
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 校验 ====================
|
||||||
|
|
||||||
|
class TestValidateNewProfile:
|
||||||
|
def test_name_rules(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
assert model_config.validate_new_profile(cm, "", "https://x", "m")
|
||||||
|
assert model_config.validate_new_profile(cm, "坏名字!", "https://x", "m")
|
||||||
|
assert model_config.validate_new_profile(cm, "deepseek", "https://x", "m") # 已存在
|
||||||
|
assert model_config.validate_new_profile(cm, "ok_name", "https://x", "m") is None
|
||||||
|
|
||||||
|
def test_url_scheme_and_model_required(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
assert model_config.validate_new_profile(cm, "kimi", "ftp://x", "m")
|
||||||
|
assert model_config.validate_new_profile(cm, "kimi", "https://x", "")
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 添加 ====================
|
||||||
|
|
||||||
|
class TestAddProfile:
|
||||||
|
def test_add_with_env_key_and_default(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
msgs = model_config.add_profile(
|
||||||
|
cm, "kimi", "https://api.moonshot.cn/v1/chat/completions",
|
||||||
|
"moonshot-v1-8k", api_key_env="KIMI_KEY",
|
||||||
|
display_name="Kimi", set_default=True)
|
||||||
|
assert any("已添加" in m for m in msgs)
|
||||||
|
assert any("设为默认" in m for m in msgs)
|
||||||
|
assert cm.get("api.default_model") == "kimi"
|
||||||
|
assert cm.get("models.kimi.api_key_env") == "KIMI_KEY"
|
||||||
|
assert cm.get("models.kimi.api_url").startswith("https://")
|
||||||
|
# 落盘文件存在且可解析
|
||||||
|
import json
|
||||||
|
disk = json.loads((tmp_path / "config.json").read_text(encoding="utf-8"))
|
||||||
|
assert disk["models"]["kimi"]["model"] == "moonshot-v1-8k"
|
||||||
|
|
||||||
|
def test_add_without_key_is_exempt_endpoint(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {})
|
||||||
|
msgs = model_config.add_profile(
|
||||||
|
cm, "local", "http://localhost:11434/v1/chat/completions",
|
||||||
|
"qwen3", api_key_env="")
|
||||||
|
assert any("免密" in m for m in msgs)
|
||||||
|
assert cm.get("models.local.api_key_env") == ""
|
||||||
|
|
||||||
|
def test_add_rejects_invalid_fallback(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
msgs = model_config.add_profile(
|
||||||
|
cm, "kimi", "https://x/v1", "m1",
|
||||||
|
fallback=["ghost", "deepseek"])
|
||||||
|
assert any("失败转移" in m for m in msgs)
|
||||||
|
assert cm.get("models.kimi") is None # 未写入
|
||||||
|
|
||||||
|
def test_add_temperature_timeout_clamped(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {})
|
||||||
|
model_config.add_profile(cm, "a", "https://x/v1", "m1",
|
||||||
|
temperature=9.9, timeout=1)
|
||||||
|
assert cm.get("models.a.temperature") == 2.0
|
||||||
|
assert cm.get("models.a.timeout") == 5
|
||||||
|
|
||||||
|
def test_add_invalid_temperature_rejected(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {})
|
||||||
|
msgs = model_config.add_profile(cm, "a", "https://x/v1", "m1",
|
||||||
|
temperature="abc")
|
||||||
|
assert any("温度" in m for m in msgs)
|
||||||
|
assert cm.get("models.a") is None
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 删除 ====================
|
||||||
|
|
||||||
|
class TestDeleteProfile:
|
||||||
|
def test_delete_cleans_fallback_refs(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {
|
||||||
|
"models": {
|
||||||
|
"main": {"api_url": "https://a", "model": "m1",
|
||||||
|
"fallback": ["backup"]},
|
||||||
|
"backup": {"api_url": "https://b", "model": "m2"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
msgs = model_config.delete_profile(cm, "backup")
|
||||||
|
assert any("已删除" in m for m in msgs)
|
||||||
|
assert any("失败转移引用" in m for m in msgs)
|
||||||
|
assert cm.get("models.main.fallback") == []
|
||||||
|
|
||||||
|
def test_delete_default_clears_default_flag(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
model_config.delete_profile(cm, "deepseek")
|
||||||
|
assert cm.get("api.default_model") is None
|
||||||
|
|
||||||
|
def test_delete_nonexistent(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
msgs = model_config.delete_profile(cm, "ghost")
|
||||||
|
assert any("不存在" in m for m in msgs)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 连通性验证 ====================
|
||||||
|
|
||||||
|
class TestTestProfile:
|
||||||
|
def test_connectivity_ok(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
fake = mock.MagicMock()
|
||||||
|
fake.chat_completion.return_value = "pong"
|
||||||
|
with mock.patch("llm_client.create_llm_client", return_value=fake):
|
||||||
|
r = model_config.test_profile(cm, "deepseek")
|
||||||
|
assert "连通正常" in r and "pong" in r
|
||||||
|
|
||||||
|
def test_connectivity_failure(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
fake = mock.MagicMock()
|
||||||
|
fake.chat_completion.side_effect = RuntimeError("连接被拒绝")
|
||||||
|
with mock.patch("llm_client.create_llm_client", return_value=fake):
|
||||||
|
r = model_config.test_profile(cm, "deepseek")
|
||||||
|
assert "验证失败" in r and "连接被拒绝" in r
|
||||||
|
|
||||||
|
def test_connectivity_missing_profile(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, _base_models())
|
||||||
|
assert "不存在" in model_config.test_profile(cm, "ghost")
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 落盘失败感知 ====================
|
||||||
|
|
||||||
|
class TestSaveFailure:
|
||||||
|
def test_add_reports_save_failure(self, tmp_path):
|
||||||
|
cm = _cfg(tmp_path, {})
|
||||||
|
# J5:写路径改为原子写,磁盘失败经原子写函数模拟
|
||||||
|
with mock.patch("common.config.atomic_write_text",
|
||||||
|
side_effect=PermissionError("disk ro")):
|
||||||
|
msgs = model_config.add_profile(
|
||||||
|
cm, "a", "https://x/v1", "m1", api_key_env="K")
|
||||||
|
assert any("保存失败" in m for m in msgs)
|
||||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,829 @@
|
|||||||
|
"""tests.test_n_fixes - Batch N 修复单元测试
|
||||||
|
|
||||||
|
波次 3 平台兼容(2026-08-21,审计第八轮 T 系列):
|
||||||
|
- T3 basic_tools 平台感知引用 + Windows PowerShell/curl.exe 指引
|
||||||
|
- T11 Windows 进程树收割(taskkill /T /F,executors + process_manager)
|
||||||
|
- T25 MemoryDatabase 统一连接工厂(busy_timeout)
|
||||||
|
- T26 SkillManager 原子 reload(中途不清空 self.skills)
|
||||||
|
- T27 atomic_write_text 对 Windows PermissionError 短退避重试
|
||||||
|
|
||||||
|
波次 4 LLM 健壮性(2026-08-21):
|
||||||
|
- T10 调度器 wait_for 硬超时兜底
|
||||||
|
- T19 FailoverClient 异常收窄(永久配置错误不转移不冷却)
|
||||||
|
- T20 流式响应未收到 [DONE] 抛 TransientLLMError(不再静默截断)
|
||||||
|
- T21 env 覆盖黑名单(LD_*/DYLD_*/PATH/PYTHONPATH 等劫持面)
|
||||||
|
- T23 压缩后 tool 配对双向修复(assistant.tool_calls 失配剥离)
|
||||||
|
(T14 FTS5 经复现证伪:JOIN 别名场景下表名 MATCH 正常命中,
|
||||||
|
全链路中文检索长词走 trigram、短词走 LIKE 兜底,无需修复)
|
||||||
|
|
||||||
|
波次 5 权限隔离(2026-08-21,产品决策:T1 仅告警不加鉴权、T2 维持现状):
|
||||||
|
- T22 定时任务精简 profile(task:无 update_identity/update_soul/set_proxy)
|
||||||
|
+ ProcessManager owner 隔离(写/杀权限,观察面共享)
|
||||||
|
- T-L17 model_config 按域名推荐 env 名(防 DeepSeek 凭据发往他厂)
|
||||||
|
- T-L18 list_profiles 输出过滤控制字符(防手写 config ANSI 注入)
|
||||||
|
|
||||||
|
Windows 分支在 Linux 上以 mock.patch("os.name", "nt") 模拟验证命令生成
|
||||||
|
逻辑(不实际执行生成的命令)。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
import sqlite3
|
||||||
|
import threading
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import requests
|
||||||
|
|
||||||
|
from common.atomic import atomic_write_text
|
||||||
|
from subsystems import MemoryDatabase, SkillManager
|
||||||
|
from subsystems.executors import SubprocessCommandExecutor, filter_env_overrides
|
||||||
|
from tools.basic_tools import (
|
||||||
|
ListFilesTool,
|
||||||
|
ReadFileTool,
|
||||||
|
SimpleWebCrawlerTool,
|
||||||
|
WriteFileTool,
|
||||||
|
_quote_arg,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T3:平台感知引用 ====================
|
||||||
|
|
||||||
|
class TestQuoteArg:
|
||||||
|
def test_posix_uses_shlex(self):
|
||||||
|
quoted, err = _quote_arg("/tmp/my file.txt")
|
||||||
|
assert err is None
|
||||||
|
assert quoted == "'/tmp/my file.txt'"
|
||||||
|
|
||||||
|
@pytest.mark.skipif(os.name != "posix", reason="POSIX 引用语义")
|
||||||
|
def test_posix_specials_escaped(self):
|
||||||
|
quoted, err = _quote_arg("a`b$c$(d)\\e")
|
||||||
|
assert err is None
|
||||||
|
assert quoted == "'a`b$c$(d)\\e'"
|
||||||
|
|
||||||
|
def test_win_wraps_double_quotes(self):
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
quoted, err = _quote_arg(r"C:\Program Files\my app\doc.txt")
|
||||||
|
assert err is None
|
||||||
|
assert quoted == r'"C:\Program Files\my app\doc.txt"'
|
||||||
|
|
||||||
|
def test_win_rejects_double_quote(self):
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
quoted, err = _quote_arg('evil" & del /f /s /q C:\\')
|
||||||
|
assert quoted == ""
|
||||||
|
assert err is not None and "双引号" in err
|
||||||
|
|
||||||
|
def test_win_rejects_cmd_var_expansion(self):
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
quoted, err = _quote_arg("https://evil.com/?x=%PATH%")
|
||||||
|
assert quoted == ""
|
||||||
|
assert err is not None and "注入" in err
|
||||||
|
|
||||||
|
def test_win_allows_url_percent_encoding(self):
|
||||||
|
"""URL 编码(%20/%E4 等)非 %变量名% 形态,不拒绝。"""
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
quoted, err = _quote_arg("https://example.com/a%20b?x=1&y=2")
|
||||||
|
assert err is None
|
||||||
|
assert quoted == '"https://example.com/a%20b?x=1&y=2"'
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T3:Windows 指引命令生成 ====================
|
||||||
|
|
||||||
|
class TestReadFileWindows:
|
||||||
|
def test_head_powershell(self, tmp_path):
|
||||||
|
f = tmp_path / "data.txt"
|
||||||
|
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = ReadFileTool().execute(filepath=str(f), head=True, max_lines=10)
|
||||||
|
assert "powershell" in r and "Get-Content" in r
|
||||||
|
assert "run_command" in r
|
||||||
|
assert "sed -n" not in r and "head -n" not in r
|
||||||
|
|
||||||
|
def test_tail_powershell(self, tmp_path):
|
||||||
|
f = tmp_path / "data.txt"
|
||||||
|
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = ReadFileTool().execute(filepath=str(f), tail=True, max_lines=5)
|
||||||
|
assert "Get-Content" in r and "-Tail 5" in r
|
||||||
|
|
||||||
|
def test_range_powershell(self, tmp_path):
|
||||||
|
f = tmp_path / "data.txt"
|
||||||
|
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = ReadFileTool().execute(filepath=str(f), start_line=5, end_line=15)
|
||||||
|
assert "Select-Object -Skip 4 -First 11" in r
|
||||||
|
|
||||||
|
def test_posix_unchanged(self, tmp_path):
|
||||||
|
f = tmp_path / "data.txt"
|
||||||
|
f.write_text("hello\n" * 30, encoding="utf-8")
|
||||||
|
r = ReadFileTool().execute(filepath=str(f), start_line=1, end_line=5)
|
||||||
|
assert "sed -n '1,5p'" in r
|
||||||
|
|
||||||
|
|
||||||
|
class TestWriteFileWindows:
|
||||||
|
def test_base64_set_content(self, tmp_path):
|
||||||
|
target = tmp_path / "sub" / "out.txt"
|
||||||
|
content = "line1\nline2\n含中文 \"quoted\" & special"
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = WriteFileTool().execute(filepath=str(target), content=content, mode="w")
|
||||||
|
assert "Set-Content" in r
|
||||||
|
assert "FromBase64String" in r
|
||||||
|
# 明文内容(含元字符)不得出现在命令行中
|
||||||
|
assert "quoted" not in r and "special" not in r
|
||||||
|
|
||||||
|
def test_append_uses_add_content(self, tmp_path):
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = WriteFileTool().execute(
|
||||||
|
filepath=str(tmp_path / "a.txt"), content="more", mode="a")
|
||||||
|
assert "Add-Content" in r
|
||||||
|
|
||||||
|
def test_oversize_rejected(self, tmp_path):
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = WriteFileTool().execute(
|
||||||
|
filepath=str(tmp_path / "big.txt"), content="x" * 10000, mode="w")
|
||||||
|
assert "过大" in r and "追加" in r
|
||||||
|
|
||||||
|
|
||||||
|
class TestListFilesWindows:
|
||||||
|
def test_get_child_item(self, tmp_path):
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = ListFilesTool().execute(directory=str(tmp_path), show_hidden=True)
|
||||||
|
assert "Get-ChildItem" in r and "-Force" in r
|
||||||
|
assert "ls -l" not in r
|
||||||
|
|
||||||
|
def test_posix_unchanged(self, tmp_path):
|
||||||
|
r = ListFilesTool().execute(directory=str(tmp_path))
|
||||||
|
assert "ls -l " in r
|
||||||
|
|
||||||
|
|
||||||
|
class TestWebWindows:
|
||||||
|
def test_get_curl_exe_amp_url_quoted(self):
|
||||||
|
tool = SimpleWebCrawlerTool()
|
||||||
|
url = "https://example.com/?a=1&b=2"
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = tool.execute(action="get", url=url)
|
||||||
|
assert "curl.exe" in r
|
||||||
|
# & 必须在双引号内(cmd 双引号内为字面量)
|
||||||
|
assert '"https://example.com/?a=1&b=2"' in r
|
||||||
|
|
||||||
|
def test_cmd_var_url_rejected(self):
|
||||||
|
tool = SimpleWebCrawlerTool()
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = tool.execute(action="get", url="https://evil.com/?x=%PATH%")
|
||||||
|
assert "拒绝" in r
|
||||||
|
|
||||||
|
def test_text_action_powershell_pipe(self):
|
||||||
|
tool = SimpleWebCrawlerTool()
|
||||||
|
with mock.patch("os.name", "nt"):
|
||||||
|
r = tool.execute(action="text", url="https://example.com/")
|
||||||
|
assert "curl.exe" in r and "powershell" in r
|
||||||
|
|
||||||
|
def test_posix_unchanged(self):
|
||||||
|
tool = SimpleWebCrawlerTool()
|
||||||
|
r = tool.execute(action="get", url="https://example.com/")
|
||||||
|
assert "curl -sL" in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T11:Windows 进程树收割 ====================
|
||||||
|
|
||||||
|
class TestTerminateTreeWindows:
|
||||||
|
def test_executors_taskkill_tree(self):
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
proc = mock.Mock()
|
||||||
|
proc.pid = 12345
|
||||||
|
proc.poll.return_value = None
|
||||||
|
with mock.patch("os.name", "nt"), \
|
||||||
|
mock.patch("subprocess.run") as run:
|
||||||
|
ex._terminate_tree(proc, force=True)
|
||||||
|
run.assert_called_once()
|
||||||
|
assert run.call_args[0][0] == ["taskkill", "/PID", "12345", "/T", "/F"]
|
||||||
|
proc.kill.assert_called_once() # 兜底单进程 kill
|
||||||
|
|
||||||
|
def test_executors_taskkill_process_gone(self):
|
||||||
|
"""taskkill 目标已退出:兜底 kill 不执行(poll 非 None)。"""
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
proc = mock.Mock()
|
||||||
|
proc.pid = 12346
|
||||||
|
proc.poll.return_value = 0
|
||||||
|
with mock.patch("os.name", "nt"), \
|
||||||
|
mock.patch("subprocess.run") as run:
|
||||||
|
ex._terminate_tree(proc, force=False)
|
||||||
|
run.assert_called_once()
|
||||||
|
proc.kill.assert_not_called()
|
||||||
|
|
||||||
|
def test_process_manager_taskkill_tree(self):
|
||||||
|
from subsystems.process_manager import ProcessManager
|
||||||
|
pm = ProcessManager()
|
||||||
|
entry = mock.Mock()
|
||||||
|
entry.popen.pid = 999
|
||||||
|
entry.popen.poll.return_value = None
|
||||||
|
with mock.patch("os.name", "nt"), \
|
||||||
|
mock.patch("subprocess.run") as run:
|
||||||
|
pm._terminate_tree(entry, force=True)
|
||||||
|
assert run.call_args[0][0] == ["taskkill", "/PID", "999", "/T", "/F"]
|
||||||
|
|
||||||
|
@pytest.mark.skipif(os.name != "posix", reason="POSIX 进程组语义")
|
||||||
|
def test_posix_no_taskkill(self):
|
||||||
|
"""POSIX 分支不调用 taskkill(走 killpg)。"""
|
||||||
|
import subprocess as _sp
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
proc = _sp.Popen(["sleep", "30"], start_new_session=True)
|
||||||
|
try:
|
||||||
|
with mock.patch("subprocess.run") as run:
|
||||||
|
ex._terminate_tree(proc, force=False)
|
||||||
|
run.assert_not_called()
|
||||||
|
finally:
|
||||||
|
try:
|
||||||
|
import signal as _sig
|
||||||
|
os.killpg(os.getpgid(proc.pid), _sig.SIGKILL)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T25:MemoryDatabase busy_timeout ====================
|
||||||
|
|
||||||
|
class TestMemoryDbBusyTimeout:
|
||||||
|
def test_connect_sets_busy_timeout(self, tmp_path):
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||||
|
conn = db._connect()
|
||||||
|
try:
|
||||||
|
value = conn.execute("PRAGMA busy_timeout").fetchone()[0]
|
||||||
|
finally:
|
||||||
|
conn.close()
|
||||||
|
assert value == 10000
|
||||||
|
|
||||||
|
def test_all_paths_via_factory(self, tmp_path):
|
||||||
|
"""写/读路径不再直连 sqlite3.connect(统一走 _connect 工厂)。"""
|
||||||
|
db = MemoryDatabase(db_path=str(tmp_path / "mem.db"))
|
||||||
|
mid = db.add_memory(topic="t", summary="s", content="c",
|
||||||
|
start_time="2026-01-01 00:00:00",
|
||||||
|
end_time="2026-01-01 00:01:00")
|
||||||
|
assert mid is not None
|
||||||
|
rows = db.get_recent_memories(limit=5)
|
||||||
|
assert any(r["id"] == mid for r in rows)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T26:SkillManager 原子 reload ====================
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def isolated_skill_dir(tmp_path, monkeypatch):
|
||||||
|
from common import constants
|
||||||
|
skills_dir = tmp_path / "skills"
|
||||||
|
skills_dir.mkdir()
|
||||||
|
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||||
|
property(lambda self: str(skills_dir)))
|
||||||
|
return skills_dir
|
||||||
|
|
||||||
|
|
||||||
|
class TestSkillReloadAtomic:
|
||||||
|
def test_skills_visible_during_reload(self, isolated_skill_dir):
|
||||||
|
"""reload 过程中 self.skills 保持旧表(T26:不再先清空)。"""
|
||||||
|
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||||
|
sm = SkillManager()
|
||||||
|
assert sm.count() == 1
|
||||||
|
|
||||||
|
(isolated_skill_dir / "b.md").write_text("# B\n内容B", encoding="utf-8")
|
||||||
|
original = sm._read_markdown_skill
|
||||||
|
observed = []
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def spy(path):
|
||||||
|
calls.append(path)
|
||||||
|
if len(calls) == 2:
|
||||||
|
observed.append(dict(sm.skills))
|
||||||
|
return original(path)
|
||||||
|
|
||||||
|
sm._read_markdown_skill = spy
|
||||||
|
sm.reload()
|
||||||
|
|
||||||
|
# 加载第 2 个技能时,旧表完整可见(未被清空/半截)
|
||||||
|
assert len(observed) == 1
|
||||||
|
assert set(observed[0].keys()) == {"a"}
|
||||||
|
# reload 完成后一次性替换为新表
|
||||||
|
assert set(sm.skills.keys()) == {"a", "b"}
|
||||||
|
|
||||||
|
def test_reload_missing_dir_keeps_old(self, isolated_skill_dir, monkeypatch):
|
||||||
|
"""目录消失时清空(旧行为保留),正常 reload 不经过空表中间态。"""
|
||||||
|
from common import constants
|
||||||
|
(isolated_skill_dir / "a.md").write_text("# A\n内容A", encoding="utf-8")
|
||||||
|
sm = SkillManager()
|
||||||
|
assert sm.count() == 1
|
||||||
|
monkeypatch.setattr(type(constants), "SKILLS_DIR",
|
||||||
|
property(lambda self: str(isolated_skill_dir / "nope")))
|
||||||
|
sm.reload()
|
||||||
|
assert sm.count() == 0
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T27:atomic_write_text 重试 ====================
|
||||||
|
|
||||||
|
class TestAtomicWriteRetry:
|
||||||
|
def test_retries_then_succeeds(self, tmp_path):
|
||||||
|
target = tmp_path / "cfg.json"
|
||||||
|
real_replace = os.replace
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def flaky_replace(src, dst):
|
||||||
|
calls.append((src, dst))
|
||||||
|
if len(calls) < 3:
|
||||||
|
raise PermissionError(5, "Access is denied")
|
||||||
|
return real_replace(src, dst)
|
||||||
|
|
||||||
|
with mock.patch("os.replace", side_effect=flaky_replace):
|
||||||
|
atomic_write_text(str(target), '{"k": 1}')
|
||||||
|
assert len(calls) == 3
|
||||||
|
assert target.read_text(encoding="utf-8") == '{"k": 1}'
|
||||||
|
|
||||||
|
def test_raises_after_exhausted(self, tmp_path):
|
||||||
|
target = tmp_path / "cfg.json"
|
||||||
|
with mock.patch("os.replace",
|
||||||
|
side_effect=PermissionError(5, "Access is denied")):
|
||||||
|
with pytest.raises(PermissionError):
|
||||||
|
atomic_write_text(str(target), "x")
|
||||||
|
assert not target.exists()
|
||||||
|
# 临时文件已清理
|
||||||
|
leftovers = [p for p in tmp_path.iterdir() if p.name.startswith(".aibox_tmp_")]
|
||||||
|
assert leftovers == []
|
||||||
|
|
||||||
|
def test_normal_write_no_retry(self, tmp_path):
|
||||||
|
target = tmp_path / "ok.txt"
|
||||||
|
with mock.patch("os.replace", wraps=os.replace) as spy:
|
||||||
|
atomic_write_text(str(target), "hello")
|
||||||
|
spy.assert_called_once()
|
||||||
|
assert target.read_text(encoding="utf-8") == "hello"
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T19:FailoverClient 异常收窄 ====================
|
||||||
|
|
||||||
|
def _make_failover_client():
|
||||||
|
"""deepseek(→kimi fallback) 链,Key 走 mock 环境变量。"""
|
||||||
|
from llm_client import create_llm_client
|
||||||
|
from test_models_profile import MODELS, _cfg
|
||||||
|
models = {
|
||||||
|
"deepseek": dict(MODELS["deepseek"], fallback=["kimi"]),
|
||||||
|
"kimi": dict(MODELS["kimi"]),
|
||||||
|
}
|
||||||
|
env = {"DEEPSEEK_API_KEY": "sk-ds", "MOONSHOT_API_KEY": "sk-moon"}
|
||||||
|
with mock.patch.dict(os.environ, env):
|
||||||
|
return create_llm_client(_cfg(models=models))
|
||||||
|
|
||||||
|
|
||||||
|
class TestT19TransientExceptionScope:
|
||||||
|
def test_permanent_schema_error_not_failed_over(self):
|
||||||
|
"""MissingSchema(api_url 手写错等永久配置错误)→ 流式路径直接上抛,
|
||||||
|
不转移(只打一次请求)、不计 failover、不冷却。"""
|
||||||
|
fc = _make_failover_client()
|
||||||
|
with mock.patch("llm_client.requests.post",
|
||||||
|
side_effect=requests.exceptions.MissingSchema("bad url")) as mp:
|
||||||
|
with pytest.raises(requests.exceptions.MissingSchema):
|
||||||
|
list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
assert len(mp.call_args_list) == 1
|
||||||
|
assert fc.get_stats()["failovers"] == 0
|
||||||
|
|
||||||
|
def test_permanent_error_no_cooldown_pollution(self):
|
||||||
|
"""永久错误后档案不进冷却——下次调用仍打主档案。"""
|
||||||
|
fc = _make_failover_client()
|
||||||
|
ok_stream = _stream_response([
|
||||||
|
'data: {"choices":[{"delta":{"content":"recovered"}}]}',
|
||||||
|
'data: [DONE]'])
|
||||||
|
with mock.patch("llm_client.requests.post",
|
||||||
|
side_effect=[requests.exceptions.MissingSchema("bad"),
|
||||||
|
ok_stream]) as mp:
|
||||||
|
with pytest.raises(requests.exceptions.MissingSchema):
|
||||||
|
list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
events = list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
# 第二次仍先打主档案(deepseek),未进冷却跳过
|
||||||
|
assert mp.call_args_list[1].args[0].startswith("https://api.deepseek.com")
|
||||||
|
assert any(e["type"] == "content" and e["text"] == "recovered" for e in events)
|
||||||
|
assert fc.get_stats()["cooldown_skips"] == 0
|
||||||
|
|
||||||
|
def test_connection_error_still_fails_over(self):
|
||||||
|
"""瞬时连接错误仍正常转移(收窄不误伤)。"""
|
||||||
|
fc = _make_failover_client()
|
||||||
|
from test_models_profile import _resp
|
||||||
|
with mock.patch("llm_client.requests.post",
|
||||||
|
side_effect=[requests.exceptions.ConnectionError("boom"),
|
||||||
|
_resp(200, content="ok-conn")]):
|
||||||
|
assert fc.chat_completion([{"role": "user", "content": "hi"}]) == "ok-conn"
|
||||||
|
|
||||||
|
def test_chunked_encoding_error_fails_over(self):
|
||||||
|
"""ChunkedEncodingError(响应体中途损坏)属瞬时故障 → 转移。"""
|
||||||
|
fc = _make_failover_client()
|
||||||
|
from test_models_profile import _resp
|
||||||
|
with mock.patch("llm_client.requests.post",
|
||||||
|
side_effect=[requests.exceptions.ChunkedEncodingError("broken"),
|
||||||
|
_resp(200, content="ok-chunked")]):
|
||||||
|
assert fc.chat_completion([{"role": "user", "content": "hi"}]) == "ok-chunked"
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T20:不完整流信号 ====================
|
||||||
|
|
||||||
|
def _stream_response(lines):
|
||||||
|
m = mock.MagicMock()
|
||||||
|
m.status_code = 200
|
||||||
|
m.iter_lines.return_value = [l.encode("utf-8") for l in lines]
|
||||||
|
return m
|
||||||
|
|
||||||
|
|
||||||
|
def _make_deepseek_client():
|
||||||
|
from llm_client import DeepSeekClient
|
||||||
|
return DeepSeekClient(
|
||||||
|
api_key="sk-test", model="test-model",
|
||||||
|
base_url="https://api.example/v1/chat/completions")
|
||||||
|
|
||||||
|
|
||||||
|
class TestT20IncompleteStream:
|
||||||
|
def test_missing_done_raises_transient(self):
|
||||||
|
"""流中途断开(iter_lines 耗尽无 [DONE])→ TransientLLMError,
|
||||||
|
不再静默视为成功导致回答无声截断。"""
|
||||||
|
from llm_client import TransientLLMError
|
||||||
|
c = _make_deepseek_client()
|
||||||
|
lines = [
|
||||||
|
'data: {"choices":[{"delta":{"content":"he"}}]}',
|
||||||
|
'data: {"choices":[{"delta":{"content":"llo"}}]}',
|
||||||
|
# 连接断开:无 [DONE],末行可能还是半截 JSON
|
||||||
|
'data: {"choices":[{"delta":{"cont',
|
||||||
|
]
|
||||||
|
with mock.patch("llm_client.requests.post",
|
||||||
|
return_value=_stream_response(lines)):
|
||||||
|
with pytest.raises(TransientLLMError, match="不完整"):
|
||||||
|
list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
|
||||||
|
def test_with_done_completes_normally(self):
|
||||||
|
from llm_client import TransientLLMError
|
||||||
|
c = _make_deepseek_client()
|
||||||
|
lines = [
|
||||||
|
'data: {"choices":[{"delta":{"content":"hi"}}]}',
|
||||||
|
'data: [DONE]',
|
||||||
|
]
|
||||||
|
with mock.patch("llm_client.requests.post",
|
||||||
|
return_value=_stream_response(lines)):
|
||||||
|
events = list(c.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
assert events[-1] == {"type": "done"}
|
||||||
|
assert any(e["type"] == "content" and e["text"] == "hi" for e in events)
|
||||||
|
|
||||||
|
def test_failover_first_event_incomplete_stream(self):
|
||||||
|
"""FailoverClient:首事件前不完整流(0 事件 + 无 [DONE])→ 可转移。"""
|
||||||
|
fc = _make_failover_client()
|
||||||
|
# 主档案:断流无任何事件;fallback(kimi):正常完整流
|
||||||
|
broken = _stream_response(['data: {"choices":[{"delta'])
|
||||||
|
ok_stream = _stream_response([
|
||||||
|
'data: {"choices":[{"delta":{"content":"from-kimi"}}]}',
|
||||||
|
'data: [DONE]'])
|
||||||
|
with mock.patch("llm_client.requests.post",
|
||||||
|
side_effect=[broken, ok_stream]):
|
||||||
|
events = list(fc.stream_chat([{"role": "user", "content": "hi"}]))
|
||||||
|
assert any(e["type"] == "content" and e["text"] == "from-kimi" for e in events)
|
||||||
|
assert events[-1]["type"] == "done"
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T21:env 覆盖黑名单 ====================
|
||||||
|
|
||||||
|
class TestT21EnvDenylist:
|
||||||
|
def test_denies_hijack_keys(self):
|
||||||
|
safe, rejected = filter_env_overrides({
|
||||||
|
"LD_PRELOAD": "/evil", "PATH": "/bin", "PYTHONPATH": "/p",
|
||||||
|
"MY_VAR": "ok"})
|
||||||
|
assert safe == {"MY_VAR": "ok"}
|
||||||
|
assert set(rejected) == {"LD_PRELOAD", "PATH", "PYTHONPATH"}
|
||||||
|
|
||||||
|
def test_case_insensitive_and_prefixes(self):
|
||||||
|
safe, rejected = filter_env_overrides({
|
||||||
|
"ld_preload": "x", "DYLD_INSERT_LIBRARIES": "y", "path": "z",
|
||||||
|
"PSMODULEPATH": "w", "OK1": "1"})
|
||||||
|
assert safe == {"OK1": "1"}
|
||||||
|
assert len(rejected) == 4
|
||||||
|
|
||||||
|
def test_empty_and_none(self):
|
||||||
|
assert filter_env_overrides(None) == ({}, [])
|
||||||
|
assert filter_env_overrides({}) == ({}, [])
|
||||||
|
|
||||||
|
def test_executor_blocks_denied_env(self):
|
||||||
|
"""执行路径集成:正常 env 注入生效,黑名单键被拒。"""
|
||||||
|
ex = SubprocessCommandExecutor()
|
||||||
|
r = ex.execute(
|
||||||
|
"echo $T21VAR",
|
||||||
|
env={"T21VAR": "visible", "LD_PRELOAD": "/nonexistent/libevil.so"})
|
||||||
|
assert "visible" in r
|
||||||
|
# LD_PRELOAD 未注入(若注入且文件不存在,loader 报错会进合并输出)
|
||||||
|
assert "/nonexistent" not in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T23:压缩后 tool 配对双向修复 ====================
|
||||||
|
|
||||||
|
def _make_compressor():
|
||||||
|
from subsystems.conversation_compressor import ConversationCompressor
|
||||||
|
return ConversationCompressor(llm_client=None)
|
||||||
|
|
||||||
|
|
||||||
|
class TestT23ToolPairValidation:
|
||||||
|
def test_orphan_tool_calls_stripped(self):
|
||||||
|
"""assistant.tool_calls 的 tool 结果被压缩掉 → 剥离 tool_calls。"""
|
||||||
|
comp = _make_compressor()
|
||||||
|
msgs = [
|
||||||
|
{"role": "user", "content": "q"},
|
||||||
|
{"role": "assistant", "content": "看结果",
|
||||||
|
"tool_calls": [{"id": "t1", "type": "function",
|
||||||
|
"function": {"name": "read", "arguments": "{}"}}]},
|
||||||
|
{"role": "user", "content": "next"},
|
||||||
|
]
|
||||||
|
out = comp._validate_message_structure(msgs)
|
||||||
|
assert len(out) == 3
|
||||||
|
assert "tool_calls" not in out[1]
|
||||||
|
assert out[1]["content"] == "看结果" # 原有内容保留
|
||||||
|
|
||||||
|
def test_empty_content_filled_when_stripped(self):
|
||||||
|
comp = _make_compressor()
|
||||||
|
msgs = [
|
||||||
|
{"role": "assistant", "content": None,
|
||||||
|
"tool_calls": [{"id": "t1", "type": "function",
|
||||||
|
"function": {"name": "read", "arguments": "{}"}}]},
|
||||||
|
]
|
||||||
|
out = comp._validate_message_structure(msgs)
|
||||||
|
assert "tool_calls" not in out[0]
|
||||||
|
assert out[0]["content"] == "" # 防产生非法空消息
|
||||||
|
|
||||||
|
def test_partial_pair_keeps_matched_calls(self):
|
||||||
|
"""两个 tool_calls 只剩一个配对 → 保留配对的那个。"""
|
||||||
|
comp = _make_compressor()
|
||||||
|
msgs = [
|
||||||
|
{"role": "assistant", "content": "run",
|
||||||
|
"tool_calls": [
|
||||||
|
{"id": "t1", "type": "function",
|
||||||
|
"function": {"name": "a", "arguments": "{}"}},
|
||||||
|
{"id": "t2", "type": "function",
|
||||||
|
"function": {"name": "b", "arguments": "{}"}}]},
|
||||||
|
{"role": "tool", "tool_call_id": "t2", "content": "result-b"},
|
||||||
|
]
|
||||||
|
out = comp._validate_message_structure(msgs)
|
||||||
|
assert len(out) == 2
|
||||||
|
assert [tc["id"] for tc in out[0]["tool_calls"]] == ["t2"]
|
||||||
|
assert out[1]["tool_call_id"] == "t2"
|
||||||
|
|
||||||
|
def test_orphan_tool_message_removed(self):
|
||||||
|
"""既有行为不回归:孤立 tool 消息(无 assistant)移除。"""
|
||||||
|
comp = _make_compressor()
|
||||||
|
msgs = [
|
||||||
|
{"role": "user", "content": "q"},
|
||||||
|
{"role": "tool", "tool_call_id": "ghost", "content": "x"},
|
||||||
|
]
|
||||||
|
out = comp._validate_message_structure(msgs)
|
||||||
|
assert [m["role"] for m in out] == ["user"]
|
||||||
|
|
||||||
|
def test_intact_pair_untouched(self):
|
||||||
|
"""完整配对原样保留(同一对象引用,零开销路径)。"""
|
||||||
|
comp = _make_compressor()
|
||||||
|
assistant = {"role": "assistant", "content": "ok",
|
||||||
|
"tool_calls": [{"id": "t1", "type": "function",
|
||||||
|
"function": {"name": "a", "arguments": "{}"}}]}
|
||||||
|
tool = {"role": "tool", "tool_call_id": "t1", "content": "r"}
|
||||||
|
out = comp._validate_message_structure([assistant, tool])
|
||||||
|
assert out == [assistant, tool]
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T10:调度器硬超时兜底 ====================
|
||||||
|
|
||||||
|
class _UnkillableAI:
|
||||||
|
"""软取消无法中断的卡死 AI(模拟 requests 读阻塞被滴流输出规避)。"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.cancel_called = 0
|
||||||
|
self._release = threading.Event()
|
||||||
|
|
||||||
|
def cancel(self):
|
||||||
|
self.cancel_called += 1
|
||||||
|
# 无效:think_and_respond 仍不返回
|
||||||
|
|
||||||
|
def think_and_respond(self, input_text, output_callback=None,
|
||||||
|
is_continuation=False):
|
||||||
|
self._release.wait(timeout=30)
|
||||||
|
return "never-reached"
|
||||||
|
|
||||||
|
|
||||||
|
class TestT10SchedulerHardTimeout:
|
||||||
|
def test_stuck_ai_hard_timeout(self):
|
||||||
|
"""软取消无效的卡死任务 → wait_for 硬超时放行,调度循环继续。"""
|
||||||
|
from aibox import ConfigManager, TaskScheduler
|
||||||
|
|
||||||
|
class _FakeMemoDB:
|
||||||
|
def get_memo(self, memo_id):
|
||||||
|
return None
|
||||||
|
|
||||||
|
loop = asyncio.new_event_loop()
|
||||||
|
fake_ai = _UnkillableAI()
|
||||||
|
try:
|
||||||
|
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||||
|
memo = {"id": 7, "title": "卡死任务", "content": "内容", "trigger_count": 0}
|
||||||
|
orig_get = ConfigManager.get
|
||||||
|
|
||||||
|
def fake_get(self, key, default=None):
|
||||||
|
# task_timeout 设为 -29.5:软取消立即触发(call_later 负延迟),
|
||||||
|
# wait_for 硬超时 = -29.5 + 30 = 0.5s 后放行
|
||||||
|
if key == "scheduler.task_timeout":
|
||||||
|
return -29.5
|
||||||
|
return orig_get(self, key, default)
|
||||||
|
|
||||||
|
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||||
|
mock.patch.object(scheduler, "_create_task_ai", return_value=fake_ai):
|
||||||
|
result = loop.run_until_complete(
|
||||||
|
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||||
|
assert "硬超时" in result
|
||||||
|
assert fake_ai.cancel_called >= 1 # 软取消先于硬超时触发
|
||||||
|
finally:
|
||||||
|
fake_ai._release.set() # 释放卡死线程
|
||||||
|
loop.close()
|
||||||
|
|
||||||
|
def test_normal_task_unaffected(self):
|
||||||
|
"""正常返回的任务不受硬超时包装影响(结果正常透传)。"""
|
||||||
|
from aibox import ConfigManager, TaskScheduler
|
||||||
|
|
||||||
|
class _FakeMemoDB:
|
||||||
|
def get_memo(self, memo_id):
|
||||||
|
return None
|
||||||
|
|
||||||
|
class _FastAI:
|
||||||
|
def cancel(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def think_and_respond(self, input_text, output_callback=None,
|
||||||
|
is_continuation=False):
|
||||||
|
return "任务完成"
|
||||||
|
|
||||||
|
loop = asyncio.new_event_loop()
|
||||||
|
try:
|
||||||
|
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||||
|
memo = {"id": 8, "title": "正常任务", "content": "内容", "trigger_count": 0}
|
||||||
|
orig_get = ConfigManager.get
|
||||||
|
|
||||||
|
def fake_get(self, key, default=None):
|
||||||
|
if key == "scheduler.task_timeout":
|
||||||
|
return 600
|
||||||
|
return orig_get(self, key, default)
|
||||||
|
|
||||||
|
with mock.patch.object(ConfigManager, "get", fake_get), \
|
||||||
|
mock.patch.object(scheduler, "_create_task_ai", return_value=_FastAI()):
|
||||||
|
result = loop.run_until_complete(
|
||||||
|
scheduler.execute_task_as_conversation(memo, _FakeMemoDB()))
|
||||||
|
assert "任务完成" in result
|
||||||
|
finally:
|
||||||
|
loop.close()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T22:owner 隔离(ProcessManager) ====================
|
||||||
|
|
||||||
|
class TestT22OwnerIsolation:
|
||||||
|
def test_kill_denied_for_other_owner(self):
|
||||||
|
from subsystems.process_manager import ProcessManager
|
||||||
|
pm = ProcessManager()
|
||||||
|
r = pm.start("sleep 30", owner="session-a")
|
||||||
|
handle = r["handle"]
|
||||||
|
try:
|
||||||
|
denied = pm.kill(handle, owner="session-b")
|
||||||
|
assert not denied["ok"]
|
||||||
|
assert "属主" in denied["error"]
|
||||||
|
# 进程仍在运行(未被误杀)
|
||||||
|
assert pm._lookup(handle).status == "running"
|
||||||
|
finally:
|
||||||
|
pm.kill(handle)
|
||||||
|
|
||||||
|
def test_kill_allowed_for_owner(self):
|
||||||
|
from subsystems.process_manager import ProcessManager
|
||||||
|
pm = ProcessManager()
|
||||||
|
r = pm.start("sleep 30", owner="session-a")
|
||||||
|
result = pm.kill(r["handle"], owner="session-a")
|
||||||
|
assert result["ok"]
|
||||||
|
|
||||||
|
def test_write_denied_for_other_owner(self):
|
||||||
|
from subsystems.process_manager import ProcessManager
|
||||||
|
pm = ProcessManager()
|
||||||
|
r = pm.start("cat", owner="session-a")
|
||||||
|
handle = r["handle"]
|
||||||
|
try:
|
||||||
|
denied = pm.write(handle, "hi", owner="session-b")
|
||||||
|
assert not denied["ok"]
|
||||||
|
assert "属主" in denied["error"]
|
||||||
|
ok = pm.write(handle, "hi", owner="session-a")
|
||||||
|
assert ok["ok"]
|
||||||
|
finally:
|
||||||
|
pm.kill(handle)
|
||||||
|
|
||||||
|
def test_backward_compatible_without_owner(self):
|
||||||
|
"""任一侧无 owner 标识 → 不隔离(测试/管理面直调兼容)。"""
|
||||||
|
from subsystems.process_manager import ProcessManager
|
||||||
|
pm = ProcessManager()
|
||||||
|
r1 = pm.start("sleep 5", owner="a")
|
||||||
|
assert pm.kill(r1["handle"])["ok"] # caller 无 owner
|
||||||
|
r2 = pm.start("sleep 5") # entry 无 owner
|
||||||
|
assert pm.kill(r2["handle"], owner="b")["ok"]
|
||||||
|
|
||||||
|
def test_list_shows_owner_field(self):
|
||||||
|
from subsystems.process_manager import ProcessManager
|
||||||
|
pm = ProcessManager()
|
||||||
|
r = pm.start("sleep 5", owner="whoami")
|
||||||
|
try:
|
||||||
|
infos = pm.list_procs()
|
||||||
|
assert any(e.get("owner") == "whoami" for e in infos)
|
||||||
|
finally:
|
||||||
|
pm.kill(r["handle"])
|
||||||
|
|
||||||
|
def test_output_not_isolated(self):
|
||||||
|
"""观察面(get_output)不受 owner 限制。"""
|
||||||
|
from subsystems.process_manager import ProcessManager
|
||||||
|
pm = ProcessManager()
|
||||||
|
r = pm.start("echo hello", owner="session-a")
|
||||||
|
handle = r["handle"]
|
||||||
|
import time as _t
|
||||||
|
_t.sleep(0.3)
|
||||||
|
result = pm.get_output(handle) # 观察面无 owner 概念
|
||||||
|
assert result["ok"]
|
||||||
|
assert "hello" in result.get("output", "")
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T22:task profile(定时任务精简工具面) ====================
|
||||||
|
|
||||||
|
def _make_chat(profile, owner=None):
|
||||||
|
from aibox import DeepSeekChat
|
||||||
|
with mock.patch("aibox.create_llm_client") as mclc:
|
||||||
|
mclc.return_value = mock.MagicMock(
|
||||||
|
api_key="sk-test", name="DS", default_temperature=None)
|
||||||
|
return DeepSeekChat(tool_profile=profile, owner_id=owner)
|
||||||
|
|
||||||
|
|
||||||
|
class TestT22TaskProfile:
|
||||||
|
def test_task_profile_excludes_global_side_effect_tools(self):
|
||||||
|
ai = _make_chat("task")
|
||||||
|
names = set(ai.tool_registry._tools.keys())
|
||||||
|
assert "update_identity" not in names
|
||||||
|
assert "update_soul" not in names
|
||||||
|
assert "set_proxy" not in names
|
||||||
|
# 管理面其余工具保留
|
||||||
|
assert "get_proxy" in names
|
||||||
|
assert "create_subagent" in names
|
||||||
|
|
||||||
|
def test_orchestrator_profile_still_has_them(self):
|
||||||
|
ai = _make_chat("orchestrator")
|
||||||
|
names = set(ai.tool_registry._tools.keys())
|
||||||
|
assert "update_identity" in names
|
||||||
|
assert "update_soul" in names
|
||||||
|
assert "set_proxy" in names
|
||||||
|
|
||||||
|
def test_worker_bg_tools_carry_owner(self):
|
||||||
|
ai = _make_chat("worker", owner="sess-1")
|
||||||
|
assert ai.tool_registry.get_tool("bg_run").owner == "sess-1"
|
||||||
|
assert ai.tool_registry.get_tool("bg_kill").owner == "sess-1"
|
||||||
|
assert ai.tool_registry.get_tool("bg_write").owner == "sess-1"
|
||||||
|
# 观察面工具不带 owner
|
||||||
|
assert getattr(ai.tool_registry.get_tool("bg_list"), "owner", "unset") == "unset"
|
||||||
|
|
||||||
|
def test_create_task_ai_uses_task_profile(self):
|
||||||
|
from aibox import TaskScheduler
|
||||||
|
# 显式新建 loop:全量运行时主线程可能已无可用的当前事件循环
|
||||||
|
loop = asyncio.new_event_loop()
|
||||||
|
try:
|
||||||
|
scheduler = TaskScheduler(loop=loop, logger=None)
|
||||||
|
with mock.patch("aibox.create_llm_client") as mclc:
|
||||||
|
mclc.return_value = mock.MagicMock(
|
||||||
|
api_key="sk-test", name="DS", default_temperature=None)
|
||||||
|
ai = scheduler._create_task_ai()
|
||||||
|
assert ai.tool_profile == "task"
|
||||||
|
assert ai.owner_id and ai.owner_id.startswith("task-")
|
||||||
|
finally:
|
||||||
|
loop.close()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== T-L17/T-L18:model_config 展示与推荐 ====================
|
||||||
|
|
||||||
|
class TestTL17EnvSuggestion:
|
||||||
|
def test_suggest_by_domain(self):
|
||||||
|
from model_config import _suggest_env_name
|
||||||
|
assert _suggest_env_name("https://api.moonshot.cn/v1") == "MOONSHOT_API_KEY"
|
||||||
|
assert _suggest_env_name("https://api.deepseek.com/chat") == "DEEPSEEK_API_KEY"
|
||||||
|
assert _suggest_env_name("https://dashscope.aliyuncs.com/x") == "DASHSCOPE_API_KEY"
|
||||||
|
|
||||||
|
def test_unknown_domain_empty(self):
|
||||||
|
from model_config import _suggest_env_name
|
||||||
|
assert _suggest_env_name("http://localhost:11434/v1") == ""
|
||||||
|
assert _suggest_env_name("") == ""
|
||||||
|
assert _suggest_env_name(None) == ""
|
||||||
|
|
||||||
|
|
||||||
|
class TestTL18DisplaySanitize:
|
||||||
|
def test_sanitize_strips_control_chars(self):
|
||||||
|
from model_config import _sanitize_display
|
||||||
|
out = _sanitize_display("a\x1b]0;pwned\x07b")
|
||||||
|
assert "\x1b" not in out and "\x07" not in out
|
||||||
|
assert out == "a]0;pwnedb" # 控制符剥离,可打印部分保留
|
||||||
|
assert _sanitize_display("正常文本") == "正常文本"
|
||||||
|
assert _sanitize_display(None) == ""
|
||||||
|
|
||||||
|
def test_print_profiles_filters_escape_sequences(self, capsys):
|
||||||
|
from model_config import _print_profiles
|
||||||
|
_print_profiles([{
|
||||||
|
"name": "evil\x1b[31m", "is_default": False,
|
||||||
|
"display_name": "d\x07x", "api_url": "https://x\x1b[2J",
|
||||||
|
"model": "m", "api_key_env": "K\x1bZ",
|
||||||
|
"temperature": None, "timeout": None, "fallback": ["fb\x07"],
|
||||||
|
}])
|
||||||
|
out = capsys.readouterr().out
|
||||||
|
assert "\x1b" not in out
|
||||||
|
assert "\x07" not in out
|
||||||
|
assert "evil" in out # 可打印部分保留
|
||||||
@@ -0,0 +1,448 @@
|
|||||||
|
"""后台进程管理器(subsystems.process_manager)与 bg_* 工具回归测试
|
||||||
|
|
||||||
|
覆盖三决策语义与安全边界:
|
||||||
|
- bg 生命周期:启动立即返回 / 惰性收割 / kill 幂等 / TTL 过期
|
||||||
|
- 输出缓冲:tail/since 游标续读 / 环形淘汰提示 / 单行截断 / 拉取字节上限
|
||||||
|
- 上限:运行中进程数上限 / 总条目 LRU 淘汰
|
||||||
|
- 收尾:shutdown 杀非 detach、保留 detach;atexit 单例共享
|
||||||
|
- 工具层:参数规整(字符串句柄)/ 双闸禁用 / 注册表分发
|
||||||
|
|
||||||
|
跨平台:长驻进程用 sys.executable -c "time.sleep(...)"(echo/ping 不可移植)。
|
||||||
|
Windows 注意:shell=True 下 terminate 只杀 shell 包装进程(与执行器 L9
|
||||||
|
回退语义一致),被杀判定以 popen.poll() 为准。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import queue
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
import types
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
from subsystems import ProcessManager, get_process_manager, reset_process_manager # noqa: E402
|
||||||
|
from subsystems.process_manager import MAX_LINE_BYTES # noqa: E402
|
||||||
|
from tools import ( # noqa: E402
|
||||||
|
ToolRegistry,
|
||||||
|
BgRunTool, BgWriteTool, BgListTool, BgOutputTool, BgKillTool,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
PY = sys.executable
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def echo_srv(tmp_path):
|
||||||
|
"""行协议回环服务器:stdin 读一行 → stdout 回 ECHO:行,EOF 退出。
|
||||||
|
交互式 I/O(bg_write/watermark/优雅终止)测试的标准对端。"""
|
||||||
|
script = (
|
||||||
|
"import sys\n"
|
||||||
|
"while True:\n"
|
||||||
|
" line = sys.stdin.readline()\n"
|
||||||
|
" if not line:\n"
|
||||||
|
" break\n"
|
||||||
|
" sys.stdout.write('ECHO:' + line.strip() + '\\n')\n"
|
||||||
|
" sys.stdout.flush()\n"
|
||||||
|
)
|
||||||
|
f = tmp_path / "echo_srv.py"
|
||||||
|
f.write_text(script, encoding="utf-8")
|
||||||
|
return f'"{PY}" "{f}"'
|
||||||
|
|
||||||
|
|
||||||
|
def _py(code: str) -> str:
|
||||||
|
"""跨平台 python 命令(code 内只用单引号)"""
|
||||||
|
return f'"{PY}" -c "{code}"'
|
||||||
|
|
||||||
|
|
||||||
|
ECHO_10 = _py("for i in range(10): print('line-%d' % i)")
|
||||||
|
SLEEP_10 = _py("import time; time.sleep(10)")
|
||||||
|
HELLO = _py("print('hello-bg')")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def mgr():
|
||||||
|
"""每测试独立管理器(不走单例),收尾兜底清进程"""
|
||||||
|
m = ProcessManager()
|
||||||
|
try:
|
||||||
|
yield m
|
||||||
|
finally:
|
||||||
|
m.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
def _wait_status(m: ProcessManager, handle: int, status: str, timeout: float = 10.0):
|
||||||
|
"""轮询等待句柄达到指定状态"""
|
||||||
|
deadline = time.time() + timeout
|
||||||
|
while time.time() < deadline:
|
||||||
|
for p in m.list_procs():
|
||||||
|
if p["handle"] == handle and p["status"] == status:
|
||||||
|
return p
|
||||||
|
time.sleep(0.1)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 启动与生命周期 ====================
|
||||||
|
|
||||||
|
class TestStart:
|
||||||
|
def test_start_returns_handle_immediately(self, mgr):
|
||||||
|
t0 = time.time()
|
||||||
|
r = mgr.start(SLEEP_10, name="sleeper")
|
||||||
|
assert r["ok"] is True
|
||||||
|
assert time.time() - t0 < 2 # 立即返回,不等进程结束
|
||||||
|
assert isinstance(r["handle"], int) and r["pid"] > 0
|
||||||
|
assert r["name"] == "sleeper"
|
||||||
|
|
||||||
|
def test_start_default_name_from_command(self, mgr):
|
||||||
|
r = mgr.start(HELLO)
|
||||||
|
assert r["ok"] is True
|
||||||
|
# 默认名称 = 命令前 40 字符
|
||||||
|
assert r["name"] == HELLO.strip()[:40]
|
||||||
|
|
||||||
|
def test_start_rejects_empty_command(self, mgr):
|
||||||
|
assert mgr.start(" ")["ok"] is False
|
||||||
|
assert mgr.start(None)["ok"] is False
|
||||||
|
|
||||||
|
def test_start_rejects_bad_name(self, mgr):
|
||||||
|
# 名称含控制字符(M5 惯例:防终端/前端注入)
|
||||||
|
assert mgr.start(HELLO, name="bad\x1b[31m")["ok"] is False
|
||||||
|
assert mgr.start(HELLO, name="x" * 65)["ok"] is False
|
||||||
|
|
||||||
|
def test_start_missing_binary_returns_error(self, mgr):
|
||||||
|
r = mgr.start("definitely-not-a-command-xyz --version")
|
||||||
|
# shell 本身能启动(rc=127/9009),故启动成功、退出码非零
|
||||||
|
assert r["ok"] is True
|
||||||
|
p = _wait_status(mgr, r["handle"], "exited")
|
||||||
|
assert p is not None
|
||||||
|
assert p["exit_code"] != 0
|
||||||
|
|
||||||
|
|
||||||
|
class TestLifecycle:
|
||||||
|
def test_exited_status_and_exit_code(self, mgr):
|
||||||
|
r = mgr.start(HELLO)
|
||||||
|
p = _wait_status(mgr, r["handle"], "exited")
|
||||||
|
assert p is not None
|
||||||
|
assert p["exit_code"] == 0
|
||||||
|
|
||||||
|
def test_kill_running_process(self, mgr):
|
||||||
|
r = mgr.start(SLEEP_10)
|
||||||
|
time.sleep(0.5) # 等进程真正起来
|
||||||
|
k = mgr.kill(r["handle"])
|
||||||
|
assert k["ok"] is True and k["status"] == "killed"
|
||||||
|
assert mgr._procs[r["handle"]].popen.poll() is not None
|
||||||
|
|
||||||
|
def test_kill_is_idempotent_on_exited(self, mgr):
|
||||||
|
r = mgr.start(HELLO)
|
||||||
|
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||||
|
k = mgr.kill(r["handle"])
|
||||||
|
assert k["ok"] is True and k["already_finished"] is True
|
||||||
|
|
||||||
|
def test_kill_missing_handle(self, mgr):
|
||||||
|
assert mgr.kill(999)["ok"] is False
|
||||||
|
|
||||||
|
def test_ttl_expiry(self):
|
||||||
|
m = ProcessManager(max_lifetime=0.5)
|
||||||
|
try:
|
||||||
|
r = m.start(SLEEP_10)
|
||||||
|
time.sleep(0.8)
|
||||||
|
# list 触发惰性 sweep
|
||||||
|
assert _wait_status(m, r["handle"], "expired", timeout=2) is not None
|
||||||
|
finally:
|
||||||
|
m.shutdown()
|
||||||
|
|
||||||
|
def test_max_processes_limit(self):
|
||||||
|
m = ProcessManager(max_processes=2)
|
||||||
|
try:
|
||||||
|
assert m.start(SLEEP_10)["ok"] is True
|
||||||
|
assert m.start(SLEEP_10)["ok"] is True
|
||||||
|
r3 = m.start(SLEEP_10)
|
||||||
|
assert r3["ok"] is False and "上限" in r3["error"]
|
||||||
|
finally:
|
||||||
|
m.shutdown()
|
||||||
|
|
||||||
|
def test_max_entries_evicts_oldest_exited(self):
|
||||||
|
m = ProcessManager(max_entries=2)
|
||||||
|
try:
|
||||||
|
h1 = m.start(HELLO)["handle"]
|
||||||
|
h2 = m.start(HELLO)["handle"]
|
||||||
|
assert _wait_status(m, h1, "exited") is not None
|
||||||
|
assert _wait_status(m, h2, "exited") is not None
|
||||||
|
# 第三条触发最老已退出条目(#1)LRU 淘汰
|
||||||
|
h3 = m.start(HELLO)["handle"]
|
||||||
|
assert h3 > 0
|
||||||
|
handles = {p["handle"] for p in m.list_procs()}
|
||||||
|
assert h1 not in handles and h2 in handles and h3 in handles
|
||||||
|
finally:
|
||||||
|
m.shutdown()
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 输出缓冲与游标 ====================
|
||||||
|
|
||||||
|
class TestOutput:
|
||||||
|
def test_tail_and_since_cursor(self, mgr):
|
||||||
|
r = mgr.start(ECHO_10)
|
||||||
|
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||||
|
time.sleep(0.3) # reader 排空末尾行
|
||||||
|
|
||||||
|
out = mgr.get_output(r["handle"], tail=3)
|
||||||
|
assert out["ok"] is True
|
||||||
|
assert out["lines"] == 3
|
||||||
|
assert "line-9" in out["output"] and "line-6" not in out["output"]
|
||||||
|
|
||||||
|
full = mgr.get_output(r["handle"], since=0)
|
||||||
|
assert full["lines"] == 10
|
||||||
|
assert full["next_cursor"] == 10
|
||||||
|
|
||||||
|
inc = mgr.get_output(r["handle"], since=full["next_cursor"])
|
||||||
|
assert inc["lines"] == 0 and inc["next_cursor"] == 10
|
||||||
|
|
||||||
|
def test_since_before_eviction_reports_gap(self):
|
||||||
|
m = ProcessManager(buffer_per_process=2048)
|
||||||
|
try:
|
||||||
|
r = m.start(_py("for i in range(50): print('pad-%03d-%s' % (i, 'x'*80))"))
|
||||||
|
assert _wait_status(m, r["handle"], "exited") is not None
|
||||||
|
time.sleep(0.3)
|
||||||
|
out = m.get_output(r["handle"], since=0)
|
||||||
|
assert out["ok"] is True
|
||||||
|
assert out["evicted"] > 0 # 早期行被环形淘汰
|
||||||
|
assert out["lines"] < 50
|
||||||
|
finally:
|
||||||
|
m.shutdown()
|
||||||
|
|
||||||
|
def test_single_line_truncated(self, mgr):
|
||||||
|
r = mgr.start(_py("print('A' * (1024*64))"))
|
||||||
|
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||||
|
time.sleep(0.3)
|
||||||
|
out = mgr.get_output(r["handle"], tail=5)
|
||||||
|
assert "单行截断" in out["output"]
|
||||||
|
# 缓冲内存预算不被超长单行击穿
|
||||||
|
entry = mgr._procs[r["handle"]]
|
||||||
|
assert entry.buffer_bytes <= MAX_LINE_BYTES + 1024
|
||||||
|
|
||||||
|
def test_fetch_byte_limit_truncates(self, mgr):
|
||||||
|
r = mgr.start(_py("for i in range(100): print('y' * 200)"))
|
||||||
|
assert _wait_status(mgr, r["handle"], "exited") is not None
|
||||||
|
time.sleep(0.3)
|
||||||
|
out = mgr.get_output(r["handle"], since=0, max_bytes=4096)
|
||||||
|
assert out["truncated"] is True
|
||||||
|
|
||||||
|
def test_missing_handle(self, mgr):
|
||||||
|
assert mgr.get_output(999)["ok"] is False
|
||||||
|
|
||||||
|
def test_invalid_since(self, mgr):
|
||||||
|
r = mgr.start(HELLO)
|
||||||
|
assert mgr.get_output(r["handle"], since="abc")["ok"] is False
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 交互式 I/O(stdin 管道 + watermark 对齐) ====================
|
||||||
|
|
||||||
|
class TestInteractive:
|
||||||
|
def test_write_roundtrip_watermark_alignment(self, mgr, echo_srv):
|
||||||
|
"""核心对齐语义:bg_output(since=watermark) = 本条输入的响应且不含回显"""
|
||||||
|
h = mgr.start(echo_srv)["handle"]
|
||||||
|
r = mgr.write(h, "hello", wait_lines=1, wait_timeout=10)
|
||||||
|
assert r["ok"] is True
|
||||||
|
assert "ECHO:hello" in r["wait"]["output"]
|
||||||
|
assert r["wait"]["timed_out"] is False
|
||||||
|
|
||||||
|
out = mgr.get_output(h, since=r["watermark"])
|
||||||
|
assert "ECHO:hello" in out["output"]
|
||||||
|
assert "[stdin]" not in out["output"] # 回显行序号 == watermark,不包含
|
||||||
|
|
||||||
|
def test_echo_line_visible_in_tail(self, mgr, echo_srv):
|
||||||
|
"""[stdin] 回显进环形缓冲,tail 视图可见完整交互时序"""
|
||||||
|
h = mgr.start(echo_srv)["handle"]
|
||||||
|
assert mgr.write(h, "abc", wait_lines=1)["ok"] is True
|
||||||
|
out = mgr.get_output(h, tail=10)
|
||||||
|
assert "[stdin] abc" in out["output"]
|
||||||
|
assert "ECHO:abc" in out["output"]
|
||||||
|
|
||||||
|
def test_write_wait_sync_roundtrip(self, mgr, echo_srv):
|
||||||
|
"""wait_lines 同步等待:一步问答(expect 式)"""
|
||||||
|
h = mgr.start(echo_srv)["handle"]
|
||||||
|
r = mgr.write(h, "1+1", wait_lines=1)
|
||||||
|
assert r["wait"]["lines"] == 1
|
||||||
|
assert r["wait"]["output"].strip() == "ECHO:1+1"
|
||||||
|
assert r["wait"]["next_cursor"] == r["watermark"] + 1
|
||||||
|
|
||||||
|
def test_write_wait_timeout_on_silent_process(self, mgr):
|
||||||
|
"""不读不响应的进程:写入成功(管道缓冲),等待如实超时"""
|
||||||
|
h = mgr.start(SLEEP_10)["handle"]
|
||||||
|
r = mgr.write(h, "ping", wait_lines=1, wait_timeout=0.5)
|
||||||
|
assert r["ok"] is True
|
||||||
|
assert r["wait"]["timed_out"] is True
|
||||||
|
assert r["wait"]["lines"] == 0
|
||||||
|
|
||||||
|
def test_write_to_exited_process(self, mgr):
|
||||||
|
h = mgr.start(HELLO)["handle"]
|
||||||
|
assert _wait_status(mgr, h, "exited") is not None
|
||||||
|
assert mgr.write(h, "x")["ok"] is False
|
||||||
|
|
||||||
|
def test_write_rejects_empty_and_oversize(self, mgr, echo_srv):
|
||||||
|
h = mgr.start(echo_srv)["handle"]
|
||||||
|
assert mgr.write(h, "")["ok"] is False
|
||||||
|
assert mgr.write(h, None)["ok"] is False
|
||||||
|
assert mgr.write(h, "A" * (64 * 1024 + 1))["ok"] is False
|
||||||
|
|
||||||
|
def test_stdin_eof_graceful_kill(self, mgr, echo_srv):
|
||||||
|
"""三级优雅终止第一级:关 stdin → 回环服务器见 EOF 自行退出"""
|
||||||
|
h = mgr.start(echo_srv)["handle"]
|
||||||
|
time.sleep(0.5)
|
||||||
|
k = mgr.kill(h) # 不 force:先关 stdin
|
||||||
|
assert k["ok"] is True
|
||||||
|
assert mgr._procs[h].popen.poll() is not None
|
||||||
|
assert _wait_status(mgr, h, "killed") is not None
|
||||||
|
|
||||||
|
def test_write_after_kill_rejected(self, mgr, echo_srv):
|
||||||
|
h = mgr.start(echo_srv)["handle"]
|
||||||
|
assert mgr.kill(h)["ok"] is True
|
||||||
|
assert mgr.write(h, "late")["ok"] is False
|
||||||
|
|
||||||
|
def test_log_file_tee(self, mgr, monkeypatch, tmp_path):
|
||||||
|
"""log_file 全档镜像:输出与 [stdin] 回显均落盘,可被原生工具读取"""
|
||||||
|
from common import constants as cm
|
||||||
|
monkeypatch.setattr(type(cm), "WORKSPACE_DIR",
|
||||||
|
property(lambda self: str(tmp_path)))
|
||||||
|
script = tmp_path / "srv.py"
|
||||||
|
script.write_text(
|
||||||
|
"import sys\n"
|
||||||
|
"for line in sys.stdin:\n"
|
||||||
|
" sys.stdout.write('ECHO:' + line.strip() + '\\n')\n"
|
||||||
|
" sys.stdout.flush()\n",
|
||||||
|
encoding="utf-8")
|
||||||
|
cmd = f'"{PY}" "{script}"'
|
||||||
|
r = mgr.start(cmd, log_file=True)
|
||||||
|
assert r["ok"] is True and r["log_file"] is not None
|
||||||
|
assert mgr.write(r["handle"], "tee-me", wait_lines=1)["ok"] is True
|
||||||
|
mgr.kill(r["handle"])
|
||||||
|
log = (tmp_path / f"bg_{r['handle']}.log").read_text(encoding="utf-8")
|
||||||
|
assert "[stdin] tee-me" in log
|
||||||
|
assert "ECHO:tee-me" in log
|
||||||
|
|
||||||
|
def test_bg_write_tool_flow(self, mgr, echo_srv):
|
||||||
|
"""工具层冒烟:BgWriteTool 参数规整 + 对齐读取提示"""
|
||||||
|
assert BgRunTool(manager=mgr).execute(command=echo_srv).startswith("✅")
|
||||||
|
out = BgWriteTool(manager=mgr).execute(handle="1", text="hi",
|
||||||
|
wait_lines="1", wait_timeout=10)
|
||||||
|
assert "ECHO:hi" in out
|
||||||
|
assert "watermark" in out
|
||||||
|
# 描述里的行缓冲指引(给模型看的适用前提)
|
||||||
|
assert "PYTHONUNBUFFERED" in BgWriteTool(manager=mgr).get_description()
|
||||||
|
|
||||||
|
def test_registry_dispatch_bg_write(self, mgr, echo_srv):
|
||||||
|
reg = ToolRegistry()
|
||||||
|
reg.register_many([BgRunTool(manager=mgr), BgWriteTool(manager=mgr)],
|
||||||
|
is_builtin=True)
|
||||||
|
assert reg.execute_tool("bg_run", command=echo_srv).startswith("✅")
|
||||||
|
r = reg.execute_tool("bg_write", handle=1, text="via-reg",
|
||||||
|
wait_lines=1, wait_timeout=10)
|
||||||
|
assert "ECHO:via-reg" in r
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 退出收尾 ====================
|
||||||
|
|
||||||
|
class TestShutdown:
|
||||||
|
def test_shutdown_kills_non_detach_keeps_detach(self):
|
||||||
|
m = ProcessManager()
|
||||||
|
try:
|
||||||
|
h1 = m.start(SLEEP_10)["handle"] # 常规:退出时终止
|
||||||
|
h2 = m.start(SLEEP_10, detach=True)["handle"] # detach:保留
|
||||||
|
time.sleep(0.5)
|
||||||
|
e1, e2 = m._procs[h1], m._procs[h2]
|
||||||
|
|
||||||
|
r = m.shutdown()
|
||||||
|
assert r["killed"] == 1 and r["detached"] == 1
|
||||||
|
assert e1.popen.poll() is not None # 常规进程已终止
|
||||||
|
assert e2.popen.poll() is None # detach 进程仍存活
|
||||||
|
assert m.list_procs() == [] # 注册表已清空(detach 仅注销)
|
||||||
|
|
||||||
|
# 幂等
|
||||||
|
assert m.shutdown() == {"killed": 0, "detached": 0}
|
||||||
|
# 收尾后拒绝新进程
|
||||||
|
assert m.start(HELLO)["ok"] is False
|
||||||
|
# 清理 detach 测试进程
|
||||||
|
e2.popen.kill()
|
||||||
|
finally:
|
||||||
|
m.shutdown()
|
||||||
|
|
||||||
|
def test_singleton_shared_across_callers(self):
|
||||||
|
m1 = get_process_manager()
|
||||||
|
m2 = get_process_manager()
|
||||||
|
assert m1 is m2 # 跨会话/跨子代理共享同一注册表
|
||||||
|
reset_process_manager(None)
|
||||||
|
m3 = get_process_manager()
|
||||||
|
assert m3 is not m1
|
||||||
|
m3.shutdown()
|
||||||
|
reset_process_manager(None)
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 工具层 ====================
|
||||||
|
|
||||||
|
class TestTools:
|
||||||
|
def test_bg_run_and_output_flow(self, mgr):
|
||||||
|
run = BgRunTool(manager=mgr)
|
||||||
|
r = run.execute(command=HELLO, proc_name="hello")
|
||||||
|
assert r.startswith("✅") and "#" in r
|
||||||
|
handle = 1
|
||||||
|
assert _wait_status(mgr, handle, "exited") is not None
|
||||||
|
time.sleep(0.3)
|
||||||
|
|
||||||
|
out = BgOutputTool(manager=mgr).execute(handle=handle, since=0)
|
||||||
|
assert "hello-bg" in out
|
||||||
|
assert "next_cursor" in out
|
||||||
|
|
||||||
|
def test_bg_run_string_handle_coercion(self, mgr):
|
||||||
|
"""LLM 可能传字符串数字句柄(W 批输入规整惯例)"""
|
||||||
|
run = BgRunTool(manager=mgr)
|
||||||
|
assert run.execute(command=HELLO).startswith("✅")
|
||||||
|
assert _wait_status(mgr, 1, "exited") is not None
|
||||||
|
out = BgOutputTool(manager=mgr).execute(handle="1", tail=5)
|
||||||
|
assert "hello-bg" in out
|
||||||
|
|
||||||
|
def test_bg_run_missing_command(self, mgr):
|
||||||
|
r = BgRunTool(manager=mgr).execute()
|
||||||
|
assert "缺少" in r
|
||||||
|
|
||||||
|
def test_bg_list_empty_and_filled(self, mgr):
|
||||||
|
lst = BgListTool(manager=mgr)
|
||||||
|
assert "没有后台进程" in lst.execute()
|
||||||
|
mgr.start(SLEEP_10, name="sleeper")
|
||||||
|
r = lst.execute()
|
||||||
|
assert "sleeper" in r and "运行中 1" in r
|
||||||
|
|
||||||
|
def test_bg_kill_flow(self, mgr):
|
||||||
|
mgr.start(SLEEP_10, name="sleeper")
|
||||||
|
time.sleep(0.5)
|
||||||
|
r = BgKillTool(manager=mgr).execute(handle=1)
|
||||||
|
assert r.startswith("✅") and "已终止" in r
|
||||||
|
|
||||||
|
def test_disabled_gate_blocks_all(self, mgr, monkeypatch):
|
||||||
|
from common import constants
|
||||||
|
monkeypatch.setattr(type(constants), "BG_PROCESSES_ENABLED",
|
||||||
|
property(lambda self: False))
|
||||||
|
assert "禁用" in BgRunTool(manager=mgr).execute(command=HELLO)
|
||||||
|
assert "禁用" in BgListTool(manager=mgr).execute()
|
||||||
|
assert "禁用" in BgOutputTool(manager=mgr).execute(handle=1)
|
||||||
|
assert "禁用" in BgKillTool(manager=mgr).execute(handle=1)
|
||||||
|
|
||||||
|
def test_registry_dispatch(self, mgr):
|
||||||
|
"""4 个工具经 ToolRegistry 正常分发(worker 注册链路冒烟)"""
|
||||||
|
reg = ToolRegistry()
|
||||||
|
reg.register_many([
|
||||||
|
BgRunTool(manager=mgr), BgListTool(manager=mgr),
|
||||||
|
BgOutputTool(manager=mgr), BgKillTool(manager=mgr),
|
||||||
|
], is_builtin=True)
|
||||||
|
r = reg.execute_tool("bg_run", command=HELLO, proc_name="smoke")
|
||||||
|
assert r.startswith("✅")
|
||||||
|
assert _wait_status(mgr, 1, "exited") is not None
|
||||||
|
time.sleep(0.3)
|
||||||
|
r = reg.execute_tool("bg_output", handle=1, since=0)
|
||||||
|
assert "hello-bg" in r
|
||||||
|
r = reg.execute_tool("bg_list")
|
||||||
|
assert "smoke" in r
|
||||||
|
|
||||||
|
|
||||||
|
def test_worker_only_tools_includes_bg():
|
||||||
|
"""编排者误调拦截表包含 bg 工具(worker 执行面边界)"""
|
||||||
|
from aibox import WORKER_ONLY_TOOLS
|
||||||
|
assert {"bg_run", "bg_write", "bg_list", "bg_output", "bg_kill"} <= WORKER_ONLY_TOOLS
|
||||||
@@ -0,0 +1,178 @@
|
|||||||
|
"""run_full_subagent 最大轮次(max_turns)行为回归测试(2026-08-19)
|
||||||
|
|
||||||
|
覆盖 aibox.py run_full_subagent 的轮次限制链路(think_with_limit 补丁):
|
||||||
|
- 超限拦截:任务轮 LLM 调用被限制在 max_turns 次,超限递归返回哨兵串
|
||||||
|
- 哨兵串文本契约("已达到最大对话轮次,任务结束。")
|
||||||
|
- 收尾报告:超限/正常结束均由 original_think 绕过轮次限制生成,作为返回值
|
||||||
|
- 入参钳制:max_turns 为 0/None 回退默认 10,负数钳制下限 1
|
||||||
|
- 超限不置 conversation_active=False(收尾保留完整上下文,aibox.py 763-765 注释)
|
||||||
|
- 收尾失败兜底链:最后一条 assistant 消息 → "子代理未产生有效对话内容。"
|
||||||
|
- SubAgentTool.execute 的 max_turns 参数透传(默认 10)
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import threading
|
||||||
|
import types
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
# ---- 环境桩:websockets 未安装时保证 import aibox.py 不失败 ----
|
||||||
|
if "websockets" not in sys.modules:
|
||||||
|
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||||
|
|
||||||
|
from aibox import SubAgentTool, run_full_subagent # noqa: E402
|
||||||
|
|
||||||
|
# 与 aibox.py think_with_limit 超限返回值一致的哨兵串(文本契约)
|
||||||
|
SENTINEL = "已达到最大对话轮次,任务结束。"
|
||||||
|
WRAP_UP_MARKER = "【收尾】"
|
||||||
|
|
||||||
|
|
||||||
|
def _make_fake_chat(task_rounds=999, final_answer=None, wrap_report="收尾报告"):
|
||||||
|
"""模拟 DeepSeekChat:think_and_respond 的"原始实现"。
|
||||||
|
|
||||||
|
每次任务轮调用递归一次(模拟 aibox.py:1801 工具调用后递归继续,
|
||||||
|
经由补丁后的 think_and_respond 走轮次计数);累计 task_rounds 轮后
|
||||||
|
若给出 final_answer 则正常结束(模拟 AI 返回纯文本)。task_rounds
|
||||||
|
设为大数且不给 final_answer 时永不自行收敛——只有轮次限制能终止它。
|
||||||
|
检测到收尾 prompt(含【收尾】标记)时返回 wrap_report。
|
||||||
|
"""
|
||||||
|
fake = mock.MagicMock()
|
||||||
|
state = {"n": 0, "task_inputs": [], "recurse_returns": []}
|
||||||
|
|
||||||
|
def original_think(input_text, output_callback=None, is_continuation=False):
|
||||||
|
if WRAP_UP_MARKER in input_text:
|
||||||
|
return wrap_report
|
||||||
|
state["task_inputs"].append(input_text)
|
||||||
|
state["n"] += 1
|
||||||
|
if final_answer is not None and state["n"] >= task_rounds:
|
||||||
|
return final_answer
|
||||||
|
ret = fake.think_and_respond("请基于以上工具执行结果继续回答。",
|
||||||
|
output_callback, is_continuation=True)
|
||||||
|
state["recurse_returns"].append(ret)
|
||||||
|
return ret
|
||||||
|
|
||||||
|
fake.think_and_respond = original_think
|
||||||
|
fake.stats = {}
|
||||||
|
fake.messages = []
|
||||||
|
fake._cancel_event = threading.Event()
|
||||||
|
fake.conversation_active = True
|
||||||
|
fake._state = state
|
||||||
|
return fake
|
||||||
|
|
||||||
|
|
||||||
|
def _run(fake, **kwargs):
|
||||||
|
with mock.patch("aibox.DeepSeekChat", return_value=fake), \
|
||||||
|
mock.patch("aibox.create_llm_client"), \
|
||||||
|
mock.patch("aibox.SubprocessCommandExecutor"):
|
||||||
|
return run_full_subagent("角色", "任务", None, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
class TestMaxTurnsCutoff:
|
||||||
|
def test_task_turns_capped_at_max_turns(self):
|
||||||
|
"""永不收敛的任务:任务轮 LLM 调用恰为 max_turns 次,第 max_turns+1 次被拦截"""
|
||||||
|
fake = _make_fake_chat(task_rounds=999)
|
||||||
|
result = _run(fake, max_turns=3)
|
||||||
|
assert len(fake._state["task_inputs"]) == 3
|
||||||
|
assert result == "收尾报告"
|
||||||
|
|
||||||
|
def test_overrun_recursion_returns_sentinel(self):
|
||||||
|
"""超限那轮递归返回哨兵串(不调 LLM),并沿递归链上抛"""
|
||||||
|
fake = _make_fake_chat(task_rounds=999)
|
||||||
|
_run(fake, max_turns=3)
|
||||||
|
assert fake._state["recurse_returns"]
|
||||||
|
assert all(r == SENTINEL for r in fake._state["recurse_returns"])
|
||||||
|
|
||||||
|
def test_wrap_up_runs_after_overrun(self):
|
||||||
|
"""超限后收尾报告仍被请求(original_think 绕过轮次限制),且为最终返回值"""
|
||||||
|
fake = _make_fake_chat(task_rounds=999, wrap_report="最终报告")
|
||||||
|
result = _run(fake, max_turns=2)
|
||||||
|
assert result == "最终报告"
|
||||||
|
|
||||||
|
def test_max_turns_zero_clamped_to_default(self):
|
||||||
|
"""max_turns=0 → 回退默认 10(对齐 max_tokens 的 0/空回退惯例)"""
|
||||||
|
fake = _make_fake_chat(task_rounds=999)
|
||||||
|
_run(fake, max_turns=0)
|
||||||
|
assert len(fake._state["task_inputs"]) == 10
|
||||||
|
|
||||||
|
def test_negative_max_turns_clamped_to_one(self):
|
||||||
|
"""max_turns 为负 → 钳制下限 1(至少执行一个任务轮,不再首轮即拦截)"""
|
||||||
|
fake = _make_fake_chat(task_rounds=999)
|
||||||
|
_run(fake, max_turns=-3)
|
||||||
|
assert len(fake._state["task_inputs"]) == 1
|
||||||
|
|
||||||
|
def test_max_turns_none_falls_back_to_default(self):
|
||||||
|
"""max_turns=None → 默认 10"""
|
||||||
|
fake = _make_fake_chat(task_rounds=999)
|
||||||
|
_run(fake, max_turns=None)
|
||||||
|
assert len(fake._state["task_inputs"]) == 10
|
||||||
|
|
||||||
|
def test_overrun_does_not_deactivate_conversation(self):
|
||||||
|
"""超限不置 conversation_active=False(否则收尾前上下文被重置清空)"""
|
||||||
|
fake = _make_fake_chat(task_rounds=999)
|
||||||
|
_run(fake, max_turns=2)
|
||||||
|
assert fake.conversation_active is True
|
||||||
|
|
||||||
|
|
||||||
|
class TestNormalCompletion:
|
||||||
|
def test_result_is_wrap_up_report_not_final_answer(self):
|
||||||
|
"""正常收敛时返回值同样是收尾报告(最终回答不作为工具结果返回)"""
|
||||||
|
fake = _make_fake_chat(task_rounds=2, final_answer="最终答案")
|
||||||
|
result = _run(fake, max_turns=10)
|
||||||
|
assert len(fake._state["task_inputs"]) == 2
|
||||||
|
# 第一轮递归的返回值即第 2 轮的最终回答(无哨兵串、无超限拦截)
|
||||||
|
assert fake._state["recurse_returns"] == ["最终答案"]
|
||||||
|
assert result == "收尾报告"
|
||||||
|
|
||||||
|
def test_sentinel_never_appears_when_within_limit(self):
|
||||||
|
"""未达上限:任何递归都不返回哨兵串"""
|
||||||
|
fake = _make_fake_chat(task_rounds=3, final_answer="最终答案")
|
||||||
|
_run(fake, max_turns=10)
|
||||||
|
assert SENTINEL not in fake._state["recurse_returns"]
|
||||||
|
|
||||||
|
|
||||||
|
class TestWrapUpFallback:
|
||||||
|
def test_empty_wrap_up_falls_back_to_last_assistant_message(self):
|
||||||
|
"""收尾返回空 → 兜底取最后一条有内容的 assistant 消息"""
|
||||||
|
fake = _make_fake_chat(task_rounds=1, final_answer="x", wrap_report="")
|
||||||
|
fake.messages = [
|
||||||
|
{"role": "user", "content": "hi"},
|
||||||
|
{"role": "assistant", "content": "中间回答"},
|
||||||
|
{"role": "assistant", "content": ""},
|
||||||
|
{"role": "assistant", "content": "最终回答"},
|
||||||
|
]
|
||||||
|
assert _run(fake, max_turns=10) == "【子代理最终回答】\n最终回答"
|
||||||
|
|
||||||
|
def test_no_wrap_up_and_no_messages(self):
|
||||||
|
"""收尾为空且无任何 assistant 消息 → 固定兜底文案"""
|
||||||
|
fake = _make_fake_chat(task_rounds=1, final_answer="x", wrap_report="")
|
||||||
|
assert _run(fake, max_turns=10) == "子代理未产生有效对话内容。"
|
||||||
|
|
||||||
|
def test_wrap_up_exception_falls_back(self):
|
||||||
|
"""收尾调用抛异常 → 不中断,走 assistant 消息兜底"""
|
||||||
|
fake = _make_fake_chat(task_rounds=1, final_answer="x")
|
||||||
|
original = fake.think_and_respond
|
||||||
|
|
||||||
|
def boom_on_wrap_up(input_text, output_callback=None, is_continuation=False):
|
||||||
|
if WRAP_UP_MARKER in input_text:
|
||||||
|
raise RuntimeError("wrap-up boom")
|
||||||
|
return original(input_text, output_callback, is_continuation)
|
||||||
|
|
||||||
|
fake.think_and_respond = boom_on_wrap_up
|
||||||
|
fake.messages = [{"role": "assistant", "content": "最终回答"}]
|
||||||
|
assert _run(fake, max_turns=10) == "【子代理最终回答】\n最终回答"
|
||||||
|
|
||||||
|
|
||||||
|
class TestSubAgentToolMaxTurns:
|
||||||
|
def test_max_turns_passed_through(self):
|
||||||
|
"""SubAgentTool.execute 的 max_turns 透传给 run_full_subagent"""
|
||||||
|
tool = SubAgentTool(logger=None)
|
||||||
|
with mock.patch("aibox.run_full_subagent", return_value="ok") as m:
|
||||||
|
tool.execute(role="r", task="t", max_turns=4)
|
||||||
|
assert m.call_args.kwargs["max_turns"] == 4
|
||||||
|
|
||||||
|
def test_max_turns_default_is_10(self):
|
||||||
|
tool = SubAgentTool(logger=None)
|
||||||
|
with mock.patch("aibox.run_full_subagent", return_value="ok") as m:
|
||||||
|
tool.execute(role="r", task="t")
|
||||||
|
assert m.call_args.kwargs["max_turns"] == 10
|
||||||
@@ -45,6 +45,13 @@ from .prompt_tools import (
|
|||||||
)
|
)
|
||||||
from .proxy_tools import SetProxyTool, GetProxyTool
|
from .proxy_tools import SetProxyTool, GetProxyTool
|
||||||
from .misc_tools import AskUserTool
|
from .misc_tools import AskUserTool
|
||||||
|
from .process_tools import (
|
||||||
|
BgRunTool,
|
||||||
|
BgWriteTool,
|
||||||
|
BgListTool,
|
||||||
|
BgOutputTool,
|
||||||
|
BgKillTool,
|
||||||
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"Tool",
|
"Tool",
|
||||||
@@ -71,4 +78,10 @@ __all__ = [
|
|||||||
"SetProxyTool",
|
"SetProxyTool",
|
||||||
"GetProxyTool",
|
"GetProxyTool",
|
||||||
"AskUserTool",
|
"AskUserTool",
|
||||||
|
# 后台进程管理工具(bg_run/bg_write/bg_list/bg_output/bg_kill)
|
||||||
|
"BgRunTool",
|
||||||
|
"BgWriteTool",
|
||||||
|
"BgListTool",
|
||||||
|
"BgOutputTool",
|
||||||
|
"BgKillTool",
|
||||||
]
|
]
|
||||||
+204
-9
@@ -15,8 +15,10 @@ tools.basic_tools - 基础工具集(A 批抽离)
|
|||||||
依赖方向:common <- tools.basic_tools,不反向依赖 aibox 任何符号。
|
依赖方向:common <- tools.basic_tools,不反向依赖 aibox 任何符号。
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import base64
|
||||||
import os
|
import os
|
||||||
import queue
|
import queue
|
||||||
|
import re
|
||||||
import shlex
|
import shlex
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Dict, Optional, Tuple
|
from typing import Dict, Optional, Tuple
|
||||||
@@ -25,6 +27,40 @@ from common import ConfigManager, constants
|
|||||||
from .base import Tool
|
from .base import Tool
|
||||||
|
|
||||||
|
|
||||||
|
# T3(审计 2026-08-21):cmd.exe 变量展开形态——%VAR% 匹配已定义环境变量时
|
||||||
|
# 展开为变量值,值中的 & | 会成为命令分隔符(经典 cmd 注入)。URL 编码
|
||||||
|
# (%20/%E4 等)与 PowerShell 的 %{http_code} 均不匹配该模式,不受影响。
|
||||||
|
_WIN_VAR_PATTERN = re.compile(r"%[A-Za-z_][A-Za-z0-9_]*%")
|
||||||
|
|
||||||
|
|
||||||
|
def _quote_arg(value: str) -> Tuple[str, Optional[str]]:
|
||||||
|
"""平台感知的 shell 参数引用,返回 (quoted, error)。
|
||||||
|
|
||||||
|
POSIX:shlex.quote(sh 单引号语义)。
|
||||||
|
Windows(run_command 经 cmd.exe / shell=True 执行):
|
||||||
|
- 双引号包裹(& | < > ( ) ^ 在双引号内为字面量)
|
||||||
|
- 含 " → 拒绝(闭合外层双引号,其后内容暴露给 cmd 解析)
|
||||||
|
- 匹配 %VAR% 形态 → 拒绝(cmd 变量展开注入面)
|
||||||
|
- 含换行/控制字符 → 拒绝(cmd 命令行不可含换行)
|
||||||
|
"""
|
||||||
|
value = str(value)
|
||||||
|
if os.name != "nt":
|
||||||
|
return shlex.quote(value), None
|
||||||
|
if '"' in value:
|
||||||
|
return "", f"❌ 参数含双引号,无法安全生成 Windows 命令: {value[:100]}"
|
||||||
|
if "\n" in value or "\r" in value:
|
||||||
|
return "", f"❌ 参数含换行符,无法生成命令: {value[:100]}"
|
||||||
|
if _WIN_VAR_PATTERN.search(value):
|
||||||
|
return "", (f"❌ 参数疑似含 cmd 变量引用(%变量名%),已拒绝以防命令注入: "
|
||||||
|
f"{value[:100]}")
|
||||||
|
return f'"{value}"', None
|
||||||
|
|
||||||
|
|
||||||
|
def _ps_quote(value: str) -> str:
|
||||||
|
"""PowerShell 单引号字面量(内部单引号翻倍),用于 -Command 内层参数。"""
|
||||||
|
return "'" + str(value).replace("'", "''") + "'"
|
||||||
|
|
||||||
|
|
||||||
# ==================== 文件操作工具基类 ====================
|
# ==================== 文件操作工具基类 ====================
|
||||||
|
|
||||||
class FileToolBase(Tool):
|
class FileToolBase(Tool):
|
||||||
@@ -96,9 +132,64 @@ class ReadFileTool(FileToolBase):
|
|||||||
|
|
||||||
# 限制max_lines上限,防止请求超大范围
|
# 限制max_lines上限,防止请求超大范围
|
||||||
max_lines = min(max_lines, 500)
|
max_lines = min(max_lines, 500)
|
||||||
quoted = shlex.quote(abs_path)
|
# L1(审计 2026-08-18):start/end_line 为 LLM 可控参数,此前未强转
|
||||||
|
# 直拼 sed 单引号模板——LLM 传含引号的字符串可闭合引号注入命令。
|
||||||
|
# 强转 int(非法输入在此拦截)并夹取合理范围。
|
||||||
|
try:
|
||||||
|
start_line = int(start_line)
|
||||||
|
end_line = int(end_line)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return "❌ 参数错误: start_line/end_line 必须为整数"
|
||||||
|
start_line = max(1, start_line)
|
||||||
|
end_line = max(0, min(end_line, 10_000_000))
|
||||||
|
|
||||||
# 构造等价的命令行读取命令
|
# T3(审计 2026-08-21):shlex.quote 是 POSIX 单引号语义,Windows 上
|
||||||
|
# run_command 经 cmd.exe 执行——单引号不被识别,& | % < > 仍是元字符
|
||||||
|
# (URL 含 & 即命令被截断;%VAR% 变量展开可注入)。按平台生成指引:
|
||||||
|
# POSIX 保持 head/tail/sed,Windows 走 PowerShell 等价命令。
|
||||||
|
quoted, qerr = _quote_arg(abs_path)
|
||||||
|
if qerr:
|
||||||
|
return qerr
|
||||||
|
|
||||||
|
if os.name == "nt":
|
||||||
|
# PowerShell 等价命令(经 cmd 调 powershell,-Command 整体双引号
|
||||||
|
# 包裹,内层路径用 _ps_quote 单引号字面量)
|
||||||
|
if head:
|
||||||
|
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||||
|
f'{_ps_quote(abs_path)} -TotalCount {max_lines}"')
|
||||||
|
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||||
|
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||||
|
elif tail:
|
||||||
|
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||||
|
f'{_ps_quote(abs_path)} -Tail {max_lines}"')
|
||||||
|
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||||
|
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||||
|
else:
|
||||||
|
if end_line > 0:
|
||||||
|
first = start_line - 1
|
||||||
|
take = end_line - start_line + 1
|
||||||
|
elif max_lines > 0:
|
||||||
|
first = start_line - 1
|
||||||
|
take = max_lines
|
||||||
|
else:
|
||||||
|
first = start_line - 1
|
||||||
|
take = 10_000_000
|
||||||
|
cmd = (f'powershell -NoProfile -Command "Get-Content -Path '
|
||||||
|
f'{_ps_quote(abs_path)} | Select-Object -Skip {first} '
|
||||||
|
f'-First {take}"')
|
||||||
|
wc_cmd = (f'powershell -NoProfile -Command "(Get-Content '
|
||||||
|
f'{_ps_quote(abs_path)} | Measure-Object -Line).Lines"')
|
||||||
|
prompt = f"""请使用 `run_command` 工具执行以下命令来读取文件内容,不要直接返回文件内容:
|
||||||
|
|
||||||
|
```
|
||||||
|
{cmd}
|
||||||
|
```
|
||||||
|
|
||||||
|
- 若需统计总行数:`{wc_cmd}`
|
||||||
|
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节,单次最多读取 500 行"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
# 构造等价的命令行读取命令(POSIX)
|
||||||
if head:
|
if head:
|
||||||
cmd = f"head -n {max_lines} {quoted}"
|
cmd = f"head -n {max_lines} {quoted}"
|
||||||
elif tail:
|
elif tail:
|
||||||
@@ -170,13 +261,53 @@ class WriteFileTool(FileToolBase):
|
|||||||
if not is_safe:
|
if not is_safe:
|
||||||
return error
|
return error
|
||||||
|
|
||||||
quoted = shlex.quote(abs_path)
|
quoted, qerr = _quote_arg(abs_path)
|
||||||
|
if qerr:
|
||||||
|
return qerr
|
||||||
|
|
||||||
|
if os.name == "nt":
|
||||||
|
# T3:Windows 下 heredoc 不可用(cmd 命令行不能含换行,且引号
|
||||||
|
# 转义跨 cmd/PowerShell 两层极易出错)。content 整体 Base64 编码
|
||||||
|
# 后经 PowerShell 解码写入——Base64 字符集 [A-Za-z0-9+/=] 不含
|
||||||
|
# 任何 cmd/PowerShell 元字符,多行/引号/特殊字符内容全覆盖。
|
||||||
|
b64 = base64.b64encode(content.encode("utf-8")).decode("ascii")
|
||||||
|
if len(b64) > 6000:
|
||||||
|
return ("❌ 写入内容过大:Windows 命令行长度上限约 8191 字符,"
|
||||||
|
"Base64 编码后超出。请将内容分成多次调用(每次 ≤4KB),"
|
||||||
|
"首次 mode=w、后续 mode=a 追加写入")
|
||||||
|
writer = "Add-Content" if mode == "a" else "Set-Content"
|
||||||
|
write_cmd = (f'powershell -NoProfile -Command "{writer} -Path '
|
||||||
|
f'{_ps_quote(abs_path)} -Value ([Text.Encoding]::UTF8.'
|
||||||
|
f'GetString([Convert]::FromBase64String(\'{b64}\')))"')
|
||||||
|
directory = os.path.dirname(abs_path)
|
||||||
|
if directory:
|
||||||
|
write_cmd = (f'powershell -NoProfile -Command "New-Item '
|
||||||
|
f'-ItemType Directory -Force -Path '
|
||||||
|
f'{_ps_quote(directory)}" & {write_cmd}')
|
||||||
|
prompt = f"""请使用 `run_command` 工具执行以下命令来写入文件,不要在此工具中直接写入:
|
||||||
|
|
||||||
|
```
|
||||||
|
{write_cmd}
|
||||||
|
```
|
||||||
|
|
||||||
|
- 写入内容已编码为 Base64(避免换行/引号/特殊字符的跨 cmd+PowerShell 转义问题),命令会自动解码并以 UTF-8 写入
|
||||||
|
- 若需检查文件大小:`powershell -NoProfile -Command "(Get-Item {_ps_quote(abs_path)}).Length"`
|
||||||
|
- 文件大小上限 {constants.MAX_FILE_SIZE} 字节"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
directory = os.path.dirname(abs_path)
|
directory = os.path.dirname(abs_path)
|
||||||
dir_cmd = f"mkdir -p {shlex.quote(directory)}" if directory else ""
|
dir_cmd = f"mkdir -p {shlex.quote(directory)}" if directory else ""
|
||||||
|
|
||||||
# 内容写入使用 heredoc,避免引号/特殊字符转义问题
|
# 内容写入使用 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 ">>"
|
redirect = ">" if mode == "w" else ">>"
|
||||||
write_cmd = f"cat {redirect} {quoted} <<'AIBOX_EOF'\n{content}\nAIBOX_EOF"
|
write_cmd = f"cat {redirect} {quoted} <<'{delimiter}'\n{content}\n{delimiter}"
|
||||||
if dir_cmd:
|
if dir_cmd:
|
||||||
write_cmd = f"{dir_cmd} && {write_cmd}"
|
write_cmd = f"{dir_cmd} && {write_cmd}"
|
||||||
|
|
||||||
@@ -237,7 +368,26 @@ class ListFilesTool(FileToolBase):
|
|||||||
if not os.path.isdir(abs_path):
|
if not os.path.isdir(abs_path):
|
||||||
return f"❌ 路径不是目录: {abs_path}"
|
return f"❌ 路径不是目录: {abs_path}"
|
||||||
|
|
||||||
quoted = shlex.quote(abs_path)
|
quoted, qerr = _quote_arg(abs_path)
|
||||||
|
if qerr:
|
||||||
|
return qerr
|
||||||
|
|
||||||
|
if os.name == "nt":
|
||||||
|
# T3:Windows 走 PowerShell 等价命令(ls/du/find/file 无原生等价物)
|
||||||
|
force_flag = " -Force" if show_hidden else ""
|
||||||
|
cmd = (f'powershell -NoProfile -Command "Get-ChildItem{force_flag} '
|
||||||
|
f'{_ps_quote(abs_path)} | Select-Object '
|
||||||
|
f'Mode,LastWriteTime,Length,Name"')
|
||||||
|
prompt = f"""请使用 `run_command` 工具执行以下命令来列出目录内容,不要直接在此工具中列出:
|
||||||
|
|
||||||
|
```
|
||||||
|
{cmd}
|
||||||
|
```
|
||||||
|
|
||||||
|
- 若需统计目录总大小:`powershell -NoProfile -Command "(Get-ChildItem -Recurse {_ps_quote(abs_path)} | Measure-Object Length -Sum).Sum"`
|
||||||
|
- 若需递归列出:`powershell -NoProfile -Command "Get-ChildItem -Recurse -Depth 2 {_ps_quote(abs_path)} | Select-Object FullName"`"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
flag = "-la" if show_hidden else "-l"
|
flag = "-la" if show_hidden else "-l"
|
||||||
|
|
||||||
# 返回 markdown 提示词,指导智能体用命令行工具列目录
|
# 返回 markdown 提示词,指导智能体用命令行工具列目录
|
||||||
@@ -307,16 +457,58 @@ class SimpleWebCrawlerTool(Tool):
|
|||||||
host_port = proxy_url.split("://", 1)[-1]
|
host_port = proxy_url.split("://", 1)[-1]
|
||||||
if "@" in host_port:
|
if "@" in host_port:
|
||||||
host_port = host_port.split("@", 1)[-1]
|
host_port = host_port.split("@", 1)[-1]
|
||||||
return f"--socks5 {shlex.quote(host_port)}"
|
# T3:平台感知引用(原 shlex.quote 在 Windows cmd.exe 下无效)
|
||||||
|
quoted, _err = _quote_arg(host_port)
|
||||||
|
return f"--socks5 {quoted}" if quoted else ""
|
||||||
|
|
||||||
def execute(self, action: str, url: str, **kwargs) -> str:
|
def execute(self, action: str, url: str, **kwargs) -> str:
|
||||||
if action not in ("get", "text", "title", "links", "search"):
|
if action not in ("get", "text", "title", "links", "search"):
|
||||||
return f"❌ 未知操作: {action}"
|
return f"❌ 未知操作: {action}"
|
||||||
|
|
||||||
quoted_url = shlex.quote(url)
|
# L10(审计 2026-08-18):url 为 LLM 可控输入(可被网页内容诱导),
|
||||||
|
# 此前不限 scheme——file:// 可经生成的 curl 指引读取本地文件(SSRF)。
|
||||||
|
# 白名单 http/https 后其余协议直接拒绝。
|
||||||
|
if not url or not str(url).strip().lower().startswith(("http://", "https://")):
|
||||||
|
return "❌ 仅支持 http/https 协议地址(已拦截 file:// 等非网络协议)"
|
||||||
|
|
||||||
|
quoted_url, qerr = _quote_arg(url)
|
||||||
|
if qerr:
|
||||||
|
return qerr
|
||||||
timeout = constants.WEB_TIMEOUT
|
timeout = constants.WEB_TIMEOUT
|
||||||
user_agent = shlex.quote(constants.WEB_USER_AGENT)
|
user_agent, ua_err = _quote_arg(constants.WEB_USER_AGENT)
|
||||||
|
if ua_err:
|
||||||
|
return ua_err
|
||||||
proxy = self._proxy_curl_args()
|
proxy = self._proxy_curl_args()
|
||||||
|
|
||||||
|
if os.name == "nt":
|
||||||
|
# T3:Windows 用系统内置 curl.exe(Win10 1803+);sed/grep 等
|
||||||
|
# POSIX 管道无原生等价,分析类动作(text/title/links/search)
|
||||||
|
# 指引先取原始 HTML 再由 AI 自行解析。URL 含 & 时靠双引号
|
||||||
|
# 保护(cmd 双引号内为字面量),%VAR% 形态已在 _quote_arg 拒绝。
|
||||||
|
base = f"curl.exe -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
||||||
|
if action == "get":
|
||||||
|
cmd = f'{base} -w "\\nHTTP_CODE:%{{http_code}}" {quoted_url}'
|
||||||
|
extra = "如需查看响应头:`curl.exe -sIL ...`;如需保存到文件:`curl.exe -sL ... -o 文件名`"
|
||||||
|
elif action == "text":
|
||||||
|
cmd = (f'{base} {quoted_url} | powershell -NoProfile -Command '
|
||||||
|
f'"$input -replace \'<[^>]+>\',\'\'"')
|
||||||
|
extra = "去除 HTML 标签提取纯文本"
|
||||||
|
else:
|
||||||
|
cmd = f"{base} {quoted_url}"
|
||||||
|
extra = ("Windows 下无 grep/sed 管道,请取回原始 HTML 后自行分析"
|
||||||
|
"(提取标题/链接/检索关键词)")
|
||||||
|
prompt = f"""请使用 `run_command` 工具执行以下命令来完成网页操作,不要直接返回网页内容:
|
||||||
|
|
||||||
|
```
|
||||||
|
{cmd}
|
||||||
|
```
|
||||||
|
|
||||||
|
{extra}
|
||||||
|
- 下载大小限制 {constants.WEB_MAX_DOWNLOAD_SIZE} 字节,可在命令后追加 `| powershell -NoProfile -Command "$input | Select-Object -First 100"` 限制
|
||||||
|
- 当前已按配置自动附带代理:{proxy or '无(如需代理请先调用 set_proxy 工具)'}"""
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
# POSIX 分支:_quote_arg 在 POSIX 下即 shlex.quote,直接沿用
|
||||||
base = f"curl -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
base = f"curl -sL --max-time {timeout} -A {user_agent} {proxy}".strip()
|
||||||
|
|
||||||
if action == "get":
|
if action == "get":
|
||||||
@@ -335,7 +527,10 @@ class SimpleWebCrawlerTool(Tool):
|
|||||||
keyword = kwargs.get("keyword")
|
keyword = kwargs.get("keyword")
|
||||||
if not keyword:
|
if not keyword:
|
||||||
return "❌ 需要提供 keyword 参数"
|
return "❌ 需要提供 keyword 参数"
|
||||||
cmd = f"{base} {quoted_url} | grep -c {shlex.quote(keyword)}"
|
quoted_kw, kw_err = _quote_arg(keyword)
|
||||||
|
if kw_err:
|
||||||
|
return kw_err
|
||||||
|
cmd = f"{base} {quoted_url} | grep -c {quoted_kw}"
|
||||||
extra = f"返回匹配行数,大于 0 即表示找到关键词 '{keyword}'"
|
extra = f"返回匹配行数,大于 0 即表示找到关键词 '{keyword}'"
|
||||||
|
|
||||||
# 返回 markdown 提示词,指导智能体用 curl 完成网页访问
|
# 返回 markdown 提示词,指导智能体用 curl 完成网页访问
|
||||||
|
|||||||
+12
-153
@@ -5,14 +5,15 @@ tools.command_tool - 命令行执行工具
|
|||||||
抽离自 aibox.py,使用抽象的 CommandExecutor 实现命令执行(流式 + 超时控制)。
|
抽离自 aibox.py,使用抽象的 CommandExecutor 实现命令执行(流式 + 超时控制)。
|
||||||
工具本身不持有 subprocess 逻辑,而是委托给 SubprocessCommandExecutor。
|
工具本身不持有 subprocess 逻辑,而是委托给 SubprocessCommandExecutor。
|
||||||
|
|
||||||
|
H1 修复(审计 2026-08-17):此前本类内置一套重复的 Popen 读取循环,
|
||||||
|
完全绕过注入的 executor,导致 1MB 输出上限、cancel() 三层中断、
|
||||||
|
超时机制三项保护同时失效。现改为纯委托,行为与 executor 完全一致。
|
||||||
|
|
||||||
依赖:common (ConfigManager, constants) + tools.base.Tool + subsystems.executors
|
依赖:common (ConfigManager, constants) + tools.base.Tool + subsystems.executors
|
||||||
不反向依赖 aibox 任何符号。
|
不反向依赖 aibox 任何符号。
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
|
||||||
import queue
|
import queue
|
||||||
import subprocess
|
|
||||||
import time
|
|
||||||
from typing import Dict
|
from typing import Dict
|
||||||
|
|
||||||
from common import ConfigManager
|
from common import ConfigManager
|
||||||
@@ -60,163 +61,21 @@ class CommandLineTool(Tool):
|
|||||||
}
|
}
|
||||||
|
|
||||||
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
def execute_streaming(self, output_queue: queue.Queue, **kwargs):
|
||||||
"""流式执行命令,支持超时控制"""
|
"""流式执行命令:纯委托给 executor(1MB 上限/超时/cancel 全部生效)"""
|
||||||
command = kwargs.get("command")
|
command = kwargs.get("command")
|
||||||
timeout = kwargs.get("timeout", 30)
|
|
||||||
work_dir = kwargs.get("work_dir", ".")
|
|
||||||
extra_env = kwargs.get("env", {})
|
|
||||||
|
|
||||||
if not command:
|
if not command:
|
||||||
output_queue.put(("line", "❌ 错误:缺少命令参数\n"))
|
output_queue.put(("line", "❌ 错误:缺少命令参数\n"))
|
||||||
output_queue.put(("complete", "❌ 错误:缺少命令参数"))
|
output_queue.put(("complete", "❌ 错误:缺少命令参数"))
|
||||||
return
|
return
|
||||||
|
|
||||||
# 发送开始执行的消息
|
self.executor.execute_streaming(
|
||||||
output_queue.put(("line", f"\n🔧 执行命令: {command}\n"))
|
command,
|
||||||
output_queue.put(("line", f"⏱️ 超时设置: {timeout} 秒\n"))
|
output_queue,
|
||||||
if work_dir != ".":
|
timeout=kwargs.get("timeout", 30),
|
||||||
output_queue.put(("line", f"📁 工作目录: {work_dir}\n"))
|
work_dir=kwargs.get("work_dir", "."),
|
||||||
output_queue.put(("line", "-" * 50 + "\n"))
|
env=kwargs.get("env"),
|
||||||
|
)
|
||||||
# 准备环境变量
|
|
||||||
env = os.environ.copy()
|
|
||||||
if extra_env:
|
|
||||||
env.update(extra_env)
|
|
||||||
|
|
||||||
try:
|
|
||||||
# 使用 subprocess.Popen 实现流式输出,支持超时
|
|
||||||
process = subprocess.Popen(
|
|
||||||
command,
|
|
||||||
shell=True,
|
|
||||||
stdout=subprocess.PIPE,
|
|
||||||
stderr=subprocess.STDOUT,
|
|
||||||
stdin=subprocess.PIPE,
|
|
||||||
text=True,
|
|
||||||
bufsize=1,
|
|
||||||
universal_newlines=True,
|
|
||||||
env=env,
|
|
||||||
cwd=os.path.expanduser(work_dir)
|
|
||||||
)
|
|
||||||
|
|
||||||
start_time = time.time()
|
|
||||||
output_lines = []
|
|
||||||
timeout_occurred = False
|
|
||||||
|
|
||||||
# 实时读取输出
|
|
||||||
while True:
|
|
||||||
# 检查超时
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
if timeout > 0 and elapsed > timeout:
|
|
||||||
process.terminate()
|
|
||||||
timeout_occurred = True
|
|
||||||
output_queue.put(("line", f"\n⏰ 命令执行超时({timeout}秒),已终止\n"))
|
|
||||||
break
|
|
||||||
|
|
||||||
# 使用非阻塞方式读取一行
|
|
||||||
import select
|
|
||||||
import sys
|
|
||||||
|
|
||||||
# 检查是否有输出可读
|
|
||||||
if sys.platform == 'win32':
|
|
||||||
# Windows 下使用 readline 带超时
|
|
||||||
try:
|
|
||||||
import msvcrt
|
|
||||||
if msvcrt.kbhit():
|
|
||||||
pass
|
|
||||||
line = process.stdout.readline()
|
|
||||||
if not line and process.poll() is not None:
|
|
||||||
break
|
|
||||||
if line:
|
|
||||||
line = line.rstrip('\n')
|
|
||||||
output_queue.put(("line", line + "\n"))
|
|
||||||
output_lines.append(line)
|
|
||||||
else:
|
|
||||||
time.sleep(0.05)
|
|
||||||
except:
|
|
||||||
line = process.stdout.readline()
|
|
||||||
if not line and process.poll() is not None:
|
|
||||||
break
|
|
||||||
if line:
|
|
||||||
line = line.rstrip('\n')
|
|
||||||
output_queue.put(("line", line + "\n"))
|
|
||||||
output_lines.append(line)
|
|
||||||
else:
|
|
||||||
# Unix 系统可以使用 select
|
|
||||||
import select
|
|
||||||
rlist, _, _ = select.select([process.stdout], [], [], 0.5)
|
|
||||||
if rlist:
|
|
||||||
line = process.stdout.readline()
|
|
||||||
if not line and process.poll() is not None:
|
|
||||||
break
|
|
||||||
if line:
|
|
||||||
line = line.rstrip('\n')
|
|
||||||
output_queue.put(("line", line + "\n"))
|
|
||||||
output_lines.append(line)
|
|
||||||
else:
|
|
||||||
# 检查进程是否结束
|
|
||||||
if process.poll() is not None:
|
|
||||||
# 读取剩余输出
|
|
||||||
remaining = process.stdout.read()
|
|
||||||
if remaining:
|
|
||||||
remaining = remaining.rstrip('\n')
|
|
||||||
if remaining:
|
|
||||||
output_queue.put(("line", remaining + "\n"))
|
|
||||||
output_lines.append(remaining)
|
|
||||||
break
|
|
||||||
|
|
||||||
# 如果超时,确保进程被终止
|
|
||||||
if timeout_occurred and process.poll() is None:
|
|
||||||
process.kill()
|
|
||||||
process.wait(timeout=2)
|
|
||||||
|
|
||||||
# 获取返回码
|
|
||||||
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"))
|
|
||||||
else:
|
|
||||||
output_queue.put(("line", f"✅ 命令执行完成,返回码: {return_code}\n"))
|
|
||||||
|
|
||||||
# 准备完整结果
|
|
||||||
full_output = "\n".join(output_lines)
|
|
||||||
data = {
|
|
||||||
"输出": full_output if full_output else "(无输出)",
|
|
||||||
"返回码": return_code
|
|
||||||
}
|
|
||||||
|
|
||||||
if timeout_occurred:
|
|
||||||
result = self.format_result(False, f"命令执行超时({timeout}秒)", data)
|
|
||||||
else:
|
|
||||||
result = self.format_result(return_code == 0, "命令执行完成", data)
|
|
||||||
|
|
||||||
output_queue.put(("complete", result))
|
|
||||||
|
|
||||||
except subprocess.TimeoutExpired:
|
|
||||||
error_msg = f"❌ 命令执行超时({timeout}秒)"
|
|
||||||
output_queue.put(("line", error_msg + "\n"))
|
|
||||||
output_queue.put(("complete", error_msg))
|
|
||||||
except FileNotFoundError as e:
|
|
||||||
error_msg = f"❌ 命令执行错误: 命令未找到 - {str(e)}"
|
|
||||||
output_queue.put(("line", error_msg + "\n"))
|
|
||||||
output_queue.put(("complete", error_msg))
|
|
||||||
except PermissionError as e:
|
|
||||||
error_msg = f"❌ 命令执行错误: 权限不足 - {str(e)}"
|
|
||||||
output_queue.put(("line", error_msg + "\n"))
|
|
||||||
output_queue.put(("complete", error_msg))
|
|
||||||
except Exception as e:
|
|
||||||
error_msg = f"❌ 执行错误: {str(e)}"
|
|
||||||
output_queue.put(("line", error_msg + "\n"))
|
|
||||||
output_queue.put(("complete", error_msg))
|
|
||||||
finally:
|
|
||||||
# 确保进程被清理
|
|
||||||
if 'process' in locals() and process and process.poll() is None:
|
|
||||||
try:
|
|
||||||
process.terminate()
|
|
||||||
process.wait(timeout=2)
|
|
||||||
except:
|
|
||||||
process.kill()
|
|
||||||
|
|
||||||
def execute(self, command: str, timeout: int = 30,
|
def execute(self, command: str, timeout: int = 30,
|
||||||
work_dir: str = ".", env: Dict = None, **kwargs) -> str:
|
work_dir: str = ".", env: Dict = None, **kwargs) -> str:
|
||||||
|
|||||||
+14
-3
@@ -745,19 +745,30 @@ class MemoTool(Tool):
|
|||||||
return f"❌ 无法完成任务 #{memo_id}"
|
return f"❌ 无法完成任务 #{memo_id}"
|
||||||
|
|
||||||
def _delete_memo(self, kwargs: Dict) -> str:
|
def _delete_memo(self, kwargs: Dict) -> str:
|
||||||
"""删除任务"""
|
"""删除任务(单条 memo_id 或批量 memo_ids)"""
|
||||||
confirm = kwargs.get("confirm", False)
|
confirm = kwargs.get("confirm", False)
|
||||||
if not confirm:
|
if not confirm:
|
||||||
return "❌ 需要设置 confirm=True 来确认删除操作"
|
return "❌ 需要设置 confirm=True 来确认删除操作"
|
||||||
|
|
||||||
|
memo_ids = kwargs.get("memo_ids")
|
||||||
|
if memo_ids:
|
||||||
|
# J10(审计 2026-08-19):schema 声明了 memo_ids 参数却从未实现,
|
||||||
|
# 批量请求此前被静默忽略
|
||||||
|
if not isinstance(memo_ids, (list, tuple)):
|
||||||
|
return "❌ memo_ids 必须是列表"
|
||||||
|
deleted = sum(1 for mid in memo_ids if self.memo_db.delete_memo(mid))
|
||||||
|
return self.format_result(True, f"已删除 {deleted}/{len(memo_ids)} 条任务")
|
||||||
|
|
||||||
memo_id = kwargs.get("memo_id")
|
memo_id = kwargs.get("memo_id")
|
||||||
if memo_id:
|
if memo_id:
|
||||||
if self.memo_db.complete_memo(memo_id, force_complete=True):
|
# M5:真实物理删除(此前调 complete_memo 只标记完成,行不删除,
|
||||||
|
# 重复任务被推进而非移除,stats 计数失真)
|
||||||
|
if self.memo_db.delete_memo(memo_id):
|
||||||
return self.format_result(True, f"任务 #{memo_id} 已删除")
|
return self.format_result(True, f"任务 #{memo_id} 已删除")
|
||||||
else:
|
else:
|
||||||
return f"❌ 无法删除任务 #{memo_id}"
|
return f"❌ 无法删除任务 #{memo_id}"
|
||||||
else:
|
else:
|
||||||
return "❌ 需要提供 memo_id"
|
return "❌ 需要提供 memo_id 或 memo_ids"
|
||||||
|
|
||||||
def _search_memos(self, kwargs: Dict) -> str:
|
def _search_memos(self, kwargs: Dict) -> str:
|
||||||
"""搜索任务"""
|
"""搜索任务"""
|
||||||
|
|||||||
@@ -39,10 +39,13 @@ class AskUserTool(Tool):
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"required": []
|
# L8:此前 required 为空,LLM 不传 question 时 str + None 抛 TypeError
|
||||||
|
"required": ["question"]
|
||||||
}
|
}
|
||||||
|
|
||||||
def execute(self, question: str, **kwargs) -> str:
|
def execute(self, question: str = None, **kwargs) -> str:
|
||||||
|
if not question:
|
||||||
|
return "❌ 需要提供问题内容(question 参数)"
|
||||||
return (
|
return (
|
||||||
"你遇到了想搞清楚的事情: " + question + "\n你现在可以询问用户。"
|
"你遇到了想搞清楚的事情: " + question + "\n你现在可以询问用户。"
|
||||||
)
|
)
|
||||||
@@ -0,0 +1,424 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
tools.process_tools - 后台进程管理工具
|
||||||
|
|
||||||
|
五个 worker 执行面工具,配合 subsystems.process_manager.ProcessManager:
|
||||||
|
- BgRunTool : bg_run 后台启动长驻进程,立即返回句柄
|
||||||
|
- BgWriteTool : bg_write 向后台进程 stdin 写入(watermark 写读对齐)
|
||||||
|
- BgListTool : bg_list 列出全部后台进程(运行中/已退出)
|
||||||
|
- BgOutputTool : bg_output 拉取输出(tail 尾部 / since 游标增量)
|
||||||
|
- BgKillTool : bg_kill 终止指定后台进程(关stdin→SIGTERM→SIGKILL)
|
||||||
|
|
||||||
|
设计语义(与用户约定的三决策):
|
||||||
|
- 用户中断(Ctrl+C/取消事件)不影响后台进程
|
||||||
|
- aibox 进程退出时默认终止全部,bg_run detach=true 豁免
|
||||||
|
- 注册表全局共享:任何子代理可管理其他子代理启动的进程
|
||||||
|
|
||||||
|
交互语义:
|
||||||
|
- stdin 为管道:行协议程序(REPL/问答 CLI)可 bg_write 写入;
|
||||||
|
watermark = [stdin] 回显行序号,bg_output(since=watermark) 即本条
|
||||||
|
输入的响应(写读对齐)
|
||||||
|
- bg_run log_file=true 时输出 tee 到 workspace/bg_<句柄>.log(无界
|
||||||
|
全档),可被原生工具(run_command tail/cat)读取
|
||||||
|
|
||||||
|
依赖:common (ConfigManager, constants) + tools.base.Tool +
|
||||||
|
subsystems.process_manager(get_process_manager 单例)
|
||||||
|
不反向依赖 aibox 任何符号。
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Dict
|
||||||
|
|
||||||
|
from common import ConfigManager, constants
|
||||||
|
from subsystems.process_manager import get_process_manager, ProcessManager
|
||||||
|
from .base import Tool
|
||||||
|
|
||||||
|
|
||||||
|
def _enabled() -> bool:
|
||||||
|
"""后台进程功能开关:commands.enabled 与 background_processes.enabled
|
||||||
|
双闸(后台执行本质仍是命令执行,前者关闭时一并拒绝)"""
|
||||||
|
return constants.COMMANDS_ENABLED and constants.BG_PROCESSES_ENABLED
|
||||||
|
|
||||||
|
|
||||||
|
def _fmt_duration(seconds: float) -> str:
|
||||||
|
"""秒数 → 人读时长(如 3h05m / 12.5s)"""
|
||||||
|
try:
|
||||||
|
seconds = float(seconds)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return "?"
|
||||||
|
if seconds < 60:
|
||||||
|
return f"{seconds:.1f}s"
|
||||||
|
m, s = divmod(int(seconds), 60)
|
||||||
|
if m < 60:
|
||||||
|
return f"{m}m{s:02d}s"
|
||||||
|
h, m = divmod(m, 60)
|
||||||
|
return f"{h}h{m:02d}m{s:02d}s"
|
||||||
|
|
||||||
|
|
||||||
|
def _coerce_int(value, default=None):
|
||||||
|
"""命令面输入规整(W 批惯例):LLM 可能传字符串数字,统一安全转 int"""
|
||||||
|
if value is None:
|
||||||
|
return default
|
||||||
|
try:
|
||||||
|
return int(value)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return default
|
||||||
|
|
||||||
|
|
||||||
|
def _coerce_bool(value, default: bool) -> bool:
|
||||||
|
"""命令面输入规整:LLM 可能传字符串 "false"/"0",统一安全转 bool"""
|
||||||
|
if value is None:
|
||||||
|
return default
|
||||||
|
if isinstance(value, bool):
|
||||||
|
return value
|
||||||
|
if isinstance(value, str):
|
||||||
|
return value.strip().lower() not in ("false", "0", "no", "off", "")
|
||||||
|
return bool(value)
|
||||||
|
|
||||||
|
|
||||||
|
class BgRunTool(Tool):
|
||||||
|
"""后台启动长驻进程"""
|
||||||
|
|
||||||
|
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||||
|
super().__init__(logger)
|
||||||
|
self.manager = manager or get_process_manager(logger)
|
||||||
|
# T22:属主标识——写入/终止权限隔离(观察面 list/output 不限)
|
||||||
|
self.owner = owner
|
||||||
|
|
||||||
|
def get_name(self) -> str:
|
||||||
|
return "bg_run"
|
||||||
|
|
||||||
|
def get_description(self) -> str:
|
||||||
|
return ("后台启动长驻进程(如 dev server、训练任务、watch 监控),立即返回句柄不阻塞。"
|
||||||
|
"与 run_command 的区别:run_command 同步等待命令结束(超时即杀),适合短命令;"
|
||||||
|
"bg_run 适合需要持续运行的进程,启动后用 bg_list 查看状态、bg_output 拉取输出、"
|
||||||
|
"bg_kill 终止。stdin 为管道:行协议交互程序(REPL/问答式 CLI)可用 bg_write 写入。"
|
||||||
|
"注意:用户中断不影响后台进程;aibox 退出时后台进程默认被终止,"
|
||||||
|
"需要跨重启存活的进程用 detach=true(随后无法再经句柄管理)。")
|
||||||
|
|
||||||
|
def get_parameters(self) -> Dict:
|
||||||
|
return {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"command": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "要执行的命令(shell 命令行)",
|
||||||
|
},
|
||||||
|
"proc_name": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "进程标签(可选,便于识别,≤64 字符,默认取命令前 40 字符)",
|
||||||
|
},
|
||||||
|
"work_dir": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "工作目录(默认 ~/.aibox/workspace)",
|
||||||
|
},
|
||||||
|
"env": {
|
||||||
|
"type": "object",
|
||||||
|
"description": "额外的环境变量",
|
||||||
|
},
|
||||||
|
"detach": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "aibox 退出时不终止该进程(默认 false;detach 后进程脱离句柄管理)",
|
||||||
|
},
|
||||||
|
"log_file": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "同时镜像全程输出到 workspace/bg_<句柄>.log(无界全档,可被 run_command 的 tail/cat 读取;默认 false 仅环形缓冲)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["command"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def execute(self, command: str = None, proc_name: str = None,
|
||||||
|
work_dir: str = None, env: Dict = None,
|
||||||
|
detach: bool = False, log_file: bool = False, **kwargs) -> str:
|
||||||
|
if error := self.validate_params(["command"], command=command):
|
||||||
|
return error
|
||||||
|
if not _enabled():
|
||||||
|
return "❌ 后台进程功能已禁用(commands.enabled 或 background_processes.enabled 为 false)"
|
||||||
|
|
||||||
|
# T18:LLM 可能传字符串 "false",裸 bool("false") is True——
|
||||||
|
# detach 误开致进程脱离管理、log_file 误开,统一规整
|
||||||
|
detach = _coerce_bool(detach, False)
|
||||||
|
log_file = _coerce_bool(log_file, False)
|
||||||
|
|
||||||
|
result = self.manager.start(command, name=proc_name, work_dir=work_dir,
|
||||||
|
env=env, detach=detach, log_file=log_file,
|
||||||
|
owner=self.owner)
|
||||||
|
if not result.get("ok"):
|
||||||
|
return self.format_result(False, "后台进程启动失败",
|
||||||
|
{"错误": result.get("error", "未知错误")})
|
||||||
|
data = {
|
||||||
|
"句柄": result["handle"],
|
||||||
|
"PID": result["pid"],
|
||||||
|
"名称": result["name"],
|
||||||
|
"detach": "是(aibox 退出不终止)" if result.get("detach") else "否(aibox 退出时终止)",
|
||||||
|
"后续": f"bg_output 拉取输出 / bg_list 查看状态 / bg_kill 终止(句柄 #{result['handle']})",
|
||||||
|
}
|
||||||
|
if result.get("log_file"):
|
||||||
|
data["日志镜像"] = result["log_file"]
|
||||||
|
return self.format_result(True, f"后台进程已启动(#{result['handle']})", data)
|
||||||
|
|
||||||
|
|
||||||
|
class BgWriteTool(Tool):
|
||||||
|
"""向后台进程 stdin 写入(交互式问答)"""
|
||||||
|
|
||||||
|
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||||
|
super().__init__(logger)
|
||||||
|
self.manager = manager or get_process_manager(logger)
|
||||||
|
# T22:属主标识——仅属主可写(他人 stdin 不容注入)
|
||||||
|
self.owner = owner
|
||||||
|
|
||||||
|
def get_name(self) -> str:
|
||||||
|
return "bg_write"
|
||||||
|
|
||||||
|
def get_description(self) -> str:
|
||||||
|
return ("向后台进程的标准输入写入一行(行协议交互:REPL / 问答式 CLI)。"
|
||||||
|
"返回 watermark 水位:bg_output(handle=..., since=watermark) 得到的必然且仅是"
|
||||||
|
"本条输入引发的进程响应(写读对齐,不含 [stdin] 回显行)。"
|
||||||
|
"适用前提:进程按行读 stdin 并及时输出——建议启动时强制行缓冲"
|
||||||
|
"(python -i -u、PYTHONUNBUFFERED=1、stdbuf -oL),否则响应可能滞留子进程缓冲。"
|
||||||
|
"限制:无 TTY,全屏程序(vim/top 等)不可用;喂密码类交互(如 sudo -S)会留在"
|
||||||
|
"对话记录中,敏感凭据请改用环境变量注入。")
|
||||||
|
|
||||||
|
def get_parameters(self) -> Dict:
|
||||||
|
return {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"handle": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "bg_run 返回的句柄编号",
|
||||||
|
},
|
||||||
|
"text": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "要写入的文本(默认追加换行;多行输入会折叠为单行回显)",
|
||||||
|
},
|
||||||
|
"append_newline": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "写入后是否自动补换行(默认 true;false 发送原始字节)",
|
||||||
|
},
|
||||||
|
"wait_lines": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "写完同步等待 N 行响应再返回(默认 0 立即返回,用 bg_output 轮询)",
|
||||||
|
},
|
||||||
|
"wait_timeout": {
|
||||||
|
"type": "number",
|
||||||
|
"description": "同步等待响应的超时秒数(默认 5,上限 60)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["handle", "text"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def execute(self, handle=None, text: str = None, append_newline=True,
|
||||||
|
wait_lines=None, wait_timeout=None, **kwargs) -> str:
|
||||||
|
if error := self.validate_params(["handle", "text"], handle=handle, text=text):
|
||||||
|
return error
|
||||||
|
if not _enabled():
|
||||||
|
return "❌ 后台进程功能已禁用"
|
||||||
|
|
||||||
|
handle = _coerce_int(handle)
|
||||||
|
if handle is None:
|
||||||
|
return "❌ 错误:handle 必须是整数句柄"
|
||||||
|
wait_i = _coerce_int(wait_lines, 0)
|
||||||
|
wait_t = _coerce_int(wait_timeout, 5)
|
||||||
|
|
||||||
|
result = self.manager.write(
|
||||||
|
handle, text,
|
||||||
|
append_newline=_coerce_bool(append_newline, True),
|
||||||
|
wait_lines=max(0, wait_i),
|
||||||
|
wait_timeout=wait_t,
|
||||||
|
owner=self.owner,
|
||||||
|
)
|
||||||
|
if not result.get("ok"):
|
||||||
|
return self.format_result(False, f"写入后台进程 #{handle} 失败",
|
||||||
|
{"错误": result.get("error", "未知错误")})
|
||||||
|
|
||||||
|
data = {
|
||||||
|
"句柄": handle,
|
||||||
|
"已写入": f"{result['sent_bytes']} 字节",
|
||||||
|
"watermark": result["watermark"],
|
||||||
|
"对齐读取": f"bg_output(handle={handle}, since={result['watermark']}) → 本条输入的响应",
|
||||||
|
}
|
||||||
|
if "wait" in result:
|
||||||
|
w = result["wait"]
|
||||||
|
data["等待输出"] = w.get("output") or "(超时未收到响应)"
|
||||||
|
data["next_cursor"] = w.get("next_cursor")
|
||||||
|
if w.get("timed_out"):
|
||||||
|
data["警告"] = "等待响应超时(进程可能无输出或缓冲滞留),可用 bg_output 轮询"
|
||||||
|
return self.format_result(True, f"已写入后台进程 #{handle}", data)
|
||||||
|
|
||||||
|
|
||||||
|
class BgListTool(Tool):
|
||||||
|
"""列出后台进程"""
|
||||||
|
|
||||||
|
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||||
|
super().__init__(logger)
|
||||||
|
self.manager = manager or get_process_manager(logger)
|
||||||
|
|
||||||
|
def get_name(self) -> str:
|
||||||
|
return "bg_list"
|
||||||
|
|
||||||
|
def get_description(self) -> str:
|
||||||
|
return ("列出全部后台进程(含已退出的,保留尾部输出供验尸)。运行中在前。"
|
||||||
|
"输出包含句柄/名称/PID/状态/运行时长,用句柄配合 bg_output/bg_kill 管理。")
|
||||||
|
|
||||||
|
def get_parameters(self) -> Dict:
|
||||||
|
return {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
def execute(self, **kwargs) -> str:
|
||||||
|
if not _enabled():
|
||||||
|
return "❌ 后台进程功能已禁用"
|
||||||
|
procs = self.manager.list_procs()
|
||||||
|
if not procs:
|
||||||
|
return "✅ 当前没有后台进程(bg_run 启动)"
|
||||||
|
lines = []
|
||||||
|
status_emoji = {"running": "🟢", "exited": "⚪", "killed": "🔴", "expired": "⏹️"}
|
||||||
|
for p in procs:
|
||||||
|
emoji = status_emoji.get(p["status"], "❔")
|
||||||
|
extra = ""
|
||||||
|
if p["status"] == "running":
|
||||||
|
extra = f"运行 {_fmt_duration(p['runtime_sec'])}"
|
||||||
|
else:
|
||||||
|
extra = f"已结束({_fmt_duration(p['runtime_sec'])},返回码 {p['exit_code']})"
|
||||||
|
detach_tag = " [detach]" if p["detach"] else ""
|
||||||
|
lines.append(f"{emoji} #{p['handle']} {p['name']}{detach_tag} "
|
||||||
|
f"(pid {p['pid']}, {extra})")
|
||||||
|
lines.append(f" 命令: {p['command']}")
|
||||||
|
running = sum(1 for p in procs if p["status"] == "running")
|
||||||
|
header = f"共 {len(procs)} 个后台进程(运行中 {running})"
|
||||||
|
return "✅ " + header + "\n" + "\n".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
class BgOutputTool(Tool):
|
||||||
|
"""拉取后台进程输出"""
|
||||||
|
|
||||||
|
def __init__(self, manager: ProcessManager = None, logger=None):
|
||||||
|
super().__init__(logger)
|
||||||
|
self.manager = manager or get_process_manager(logger)
|
||||||
|
|
||||||
|
def get_name(self) -> str:
|
||||||
|
return "bg_output"
|
||||||
|
|
||||||
|
def get_description(self) -> str:
|
||||||
|
return ("拉取后台进程输出。两种模式:tail=N 取最近 N 行(默认 50);"
|
||||||
|
"since=游标 增量拉取上次读取之后的新输出(游标用上次返回的 next_cursor)。"
|
||||||
|
"长驻进程轮询建议用 since 模式避免重复传输旧输出。")
|
||||||
|
|
||||||
|
def get_parameters(self) -> Dict:
|
||||||
|
return {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"handle": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "bg_run 返回的句柄编号",
|
||||||
|
},
|
||||||
|
"tail": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "取最近 N 行(默认 50;与 since 二选一)",
|
||||||
|
},
|
||||||
|
"since": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "增量游标:上次返回的 next_cursor,取此序号之后的新行(与 tail 二选一)",
|
||||||
|
},
|
||||||
|
"max_bytes": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "本次拉取最大字节数(默认 64KB,上限 1MB)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["handle"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def execute(self, handle=None, tail=None, since=None, max_bytes=None, **kwargs) -> str:
|
||||||
|
if error := self.validate_params(["handle"], handle=handle):
|
||||||
|
return error
|
||||||
|
if not _enabled():
|
||||||
|
return "❌ 后台进程功能已禁用"
|
||||||
|
|
||||||
|
handle = _coerce_int(handle)
|
||||||
|
if handle is None:
|
||||||
|
return "❌ 错误:handle 必须是整数句柄"
|
||||||
|
since_i = _coerce_int(since)
|
||||||
|
tail_i = _coerce_int(tail)
|
||||||
|
|
||||||
|
result = self.manager.get_output(handle, tail=tail_i, since=since_i,
|
||||||
|
max_bytes=_coerce_int(max_bytes, 65536))
|
||||||
|
if not result.get("ok"):
|
||||||
|
return self.format_result(False, "拉取输出失败",
|
||||||
|
{"错误": result.get("error", "未知错误")})
|
||||||
|
|
||||||
|
data = {
|
||||||
|
"状态": result["status"] + (f"(返回码 {result['exit_code']})"
|
||||||
|
if result["exit_code"] is not None else ""),
|
||||||
|
"本次行数": result["lines"],
|
||||||
|
"next_cursor": result["next_cursor"],
|
||||||
|
}
|
||||||
|
if result.get("truncated"):
|
||||||
|
data["警告"] = "输出超过本次字节上限,已截断(提高 max_bytes 或用 since 分段拉取)"
|
||||||
|
if result.get("evicted"):
|
||||||
|
data["提示"] = (f"有 {result['evicted']} 行早期输出因环形缓冲淘汰不可读取"
|
||||||
|
"(需要完整日志请让进程重定向到文件)")
|
||||||
|
output = result.get("output") or "(本次无新输出)"
|
||||||
|
data["输出"] = output
|
||||||
|
return self.format_result(True, f"后台进程 #{handle} 输出", data)
|
||||||
|
|
||||||
|
|
||||||
|
class BgKillTool(Tool):
|
||||||
|
"""终止后台进程"""
|
||||||
|
|
||||||
|
def __init__(self, manager: ProcessManager = None, logger=None, owner: str = None):
|
||||||
|
super().__init__(logger)
|
||||||
|
self.manager = manager or get_process_manager(logger)
|
||||||
|
# T22:属主标识——仅属主可终止(防提示注入跨会话杀 dev server)
|
||||||
|
self.owner = owner
|
||||||
|
|
||||||
|
def get_name(self) -> str:
|
||||||
|
return "bg_kill"
|
||||||
|
|
||||||
|
def get_description(self) -> str:
|
||||||
|
return ("终止后台进程(整个进程组,含子进程)。默认先 SIGTERM 优雅终止,"
|
||||||
|
"宽限期内未退出再 SIGKILL;force=true 直接 SIGKILL。对已退出进程幂等。")
|
||||||
|
|
||||||
|
def get_parameters(self) -> Dict:
|
||||||
|
return {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"handle": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "bg_run 返回的句柄编号",
|
||||||
|
},
|
||||||
|
"force": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "跳过优雅终止直接 SIGKILL(默认 false)",
|
||||||
|
},
|
||||||
|
"grace": {
|
||||||
|
"type": "number",
|
||||||
|
"description": "SIGTERM 后等待进程退出的宽限秒数(默认 3,上限 30)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["handle"],
|
||||||
|
}
|
||||||
|
|
||||||
|
def execute(self, handle=None, force: bool = False, grace: float = 3.0, **kwargs) -> str:
|
||||||
|
if error := self.validate_params(["handle"], handle=handle):
|
||||||
|
return error
|
||||||
|
if not _enabled():
|
||||||
|
return "❌ 后台进程功能已禁用"
|
||||||
|
|
||||||
|
handle = _coerce_int(handle)
|
||||||
|
if handle is None:
|
||||||
|
return "❌ 错误:handle 必须是整数句柄"
|
||||||
|
|
||||||
|
# T18:force 字符串 "false" 会被裸 bool 判为 True,跳过宽限直接 SIGKILL
|
||||||
|
force = _coerce_bool(force, False)
|
||||||
|
|
||||||
|
result = self.manager.kill(handle, force=force, grace=grace, owner=self.owner)
|
||||||
|
if not result.get("ok"):
|
||||||
|
return self.format_result(False, f"终止后台进程 #{handle} 失败",
|
||||||
|
{"错误": result.get("error", "未知错误")})
|
||||||
|
data = {
|
||||||
|
"句柄": handle,
|
||||||
|
"状态": result["status"],
|
||||||
|
}
|
||||||
|
if result.get("exit_code") is not None:
|
||||||
|
data["返回码"] = result["exit_code"]
|
||||||
|
if result.get("already_finished"):
|
||||||
|
return self.format_result(True, f"后台进程 #{handle} 此前已结束", data)
|
||||||
|
return self.format_result(True, f"后台进程 #{handle} 已终止", data)
|
||||||
@@ -56,6 +56,11 @@ class UpdateIdentityTool(Tool):
|
|||||||
if not content:
|
if not content:
|
||||||
return "❌ 需要提供身份信息内容"
|
return "❌ 需要提供身份信息内容"
|
||||||
|
|
||||||
|
# N9:持久化注入面约束——内容超限拒绝写入(管理端另有兜底)
|
||||||
|
if len(content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||||
|
return (f"❌ 内容过长({len(content)} > {PromptManager.MAX_CONTENT_LENGTH} 字符),"
|
||||||
|
f"身份文件作为系统提示词加载,请精简后重试")
|
||||||
|
|
||||||
if not confirm:
|
if not confirm:
|
||||||
return "❌ 身份更新已取消,如需更新请设置 confirm=true"
|
return "❌ 身份更新已取消,如需更新请设置 confirm=true"
|
||||||
|
|
||||||
@@ -119,6 +124,11 @@ class UpdateSoulTool(Tool):
|
|||||||
if not content:
|
if not content:
|
||||||
return "❌ 需要提供智能体自我认识定义内容"
|
return "❌ 需要提供智能体自我认识定义内容"
|
||||||
|
|
||||||
|
# N9:持久化注入面约束——内容超限拒绝写入(管理端另有兜底)
|
||||||
|
if len(content) > PromptManager.MAX_CONTENT_LENGTH:
|
||||||
|
return (f"❌ 内容过长({len(content)} > {PromptManager.MAX_CONTENT_LENGTH} 字符),"
|
||||||
|
f"灵魂文件作为系统提示词加载,请精简后重试")
|
||||||
|
|
||||||
if not confirm:
|
if not confirm:
|
||||||
return "❌ 自我认识更新已取消,如需更新请设置 confirm=true"
|
return "❌ 自我认识更新已取消,如需更新请设置 confirm=true"
|
||||||
|
|
||||||
@@ -223,6 +233,11 @@ class SaveMemoryAndEndConversationTool(Tool):
|
|||||||
status_content += f"【下一步计划】\n{next_goals}\n\n"
|
status_content += f"【下一步计划】\n{next_goals}\n\n"
|
||||||
status_content += f"【对话总结】\n{summary}"
|
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)
|
status_updated = self.prompt_manager.update_status(status_content)
|
||||||
|
|
||||||
|
|||||||
+121
-7
@@ -3,18 +3,26 @@
|
|||||||
tools.proxy_tools - 代理配置工具
|
tools.proxy_tools - 代理配置工具
|
||||||
|
|
||||||
抽离自 aibox.py,包含 2 个工具:
|
抽离自 aibox.py,包含 2 个工具:
|
||||||
- SetProxyTool : 设置 SOCKS5 代理
|
- SetProxyTool : 设置 SOCKS5 代理(设置前验证连通性)
|
||||||
- GetProxyTool : 查看当前代理配置
|
- GetProxyTool : 查看当前代理配置
|
||||||
|
|
||||||
依赖:common (ConfigManager) + tools.base.Tool
|
依赖:common (ConfigManager) + tools.base.Tool
|
||||||
不反向依赖 aibox 任何符号。
|
不反向依赖 aibox 任何符号。
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from typing import Dict
|
import os
|
||||||
|
from typing import Dict, Optional, Tuple
|
||||||
|
|
||||||
|
import requests
|
||||||
|
|
||||||
from common import ConfigManager
|
from common import ConfigManager
|
||||||
from .base import Tool
|
from .base import Tool
|
||||||
|
|
||||||
|
# 探测端点:无内容的轻量响应,验证完整 SOCKS5 握手(含认证)与出网连通性
|
||||||
|
_PROBE_URL = "http://www.gstatic.com/generate_204"
|
||||||
|
# (连接超时, 读超时):验证总耗时上限约 10s,避免阻塞工具调用过久
|
||||||
|
_PROBE_TIMEOUT = (4, 6)
|
||||||
|
|
||||||
|
|
||||||
class SetProxyTool(Tool):
|
class SetProxyTool(Tool):
|
||||||
"""设置 SOCKS5 代理的工具"""
|
"""设置 SOCKS5 代理的工具"""
|
||||||
@@ -25,6 +33,8 @@ class SetProxyTool(Tool):
|
|||||||
def get_description(self) -> str:
|
def get_description(self) -> str:
|
||||||
return """设置智能体的 SOCKS5 代理配置。
|
return """设置智能体的 SOCKS5 代理配置。
|
||||||
调用后所有 HTTP/HTTPS 请求(包括网页访问、API 调用)都会通过该代理。
|
调用后所有 HTTP/HTTPS 请求(包括网页访问、API 调用)都会通过该代理。
|
||||||
|
启用代理时会先验证连通性(完整 SOCKS5 握手 + 探测请求),验证失败不保存配置;
|
||||||
|
代理暂未启动等场景可用 skip_check=true 跳过验证强制保存。
|
||||||
可随时修改,即时生效。"""
|
可随时修改,即时生效。"""
|
||||||
|
|
||||||
def get_parameters(self) -> Dict:
|
def get_parameters(self) -> Dict:
|
||||||
@@ -34,24 +44,128 @@ class SetProxyTool(Tool):
|
|||||||
"host": {"type": "string", "description": "代理服务器地址,如 127.0.0.1"},
|
"host": {"type": "string", "description": "代理服务器地址,如 127.0.0.1"},
|
||||||
"port": {"type": "integer", "description": "代理端口,如 1080"},
|
"port": {"type": "integer", "description": "代理端口,如 1080"},
|
||||||
"username": {"type": "string", "description": "用户名(可选)"},
|
"username": {"type": "string", "description": "用户名(可选)"},
|
||||||
"password": {"type": "string", "description": "密码(可选)"},
|
"password": {"type": "string", "description": "密码(可选,明文仅本次会话有效,不落盘)"},
|
||||||
|
"password_env": {"type": "string", "description": "存有密码的环境变量名(可选,跨重启持久的推荐方式)"},
|
||||||
"enabled": {"type": "boolean", "description": "是否启用代理,默认 true"},
|
"enabled": {"type": "boolean", "description": "是否启用代理,默认 true"},
|
||||||
"save": {"type": "boolean", "description": "是否保存到配置文件(持久化),默认 true"}
|
"save": {"type": "boolean", "description": "是否保存到配置文件(持久化),默认 true"},
|
||||||
|
"skip_check": {"type": "boolean", "description": "跳过连通性验证强制保存(默认 false)", "default": False}
|
||||||
},
|
},
|
||||||
"required": ["host", "port"]
|
"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,
|
def execute(self, host: str, port: int, username: str = None, password: str = None,
|
||||||
enabled: bool = True, save: bool = True, **kwargs) -> str:
|
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()
|
config = ConfigManager()
|
||||||
config.set_proxy(host, port, username, password, enabled, save)
|
saved = config.set_proxy(host, port, username, password, password_env, enabled, save)
|
||||||
proxy_info = config.get_proxy_info()
|
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 = {
|
data = {
|
||||||
"代理状态": "启用" if proxy_info["enabled"] else "禁用",
|
"代理状态": "启用" if proxy_info["enabled"] else "禁用",
|
||||||
"服务器": f"{proxy_info['host']}:{proxy_info['port']}",
|
"服务器": f"{proxy_info['host']}:{proxy_info['port']}",
|
||||||
"认证": "有" if proxy_info.get("username") else "无",
|
"认证": "有" if proxy_info.get("username") else "无",
|
||||||
"持久化": "已保存" if save 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)
|
return self.format_result(True, "代理设置成功", data)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
+102
-23
@@ -11,10 +11,26 @@ import threading
|
|||||||
import time
|
import time
|
||||||
from typing import Dict, List, Optional, Tuple
|
from typing import Dict, List, Optional, Tuple
|
||||||
|
|
||||||
from common import Logger, i18n
|
from common import Logger, constants, i18n
|
||||||
|
|
||||||
from .base import Tool
|
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:
|
class ToolRegistry:
|
||||||
"""工具注册器 - 支持技能热重载"""
|
"""工具注册器 - 支持技能热重载"""
|
||||||
@@ -25,6 +41,8 @@ class ToolRegistry:
|
|||||||
self.logger = logger or Logger("ToolRegistry")
|
self.logger = logger or Logger("ToolRegistry")
|
||||||
self.current_output_queue = None # 用于流式输出的队列
|
self.current_output_queue = None # 用于流式输出的队列
|
||||||
self.skill_manager: Optional['SkillManager'] = None # 技能管理器引用
|
self.skill_manager: Optional['SkillManager'] = None # 技能管理器引用
|
||||||
|
# 单工具流式执行总上限(工具自身秒级超时之外的兜底;测试可调小)
|
||||||
|
self.stream_tool_timeout = 3600
|
||||||
|
|
||||||
def register(self, tool: Tool, is_builtin: bool = False):
|
def register(self, tool: Tool, is_builtin: bool = False):
|
||||||
"""注册工具,可以标记是否为内置工具"""
|
"""注册工具,可以标记是否为内置工具"""
|
||||||
@@ -49,29 +67,42 @@ class ToolRegistry:
|
|||||||
def get_tools_schemas(self) -> List[Dict]:
|
def get_tools_schemas(self) -> List[Dict]:
|
||||||
"""获取所有工具的OpenAI函数调用格式"""
|
"""获取所有工具的OpenAI函数调用格式"""
|
||||||
return [tool.to_function_call_format() for tool in self._tools.values()]
|
return [tool.to_function_call_format() for tool in self._tools.values()]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _truncate_result(result) -> str:
|
||||||
|
"""T15:工具返回值统一长度预算(非流式与流式最终结果同口径)"""
|
||||||
|
if not isinstance(result, str):
|
||||||
|
return result
|
||||||
|
limit = constants.TOOL_MAX_RESULT_LENGTH
|
||||||
|
if len(result) > limit:
|
||||||
|
return (result[:limit] +
|
||||||
|
f"\n⚠️ 结果超过 {limit} 字符上限,已截断")
|
||||||
|
return result
|
||||||
|
|
||||||
def set_output_queue(self, queue: queue.Queue):
|
def set_output_queue(self, queue: queue.Queue):
|
||||||
"""设置输出队列"""
|
"""设置输出队列(T4:仅记注册表级引用,不再传播到共享工具实例——
|
||||||
|
实例属性由 execute_streaming 每次调用自行注入局部队列,全局传播
|
||||||
|
会让并发会话互相覆盖对方的输出通路,造成跨会话输出串扰)"""
|
||||||
self.current_output_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:
|
def execute_tool(self, name: str, **kwargs) -> str:
|
||||||
"""执行工具"""
|
"""执行工具"""
|
||||||
tool = self.get_tool(name)
|
tool = self.get_tool(name)
|
||||||
if not tool:
|
if not tool:
|
||||||
return f"❌ 错误:找不到工具 '{name}'"
|
return f"❌ 错误:找不到工具 '{name}'"
|
||||||
|
|
||||||
self.logger.debug(f"工具调用: {name}({kwargs})")
|
self.logger.debug(f"工具调用: {name}({_sanitize_kwargs(kwargs)})")
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
# T4:入口快照会话队列。转发循环可能持续数秒~小时,期间并发
|
||||||
|
# 会话的 set_output_queue 会覆盖单值引用——全程用快照,保证
|
||||||
|
# 本会话输出只进本会话队列,结束清理也不影响在途转发
|
||||||
|
session_queue = self.current_output_queue
|
||||||
# 检查是否有队列可用,并且工具支持流式输出
|
# 检查是否有队列可用,并且工具支持流式输出
|
||||||
if self.current_output_queue and hasattr(tool, 'execute_streaming'):
|
if session_queue and hasattr(tool, 'execute_streaming'):
|
||||||
# 设置工具的流式输出属性
|
# 设置工具的流式输出属性
|
||||||
tool.streaming_output = True
|
tool.streaming_output = True
|
||||||
tool.output_queue = self.current_output_queue
|
|
||||||
|
|
||||||
# 创建队列来收集流式输出
|
# 创建队列来收集流式输出
|
||||||
output_queue = queue.Queue()
|
output_queue = queue.Queue()
|
||||||
@@ -84,34 +115,82 @@ class ToolRegistry:
|
|||||||
)
|
)
|
||||||
thread.daemon = True
|
thread.daemon = True
|
||||||
thread.start()
|
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
|
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():
|
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:
|
try:
|
||||||
msg_type, content = output_queue.get(timeout=0.1)
|
msg_type, content = output_queue.get(timeout=0.1)
|
||||||
if msg_type == "line":
|
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 session_queue:
|
||||||
|
session_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发送)
|
# 发送到当前队列(会通过WebSocket发送)
|
||||||
if self.current_output_queue:
|
if session_queue:
|
||||||
self.current_output_queue.put(("line", content))
|
session_queue.put(("line", piece))
|
||||||
elif msg_type == "complete":
|
elif msg_type == "complete":
|
||||||
final_result = content
|
final_result = content
|
||||||
# 也发送完整结果
|
# 也发送完整结果
|
||||||
if self.current_output_queue:
|
if session_queue:
|
||||||
self.current_output_queue.put(("complete", content))
|
session_queue.put(("complete", content))
|
||||||
except queue.Empty:
|
except queue.Empty:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
thread.join(timeout=5)
|
if not timed_out:
|
||||||
if thread.is_alive():
|
thread.join(timeout=5)
|
||||||
return "❌ 工具执行超时"
|
if thread.is_alive():
|
||||||
|
timed_out = True
|
||||||
result = final_result or ""
|
|
||||||
|
# 统一清理(正常与超时路径都必须执行):断开实例属性,
|
||||||
|
# 僵尸 daemon 线程的 send_line 落到 None 上不再外写
|
||||||
tool.streaming_output = False
|
tool.streaming_output = False
|
||||||
tool.output_queue = None
|
tool.output_queue = None
|
||||||
|
|
||||||
|
if timed_out:
|
||||||
|
result = f"❌ 工具执行超时(>{tool_timeout}s,已放弃等待)"
|
||||||
|
else:
|
||||||
|
result = final_result or ""
|
||||||
else:
|
else:
|
||||||
# 普通执行
|
# 普通执行
|
||||||
result = tool.execute(**kwargs)
|
result = tool.execute(**kwargs)
|
||||||
|
|
||||||
|
# T15:两条分支的最终返回值统一过 TOOL_MAX_RESULT_LENGTH——
|
||||||
|
# 此前仅流式转发有 line_budget(N8),非流式结果无上限直入
|
||||||
|
# 对话上下文(memory/memo get 超长 content,token/内存放大)
|
||||||
|
result = self._truncate_result(result)
|
||||||
|
|
||||||
elapsed = time.time() - start_time
|
elapsed = time.time() - start_time
|
||||||
self.logger.debug(f"工具结果 ({elapsed:.2f}s): {result[:100]}...")
|
self.logger.debug(f"工具结果 ({elapsed:.2f}s): {result[:100]}...")
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ tools.skill_tools - 技能相关工具
|
|||||||
import os
|
import os
|
||||||
import queue
|
import queue
|
||||||
import re
|
import re
|
||||||
|
import shlex
|
||||||
from typing import Dict, Optional
|
from typing import Dict, Optional
|
||||||
|
|
||||||
from common import constants, i18n, Logger
|
from common import constants, i18n, Logger
|
||||||
@@ -253,10 +254,14 @@ class InstallSkillTool(Tool):
|
|||||||
skills_dir = constants.SKILLS_DIR
|
skills_dir = constants.SKILLS_DIR
|
||||||
|
|
||||||
# 如果给了具体路径,在说明里带入,否则用占位符
|
# 如果给了具体路径,在说明里带入,否则用占位符
|
||||||
src = skill_path.strip() if skill_path else "/path/to/your_skill.md"
|
raw_src = skill_path.strip() if skill_path else ""
|
||||||
# 去掉扩展名作为技能名称,再统一补 .md 作为目标文件名
|
# 去掉扩展名作为技能名称,再统一补 .md 作为目标文件名
|
||||||
base = os.path.splitext(os.path.basename(src))[0] if src else "your_skill"
|
base = os.path.splitext(os.path.basename(raw_src))[0] if raw_src else "your_skill"
|
||||||
dst = os.path.join(skills_dir, base + ".md")
|
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"""📦 扩展技能安装指引
|
guide = f"""📦 扩展技能安装指引
|
||||||
{"=" * 55}
|
{"=" * 55}
|
||||||
@@ -282,7 +287,7 @@ class InstallSkillTool(Tool):
|
|||||||
|
|
||||||
调用 create_subagent 派生子代理执行复制命令(目标文件名统一补 .md 扩展名):
|
调用 create_subagent 派生子代理执行复制命令(目标文件名统一补 .md 扩展名):
|
||||||
|
|
||||||
cp {src} {dst}
|
cp {src} {dst_q}
|
||||||
|
|
||||||
⚠️ 注意:目标文件名必须带 .md 扩展名,
|
⚠️ 注意:目标文件名必须带 .md 扩展名,
|
||||||
系统以 .md 文件作为技能加载,且技能名称取文件名(不含 .md)。
|
系统以 .md 文件作为技能加载,且技能名称取文件名(不含 .md)。
|
||||||
|
|||||||
+68
-19
@@ -181,6 +181,9 @@ class WebSocketClient:
|
|||||||
self.connected = False
|
self.connected = False
|
||||||
print(f"\n{Colors.ERROR}❌ 连接已断开{Colors.RESET}")
|
print(f"\n{Colors.ERROR}❌ 连接已断开{Colors.RESET}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
# J14:异常退出同样标记断开(此前 connected 恒 True,用户继续
|
||||||
|
# "成功"发送但永远收不到回复,成僵尸连接)
|
||||||
|
self.connected = False
|
||||||
print(f"\n{Colors.ERROR}❌ 接收消息错误: {e}{Colors.RESET}")
|
print(f"\n{Colors.ERROR}❌ 接收消息错误: {e}{Colors.RESET}")
|
||||||
|
|
||||||
async def handle_message(self, data: dict, tts: TextToSpeech):
|
async def handle_message(self, data: dict, tts: TextToSpeech):
|
||||||
@@ -266,12 +269,13 @@ class WebSocketClient:
|
|||||||
tts: TextToSpeech = None, task_info: dict = None):
|
tts: TextToSpeech = None, task_info: dict = None):
|
||||||
if msg_type == 'chat':
|
if msg_type == 'chat':
|
||||||
stream = self.streaming_messages['chat']
|
stream = self.streaming_messages['chat']
|
||||||
stream.add_chunk(content)
|
# J14(审计 2026-08-19):前缀判断必须在 add_chunk 之前——
|
||||||
|
# add_chunk 会置 is_active=True,原顺序使 "AI助手:" 前缀永不打印
|
||||||
if not stream.is_active:
|
if not stream.is_active:
|
||||||
print(f"\n{Colors.AI}🤖 AI助手: ", end='', flush=True)
|
print(f"\n{Colors.AI}🤖 AI助手: ", end='', flush=True)
|
||||||
stream.is_active = True
|
stream.is_active = True
|
||||||
|
|
||||||
|
stream.add_chunk(content)
|
||||||
print(content, end='', flush=True)
|
print(content, end='', flush=True)
|
||||||
|
|
||||||
elif msg_type == 'task' and task_info:
|
elif msg_type == 'task' and task_info:
|
||||||
@@ -342,14 +346,19 @@ class VoiceInterface:
|
|||||||
|
|
||||||
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
|
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
|
||||||
test_file = f.name
|
test_file = f.name
|
||||||
|
|
||||||
result = subprocess.run(
|
try:
|
||||||
['arecord', '-f', 'cd', '-d', '1', test_file],
|
result = subprocess.run(
|
||||||
capture_output=True, text=True
|
['arecord', '-f', 'cd', '-d', '1', test_file],
|
||||||
)
|
capture_output=True, text=True,
|
||||||
|
timeout=5, # M7(审计 2026-08-20):设备挂死此前永久卡死
|
||||||
os.unlink(test_file)
|
)
|
||||||
|
finally:
|
||||||
|
# M7:超时/异常路径同样清理探测临时文件(此前残留 /tmp);
|
||||||
|
# TimeoutExpired 传播至外层 except,按"麦克风不可用"降级
|
||||||
|
if os.path.exists(test_file):
|
||||||
|
os.unlink(test_file)
|
||||||
|
|
||||||
if result.returncode == 0:
|
if result.returncode == 0:
|
||||||
print(f"{Colors.INFO}🎤 arecord 检测到设备,将使用 arecord 录音{Colors.RESET}")
|
print(f"{Colors.INFO}🎤 arecord 检测到设备,将使用 arecord 录音{Colors.RESET}")
|
||||||
self.has_microphone = True
|
self.has_microphone = True
|
||||||
@@ -428,6 +437,7 @@ class VoiceInterface:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
print(f"{Colors.INFO}🎤 正在听... (按 Ctrl+C 取消){Colors.RESET}")
|
print(f"{Colors.INFO}🎤 正在听... (按 Ctrl+C 取消){Colors.RESET}")
|
||||||
|
temp_file = None # J14:录音临时文件,任何异常路径统一清理
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if self.use_arecord:
|
if self.use_arecord:
|
||||||
@@ -506,6 +516,13 @@ class VoiceInterface:
|
|||||||
return None
|
return None
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"{Colors.ERROR}❌ 语音输入错误: {e}{Colors.RESET}")
|
print(f"{Colors.ERROR}❌ 语音输入错误: {e}{Colors.RESET}")
|
||||||
|
# J14:录音临时文件(含用户语音)任何异常路径都清理——
|
||||||
|
# 此前 wave.open/whisper 解析失败、arecord 超时等滞留 /tmp
|
||||||
|
if temp_file and os.path.exists(temp_file):
|
||||||
|
try:
|
||||||
|
os.unlink(temp_file)
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
return None
|
return None
|
||||||
|
|
||||||
async def handle_command(self, command: str):
|
async def handle_command(self, command: str):
|
||||||
@@ -598,21 +615,53 @@ class VoiceInterface:
|
|||||||
if not self.client or not self.client.connected:
|
if not self.client or not self.client.connected:
|
||||||
print(f"{Colors.ERROR}❌ 未连接到服务器{Colors.RESET}")
|
print(f"{Colors.ERROR}❌ 未连接到服务器{Colors.RESET}")
|
||||||
return
|
return
|
||||||
|
|
||||||
print(f"\n{Colors.USER}👤 用户: {message}{Colors.RESET}")
|
print(f"\n{Colors.USER}👤 用户: {message}{Colors.RESET}")
|
||||||
await self.client.send_message(message)
|
await self.client.send_message(message)
|
||||||
|
|
||||||
|
def _stdin_reader(self, loop: asyncio.AbstractEventLoop):
|
||||||
|
"""M1(审计 2026-08-20):单守护线程读 stdin——替代 asyncio.to_thread
|
||||||
|
每行一个 input 线程的方案:Ctrl+C 取消 await 后旧线程仍阻塞 readline
|
||||||
|
永久泄漏,并"偷吃"后续输入整行(表现为偶发丢行)。
|
||||||
|
方案照搬 client.py:735-755(守护线程 + Queue + call_soon_threadsafe)。"""
|
||||||
|
def _put(item):
|
||||||
|
try:
|
||||||
|
loop.call_soon_threadsafe(self._input_queue.put_nowait, item)
|
||||||
|
except RuntimeError:
|
||||||
|
pass # 事件循环已关闭(程序退出中)
|
||||||
|
|
||||||
|
while self._running:
|
||||||
|
try:
|
||||||
|
line = sys.stdin.readline()
|
||||||
|
except Exception:
|
||||||
|
break
|
||||||
|
if not line: # EOF (Ctrl+D)
|
||||||
|
_put(None)
|
||||||
|
break
|
||||||
|
_put(line)
|
||||||
|
|
||||||
async def run(self):
|
async def run(self):
|
||||||
self.print_banner()
|
self.print_banner()
|
||||||
|
|
||||||
prompt = f"{Colors.USER}>>> {Colors.RESET}"
|
prompt = f"{Colors.USER}>>> {Colors.RESET}"
|
||||||
|
|
||||||
|
# M1:进程生命周期只启动一次 stdin 守护线程
|
||||||
|
self._input_queue = asyncio.Queue()
|
||||||
|
self._running = True
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
if getattr(self, '_stdin_thread', None) is None:
|
||||||
|
self._stdin_thread = threading.Thread(
|
||||||
|
target=self._stdin_reader, args=(loop,),
|
||||||
|
daemon=True, name="stdin-reader")
|
||||||
|
self._stdin_thread.start()
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
try:
|
try:
|
||||||
user_input = await asyncio.to_thread(
|
print(prompt, end='', flush=True)
|
||||||
lambda: input(prompt)
|
line = await self._input_queue.get()
|
||||||
)
|
if line is None: # EOF (Ctrl+D)
|
||||||
user_input = user_input.strip()
|
raise EOFError
|
||||||
|
user_input = line.strip()
|
||||||
|
|
||||||
if not user_input:
|
if not user_input:
|
||||||
if self.has_microphone:
|
if self.has_microphone:
|
||||||
|
|||||||
@@ -21,6 +21,7 @@
|
|||||||
spi4cj = "1.0.0"
|
spi4cj = "1.0.0"
|
||||||
#lua4cj = { git = "https://gitcode.com/qingss0/lua4cj.git", branch="main" }
|
#lua4cj = { git = "https://gitcode.com/qingss0/lua4cj.git", branch="main" }
|
||||||
lua4cj = "1.0.0"
|
lua4cj = "1.0.0"
|
||||||
|
nn4cj = { git = "https://gitcode.com/qingss0/nn4cj.git", branch="main"}
|
||||||
# 这种x86机器人开发板感觉不经常使用,我没用过,不知道里面的i2c调用是不是参数是一样的
|
# 这种x86机器人开发板感觉不经常使用,我没用过,不知道里面的i2c调用是不是参数是一样的
|
||||||
[target.x86_64-unknown-linux-gnu]
|
[target.x86_64-unknown-linux-gnu]
|
||||||
[target.x86_64-unknown-linux-gnu.dependencies]
|
[target.x86_64-unknown-linux-gnu.dependencies]
|
||||||
|
|||||||
+7
-7
@@ -98,13 +98,13 @@ class MatrixObject {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/* ---------------- 对Nn1进行扩展 ---------------- */
|
/* ---------------- 对Nn1进行扩展 ---------------- */
|
||||||
extend Nn1 {
|
//extend Nn1 {
|
||||||
/* -------- 转换为MatrixObject对象 -------- */
|
// /* -------- 转换为MatrixObject对象 -------- */
|
||||||
public func toMto(): MatrixObject {
|
// public func toMto(): MatrixObject {
|
||||||
let mto = MatrixObject()
|
// let mto = MatrixObject()
|
||||||
return mto
|
// return mto
|
||||||
}
|
// }
|
||||||
}
|
//}
|
||||||
|
|
||||||
/* -------- 矩阵式解析器 --------
|
/* -------- 矩阵式解析器 --------
|
||||||
* TODO(20260514-qing): 这个类输入流数据,输出MatrixObject对象
|
* TODO(20260514-qing): 这个类输入流数据,输出MatrixObject对象
|
||||||
|
|||||||
@@ -1,152 +0,0 @@
|
|||||||
/*
|
|
||||||
* i2c - 用户界面
|
|
||||||
*
|
|
||||||
* 这里已经提供了I2C通用接口,可以将上层代码移动到仓颉下面
|
|
||||||
*/
|
|
||||||
|
|
||||||
#include <stdio.h>
|
|
||||||
#include <stdint.h>
|
|
||||||
#include <fcntl.h>
|
|
||||||
#ifdef __linux__
|
|
||||||
#include <unistd.h>
|
|
||||||
#include <sys/ioctl.h>
|
|
||||||
#include <linux/i2c-dev.h> /* 其中宏的取值可能具有平台特异性 */
|
|
||||||
#include <i2c/smbus.h>
|
|
||||||
#endif
|
|
||||||
|
|
||||||
#include "bot4cj.h" /* 有没有可能在仓颉代码中解析C头文件取得宏 */
|
|
||||||
|
|
||||||
|
|
||||||
/*
|
|
||||||
* 初始化i2c设备
|
|
||||||
*
|
|
||||||
* 动态指定设备文件
|
|
||||||
*
|
|
||||||
* 返回打开的文件句柄或者错误码
|
|
||||||
*/
|
|
||||||
#ifdef __linux__
|
|
||||||
int16_t I2C_init(const char* filepath) {
|
|
||||||
int16_t fd;
|
|
||||||
if ((fd = open(filepath, O_RDWR)) < 0) {
|
|
||||||
perror("Failed to open I2C bus");
|
|
||||||
return BOT4CJ_I2C_OPENED_FAILED;
|
|
||||||
}
|
|
||||||
else {
|
|
||||||
return fd;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
#ifdef _WIN64
|
|
||||||
int16_t I2C_init(const char* filepath) {
|
|
||||||
return BOT4CJ_I2C_OPENED_FAILED;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
|
|
||||||
/* 释放I2C设备文件 */
|
|
||||||
#ifdef __linux__
|
|
||||||
void I2C_free(int16_t fd) {
|
|
||||||
close(fd);
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
#ifdef _WIN64
|
|
||||||
void I2C_free(int16_t fd) {
|
|
||||||
/* Windows平台下无需操作 */
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
/* 设置I2C地址 */
|
|
||||||
#ifdef __linux__
|
|
||||||
int I2C_setAddr(int16_t fd, uint8_t addr) {
|
|
||||||
if (ioctl(fd, I2C_SLAVE, addr) < 0) {
|
|
||||||
perror("Failed to set I2C address");
|
|
||||||
return BOT4CJ_I2C_ERR;
|
|
||||||
}
|
|
||||||
else {
|
|
||||||
return BOT4CJ_I2C_OK;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
#ifdef _WIN64
|
|
||||||
int I2C_setAddr(int16_t fd, uint8_t addr) {
|
|
||||||
return BOT4CJ_I2C_ERR;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
|
|
||||||
/* 写寄存器 */
|
|
||||||
#ifdef __linux__
|
|
||||||
int I2C_writeReg(int16_t fd, uint8_t reg, uint8_t val) {
|
|
||||||
/*uint8_t buf[] = {reg, val};
|
|
||||||
if (write(fd, buf, 2) != 2) {
|
|
||||||
perror("I2C write error");
|
|
||||||
return BOT4CJ_I2C_ERR;
|
|
||||||
}
|
|
||||||
else {
|
|
||||||
return BOT4CJ_I2C_OK;
|
|
||||||
}*/
|
|
||||||
|
|
||||||
return i2c_smbus_write_byte_data(fd, reg, val);
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
#ifdef _WIN64
|
|
||||||
int I2C_writeReg(int16_t fd, uint8_t reg, uint8_t val) {
|
|
||||||
return BOT4CJ_I2C_ERR;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
|
|
||||||
/* 读寄存器 失败返回0 */
|
|
||||||
#ifdef __linux__
|
|
||||||
uint8_t I2C_readReg(int16_t fd, uint8_t reg) {
|
|
||||||
|
|
||||||
uint8_t value;
|
|
||||||
|
|
||||||
/*if (write(fd, ®, 1) != 1) {
|
|
||||||
perror("I2C read error");
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
|
|
||||||
if (read(fd, &value, 1)!= 1) {
|
|
||||||
perror("I2C read error");
|
|
||||||
return 0;
|
|
||||||
}*/
|
|
||||||
|
|
||||||
//return value;
|
|
||||||
|
|
||||||
int8_t a = i2c_smbus_read_byte_data(fd, reg);
|
|
||||||
return *(uint8_t *)& a;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
#ifdef _WIN64
|
|
||||||
uint8_t I2C_readReg(int16_t fd, uint8_t reg) {
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
|
|
||||||
/* 读寄存器 - 有符号
|
|
||||||
* 失败返回0 */
|
|
||||||
#ifdef __linux__
|
|
||||||
int8_t I2C_readRegFlag(int16_t fd, uint8_t reg) {
|
|
||||||
|
|
||||||
int8_t value;
|
|
||||||
|
|
||||||
//if (write(fd, ®, 1) != 1) {
|
|
||||||
//perror("I2C read error");
|
|
||||||
//return 0;
|
|
||||||
//}
|
|
||||||
|
|
||||||
//if (read(fd, &value, 1)!= 1) {
|
|
||||||
//perror("I2C read error");
|
|
||||||
//return 0;
|
|
||||||
//}
|
|
||||||
|
|
||||||
//return value;
|
|
||||||
|
|
||||||
return i2c_smbus_read_byte_data(fd, reg);
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
#ifdef _WIN64
|
|
||||||
int8_t I2C_readRegFlag(int16_t fd, uint8_t reg) {
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
#endif
|
|
||||||
@@ -1,516 +0,0 @@
|
|||||||
#include <random>
|
|
||||||
#include <algorithm>
|
|
||||||
#include "Matrix.hh"
|
|
||||||
|
|
||||||
namespace qing {
|
|
||||||
|
|
||||||
template <typename T>
|
|
||||||
Matrx<T>::Matrx(): row(0), col(0) {}
|
|
||||||
|
|
||||||
template <typename T>
|
|
||||||
Matrx<T>::Matrx(const uint64_t row, const uint64_t col)
|
|
||||||
: row(row), col(col) { /* 该构造函数用于从头开始构造 */
|
|
||||||
arr = std::vector<T>(col * row);
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T>::Matrx(const uint64_t row, const uint64_t col, const std::vector<T>& src)
|
|
||||||
: row(row), col(col) { /* 该构造函数用于通过已有的vector进行构造 */
|
|
||||||
arr = src;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
const T& Matrx<T>::get(uint64_t row, uint64_t col) const { /* 访问器 - 访问矩阵元素 */
|
|
||||||
return arr[col * this->row + row];
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
T& Matrx<T>::witch(uint64_t row, uint64_t col) { /* 访问器 - 访问矩阵元素 */
|
|
||||||
return arr[col * this->row + row];
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
void Matrx<T>::set(uint64_t row, uint64_t col, T v) { /* 设置矩阵元素 */
|
|
||||||
arr[col * this->row + row] = v;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
const std::vector<T>& Matrx<T>::get_data() const{ /* 访问器 - 用于获取原始数据 */
|
|
||||||
return arr;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
const uint64_t Matrx<T>::get_col() const { /* 获取列数 */
|
|
||||||
return col;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
const uint64_t Matrx<T>::get_row() const { /* 获取行数 */
|
|
||||||
return row;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
bool Matrx<T>::has_same_shape(const Matrx<T>& other) { /* 两个矩阵是否具有相同形状 */
|
|
||||||
return other.get_col() == this->get_col() && other.get_row() == this->get_row();
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
bool Matrx<T>::has_dot_shape(const Matrx<T>& other) { /* 具有矩阵乘法的形状 */
|
|
||||||
return this->get_col() == other.get_row();
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
void Matrx<T>::assert_same_shape(const Matrx<T>& other) { /* 断言两个矩阵具有相同形状 */
|
|
||||||
if (!has_same_shape(other))
|
|
||||||
throw std::invalid_argument("The shapes of matrix are not same.");
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
void Matrx<T>::assert_dot_shape(const Matrx<T>& other) { /* 断言两个矩阵具有点积前置形状 */
|
|
||||||
if (!has_dot_shape(other)) {
|
|
||||||
throw std::invalid_argument("The shapes of matrix can not be dot.");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator+(const Matrx<T>& other) { /* 加法操作符重定义 - 矩阵加法 */
|
|
||||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
|
||||||
std::vector<T> v0 = this->get_data();
|
|
||||||
std::vector<T> vec = other.get_data();
|
|
||||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
|
||||||
*it += *it_or;
|
|
||||||
return Matrx<T>{this->get_row(), this->get_col(), v0};
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator+(T& item) { /* 矩阵加法广播 */
|
|
||||||
auto len = this->get_data().size();
|
|
||||||
auto v0 = std::vector<T> (len, item);
|
|
||||||
auto matrix = (*this) + Matrx(get_row(), get_col(), v0);
|
|
||||||
|
|
||||||
//auto matrix = *this;
|
|
||||||
//for (auto it = arr.begin(); it != arr.end(); ++it)
|
|
||||||
// *it += item;
|
|
||||||
return matrix;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator-(const Matrx<T>& other) { /* 操作符重载 - 矩阵减法 */
|
|
||||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
|
||||||
std::vector<T> v0 = this->get_data();
|
|
||||||
std::vector <T> vec = other.get_data();
|
|
||||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
|
||||||
*it -= *it_or;
|
|
||||||
return Matrx{this->get_row(), this->get_col(), v0};
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator-(T& item) { /* 操作符重载 - 矩阵减法广播 */
|
|
||||||
auto len = this->get_data().size();
|
|
||||||
auto v0 = std::vector<T> (len, item);
|
|
||||||
auto matrix = (*this) - Matrx(get_row(), get_col(), v0);
|
|
||||||
|
|
||||||
//Matrx<T> matrix = *this;
|
|
||||||
//for (auto it = arr.begin(); it != arr.end(); ++it)
|
|
||||||
// *it -= item;
|
|
||||||
return matrix;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::dot(const Matrx<T>& other) { /* 操作符重载 - 矩阵乘法 */
|
|
||||||
//assert_dot_shape(other); /* 它们的形状是否能够点积 */
|
|
||||||
auto matrix = Matrx<T>(this->get_row(), other.get_col());
|
|
||||||
for (uint64_t i=0; i<this->get_row(); ++i) /* 点积运算 */
|
|
||||||
for (uint64_t j=0; j<other.get_col(); ++j) {
|
|
||||||
T sum{0};
|
|
||||||
for (uint64_t k=0; k<this->get_col(); ++k)
|
|
||||||
sum += this->get(i, k) * other.get(k, j);
|
|
||||||
matrix.witch(i, j) = sum;
|
|
||||||
}
|
|
||||||
return matrix;
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator*(T& item) { /* 操作符重载 - 矩阵乘法广播 */
|
|
||||||
std::vector<T> vec = this->get_data();
|
|
||||||
for (auto it = vec.begin(); it != vec.end(); ++it)
|
|
||||||
*it *= item;
|
|
||||||
return Matrx<T>{this->get_row(), this->get_col(), vec};
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator*(const Matrx<T>& other) { /* 加法操作符重定义 - 矩阵加法 */
|
|
||||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
|
||||||
std::vector<T> v0 = this->get_data();
|
|
||||||
std::vector<T> vec = other.get_data();
|
|
||||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
|
||||||
*it *= *it_or;
|
|
||||||
return Matrx<T>{this->get_row(), this->get_col(), v0};
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator/(T& item) { /* 操作符重载 - 矩阵除法广播 */
|
|
||||||
std::vector<T> vec = this->get_data();
|
|
||||||
for (auto it = vec.begin(); it != vec.end(); ++it)
|
|
||||||
*it /= item;
|
|
||||||
return Matrx<T>(this->get_row(), this->get_col(), vec);
|
|
||||||
}
|
|
||||||
|
|
||||||
template<typename T>
|
|
||||||
Matrx<T> Matrx<T>::operator/(const Matrx<T>& other) { /* 加法操作符重定义 - 矩阵加法 */
|
|
||||||
//assert_same_shape(other); /* 检查两个矩阵形状是否一致 */
|
|
||||||
std::vector<T> v0 = this->get_data();
|
|
||||||
std::vector<T> vec = other.get_data();
|
|
||||||
for (auto it = v0.begin(), it_or = vec.begin(); it != v0.end(); ++it, ++it_or )
|
|
||||||
*it /= *it_or;
|
|
||||||
return Matrx<T>{this->get_row(), this->get_col(), v0};
|
|
||||||
}
|
|
||||||
|
|
||||||
template class Matrx<float>; /* 显式实例化 */
|
|
||||||
template class Matrx<double>;
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
/* FIXME: 形状不合法时返回空指针 */
|
|
||||||
using namespace qing;
|
|
||||||
extern "C" {
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* createMatrix(uint64_t row, uint64_t col) {
|
|
||||||
return (void *) new Matrx<float>(row, col);
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void freeMatrix(void* M) {
|
|
||||||
delete (Matrx<float>*) M;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
float* witchMatrix(void *m, uint64_t row, uint64_t col) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
return &(M->witch(row, col));
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
float getMatrix(void *m, uint64_t row, uint64_t col) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
return M->get(row, col);
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void setMatrix(void *m, uint64_t row, uint64_t col, float v) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
M->set(row, col, v);
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
uint64_t getrowMatrix(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
return M->get_row();
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
uint64_t getcolMatrix(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
return M->get_col();
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
bool sameShapeMatrix(void *m, void *s) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
Matrx<float> *S = (Matrx<float>*) s;
|
|
||||||
return M->has_same_shape(*S);
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
bool dotShapeMatrix(void *m, void *s) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
Matrx<float> *S = (Matrx<float>*) s;
|
|
||||||
return M->has_dot_shape(*S);
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* plusMatrix(void *m, void *s) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
Matrx<float> *S = (Matrx<float>*) s;
|
|
||||||
auto a = new Matrx<float>((*M) + (*S));
|
|
||||||
return (void *) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* plusBroadcastMatrix(void *m, float t) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
auto a = new Matrx<float> ((*M) + t);
|
|
||||||
return (void*) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* subMatrix(void *m, void *s) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
Matrx<float> *S = (Matrx<float>*) s;
|
|
||||||
auto a = new Matrx<float>((*M) - (*S));
|
|
||||||
return (void *) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void * subBroadcastMatrix(void *m, float t) {
|
|
||||||
Matrx<float> *M = (Matrx<float> *) m;
|
|
||||||
auto a = new Matrx<float> ((*M) - t);
|
|
||||||
return (void *) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* multiMatrix(void *m, void *s) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
Matrx<float> *S = (Matrx<float>*) s;
|
|
||||||
auto a = new Matrx<float>((*M) * (*S));
|
|
||||||
return (void *) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* multiBroadcastMatrix(void *m, float t) {
|
|
||||||
Matrx<float> *M = (Matrx<float> *) m;
|
|
||||||
auto a = new Matrx<float>((*M) * t);
|
|
||||||
return (void*) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* divMatrix(void *m, void *s) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
Matrx<float> *S = (Matrx<float>*) s;
|
|
||||||
auto a = new Matrx<float>((*M) / (*S));
|
|
||||||
return (void *) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* divBroadcastMatrix(void *m, float t) {
|
|
||||||
Matrx<float> *M = (Matrx<float> *) m;
|
|
||||||
auto a = new Matrx<float>((*M) / t);
|
|
||||||
return (void*) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* dotMatrix(void *m, void *s) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
Matrx<float> *S = (Matrx<float>*) s;
|
|
||||||
auto a = new Matrx<float>(M->dot(*S));
|
|
||||||
return (void *) a;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* transposeMatrix(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float> *)m;
|
|
||||||
auto N= new Matrx<float>(M->get_col(), M->get_row());
|
|
||||||
for (uint64_t row = 0; row< N->get_row(); row++) {
|
|
||||||
for (uint64_t col = 0; col < N->get_col(); col++) {
|
|
||||||
N->witch(row, col) = M->get(col, row);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
/* 正态分布Xavier初始化 */
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* xavierNormalInit(uint64_t fan_in, uint64_t fan_out)
|
|
||||||
{
|
|
||||||
/* 构建空数组 */
|
|
||||||
std::vector<float> weights(fan_in * fan_out);
|
|
||||||
|
|
||||||
/* 高质量的随机数生成器 */
|
|
||||||
static std::mt19937 generator(std::random_device{}());
|
|
||||||
|
|
||||||
/* 计算标准差 */
|
|
||||||
float sigma = std::sqrt(2.0f / (fan_in + fan_out));
|
|
||||||
|
|
||||||
/* 创建正态分布对象,均值为0,标准差为sigma */
|
|
||||||
std::normal_distribution<float> distribution(0.0f, sigma);
|
|
||||||
|
|
||||||
/* 遍历权重矩阵,进行赋值 */
|
|
||||||
uint64_t k=0;
|
|
||||||
for (uint64_t i=0; i<fan_in; ++i)
|
|
||||||
for (uint64_t j=0; j<fan_out; ++j)
|
|
||||||
weights[k++] = distribution(generator);
|
|
||||||
|
|
||||||
/* 构建矩阵 */
|
|
||||||
auto matrix = new Matrx<float>(fan_in, fan_out, weights);
|
|
||||||
return (void *)matrix;
|
|
||||||
}
|
|
||||||
|
|
||||||
/* 均匀分布Xavier初始化函数 */
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* xavierUniformInit(uint64_t fan_in, uint64_t fan_out)
|
|
||||||
{
|
|
||||||
/* 构建空的数组 */
|
|
||||||
auto weights = std::vector<float>(fan_in * fan_out);
|
|
||||||
|
|
||||||
/* 使用高质量的随机数生成器 */
|
|
||||||
static std::mt19937 generator(std::random_device{}());
|
|
||||||
|
|
||||||
/* 计算区间范围:sqrt(2 / (fan_in + fan_out)) */
|
|
||||||
float scale = std::sqrt(2.0f / (fan_in + fan_out));
|
|
||||||
|
|
||||||
/* 创建均匀分布对象,范围为[-scale, scale] */
|
|
||||||
std::uniform_real_distribution<float> distribution(-scale, scale);
|
|
||||||
|
|
||||||
/* 遍历权重矩阵,进行赋值 */
|
|
||||||
uint64_t k=0;
|
|
||||||
for (uint64_t i=0; i<fan_in; ++i)
|
|
||||||
for (uint64_t j=0; j<fan_out; ++j)
|
|
||||||
weights[k++] = distribution(generator);
|
|
||||||
|
|
||||||
/* 构建矩阵对象 */
|
|
||||||
auto matrix = new Matrx<float> (fan_in, fan_out, weights);
|
|
||||||
return (void *)matrix;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* Sigmoid_do(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& x: v) {
|
|
||||||
x = 1.0 / (1.0 + std::exp(-x));
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* Sigmoid_dot(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& o: v) {
|
|
||||||
o = o * (1.0 - o);
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* ReLU_do(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& x: v) {
|
|
||||||
x = x > 0 ? x : 0;
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* ReLU_dot(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& x: v) {
|
|
||||||
x = x > 0 ? 1.0f : 0.0f;
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* Leaky_ReLU_do(void *m, float alpha) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& x: v) {
|
|
||||||
x = x > 0 ? x : alpha * x;
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* Leaky_ReLU_dot(void *m, float alpha) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& x: v) {
|
|
||||||
x = x > 0 ? 1.0 : alpha;
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* Tanh_do(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& x: v) {
|
|
||||||
x = std::tanh(x);
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
_declspec(dllexport)
|
|
||||||
#endif
|
|
||||||
void* Tanh_dot(void *m) {
|
|
||||||
Matrx<float> *M = (Matrx<float>*) m;
|
|
||||||
std::vector<float> v = M->get_data();
|
|
||||||
for (float& x: v) {
|
|
||||||
if (x > 20.0) x = 0.0; // tanh(x) ≈ 1, 导数≈0
|
|
||||||
else if (x < -20.0) x = 0.0; // tanh(x) ≈ -1, 导数≈0
|
|
||||||
else {
|
|
||||||
float t = tanh(x);
|
|
||||||
x = 1.0 - t * t;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
auto N = new Matrx<float> (M->get_row(), M->get_col(), v);
|
|
||||||
return (void *)N;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,60 +0,0 @@
|
|||||||
/*
|
|
||||||
|
|
||||||
___col___
|
|
||||||
/ \
|
|
||||||
+---------+ \
|
|
||||||
| + + + + | | row - 第一维度
|
|
||||||
| + + + + | |
|
|
||||||
+---------+ /
|
|
||||||
|
|
||||||
*/
|
|
||||||
|
|
||||||
#pragma once
|
|
||||||
#include <cmath>
|
|
||||||
#include <vector>
|
|
||||||
#include <iostream>
|
|
||||||
#include <stdexcept>
|
|
||||||
#include <cstdint>
|
|
||||||
namespace qing {
|
|
||||||
template <class T>
|
|
||||||
class Matrx {
|
|
||||||
public:
|
|
||||||
Matrx();
|
|
||||||
Matrx(const uint64_t row, const uint64_t col);
|
|
||||||
Matrx(const uint64_t row, const uint64_t col, const std::vector<T>& src);
|
|
||||||
const T& get(uint64_t row, uint64_t col) const;
|
|
||||||
//-------------旧接口----------------
|
|
||||||
T& witch(uint64_t row, uint64_t col);
|
|
||||||
//-------------新接口----------------
|
|
||||||
void set(uint64_t row, uint64_t col, T v);
|
|
||||||
//-----------------------------------
|
|
||||||
const std::vector<T>& get_data() const;
|
|
||||||
const uint64_t get_col() const;
|
|
||||||
const uint64_t get_row() const;
|
|
||||||
bool has_same_shape(const Matrx<T>& other);
|
|
||||||
bool has_dot_shape(const Matrx<T>& other);
|
|
||||||
void assert_same_shape(const Matrx<T>& other);
|
|
||||||
void assert_dot_shape(const Matrx<T>& other);
|
|
||||||
Matrx<T> operator+(const Matrx<T>& other);
|
|
||||||
Matrx<T> operator+(T& item);
|
|
||||||
Matrx<T> operator-(const Matrx<T>& other);
|
|
||||||
Matrx<T> operator-(T& item);
|
|
||||||
Matrx<T> operator*(const Matrx<T>& other);
|
|
||||||
Matrx<T> operator*(T& item);
|
|
||||||
Matrx<T> operator/(T& item);
|
|
||||||
Matrx<T> operator/(const Matrx<T>& other);
|
|
||||||
Matrx<T> dot(const Matrx<T>& other);
|
|
||||||
|
|
||||||
void save(std::ostream& out) {
|
|
||||||
out << row << " " << col;
|
|
||||||
for (const auto& item : arr) {
|
|
||||||
out << " ";
|
|
||||||
out << item;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
private:
|
|
||||||
uint64_t row, col;
|
|
||||||
std::vector<T> arr;
|
|
||||||
};
|
|
||||||
}
|
|
||||||
@@ -1,21 +0,0 @@
|
|||||||
/* I2C操作库 - PCA9685 - MPU6050 */
|
|
||||||
#ifndef QING_I2C
|
|
||||||
#define QING_I2C
|
|
||||||
/* I2C */
|
|
||||||
#define BOT4CJ_I2C_OK 0
|
|
||||||
#define BOT4CJ_I2C_ERR 1
|
|
||||||
#define BOT4CJ_I2C_OPENED_FAILED -1
|
|
||||||
|
|
||||||
#ifdef _WIN64
|
|
||||||
#define DLLEXPORT __declspec(dllexport)
|
|
||||||
#else
|
|
||||||
#define DLLEXPORT
|
|
||||||
#endif
|
|
||||||
|
|
||||||
DLLEXPORT int16_t I2C_init(const char* filepath);
|
|
||||||
DLLEXPORT void I2C_free(int16_t fd);
|
|
||||||
DLLEXPORT int I2C_setAddr(int16_t fd, uint8_t addr);
|
|
||||||
DLLEXPORT int I2C_writeReg(int16_t fd, uint8_t reg, uint8_t val);
|
|
||||||
DLLEXPORT uint8_t I2C_readReg(int16_t fd, uint8_t reg);
|
|
||||||
DLLEXPORT int8_t I2C_readRegFlag(int16_t fd, uint8_t reg);
|
|
||||||
#endif
|
|
||||||
+37
-62
@@ -1,15 +1,20 @@
|
|||||||
package bot4cj
|
package bot4cj
|
||||||
|
|
||||||
|
/*
|
||||||
|
* 机器人脚本解析器 - 仓颉语言
|
||||||
|
*
|
||||||
|
*/
|
||||||
|
|
||||||
|
import std.argopt.*
|
||||||
|
import std.collection.*
|
||||||
|
import std.env.*
|
||||||
|
import std.convert.*
|
||||||
import std.io.*
|
import std.io.*
|
||||||
import std.fs.*
|
import std.fs.*
|
||||||
import std.env.*
|
|
||||||
import std.time.*
|
|
||||||
import std.argopt.*
|
|
||||||
import std.convert.*
|
|
||||||
import std.collection.*
|
|
||||||
import stdx.encoding.json.*
|
import stdx.encoding.json.*
|
||||||
import bot4cj.hd.*
|
import bot4cj.hd.*
|
||||||
import bot4cj.math.*
|
|
||||||
import ctype4cj.*
|
import ctype4cj.*
|
||||||
|
import nn4cj.*
|
||||||
|
|
||||||
/* ---- 神经网络输出路径 ---- */
|
/* ---- 神经网络输出路径 ---- */
|
||||||
var NNS_OUT: Option<String> = None
|
var NNS_OUT: Option<String> = None
|
||||||
@@ -17,10 +22,10 @@ main(args: Array<String>): Int64 {
|
|||||||
|
|
||||||
/* ---- 解析命令行参数 ---- */
|
/* ---- 解析命令行参数 ---- */
|
||||||
let specs = [
|
let specs = [
|
||||||
ArgumentSpec.Long("help", ArgumentMode.NoValue), /* ---- 拉起帮助菜单 ---- */
|
ArgumentSpec.Long("help", ArgumentMode.NoValue), /* ---- 拉起帮助 ---- */
|
||||||
ArgumentSpec.Short(r'i', ArgumentMode.RequiredValue), /* ---- 神经网络模型 - 输入 ---- */
|
ArgumentSpec.Short(r'i', ArgumentMode.RequiredValue), /* ---- 输入模型 ---- */
|
||||||
ArgumentSpec.Short(r'o', ArgumentMode.RequiredValue), /* ---- 神经网络模型 - 输出 ---- */
|
ArgumentSpec.Short(r'o', ArgumentMode.RequiredValue), /* ---- 输出模型 ---- */
|
||||||
ArgumentSpec.Long("test", ArgumentMode.RequiredValue) /* ---- 运行一个测试分支 - 需要输入测试类型 ---- */
|
ArgumentSpec.Long("test", ArgumentMode.RequiredValue) /* ---- 输入一个测试 ---- */
|
||||||
]
|
]
|
||||||
|
|
||||||
/* ---- 执行解析 ---- */
|
/* ---- 执行解析 ---- */
|
||||||
@@ -30,10 +35,10 @@ main(args: Array<String>): Int64 {
|
|||||||
//---------------------------
|
//---------------------------
|
||||||
/* 读取设置项 */
|
/* 读取设置项 */
|
||||||
|
|
||||||
/*---------------------------*/
|
/*
|
||||||
/*---------------------------*/
|
* 初始化神经网络 (init nn)
|
||||||
/* 初始化神经网络 (init nn)
|
* TODO: 可以去掉内部初始化参数,使用矩阵式模板进行初始化
|
||||||
* TODO: 可以去掉内部初始化参数,使用矩阵式模板进行初始化 */
|
*/
|
||||||
|
|
||||||
|
|
||||||
//var nns: ArrayList<Nn> =
|
//var nns: ArrayList<Nn> =
|
||||||
@@ -53,28 +58,27 @@ main(args: Array<String>): Int64 {
|
|||||||
NNS_OUT = Option<String>.Some(options['o'])
|
NNS_OUT = Option<String>.Some(options['o'])
|
||||||
}
|
}
|
||||||
|
|
||||||
/*---------------------------*/
|
|
||||||
/*---------------------------*/
|
/*---------------------------*/
|
||||||
/* 分支选择 */
|
/* 分支选择 */
|
||||||
|
|
||||||
|
|
||||||
/* -------- 拉起帮助手册
|
/* -------- 拉起帮助手册 -------- */
|
||||||
* -------- */
|
|
||||||
if (options.contains("help")) {
|
if (options.contains("help")) {
|
||||||
print_help()
|
print_help()
|
||||||
return 0
|
return 0
|
||||||
}
|
}
|
||||||
|
|
||||||
/* -------- 启动一个测试分支
|
/* -------- 启动一个测试分支 -------- */
|
||||||
* -------- */
|
|
||||||
else if (options.contains("test")) {
|
else if (options.contains("test")) {
|
||||||
run_test(options["test"])
|
run_test(options["test"])
|
||||||
return 0
|
return 0
|
||||||
}
|
}
|
||||||
|
|
||||||
/* -------- 缺省行为:从标准输入获取动作脚本
|
/*
|
||||||
* --------
|
* 缺省行为:从标准输入获取动作脚本
|
||||||
* -------- TODO(20260519-qing): 进行脚本引入(相对路径也可,从**环境变量**的路径中搜寻也可) */
|
*
|
||||||
|
* TODO(20260519-qing): 进行脚本引入(相对路径也可,从**环境变量**的路径中搜寻也可)
|
||||||
|
*/
|
||||||
else {
|
else {
|
||||||
let consoleReader = getStdIn()
|
let consoleReader = getStdIn()
|
||||||
let robot = RobotStruct(consoleReader)
|
let robot = RobotStruct(consoleReader)
|
||||||
@@ -82,12 +86,11 @@ main(args: Array<String>): Int64 {
|
|||||||
robot.active()
|
robot.active()
|
||||||
robot.run()
|
robot.run()
|
||||||
}
|
}
|
||||||
finally {
|
|
||||||
robot.deactive()
|
|
||||||
}
|
|
||||||
|
|
||||||
/* ---- NOTE(20260523-qing): 储存微调网络模型
|
/* 无论有没有成功激活,都进行清理 */
|
||||||
* ---- */
|
finally { robot.deactive() }
|
||||||
|
|
||||||
|
/* NOTE(20260523-qing): 储存微调网络模型 */
|
||||||
// let consoleReader = getStdIn()
|
// let consoleReader = getStdIn()
|
||||||
// run_action_drl(consoleReader, nns)
|
// run_action_drl(consoleReader, nns)
|
||||||
// if (options.contains("o")) {
|
// if (options.contains("o")) {
|
||||||
@@ -98,8 +101,7 @@ main(args: Array<String>): Int64 {
|
|||||||
return 0
|
return 0
|
||||||
}
|
}
|
||||||
|
|
||||||
/* -------- 打印帮助手册
|
/* -------- 打印帮助手册 -------- */
|
||||||
* -------- */
|
|
||||||
func print_help()
|
func print_help()
|
||||||
{
|
{
|
||||||
println("这是一个机器人动作脚本解析器。")
|
println("这是一个机器人动作脚本解析器。")
|
||||||
@@ -112,43 +114,16 @@ func print_help()
|
|||||||
println("-o 设置输出的模型的路径")
|
println("-o 设置输出的模型的路径")
|
||||||
}
|
}
|
||||||
|
|
||||||
/* -------- 集成多种常规测试、临时测试
|
/* -------- 集成多种常规测试、临时测试 -------- */
|
||||||
* -------- */
|
|
||||||
func run_test(key: String) {
|
func run_test(key: String) {
|
||||||
|
|
||||||
/* ----------- 神经网络与门测试
|
/*
|
||||||
* -----------
|
* 神经网络门测试
|
||||||
* ----------- INFO: 我们使用随机生成的动态数据集,这样会导致一些后果 */
|
*/
|
||||||
if (key == "nn") {
|
if (key == "nn1") {
|
||||||
let nn_t = ArrayList<Nn1>()
|
testNn()
|
||||||
nn_t.add(Nn1(2, 4, 0.05f32, Af.Leaky_ReLU(FP32(0.01))))
|
|
||||||
nn_t.add(Nn1(4, 1, 0.05f32, Af.Sigmoid))
|
|
||||||
let batch = 1
|
|
||||||
let num =2000
|
|
||||||
let prt_num = 100
|
|
||||||
for (i in 0..(num*batch)) {
|
|
||||||
let m = Matrix1(1, 2)
|
|
||||||
let t = Matrix1(1, 1)
|
|
||||||
prepare_andDoor(m, t)
|
|
||||||
let l = train1(nn_t, m, t, 1.0f32) // 这个1.0是学习率的折扣,而非学习率本身
|
|
||||||
if (i % prt_num == prt_num-1) {
|
|
||||||
print("${i}\t:")
|
|
||||||
print("${l}")
|
|
||||||
println()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
else if (key == "cnn") {
|
|
||||||
testCNN()
|
|
||||||
}
|
|
||||||
|
|
||||||
else if (key == "cnn1") {
|
|
||||||
testGeometryCNN()
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
/* -------- 测试解析矩阵,并且输出
|
/* -------- 测试解析矩阵,并且输出
|
||||||
* -------- */
|
* -------- */
|
||||||
else if (key == "matrix_o") {
|
else if (key == "matrix_o") {
|
||||||
|
|||||||
@@ -1,204 +0,0 @@
|
|||||||
// Adam.cj - Adam 优化器实现(添加全连接层支持)
|
|
||||||
package bot4cj.math
|
|
||||||
|
|
||||||
import std.math.*
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
public class Adam {
|
|
||||||
let learningRate: FP32
|
|
||||||
let beta1: FP32
|
|
||||||
let beta2: FP32
|
|
||||||
let epsilon: FP32
|
|
||||||
|
|
||||||
var m: HashMap<String, Matrix1>
|
|
||||||
var v: HashMap<String, Matrix1>
|
|
||||||
var t: Int64
|
|
||||||
|
|
||||||
// 完整参数构造函数
|
|
||||||
public init(
|
|
||||||
lr: FP32,
|
|
||||||
beta1: FP32,
|
|
||||||
beta2: FP32,
|
|
||||||
epsilon: FP32
|
|
||||||
) {
|
|
||||||
learningRate = lr
|
|
||||||
this.beta1 = beta1
|
|
||||||
this.beta2 = beta2
|
|
||||||
this.epsilon = epsilon
|
|
||||||
m = HashMap<String, Matrix1>()
|
|
||||||
v = HashMap<String, Matrix1>()
|
|
||||||
t = 0
|
|
||||||
}
|
|
||||||
|
|
||||||
// 默认参数构造函数
|
|
||||||
public init() {
|
|
||||||
learningRate = FP32(0.001)
|
|
||||||
beta1 = FP32(0.9)
|
|
||||||
beta2 = FP32(0.999)
|
|
||||||
epsilon = FP32(0.00000001)
|
|
||||||
m = HashMap<String, Matrix1>()
|
|
||||||
v = HashMap<String, Matrix1>()
|
|
||||||
t = 0
|
|
||||||
}
|
|
||||||
|
|
||||||
// 仅自定义学习率的构造函数
|
|
||||||
public init(lr: FP32) {
|
|
||||||
learningRate = lr
|
|
||||||
beta1 = FP32(0.9)
|
|
||||||
beta2 = FP32(0.999)
|
|
||||||
epsilon = FP32(0.00000001)
|
|
||||||
m = HashMap<String, Matrix1>()
|
|
||||||
v = HashMap<String, Matrix1>()
|
|
||||||
t = 0
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算幂运算的辅助函数
|
|
||||||
*/
|
|
||||||
private func power(base: FP32, exp: Int64): FP32 {
|
|
||||||
var result = FP32(1.0)
|
|
||||||
for (i in 0..exp) {
|
|
||||||
result = result * base
|
|
||||||
}
|
|
||||||
return result
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算平方根的辅助函数 - 使用牛顿迭代法
|
|
||||||
*/
|
|
||||||
private func sqrt(val: FP32): FP32 {
|
|
||||||
let x = val.getValue()
|
|
||||||
if (x <= 0.0) {
|
|
||||||
return FP32(0.0)
|
|
||||||
}
|
|
||||||
var guess = x
|
|
||||||
for (i in 0..10) {
|
|
||||||
guess = (guess + x / guess) / 2.0
|
|
||||||
}
|
|
||||||
return FP32(guess)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 更新矩阵参数(通用方法)
|
|
||||||
* @param param 要更新的参数矩阵
|
|
||||||
* @param grad 梯度矩阵
|
|
||||||
* @param key 参数在HashMap中的键名
|
|
||||||
*/
|
|
||||||
private func updateMatrix(
|
|
||||||
param: Matrix1,
|
|
||||||
grad: Matrix1,
|
|
||||||
key: String
|
|
||||||
): Unit {
|
|
||||||
// 初始化动量
|
|
||||||
let hasState = m.contains(key)
|
|
||||||
if (!hasState) {
|
|
||||||
let rows = param.get_row()
|
|
||||||
let cols = param.get_col()
|
|
||||||
m.add(key, Matrix1(rows, cols))
|
|
||||||
v.add(key, Matrix1(rows, cols))
|
|
||||||
}
|
|
||||||
|
|
||||||
let mMat = m.get(key).getOrThrow()
|
|
||||||
let vMat = v.get(key).getOrThrow()
|
|
||||||
|
|
||||||
let rows = param.get_row()
|
|
||||||
let cols = param.get_col()
|
|
||||||
|
|
||||||
for (r in 0..rows) {
|
|
||||||
for (c in 0..cols) {
|
|
||||||
let g = grad.get(r, c).getOrThrow()
|
|
||||||
|
|
||||||
// 更新一阶矩
|
|
||||||
let oldM = mMat.get(r, c).getOrThrow()
|
|
||||||
let newM = beta1 * oldM + (FP32(1.0) - beta1) * g
|
|
||||||
mMat.set(r, c, newM)
|
|
||||||
|
|
||||||
// 更新二阶矩
|
|
||||||
let oldV = vMat.get(r, c).getOrThrow()
|
|
||||||
let newV = beta2 * oldV + (FP32(1.0) - beta2) * g * g
|
|
||||||
vMat.set(r, c, newV)
|
|
||||||
|
|
||||||
// 偏差修正
|
|
||||||
let one = FP32(1.0)
|
|
||||||
let beta1Pow = power(beta1, t)
|
|
||||||
let beta2Pow = power(beta2, t)
|
|
||||||
|
|
||||||
let mHat = newM / (one - beta1Pow)
|
|
||||||
let vHat = newV / (one - beta2Pow)
|
|
||||||
|
|
||||||
// 更新参数
|
|
||||||
let oldParam = param.get(r, c).getOrThrow()
|
|
||||||
let update = learningRate * mHat / (sqrt(vHat) + epsilon)
|
|
||||||
param.set(r, c, oldParam - update)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 更新卷积层参数
|
|
||||||
*/
|
|
||||||
public func updateConvLayer(layer: ConvLayer): Unit {
|
|
||||||
t += 1
|
|
||||||
|
|
||||||
// 更新卷积核
|
|
||||||
var kGrad = match (layer.kernelsGrad) {
|
|
||||||
case Some(val) => val
|
|
||||||
case _ => throw Exception("No kernel gradients")
|
|
||||||
}
|
|
||||||
|
|
||||||
for (oc in 0..layer.nOutputChannels) {
|
|
||||||
let key = "conv_kernel_" + oc.toString()
|
|
||||||
updateMatrix(layer.kernels[oc], kGrad[oc], key)
|
|
||||||
}
|
|
||||||
|
|
||||||
// 更新偏置
|
|
||||||
var bGrad = match (layer.biasGrad) {
|
|
||||||
case Some(val) => val
|
|
||||||
case _ => throw Exception("No bias gradient")
|
|
||||||
}
|
|
||||||
|
|
||||||
let biasKey = "conv_bias"
|
|
||||||
updateMatrix(layer.biases, bGrad, biasKey)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 更新全连接层参数(Nn1)
|
|
||||||
*/
|
|
||||||
public func updateFCLayer(layer: Nn1): Unit {
|
|
||||||
t += 1
|
|
||||||
|
|
||||||
// 更新权重
|
|
||||||
var wGrad = match (layer.weights_grad) {
|
|
||||||
case Some(val) => val
|
|
||||||
case _ => throw Exception("No weights gradient")
|
|
||||||
}
|
|
||||||
|
|
||||||
let weightKey = "fc_weights"
|
|
||||||
updateMatrix(layer.weights, wGrad, weightKey)
|
|
||||||
|
|
||||||
// 更新偏置
|
|
||||||
var bGrad = match (layer.bias_grad) {
|
|
||||||
case Some(val) => val
|
|
||||||
case _ => throw Exception("No bias gradient")
|
|
||||||
}
|
|
||||||
|
|
||||||
let biasKey = "fc_bias"
|
|
||||||
updateMatrix(layer.bias, bGrad, biasKey)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 重置优化器状态
|
|
||||||
*/
|
|
||||||
public func reset(): Unit {
|
|
||||||
m.clear()
|
|
||||||
v.clear()
|
|
||||||
t = 0
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 获取当前迭代次数
|
|
||||||
*/
|
|
||||||
public func getStep(): Int64 {
|
|
||||||
t
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,273 +0,0 @@
|
|||||||
package bot4cj.math
|
|
||||||
import std.math.*
|
|
||||||
|
|
||||||
/* -------- Activation Functions 激活函数族 --------
|
|
||||||
* TODO(20260516-qing): 增加序列化、反序列化工厂函数 */
|
|
||||||
public enum Af {
|
|
||||||
|
|
||||||
None | Sigmoid | ReLU | Leaky_ReLU(FP32) | Tanh | Softmax
|
|
||||||
|
|
||||||
/* -------- 执行激活函数 -------- */
|
|
||||||
public func forward(x: Matrix1) {
|
|
||||||
match(this) {
|
|
||||||
case None => none(x)
|
|
||||||
case Sigmoid => sigmoid(x)
|
|
||||||
case ReLU => relu(x)
|
|
||||||
case Leaky_ReLU(alpha) => leaky_relu(x, alpha)
|
|
||||||
case Tanh => tanh(x)
|
|
||||||
case Softmax => softmax(x)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 激活函数的导数 -------- */
|
|
||||||
public func backward(y: Matrix1) {
|
|
||||||
match(this) {
|
|
||||||
case None => none_dot(y)
|
|
||||||
case Sigmoid => sigmoid_dot(y)
|
|
||||||
case ReLU => relu_dot(y)
|
|
||||||
case Leaky_ReLU(alpha) => leaky_relu_dot(y, alpha)
|
|
||||||
case Tanh => tanh_dot(y)
|
|
||||||
case Softmax => softmax_jacobian(y)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 序列化与反序列化 -------- */
|
|
||||||
public func get_type(): String {
|
|
||||||
return match(this) {
|
|
||||||
case None => "None"
|
|
||||||
case Sigmoid => "Sigmoid"
|
|
||||||
case ReLU => "ReLU"
|
|
||||||
case Leaky_ReLU(alpha) => "Leaky_ReLU"
|
|
||||||
case Tanh => "Tanh"
|
|
||||||
case Softmax => "Softmax"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/* */
|
|
||||||
public func get_value(): FP32 {
|
|
||||||
return match(this) {
|
|
||||||
case None => FP32(0.0f32)
|
|
||||||
case Sigmoid => FP32(0.0f32)
|
|
||||||
case ReLU => FP32(0.0f32)
|
|
||||||
case Leaky_ReLU(alpha) => alpha
|
|
||||||
case Tanh => FP32(0.0f32)
|
|
||||||
case Softmax => FP32(0.0f32)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 空的激活函数 -------- */
|
|
||||||
private static func none(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
N.set(r, c, FP32(x1)) }
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 空激活函数的导数 -------- */
|
|
||||||
private static func none_dot(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
N.set(r, c, FP32(1.0f32)) }
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- S型生长曲线 -------- */
|
|
||||||
private static func sigmoid(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 = 1.0f32 / (1.0f32 + exp(-x1))
|
|
||||||
N.set(r, c, FP32(x2)) }
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- S型生长曲线的导数 -------- */
|
|
||||||
private static func sigmoid_dot(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 = 1.0f32 / (1.0f32 + exp(-x1))
|
|
||||||
let x3 = x2 * (1.0f32 - x2)
|
|
||||||
N.set(r, c, FP32(x3)) }
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 线性整流单元 -------- */
|
|
||||||
private static func relu(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 =if (x1 > 0.0f32) {x1} else{0.0f32}
|
|
||||||
N.set(r, c, FP32(x2))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 线性整流单元的导数 -------- */
|
|
||||||
private static func relu_dot(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 =if (x1 > 0.0f32) {1.0f32} else{0.0f32}
|
|
||||||
N.set(r, c, FP32(x2))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 带泄露的线性整流单元 -------- */
|
|
||||||
private static func leaky_relu(M: Matrix1, alpha: FP32) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 =if (x1 > 0.0f32) {x1} else{x1 * alpha.getValue()}
|
|
||||||
N.set(r, c, FP32(x2))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 带泄露的整流单元导数 -------- */
|
|
||||||
private static func leaky_relu_dot(M: Matrix1, alpha: FP32) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 =if (x1 > 0.0f32) { 1.0f32 } else{alpha.getValue()}
|
|
||||||
N.set(r, c, FP32(x2)) }
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 双曲正切函数 --------*/
|
|
||||||
private static func tanh(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 = tanh(x1)
|
|
||||||
N.set(r, c, FP32(x2)) }
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 双曲正切函数的导数 -------- */
|
|
||||||
private static func tanh_dot(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let y = M.get(r, c).getOrThrow()
|
|
||||||
let y_val= y.getValue() /*---- NOTE(20260525-qing): 修复命名(x => y) ----*/
|
|
||||||
let rea = 1.0f32 - y_val * y_val
|
|
||||||
N.set(r, c, FP32(rea)) }
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 归一化指数函数 -------- */
|
|
||||||
private static func softmax(M: Matrix1) {
|
|
||||||
let row = M.get_row()
|
|
||||||
let col = M.get_col()
|
|
||||||
let N = Matrix1(row, col)
|
|
||||||
var max_val = 0.0f32
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
if (x1 > max_val) {
|
|
||||||
max_val = x1 }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
var sum = 0.0f32
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = M.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 = exp(x1 - max_val)
|
|
||||||
N.set(r, c, FP32(x2))
|
|
||||||
sum += x2
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let x = N.get(r, c).getOrThrow()
|
|
||||||
let x1= x.getValue()
|
|
||||||
let x2 = x1 / sum
|
|
||||||
N.set(r, c, FP32(x2))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 计算Softmax的雅可比矩阵 --------
|
|
||||||
* 输入:已计算好的softmax输出(概率分布)
|
|
||||||
* 输出:雅可比矩阵 J[i][j] = ∂y_i/∂x_j */
|
|
||||||
private func softmax_jacobian(softmax_output: Matrix1) {
|
|
||||||
let row = softmax_output.get_row()
|
|
||||||
let col = softmax_output.get_col()
|
|
||||||
let n = row * col
|
|
||||||
var jacobian = Matrix1(n, n)
|
|
||||||
|
|
||||||
for (i in 0..n) {
|
|
||||||
for (j in 0..n) {
|
|
||||||
|
|
||||||
let row_i = i / col /* 需要将一维索引转换为二维索引来访问Matrix1 */
|
|
||||||
let col_i = i % col
|
|
||||||
let row_j = j / col
|
|
||||||
let col_j = j % col
|
|
||||||
|
|
||||||
let val_i = softmax_output.get(row_i, col_i).getOrThrow().getValue()
|
|
||||||
let val_j = softmax_output.get(row_j, col_j).getOrThrow().getValue()
|
|
||||||
|
|
||||||
if (i == j) { /* 对角线元素: y_i * (1 - y_i) */
|
|
||||||
jacobian.set(i, j, FP32(val_i * (1.0f32 - val_i))) }
|
|
||||||
else { /* 非对角线元素: -y_i * y_j */
|
|
||||||
jacobian.set(i, j, FP32(-val_i * val_j)) }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return jacobian
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
@@ -1,37 +0,0 @@
|
|||||||
package bot4cj.math
|
|
||||||
import std.random.*
|
|
||||||
import std.time.*
|
|
||||||
import std.math.*
|
|
||||||
|
|
||||||
/* -------- 神经网络使用的权重矩阵初始化类 --------
|
|
||||||
* TODO(20260517-qing): 将该函数改为枚举类,统一一个接口,根据模式匹配执行初始化
|
|
||||||
* - 这样在矩阵式中可以定义初始化方式 */
|
|
||||||
public class Ini {
|
|
||||||
|
|
||||||
/* -------- xavier正态初始化 -------- */
|
|
||||||
public static func xavierNormalInit(target: Matrix1, fan_in: Int64, fan_out: Int64) {
|
|
||||||
let seed = DateTime.now().nanosecond
|
|
||||||
let random = Random(UInt64(seed))
|
|
||||||
let sigma = sqrt(2.0f32 / Float32(fan_in + fan_out)) /* 计算标准差 */
|
|
||||||
for (r in 0..fan_in) {
|
|
||||||
for (c in 0..fan_out) {
|
|
||||||
let v = random.nextGaussianFloat32(mean: 0.0f32, sigma:sigma)
|
|
||||||
target.set(r, c, FP32(v))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- Glorot均匀初始化 -------- */
|
|
||||||
public static func xavierUniformInit(target: Matrix1, fan_in: Int64, fan_out: Int64) {
|
|
||||||
let limit = sqrt(6.0f32 / Float32(fan_in + fan_out))
|
|
||||||
let seed = DateTime.now().nanosecond
|
|
||||||
let random = Random(UInt64(seed))
|
|
||||||
|
|
||||||
for (r in 0..fan_in) {
|
|
||||||
for (c in 0..fan_out) {
|
|
||||||
let v = (random.nextFloat32() * 2.0f32 - 1.0f32) * limit /* 均匀分布 [-limit, limit] */
|
|
||||||
target.set(r, c, FP32(v)) }
|
|
||||||
}
|
|
||||||
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,535 +0,0 @@
|
|||||||
package bot4cj.math
|
|
||||||
import std.fs.*
|
|
||||||
import std.io.*
|
|
||||||
import std.convert.*
|
|
||||||
//foreign func createMatrix(row: UInt64, col: UInt64): CPointer<Unit>
|
|
||||||
//foreign func freeMatrix(M: CPointer<Unit>): Unit
|
|
||||||
//foreign func witchMatrix(M: CPointer<Unit>, row: UInt64, col: UInt64): CPointer<Float32>
|
|
||||||
//foreign func getMatrix(M: CPointer<Unit>, row: UInt64, col: UInt64): Float32
|
|
||||||
//foreign func setMatrix(M: CPointer<Unit>, row: UInt64, col: UInt64, v: Float32): Unit
|
|
||||||
//foreign func getrowMatrix(M: CPointer<Unit>): UInt64
|
|
||||||
//foreign func getcolMatrix(M: CPointer<Unit>): UInt64
|
|
||||||
//foreign func sameShapeMatrix(M: CPointer<Unit>, S: CPointer<Unit>): Bool
|
|
||||||
//foreign func dotShapeMatrix(M: CPointer<Unit>, S: CPointer<Unit>): Bool
|
|
||||||
//foreign func plusMatrix(M: CPointer<Unit>, S:CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func plusBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
|
||||||
//foreign func subMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func subBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
|
||||||
//foreign func multiMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func multiBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
|
||||||
//foreign func divMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func divBroadcastMatrix(M: CPointer<Unit>, T: Float32): CPointer<Unit>
|
|
||||||
//foreign func dotMatrix(M: CPointer<Unit>, S: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func transposeMatrix(M: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func xavierNormalInit(fan_in: UInt64, fan_out: UInt64): CPointer<Unit>
|
|
||||||
//foreign func xavierUniformInit(fan_in: UInt64, fan_out: UInt64): CPointer<Unit>
|
|
||||||
|
|
||||||
|
|
||||||
/* -------- 32位浮点数 -------- */
|
|
||||||
public struct FP32 {
|
|
||||||
private var val: Float32
|
|
||||||
public init(t: Float32) { /* 以Float32初始化对象 */
|
|
||||||
val = t }
|
|
||||||
public func getValue() { /* 获取值 */
|
|
||||||
return val }
|
|
||||||
/* -------- 运算符重载 -------- */
|
|
||||||
//public operator func =(t: Float32) { /*-------- this can not --------*/
|
|
||||||
// val = t }
|
|
||||||
public operator func +(t: Float32) {
|
|
||||||
return FP32(val + t)}
|
|
||||||
public operator func +(S: FP32) {
|
|
||||||
return this + S.getValue() }
|
|
||||||
public operator func -(t: Float32) {
|
|
||||||
return FP32(val - t) }
|
|
||||||
public operator func -(S: FP32) {
|
|
||||||
return this - S.getValue() }
|
|
||||||
public operator func *(t: Float32) {
|
|
||||||
return FP32(val * t) }
|
|
||||||
public operator func *(S: FP32) {
|
|
||||||
return this * S.getValue() }
|
|
||||||
public operator func /(t: Float32) {
|
|
||||||
return FP32(val / t) }
|
|
||||||
public operator func /(S: FP32) {
|
|
||||||
return this / S.getValue() }
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
/*
|
|
||||||
* 新版的矩阵对象,仓颉原生实现
|
|
||||||
* TODO(20260517-qing): 设计原生的多维数组类
|
|
||||||
*
|
|
||||||
|
|
||||||
___col___
|
|
||||||
/ \
|
|
||||||
+---------+ \
|
|
||||||
| + + + + | | row - 第一维度
|
|
||||||
| + + + + | |
|
|
||||||
+---------+ /
|
|
||||||
|
|
||||||
|
|
||||||
* TODO(20260523-qing): 以某种方式统一矩阵式的字符串矩阵和这里的数值矩阵
|
|
||||||
*/
|
|
||||||
public class Matrix1 {
|
|
||||||
|
|
||||||
private var row: Int64 /* ---- 行数 ---- */
|
|
||||||
private var col: Int64 /* ---- 列数 ---- */
|
|
||||||
private var arr: Array<FP32> /* ---- 数据 ---- */
|
|
||||||
|
|
||||||
/* -------- 构造器,默认值为0 -------- */
|
|
||||||
public Matrix1(row: Int64, col: Int64) {
|
|
||||||
this.row = row
|
|
||||||
this.col = col
|
|
||||||
this.arr = Array<FP32>(row * col, {_ => FP32(0.0f32)})
|
|
||||||
}
|
|
||||||
|
|
||||||
/* --------- 这几个方法是访问器 -------- */
|
|
||||||
/* ------------------------------------- */
|
|
||||||
/* -------- TODO(20260524-qing): 使用prop成员 -------- */
|
|
||||||
|
|
||||||
public func get(r: Int64, c: Int64) {
|
|
||||||
return arr.get(r * this.col + c)
|
|
||||||
}
|
|
||||||
|
|
||||||
public func set(r: Int64, c: Int64, v: FP32) {
|
|
||||||
arr[r * this.col + c] = v
|
|
||||||
}
|
|
||||||
|
|
||||||
public func get_row() { return row }
|
|
||||||
|
|
||||||
public func get_col() { return col }
|
|
||||||
|
|
||||||
/* -------- 这两个函数是比较器 -------- */
|
|
||||||
/* ------------------------------------ */
|
|
||||||
public func has_same_shape(other: Matrix1) {
|
|
||||||
let a = this.get_row() == other.get_row()
|
|
||||||
let b = this.get_col() == other.get_col()
|
|
||||||
return a && b
|
|
||||||
}
|
|
||||||
|
|
||||||
public func has_dot_shape(other: Matrix1) {
|
|
||||||
return this.get_col() == other.get_row()
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 这些是运算符重载 -------- */
|
|
||||||
/* ---------------------------------- */
|
|
||||||
public operator func +(a: FP32) { /* ---- 矩阵广播加法 ---- */
|
|
||||||
let r =get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c) {
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v + a)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
public operator func +(S: Matrix1) { /* ---- 矩阵加法 ---- */
|
|
||||||
if (!has_same_shape(S)) {
|
|
||||||
throw IllegalArgumentException("Matrixes doesn't match")
|
|
||||||
}
|
|
||||||
let r =get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c) {
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
let vs = S.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v + vs)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
public operator func -(a: FP32) { /* ---- 矩阵广播减法 ----*/
|
|
||||||
let r = get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c){
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v - a)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
public operator func -(S: Matrix1) { /* ---- 矩阵减法 ---- */
|
|
||||||
if (!has_same_shape(S)) {
|
|
||||||
throw IllegalArgumentException("Matrixes doesn't match")
|
|
||||||
}
|
|
||||||
let r =get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c) {
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
let vs = S.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v - vs)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
public operator func *(a: FP32) { /* ---- 矩阵广播乘法 ---- */
|
|
||||||
let r = get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c){
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v * a)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
public operator func *(S: Matrix1) { /* ---- 逐元素相乘 ---- */
|
|
||||||
if (!has_same_shape(S)) {
|
|
||||||
throw IllegalArgumentException("Matrixes doesn't match")
|
|
||||||
}
|
|
||||||
let r =get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c) {
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
let vs = S.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v * vs)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
public operator func /(a: FP32) { /* ---- 矩阵广播除法 ---- */
|
|
||||||
let r = get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c){
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v / a)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
public operator func /(S: Matrix1) { /* ---- 逐元素除法 ---- */
|
|
||||||
if (!has_same_shape(S)) {
|
|
||||||
throw IllegalArgumentException("Matrixes doesn't match")
|
|
||||||
}
|
|
||||||
let r =get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c) {
|
|
||||||
let v = this.get(ir, ic).getOrThrow()
|
|
||||||
let vs = S.get(ir, ic).getOrThrow()
|
|
||||||
M.set(ir, ic, v / vs)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
/* ------------------------------ */
|
|
||||||
/* -------- 矩阵点积乘法 -------- */
|
|
||||||
public func dot(S: Matrix1) {
|
|
||||||
if (!has_dot_shape(S)) {
|
|
||||||
throw IllegalArgumentException("Matrixes doesn't match")
|
|
||||||
}
|
|
||||||
let r =get_row()
|
|
||||||
let c = S.get_col()
|
|
||||||
let M = Matrix1(r, c)
|
|
||||||
for (ir in 0..get_row()) {
|
|
||||||
for (ic in 0..S.get_col()) {
|
|
||||||
var sum = FP32(0.0f32)
|
|
||||||
for (ik in 0..get_col()) {
|
|
||||||
let a = get(ir, ik).getOrThrow()
|
|
||||||
let b = S.get(ik, ic).getOrThrow()
|
|
||||||
sum += a * b
|
|
||||||
}
|
|
||||||
M.set(ir, ic, sum)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return M
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 克隆一个一模一样的矩阵 -------- */
|
|
||||||
public func clone() {
|
|
||||||
let r = get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let N = Matrix1(r, c)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c) {
|
|
||||||
let v = get(ir, ic).getOrThrow()
|
|
||||||
N.set(ir, ic, v)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 矩阵转置,返回转置后的矩阵 -------- */
|
|
||||||
public func T() {
|
|
||||||
let r = get_row()
|
|
||||||
let c = get_col()
|
|
||||||
let N = Matrix1(c, r)
|
|
||||||
for (ir in 0..r) {
|
|
||||||
for (ic in 0..c) {
|
|
||||||
let v = get(ir, ic).getOrThrow()
|
|
||||||
N.set(ic, ir, v)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return N
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
/* -------- 旧版的矩阵类,依赖c++编译的共享动态库 --------
|
|
||||||
* TODO: 将矩阵元数据存储在仓颉部分(done at 202605xx ) */
|
|
||||||
//public class Matrix {
|
|
||||||
//
|
|
||||||
// /* -------- 主初始化器,初始化并获取资源 -------- */
|
|
||||||
// public init(row: UInt64, col: UInt64) {
|
|
||||||
// unsafe { M = createMatrix(row, col) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 初始化器,利用已经存在的元数据进行构建 -------- */
|
|
||||||
// public Matrix(m: CPointer<Unit>) { M = m }
|
|
||||||
//
|
|
||||||
// /* -------- Xavier正态分布初始化 -------- */
|
|
||||||
// public static func initWithXavierNormal(fan_in: UInt64, fan_out: UInt64) {
|
|
||||||
// unsafe {
|
|
||||||
// let m = xavierNormalInit(fan_in, fan_out);
|
|
||||||
// return Matrix(m)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- Glorot均匀初始化 -------- */
|
|
||||||
// public static func initWithXavierUniform(fan_in: UInt64, fan_out: UInt64) {
|
|
||||||
// unsafe {
|
|
||||||
// let m = xavierUniformInit(fan_in, fan_out);
|
|
||||||
// return Matrix(m)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* ======== 在析构函数中释放资源 ======== */
|
|
||||||
// ~init() {
|
|
||||||
// unsafe { freeMatrix(M) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 取下标,由于不能返回值类型的引用被弃用 -------- */
|
|
||||||
// //public operator func [](r: UInt64, c: UInt64): CPointer<Float32> {
|
|
||||||
// // unsafe { return witchMatrix(M, r, c) }
|
|
||||||
// //}
|
|
||||||
//
|
|
||||||
// /* -------- 获得某个矩阵单元的值 -------- */
|
|
||||||
// public func get(r: UInt64, c: UInt64): Float32 {
|
|
||||||
// unsafe { return getMatrix(M, r, c) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 设置某个矩阵单元的值 --------*/
|
|
||||||
// public func set(r: UInt64, c: UInt64, v: Float32) {
|
|
||||||
// unsafe { setMatrix(M, r, c, v) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* ------- 两个矩阵形状是否相同 --------
|
|
||||||
// * FIXME: 这个函数所重载的运算符不贴切 */
|
|
||||||
// public operator func ==(s: Matrix): Bool {
|
|
||||||
// unsafe { return sameShapeMatrix(M, s.getData()) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 获得行数 -------- */
|
|
||||||
// public func getRow(): UInt64 {
|
|
||||||
// unsafe { return getrowMatrix(M) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 获取列数 --------*/
|
|
||||||
// public func getCol(): UInt64 {
|
|
||||||
// unsafe { return getcolMatrix(M) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 获取指向元数据的C风格指针 -------- */
|
|
||||||
// public func getData(): CPointer<Unit> { return M; }
|
|
||||||
//
|
|
||||||
// /* -------- 是否两个矩阵具有点积和的形状 -------- */
|
|
||||||
// public func hasDotShapeWith(s: Matrix): Bool {
|
|
||||||
// unsafe { return dotShapeMatrix(M, s.getData()) }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 逐元素算数运算 -------- */
|
|
||||||
// /* -------------------------------- */
|
|
||||||
// public operator func +(s: Matrix): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// if (!sameShapeMatrix(M, s.getData())) {
|
|
||||||
// println('------------')
|
|
||||||
// print_shape()
|
|
||||||
// s.print_shape()
|
|
||||||
// println('------------')
|
|
||||||
// throw IllegalArgumentException("not same shape")
|
|
||||||
// }
|
|
||||||
// let r = plusMatrix(M, s.getData())
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// public operator func +(t: Float32): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// let r = plusBroadcastMatrix(M, t)
|
|
||||||
// return Matrix(r);
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// public operator func -(s: Matrix): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// if (!(this == s)) {
|
|
||||||
// println('------------')
|
|
||||||
// print_shape()
|
|
||||||
// s.print_shape()
|
|
||||||
// println('------------')
|
|
||||||
// throw IllegalArgumentException("not same shape")
|
|
||||||
// }
|
|
||||||
// let r = subMatrix(M, s.getData())
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// public operator func -(t: Float32): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// let r = subBroadcastMatrix(M, t)
|
|
||||||
// return Matrix(r);
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// public operator func *(s: Matrix): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// if (!(this == s)) {
|
|
||||||
// println('------------')
|
|
||||||
// print_shape()
|
|
||||||
// s.print_shape()
|
|
||||||
// println('------------')
|
|
||||||
// throw IllegalArgumentException("not same shape")
|
|
||||||
// }
|
|
||||||
// let r = multiMatrix(M, s.getData())
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// public operator func *(t: Float32): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// let r = multiBroadcastMatrix(M, t)
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// public operator func /(s: Matrix): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// if (!(this == s)) {
|
|
||||||
// println('------------')
|
|
||||||
// print_shape()
|
|
||||||
// s.print_shape()
|
|
||||||
// println('------------')
|
|
||||||
// throw IllegalArgumentException("not same shape")
|
|
||||||
// }
|
|
||||||
// let r = divMatrix(M, s.getData())
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// public operator func /(t: Float32): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// let r = divBroadcastMatrix(M, t)
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------------------------------- */
|
|
||||||
// /* -------------------------------- */
|
|
||||||
//
|
|
||||||
//
|
|
||||||
// /* -------- 矩阵点积和算法 -------- */
|
|
||||||
// public func dot (s: Matrix): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// if (!hasDotShapeWith(s)) {
|
|
||||||
// println('------------')
|
|
||||||
// print_shape()
|
|
||||||
// s.print_shape()
|
|
||||||
// println('------------')
|
|
||||||
// throw IllegalArgumentException("not dot shape")
|
|
||||||
// }
|
|
||||||
// let r = dotMatrix(M, s.getData())
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 矩阵转置 -------- */
|
|
||||||
// public func T(): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// let r = transposeMatrix(M)
|
|
||||||
// return Matrix(r)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 打印矩阵 -------- */
|
|
||||||
// public func print(): Unit {
|
|
||||||
// unsafe {
|
|
||||||
// for (i in 0..getRow()) {
|
|
||||||
// for (j in 0..getCol()) {
|
|
||||||
// let a = this.get(i, j)
|
|
||||||
// print(a)
|
|
||||||
// print(" ")
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// println()
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 打印形状 -------- */
|
|
||||||
// public func print_shape() {
|
|
||||||
// println("${getRow()}, ${getCol()}")
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 保存矩阵 -------- */
|
|
||||||
// public func save(file: File) {
|
|
||||||
// unsafe {
|
|
||||||
// for (i in 0..getRow()) {
|
|
||||||
// for (j in 0..getCol()) {
|
|
||||||
// file.write(this.get(i, j).toString().toArray())
|
|
||||||
// if (j < getCol()-1) {
|
|
||||||
// file.write(" ".toArray()) }
|
|
||||||
// else { file.write("\n".toArray()) }
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 静态方法,加载单独项 -------- */
|
|
||||||
// private static func loadItem(sr: StringReader<InputStream>) {
|
|
||||||
// let s = sr.readUntil({r => r==r' ' || r ==r'\n'})
|
|
||||||
// return Float32.parse(s.getOrThrow())
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 静态方法,加载矩阵 -------- */
|
|
||||||
// public static func load(sr: StringReader<InputStream>, row: UInt64, col: UInt64): Matrix {
|
|
||||||
// let m = Matrix(row, col)
|
|
||||||
// unsafe {
|
|
||||||
// for (i in 0..m.getRow()) {
|
|
||||||
// for (j in 0..m.getCol()) {
|
|
||||||
// let f = loadItem(sr)
|
|
||||||
// m.set(i, j, f)
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// m
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 矩阵的C风格指针 -------- */
|
|
||||||
// private var M: CPointer<Unit>
|
|
||||||
//}
|
|
||||||
@@ -1,56 +0,0 @@
|
|||||||
package bot4cj.math
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
/* -------- 张量 --------
|
|
||||||
* INFO(20260524-qing): 输入一个维度列表 */
|
|
||||||
class Tensor<T> {
|
|
||||||
private let data: Array<T>
|
|
||||||
private let dims: Array<Int64>
|
|
||||||
private let size: Int64
|
|
||||||
public init(dims: Array<Int64>, val: T) {
|
|
||||||
var size = 1
|
|
||||||
for (item in dims) { /* ---- 计算总量 ---- */
|
|
||||||
size *= item
|
|
||||||
}
|
|
||||||
this.data = Array<T>(size, {_ => val})
|
|
||||||
this.dims = dims
|
|
||||||
this.size = size
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 私有构造函数,用于内部创建张量 -------- */
|
|
||||||
private init(dims: Array<Int64>, data: Array<T>) {
|
|
||||||
this.dims = dims
|
|
||||||
this.data = data
|
|
||||||
var size = 1
|
|
||||||
for (item in dims) {
|
|
||||||
size *= item
|
|
||||||
}
|
|
||||||
this.size = size
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 获取维度列表(形状) -------- */
|
|
||||||
public func shape(): Array<Int64> {
|
|
||||||
return dims
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 逐元素进行某种运算 -------- */
|
|
||||||
public func oneByOne(S: Tensor<T>, f: (T, T) -> T): Tensor<T> {
|
|
||||||
let s1 = S.shape()
|
|
||||||
|
|
||||||
/* ---- 比较维度数量 ---- */
|
|
||||||
if (s1 != dims) {
|
|
||||||
throw IllegalArgumentException("Tensors' shape are not same")
|
|
||||||
}
|
|
||||||
|
|
||||||
/* ---- 比较维度形状 ----*/
|
|
||||||
for (i in 0..dims.size) {
|
|
||||||
if (s1[i] != dims[i]) {
|
|
||||||
throw IllegalArgumentException("Tensors' shape are not same")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/* ---- 逐个进行计算 ---- */
|
|
||||||
let res_data = Array<T>(size, {i => f(data[i], S.data[i])})
|
|
||||||
return Tensor<T>(dims, res_data)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,379 +0,0 @@
|
|||||||
// cnn.cj - 修正 forward 中的偏置加法
|
|
||||||
package bot4cj.math
|
|
||||||
|
|
||||||
import std.math.*
|
|
||||||
import std.random.*
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
public class ConvLayer {
|
|
||||||
|
|
||||||
let nInputChannels: Int64
|
|
||||||
let nOutputChannels: Int64
|
|
||||||
let kernelHeight: Int64
|
|
||||||
let kernelWidth: Int64
|
|
||||||
let stride: Int64
|
|
||||||
let padding: Int64
|
|
||||||
let learningRate: FP32
|
|
||||||
let f: Af
|
|
||||||
|
|
||||||
var connectingMatrix: Matrix1
|
|
||||||
var kernels: Array<Matrix1>
|
|
||||||
var biases: Matrix1
|
|
||||||
|
|
||||||
var prevInput: Option<Matrix1> = Option.None
|
|
||||||
var prevInputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
|
||||||
var prevCombination: Option<Matrix1> = Option.None
|
|
||||||
var prevOutput: Option<Matrix1> = Option.None
|
|
||||||
var prevOutputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
|
||||||
|
|
||||||
var kernelsGrad: Option<Array<Matrix1>> = Option.None
|
|
||||||
var biasGrad: Option<Matrix1> = Option.None
|
|
||||||
|
|
||||||
public init(
|
|
||||||
inCh: Int64,
|
|
||||||
outCh: Int64,
|
|
||||||
kernelSize: (Int64, Int64),
|
|
||||||
stride: Int64,
|
|
||||||
padding: Int64,
|
|
||||||
lr: FP32,
|
|
||||||
activation: Af,
|
|
||||||
connMatrix: ?Matrix1
|
|
||||||
) {
|
|
||||||
nInputChannels = inCh
|
|
||||||
nOutputChannels = outCh
|
|
||||||
kernelHeight = kernelSize[0]
|
|
||||||
kernelWidth = kernelSize[1]
|
|
||||||
this.stride = stride
|
|
||||||
this.padding = padding
|
|
||||||
learningRate = lr
|
|
||||||
this.f = activation
|
|
||||||
|
|
||||||
// 初始化 connectingMatrix
|
|
||||||
connectingMatrix = match (connMatrix) {
|
|
||||||
case m: Matrix1 => m
|
|
||||||
case _ => Matrix1(nInputChannels, nOutputChannels)
|
|
||||||
}
|
|
||||||
let isNone = match (connMatrix) {
|
|
||||||
case _: Matrix1 => false
|
|
||||||
case _ => true
|
|
||||||
}
|
|
||||||
if (isNone) {
|
|
||||||
for (r in 0..nInputChannels) {
|
|
||||||
for (c in 0..nOutputChannels) {
|
|
||||||
connectingMatrix.set(r, c, FP32(1.0))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 初始化 kernels
|
|
||||||
kernels = Array(nOutputChannels, { _ => Matrix1(1, 1) })
|
|
||||||
for (oc in 0..nOutputChannels) {
|
|
||||||
var nConnections = 0
|
|
||||||
for (ic in 0..nInputChannels) {
|
|
||||||
let val = connectingMatrix.get(ic, oc).getOrThrow()
|
|
||||||
if (val.getValue() > 0.0) {
|
|
||||||
nConnections += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
let kernelLen = nConnections * kernelHeight * kernelWidth
|
|
||||||
let kernel = Matrix1(1, kernelLen)
|
|
||||||
|
|
||||||
Ini.xavierUniformInit(kernel, 1, kernelLen)
|
|
||||||
|
|
||||||
kernels[oc] = kernel
|
|
||||||
}
|
|
||||||
|
|
||||||
// 初始化 biases
|
|
||||||
biases = Matrix1(1, nOutputChannels)
|
|
||||||
Ini.xavierUniformInit(biases, 1, nOutputChannels)
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 前向传播 ====================
|
|
||||||
public func forward(x: Matrix1, inputShape: (Int64, Int64, Int64, Int64)): Matrix1 {
|
|
||||||
prevInput = Option.Some(x)
|
|
||||||
prevInputShape = Option.Some(inputShape)
|
|
||||||
|
|
||||||
let batch = inputShape[0]
|
|
||||||
let inH = inputShape[1]
|
|
||||||
let inW = inputShape[2]
|
|
||||||
|
|
||||||
let outH = (inH + 2 * padding - kernelHeight) / stride + 1
|
|
||||||
let outW = (inW + 2 * padding - kernelWidth) / stride + 1
|
|
||||||
let outSize = batch * outH * outW
|
|
||||||
|
|
||||||
var combination = Matrix1(outSize, nOutputChannels)
|
|
||||||
var output = Matrix1(outSize, nOutputChannels)
|
|
||||||
|
|
||||||
for (b in 0..batch) {
|
|
||||||
for (oh in 0..outH) {
|
|
||||||
for (ow in 0..outW) {
|
|
||||||
let outRow = b * outH * outW + oh * outW + ow
|
|
||||||
|
|
||||||
for (oc in 0..nOutputChannels) {
|
|
||||||
var sum = FP32(0.0)
|
|
||||||
var kernelIdx = 0
|
|
||||||
|
|
||||||
for (ic in 0..nInputChannels) {
|
|
||||||
let connVal = connectingMatrix.get(ic, oc).getOrThrow()
|
|
||||||
if (connVal.getValue() > 0.0) {
|
|
||||||
for (kh in 0..kernelHeight) {
|
|
||||||
for (kw in 0..kernelWidth) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
let inRow = b * inH * inW + ih * inW + iw
|
|
||||||
let xVal = x.get(inRow, ic).getOrThrow()
|
|
||||||
let kVal = kernels[oc].get(0, kernelIdx).getOrThrow()
|
|
||||||
sum = sum + xVal * kVal
|
|
||||||
}
|
|
||||||
kernelIdx += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
combination.set(outRow, oc, sum)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 手动广播偏置加法(因为 Matrix1 的 + 要求形状相同)
|
|
||||||
for (r in 0..outSize) {
|
|
||||||
for (c in 0..nOutputChannels) {
|
|
||||||
let val = combination.get(r, c).getOrThrow()
|
|
||||||
let biasVal = biases.get(0, c).getOrThrow()
|
|
||||||
combination.set(r, c, val + biasVal)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
prevCombination = Option.Some(combination)
|
|
||||||
|
|
||||||
output = f.forward(combination)
|
|
||||||
prevOutput = Option.Some(output)
|
|
||||||
prevOutputShape = Option.Some((batch, outH, outW, nOutputChannels))
|
|
||||||
|
|
||||||
return output
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 反向传播 ====================
|
|
||||||
public func backward(errs: Matrix1): Matrix1 {
|
|
||||||
var x = Matrix1(1, 1)
|
|
||||||
match (prevInput) {
|
|
||||||
case Some(T) => x = T
|
|
||||||
case _ => throw Exception("No input data")
|
|
||||||
}
|
|
||||||
|
|
||||||
var shape = (0, 0, 0, 0)
|
|
||||||
match (prevInputShape) {
|
|
||||||
case Some(T) => shape = T
|
|
||||||
case _ => throw Exception("No input shape")
|
|
||||||
}
|
|
||||||
|
|
||||||
var combination = Matrix1(1, 1)
|
|
||||||
match (prevCombination) {
|
|
||||||
case Some(val) => combination = val
|
|
||||||
case _ => throw Exception("No combination")
|
|
||||||
}
|
|
||||||
|
|
||||||
var outShape = (0, 0, 0, 0)
|
|
||||||
match (prevOutputShape) {
|
|
||||||
case Some(val) => outShape = val
|
|
||||||
case _ => throw Exception("No output shape")
|
|
||||||
}
|
|
||||||
|
|
||||||
let batch = shape[0]
|
|
||||||
let inH = shape[1]
|
|
||||||
let inW = shape[2]
|
|
||||||
let outH = outShape[1]
|
|
||||||
let outW = outShape[2]
|
|
||||||
|
|
||||||
let delta = errs * f.backward(combination)
|
|
||||||
|
|
||||||
calcKernelGradient(x, delta, batch, inH, inW, outH, outW)
|
|
||||||
calcBiasGradient(delta)
|
|
||||||
|
|
||||||
let errsUp = calcUpsampledError(delta, batch, inH, inW, outH, outW)
|
|
||||||
|
|
||||||
return errsUp
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算卷积核梯度
|
|
||||||
*/
|
|
||||||
private func calcKernelGradient(
|
|
||||||
x: Matrix1,
|
|
||||||
delta: Matrix1,
|
|
||||||
batch: Int64,
|
|
||||||
inH: Int64,
|
|
||||||
inW: Int64,
|
|
||||||
outH: Int64,
|
|
||||||
outW: Int64
|
|
||||||
): Unit {
|
|
||||||
var grads = Array(nOutputChannels, { _ => Matrix1(1, 1) })
|
|
||||||
|
|
||||||
for (oc in 0..nOutputChannels) {
|
|
||||||
let kernelCols = kernels[oc].get_col()
|
|
||||||
var grad = Matrix1(1, kernelCols)
|
|
||||||
|
|
||||||
for (b in 0..batch) {
|
|
||||||
for (oh in 0..outH) {
|
|
||||||
for (ow in 0..outW) {
|
|
||||||
let deltaRow = b * outH * outW + oh * outW + ow
|
|
||||||
let deltaVal = delta.get(deltaRow, oc).getOrThrow()
|
|
||||||
|
|
||||||
var kernelIdx = 0
|
|
||||||
for (ic in 0..nInputChannels) {
|
|
||||||
let connVal = connectingMatrix.get(ic, oc).getOrThrow()
|
|
||||||
if (connVal.getValue() > 0.0) {
|
|
||||||
for (kh in 0..kernelHeight) {
|
|
||||||
for (kw in 0..kernelWidth) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
let inRow = b * inH * inW + ih * inW + iw
|
|
||||||
let xVal = x.get(inRow, ic).getOrThrow()
|
|
||||||
let oldGrad = grad.get(0, kernelIdx).getOrThrow()
|
|
||||||
grad.set(0, kernelIdx, oldGrad + xVal * deltaVal)
|
|
||||||
}
|
|
||||||
kernelIdx += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
grads[oc] = grad
|
|
||||||
}
|
|
||||||
|
|
||||||
kernelsGrad = Option.Some(grads)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算偏置梯度
|
|
||||||
*/
|
|
||||||
private func calcBiasGradient(delta: Matrix1): Unit {
|
|
||||||
let rows = delta.get_row()
|
|
||||||
var grad = Matrix1(1, nOutputChannels)
|
|
||||||
|
|
||||||
for (oc in 0..nOutputChannels) {
|
|
||||||
var sum = FP32(0.0)
|
|
||||||
for (r in 0..rows) {
|
|
||||||
sum = sum + delta.get(r, oc).getOrThrow()
|
|
||||||
}
|
|
||||||
grad.set(0, oc, sum)
|
|
||||||
}
|
|
||||||
|
|
||||||
biasGrad = Option.Some(grad)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算传递到上一层的误差
|
|
||||||
*/
|
|
||||||
private func calcUpsampledError(
|
|
||||||
delta: Matrix1,
|
|
||||||
batch: Int64,
|
|
||||||
inH: Int64,
|
|
||||||
inW: Int64,
|
|
||||||
outH: Int64,
|
|
||||||
outW: Int64
|
|
||||||
): Matrix1 {
|
|
||||||
let inSize = batch * inH * inW
|
|
||||||
var errsUp = Matrix1(inSize, nInputChannels)
|
|
||||||
|
|
||||||
for (b in 0..batch) {
|
|
||||||
for (oh in 0..outH) {
|
|
||||||
for (ow in 0..outW) {
|
|
||||||
let deltaRow = b * outH * outW + oh * outW + ow
|
|
||||||
|
|
||||||
for (oc in 0..nOutputChannels) {
|
|
||||||
let deltaVal = delta.get(deltaRow, oc).getOrThrow()
|
|
||||||
|
|
||||||
var kernelIdx = 0
|
|
||||||
for (ic in 0..nInputChannels) {
|
|
||||||
let connVal = connectingMatrix.get(ic, oc).getOrThrow()
|
|
||||||
if (connVal.getValue() > 0.0) {
|
|
||||||
for (kh in 0..kernelHeight) {
|
|
||||||
for (kw in 0..kernelWidth) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
let inRow = b * inH * inW + ih * inW + iw
|
|
||||||
let kVal = kernels[oc].get(0, kernelIdx).getOrThrow()
|
|
||||||
let oldErr = errsUp.get(inRow, ic).getOrThrow()
|
|
||||||
errsUp.set(inRow, ic, oldErr + deltaVal * kVal)
|
|
||||||
}
|
|
||||||
kernelIdx += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return errsUp
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 参数更新 ====================
|
|
||||||
public func update(discount: Float32): Unit {
|
|
||||||
var kGrad = Array(0, { _ => Matrix1(1, 1) })
|
|
||||||
match (kernelsGrad) {
|
|
||||||
case Some(val) => kGrad = val
|
|
||||||
case _ => throw Exception("No kernel gradients")
|
|
||||||
}
|
|
||||||
|
|
||||||
var bGrad = Matrix1(1, 1)
|
|
||||||
match (biasGrad) {
|
|
||||||
case Some(val) => bGrad = val
|
|
||||||
case _ => throw Exception("No bias gradient")
|
|
||||||
}
|
|
||||||
|
|
||||||
for (oc in 0..nOutputChannels) {
|
|
||||||
let rows = kernels[oc].get_row()
|
|
||||||
let cols = kernels[oc].get_col()
|
|
||||||
for (r in 0..rows) {
|
|
||||||
for (c in 0..cols) {
|
|
||||||
let oldVal = kernels[oc].get(r, c).getOrThrow()
|
|
||||||
let gradVal = kGrad[oc].get(r, c).getOrThrow()
|
|
||||||
kernels[oc].set(r, c, oldVal - gradVal * learningRate * FP32(discount))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
for (oc in 0..nOutputChannels) {
|
|
||||||
let oldBias = biases.get(0, oc).getOrThrow()
|
|
||||||
let gradBias = bGrad.get(0, oc).getOrThrow()
|
|
||||||
biases.set(0, oc, oldBias - gradBias * learningRate * FP32(discount))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 形状访问器 ====================
|
|
||||||
public func getInShape(): (Int64, Int64) {
|
|
||||||
(nInputChannels, kernelHeight * kernelWidth)
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getOutShape(): Int64 {
|
|
||||||
nOutputChannels
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getLR(): FP32 {
|
|
||||||
learningRate
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getFType(): String {
|
|
||||||
f.get_type()
|
|
||||||
}
|
|
||||||
|
|
||||||
public func calculateOutputSize(inputHeight: Int64, inputWidth: Int64): (Int64, Int64) {
|
|
||||||
let outH = (inputHeight + 2 * padding - kernelHeight) / stride + 1
|
|
||||||
let outW = (inputWidth + 2 * padding - kernelWidth) / stride + 1
|
|
||||||
(outH, outW)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,11 +0,0 @@
|
|||||||
package bot4cj.math
|
|
||||||
//foreign func start_fft(M: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func gen_wave(): CPointer<Unit>
|
|
||||||
/*-------- 生成测试用的正弦波 --------*/
|
|
||||||
//public func genWave(): Matrix {
|
|
||||||
// unsafe { return Matrix(gen_wave()) }
|
|
||||||
//}
|
|
||||||
/*-------- 执行fft算法 --------*/
|
|
||||||
//public func start_fft(M: Matrix): Matrix {
|
|
||||||
// unsafe { return Matrix(start_fft(M.getData())) }
|
|
||||||
//}
|
|
||||||
@@ -1,136 +0,0 @@
|
|||||||
// flatten.cj - 移除调试代码
|
|
||||||
package bot4cj.math
|
|
||||||
|
|
||||||
import std.math.*
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
public class FlattenLayer {
|
|
||||||
|
|
||||||
let nInputChannels: Int64
|
|
||||||
let inputHeight: Int64
|
|
||||||
let inputWidth: Int64
|
|
||||||
let outputSize: Int64
|
|
||||||
|
|
||||||
var prevInput: Option<Matrix1>
|
|
||||||
var prevInputShape: Option<(Int64, Int64, Int64, Int64)>
|
|
||||||
var prevOutput: Option<Matrix1>
|
|
||||||
var prevOutputShape: Option<(Int64, Int64)>
|
|
||||||
|
|
||||||
public init(inputChannels: Int64, inputHeight: Int64, inputWidth: Int64) {
|
|
||||||
nInputChannels = inputChannels
|
|
||||||
this.inputHeight = inputHeight
|
|
||||||
this.inputWidth = inputWidth
|
|
||||||
outputSize = inputChannels * inputHeight * inputWidth
|
|
||||||
|
|
||||||
prevInput = None
|
|
||||||
prevInputShape = None
|
|
||||||
prevOutput = None
|
|
||||||
prevOutputShape = None
|
|
||||||
}
|
|
||||||
|
|
||||||
public init(inputShape: (Int64, Int64, Int64)) {
|
|
||||||
nInputChannels = inputShape[0]
|
|
||||||
inputHeight = inputShape[1]
|
|
||||||
inputWidth = inputShape[2]
|
|
||||||
outputSize = nInputChannels * inputHeight * inputWidth
|
|
||||||
|
|
||||||
prevInput = None
|
|
||||||
prevInputShape = None
|
|
||||||
prevOutput = None
|
|
||||||
prevOutputShape = None
|
|
||||||
}
|
|
||||||
|
|
||||||
public func forward(x: Matrix1, inputShape: (Int64, Int64, Int64, Int64)): Matrix1 {
|
|
||||||
prevInput = x
|
|
||||||
prevInputShape = inputShape
|
|
||||||
|
|
||||||
let batch = inputShape[0]
|
|
||||||
let inH = inputShape[1]
|
|
||||||
let inW = inputShape[2]
|
|
||||||
let inCh = inputShape[3]
|
|
||||||
|
|
||||||
let totalRows = batch * outputSize
|
|
||||||
var output = Matrix1(totalRows, 1)
|
|
||||||
|
|
||||||
for (b in 0..batch) {
|
|
||||||
for (h in 0..inH) {
|
|
||||||
for (w in 0..inW) {
|
|
||||||
for (ch in 0..inCh) {
|
|
||||||
let inRow = b * inH * inW + h * inW + w
|
|
||||||
let outRow = b * outputSize + (ch * inH * inW + h * inW + w)
|
|
||||||
if (outRow < totalRows) {
|
|
||||||
let val = x.get(inRow, ch).getOrThrow()
|
|
||||||
output.set(outRow, 0, val)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
prevOutput = output
|
|
||||||
prevOutputShape = (totalRows, 1)
|
|
||||||
|
|
||||||
return output
|
|
||||||
}
|
|
||||||
|
|
||||||
public func backward(errs: Matrix1): Matrix1 {
|
|
||||||
var shape = (0, 0, 0, 0)
|
|
||||||
match (prevInputShape) {
|
|
||||||
case Some(T) => shape = T
|
|
||||||
case _ => throw Exception("FlattenLayer.backward: No input shape. Call forward() first.")
|
|
||||||
}
|
|
||||||
|
|
||||||
let batch = shape[0]
|
|
||||||
let inH = shape[1]
|
|
||||||
let inW = shape[2]
|
|
||||||
let inCh = shape[3]
|
|
||||||
|
|
||||||
let inSize = batch * inH * inW
|
|
||||||
var errsUp = Matrix1(inSize, nInputChannels)
|
|
||||||
|
|
||||||
let errsCols = errs.get_col()
|
|
||||||
|
|
||||||
for (b in 0..batch) {
|
|
||||||
for (h in 0..inH) {
|
|
||||||
for (w in 0..inW) {
|
|
||||||
let inRow = b * inH * inW + h * inW + w
|
|
||||||
|
|
||||||
for (ch in 0..inCh) {
|
|
||||||
let outRow = b * outputSize + (ch * inH * inW + h * inW + w)
|
|
||||||
|
|
||||||
if (outRow < errs.get_row()) {
|
|
||||||
if (errsCols == 1) {
|
|
||||||
let val = errs.get(outRow, 0).getOrThrow()
|
|
||||||
errsUp.set(inRow, ch, val)
|
|
||||||
} else {
|
|
||||||
var sum = FP32(0.0)
|
|
||||||
for (col in 0..errsCols) {
|
|
||||||
sum = sum + errs.get(outRow, col).getOrThrow()
|
|
||||||
}
|
|
||||||
errsUp.set(inRow, ch, sum)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return errsUp
|
|
||||||
}
|
|
||||||
|
|
||||||
public func update(discount: Float32): Unit {
|
|
||||||
// 无参数
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getInShape(): (Int64, Int64, Int64) {
|
|
||||||
(nInputChannels, inputHeight, inputWidth)
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getOutShape(): Int64 {
|
|
||||||
outputSize
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getOutputSize(): Int64 {
|
|
||||||
outputSize
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,331 +0,0 @@
|
|||||||
package bot4cj.math
|
|
||||||
import std.fs.*
|
|
||||||
import std.io.*
|
|
||||||
import std.convert.*
|
|
||||||
//foreign func Sigmoid_do(m: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func Sigmoid_dot(m: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func ReLU_do(m: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func ReLU_dot(m: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func Leaky_ReLU_do(m: CPointer<Unit>, alpha: Float32): CPointer<Unit>
|
|
||||||
//foreign func Leaky_ReLU_dot(m: CPointer<Unit>, alpha: Float32): CPointer<Unit>
|
|
||||||
//foreign func Tanh_do(m: CPointer<Unit>): CPointer<Unit>
|
|
||||||
//foreign func Tanh_dot(m: CPointer<Unit>): CPointer<Unit>
|
|
||||||
|
|
||||||
|
|
||||||
/* -------- 新版的神经网络,使用原生矩阵对象实现 --------
|
|
||||||
* TODO(20260517-qing): 增加batch批次处理,使用ArrayList<Matrix1>
|
|
||||||
* TODO*(20260519-qing): 将神经网络转换为MatrixObject矩阵式对象
|
|
||||||
* TODO(20260524-qing): 增加Adam优化器 */
|
|
||||||
|
|
||||||
public class Nn1 {
|
|
||||||
|
|
||||||
let inputno: Int64 /* -------- 输入 -------- */
|
|
||||||
let outputno: Int64 /* -------- 输出 -------- */
|
|
||||||
var weights: Matrix1 /* -- (inputs, outputs) - */
|
|
||||||
var bias: Matrix1 /* ----- (1, output) ---- */
|
|
||||||
let learning_rate: FP32 /* ------- 学习率 ------- */
|
|
||||||
let f: Af /* ------ 激活函数 ------ */
|
|
||||||
|
|
||||||
/*------------------*/
|
|
||||||
var prev_inputs: Option<Matrix1> = None /* ---- 之前的输入 ---- */
|
|
||||||
var prev_raw_inputs: Option<Matrix1> = None /* ---- 之前的原始输出 ---- */
|
|
||||||
var prev_outputs: Option<Matrix1> = None /* ---- 之前的输出 ---- */
|
|
||||||
/*------------------*/
|
|
||||||
var weights_grad: Option<Matrix1> = None /* ---- 权重梯度 ---- */
|
|
||||||
var bias_grad: Option<Matrix1> = None /* ---- 偏置梯度 ---- */
|
|
||||||
|
|
||||||
|
|
||||||
/* -------- 初始化器 --------
|
|
||||||
* TODO(20260516-qing): 不要把初始化方法写死,使用枚举类 */
|
|
||||||
public init(i: Int64, o: Int64, lr: Float32, f: Af) {
|
|
||||||
inputno = i
|
|
||||||
outputno =o
|
|
||||||
learning_rate=FP32(lr)
|
|
||||||
this.f = f
|
|
||||||
weights = Matrix1(inputno, outputno)
|
|
||||||
bias = Matrix1(1, outputno)
|
|
||||||
|
|
||||||
/* -------- 初始化权重 -------- */
|
|
||||||
Ini.xavierUniformInit(weights, inputno, outputno)
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 神经网络推理 -------- */
|
|
||||||
public func forward(x: Matrix1) {
|
|
||||||
prev_inputs = x
|
|
||||||
var r = x.dot(weights)
|
|
||||||
r = r + bias
|
|
||||||
prev_raw_inputs = r
|
|
||||||
let r1 = f.forward(r)
|
|
||||||
prev_outputs = r1
|
|
||||||
return r1
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 反向传播算法 -------- */
|
|
||||||
public func backward(errs: Matrix1)
|
|
||||||
{
|
|
||||||
if (is_softmax(f)) {
|
|
||||||
return softmax_backward(errs)
|
|
||||||
}
|
|
||||||
let fbak = f.backward(prev_raw_inputs.getOrThrow())
|
|
||||||
let first = errs * fbak /* 逐元素相乘 */
|
|
||||||
let mx = prev_inputs.getOrThrow().T()
|
|
||||||
weights_grad = mx.dot(first)
|
|
||||||
bias_grad = first
|
|
||||||
let errs_up = first.dot(weights.T())
|
|
||||||
return errs_up
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 更新权重值、偏爱因子 --------
|
|
||||||
* NOTE: discount(折扣)是为了方便深度强化学习给出奖励值 */
|
|
||||||
public func update(discount: Float32) {
|
|
||||||
weights = weights - weights_grad.getOrThrow() * learning_rate * FP32(discount)
|
|
||||||
bias = bias - bias_grad.getOrThrow() * learning_rate * FP32(discount)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**/
|
|
||||||
private static func is_softmax(f: Af) {
|
|
||||||
return match (f) {
|
|
||||||
case Af.Softmax => true
|
|
||||||
case _ => false
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
private func softmax_backward(errs: Matrix1) {
|
|
||||||
/* 使用雅可比矩阵的方式 */
|
|
||||||
let y = prev_outputs.getOrThrow() /* softmax 输出 */
|
|
||||||
let dy = errs /* 上游梯度 ∂L/∂y */
|
|
||||||
let jacobian = f.backward(y) /* n×n 雅可比矩阵 */
|
|
||||||
|
|
||||||
/* 将 dy 展平成列向量 n×1 */
|
|
||||||
let row = dy.get_row()
|
|
||||||
let col = dy.get_col()
|
|
||||||
let n = row * col
|
|
||||||
var dy_vec = Matrix1(n, 1)
|
|
||||||
var idx = 0
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let val = dy.get(r, c).getOrThrow()
|
|
||||||
dy_vec.set(idx, 0, val)
|
|
||||||
idx = idx + 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/* dx_vec = J · dy_vec (n×n × n×1 = n×1) */
|
|
||||||
var dx_vec = Matrix1(n, 1)
|
|
||||||
for (i in 0..n) {
|
|
||||||
var sum = 0.0f32
|
|
||||||
for (j in 0..n) {
|
|
||||||
let j_val = jacobian.get(i, j).getOrThrow().getValue()
|
|
||||||
let dy_val = dy_vec.get(j, 0).getOrThrow().getValue()
|
|
||||||
sum = sum + j_val * dy_val
|
|
||||||
}
|
|
||||||
dx_vec.set(i, 0, FP32(sum))
|
|
||||||
}
|
|
||||||
|
|
||||||
/* 恢复为原始形状,作为 first */
|
|
||||||
var first = Matrix1(row, col)
|
|
||||||
idx = 0
|
|
||||||
for (r in 0..row) {
|
|
||||||
for (c in 0..col) {
|
|
||||||
let val = dx_vec.get(idx, 0).getOrThrow()
|
|
||||||
first.set(r, c, val)
|
|
||||||
idx = idx + 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let mx = prev_inputs.getOrThrow().T()
|
|
||||||
weights_grad = mx.dot(first)
|
|
||||||
bias_grad = first
|
|
||||||
let errs_up = first.dot(weights.T())
|
|
||||||
return errs_up
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
/* -------- 形状访问器 -------- */
|
|
||||||
public func get_in_shape() {
|
|
||||||
return inputno }
|
|
||||||
|
|
||||||
public func get_out_shape() {
|
|
||||||
return outputno }
|
|
||||||
|
|
||||||
public func get_lr() {
|
|
||||||
return learning_rate }
|
|
||||||
|
|
||||||
public func get_f_type() {
|
|
||||||
return f.get_type() }
|
|
||||||
|
|
||||||
public func get_f_val() {
|
|
||||||
return f.get_value() }
|
|
||||||
/* ---------------------------- */
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
///* -------- 旧版激活函数 --------*/
|
|
||||||
//public enum ActivationFuncs {
|
|
||||||
// Sigmoid | ReLU | Leaky_ReLU | Tanh
|
|
||||||
// func fdo(m: Matrix): Matrix { // 暂时只设置ReLU
|
|
||||||
// unsafe {
|
|
||||||
// match(this) {
|
|
||||||
// case Sigmoid => Matrix(Sigmoid_do(m.getData()))
|
|
||||||
// case ReLU => Matrix(ReLU_do(m.getData()))
|
|
||||||
// case Leaky_ReLU => Matrix(Leaky_ReLU_do(m.getData(), 0.1f32))
|
|
||||||
// case Tanh => Matrix(Tanh_do(m.getData()))
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// func fdot(m: Matrix): Matrix {
|
|
||||||
// unsafe {
|
|
||||||
// match(this) {
|
|
||||||
// case Sigmoid => Matrix(Sigmoid_dot(m.getData()))
|
|
||||||
// case ReLU => Matrix(ReLU_dot(m.getData()))
|
|
||||||
// case Leaky_ReLU => Matrix(Leaky_ReLU_dot(m.getData(), 0.1f32))
|
|
||||||
// case Tanh => Matrix(Tanh_dot(m.getData()))
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// func toString(): String {
|
|
||||||
// match(this) {
|
|
||||||
// case Sigmoid => "Sigmoid"
|
|
||||||
// case ReLU => "ReLU"
|
|
||||||
// case Leaky_ReLU => "Leaky_ReLU"
|
|
||||||
// case Tanh => "Tanh"
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
// static func parse(s: String) {
|
|
||||||
// match(s) {
|
|
||||||
// case "Sigmoid" => ActivationFuncs.Sigmoid
|
|
||||||
// case "ReLU" => ActivationFuncs.ReLU
|
|
||||||
// case "Leaky_ReLU" => ActivationFuncs.Leaky_ReLU
|
|
||||||
// case "Tanh" => ActivationFuncs.Tanh
|
|
||||||
// case _ => throw IllegalArgumentException()
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
//}
|
|
||||||
|
|
||||||
///* -------- 旧版的神经网络类 --------
|
|
||||||
// * 这个实现依赖C++语言的动态分享库,特点是速度快
|
|
||||||
// * 这个实现可能是有问题的 */
|
|
||||||
//public class Nn {
|
|
||||||
//
|
|
||||||
// /* -------- 主初始化器 -------- */
|
|
||||||
// public init(
|
|
||||||
// inputno: UInt64, outputno: UInt64, f: ActivationFuncs, learning_rate: Float32) {
|
|
||||||
// this.inputno = inputno
|
|
||||||
// this.outputno = outputno
|
|
||||||
// weights = Matrix.initWithXavierUniform(inputno, outputno)
|
|
||||||
// bias = Matrix(1, outputno)
|
|
||||||
// this.f = f
|
|
||||||
// this.learning_rate = learning_rate
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 提供了权重矩阵的初始化器 -------- */
|
|
||||||
// public init(
|
|
||||||
// inputno: UInt64, outputno: UInt64, f: ActivationFuncs, learning_rate: Float32, weights: Matrix, bias: Matrix) {
|
|
||||||
// this.inputno = inputno
|
|
||||||
// this.outputno = outputno
|
|
||||||
// this.weights = weights
|
|
||||||
// this.bias = bias
|
|
||||||
// this.f = f
|
|
||||||
// this.learning_rate = learning_rate
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 神经网络推理 -------- */
|
|
||||||
// func forward(x: Matrix): Matrix {
|
|
||||||
// /* x => (1, inputs) */
|
|
||||||
// let r = x.dot(weights) + bias /* (1, outputno) */
|
|
||||||
// return r
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 实施激活函数 -------- */
|
|
||||||
// func fdo(x: Matrix): Matrix {
|
|
||||||
// return f.fdo(x) }
|
|
||||||
//
|
|
||||||
// /* -------- 激活函数导数 -------- */
|
|
||||||
// func fdot(x: Matrix): Matrix {
|
|
||||||
// return f.fdot(x) }
|
|
||||||
//
|
|
||||||
// /* This is Wrong */
|
|
||||||
// //func bak_bsgrad(prev_raw_outputs_dotf: Matrix, errs: Matrix): Matrix {
|
|
||||||
// /* prev_raw_outputs_dotf => (1, outputno)
|
|
||||||
// * errs => (1, outputno) */
|
|
||||||
//
|
|
||||||
// // return errs * prev_raw_outputs_dotf
|
|
||||||
// //}
|
|
||||||
//
|
|
||||||
// /* -------- 权重梯度的传播 -------- */
|
|
||||||
// func bak_wsgrad(prev_inputs: Matrix, delta: Matrix) {
|
|
||||||
// /* prev_inputs => (1, inputs) */
|
|
||||||
// return prev_inputs.T().dot(delta) /* (inputs, outputs) */
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 计算上一层的误差 -------- */
|
|
||||||
// func bak_uperrs(delta: Matrix, prev_dotf: Matrix) {
|
|
||||||
// let first = weights.dot(delta.T()).T() // (1, inputs)
|
|
||||||
//
|
|
||||||
// var uperrs = first * prev_dotf
|
|
||||||
// uperrs
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 更新权重、偏爱因子 -------- */
|
|
||||||
// func update(weights_grad: Matrix, bias_grad: Matrix, discount: Float32) {
|
|
||||||
// weights -= weights_grad * learning_rate * discount
|
|
||||||
// bias -= bias_grad * learning_rate * discount
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 打印权重、偏爱因子 -------- */
|
|
||||||
// func print() {
|
|
||||||
// weights.print()
|
|
||||||
// bias.print()
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 打印神经网络全连接层的形状 -------- */
|
|
||||||
// func print_shape() {
|
|
||||||
// println("${inputno}, ${outputno}") }
|
|
||||||
//
|
|
||||||
// /* -------- 储存神经网络 -------- */
|
|
||||||
// func save(file: File) {
|
|
||||||
// file.write("FullConnectLayer ".toArray())
|
|
||||||
// file.write(inputno.toString().toArray())
|
|
||||||
// file.write(" ".toArray())
|
|
||||||
// file.write(outputno.toString().toArray())
|
|
||||||
// file.write(" ".toArray())
|
|
||||||
// file.write(learning_rate.toString().toArray())
|
|
||||||
// file.write(" ".toArray())
|
|
||||||
// file.write(f.toString().toArray())
|
|
||||||
// file.write("\n".toArray())
|
|
||||||
// weights.save(file)
|
|
||||||
// //file.write("\n".toArray())
|
|
||||||
// bias.save(file)
|
|
||||||
// //file.write("\n".toArray())
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 静态方法,加载单个项目 -------- */
|
|
||||||
// static func loadItem<T>(sr: StringReader<InputStream>) where T <: Parsable<T>
|
|
||||||
// {
|
|
||||||
// let s = sr.readUntil({r => r == r" " || r == r"\n"})
|
|
||||||
// let ss = s.getOrThrow().trimAscii()
|
|
||||||
// println(ss)
|
|
||||||
// return T.parse(ss)
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 静态方法,加载神经网络 -------- */
|
|
||||||
// static func load(sr: StringReader<InputStream>)
|
|
||||||
// {
|
|
||||||
// let inputno = loadItem<UInt64>(sr)
|
|
||||||
// let outputno = loadItem<UInt64>(sr)
|
|
||||||
// let lr = loadItem<Float32>(sr)
|
|
||||||
// let fna = sr.readUntil(r'\n')
|
|
||||||
// let f = ActivationFuncs.parse(fna.getOrThrow().trimAscii())
|
|
||||||
// let weights = Matrix.load(sr, inputno, outputno)
|
|
||||||
// let bias = Matrix.load(sr, 1, outputno)
|
|
||||||
// Nn(inputno, outputno, f, lr, weights, bias)
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// let inputno: UInt64
|
|
||||||
// let outputno: UInt64
|
|
||||||
// var weights: Matrix /* (inputs, outputs) */
|
|
||||||
// var bias: Matrix /* (1, outputs) */
|
|
||||||
// var learning_rate: Float32
|
|
||||||
// let f: ActivationFuncs
|
|
||||||
//}
|
|
||||||
@@ -1,164 +0,0 @@
|
|||||||
package bot4cj.math
|
|
||||||
import std.io.*
|
|
||||||
import std.fs.*
|
|
||||||
import std.math.*
|
|
||||||
import std.time.*
|
|
||||||
import std.random.*
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
/* -------- 准备与门的测试用例 --------
|
|
||||||
* INFO(20260524-qing): 输入{2},输出{1} */
|
|
||||||
public func prepare_andDoor(m: Matrix1, t: Matrix1) {
|
|
||||||
let random = Random(UInt64(DateTime.now().nanosecond))
|
|
||||||
let a = random.nextInt8(2)
|
|
||||||
let b = random.nextInt8(2)
|
|
||||||
let t0 = if ((a == 0 && b == 0) || (a ==1 && b == 1)) { 1.0f32 } else { 0.0f32 }
|
|
||||||
//let t1 = if (a != b) { 1.0f32 } else { 0.0f32 }
|
|
||||||
m.set(0, 0, FP32(Float32(a)))
|
|
||||||
m.set(0, 1, FP32(Float32(b)))
|
|
||||||
t.set(0, 0, FP32(t0))
|
|
||||||
//t.set(0, 1, FP32(t1))
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- 新版神经网络推理,使用原生矩阵 -------- */
|
|
||||||
public func forward1(nns: ArrayList<Nn1>, m: Matrix1) {
|
|
||||||
var out_dof = Matrix1(1, 1) /* ---- 给一个值占位 ----*/
|
|
||||||
for (idx in 0..nns.size) {
|
|
||||||
let x = if (idx == 0) { m } else { out_dof }
|
|
||||||
out_dof= nns[idx].forward(x)
|
|
||||||
}
|
|
||||||
return out_dof
|
|
||||||
}
|
|
||||||
|
|
||||||
/* 【神经网络训练】
|
|
||||||
* description: 输入一组原始数据和教师数据,进行神经网络训练
|
|
||||||
* @param nns: 输入神经网络结构
|
|
||||||
* @param m: 输入一组原始数据
|
|
||||||
* @param t: 输入一组教师数据
|
|
||||||
* @param discount: 输入本次训练的折扣系数
|
|
||||||
*/
|
|
||||||
public func train1(nns: ArrayList<Nn1>, m: Matrix1, t: Matrix1, discount: Float32) {
|
|
||||||
|
|
||||||
/* ======== 这个对象用于暂时占位 ======== */
|
|
||||||
var res = Matrix1(1, 1)
|
|
||||||
|
|
||||||
/* ======== 正向传播 ======== */
|
|
||||||
for (idx in 0..nns.size) {
|
|
||||||
let x = if (idx == 0) { m } else { res }
|
|
||||||
res = nns[idx].forward(x)
|
|
||||||
}
|
|
||||||
|
|
||||||
/* ======== 计算误差 ======== */
|
|
||||||
var errs = t - res
|
|
||||||
|
|
||||||
/* ======== 反向传播 ======== */
|
|
||||||
for (idx in (nns.size-1)..=0:-1) {
|
|
||||||
errs = nns[idx].backward(errs)
|
|
||||||
}
|
|
||||||
|
|
||||||
/* ======== 梯度更新 ======== */
|
|
||||||
for (idx in 0..nns.size) {
|
|
||||||
nns[idx].update(discount)
|
|
||||||
}
|
|
||||||
|
|
||||||
/* ======== 计算MSE损失 ========
|
|
||||||
FIXME*(20260517-qing): 对于多个结果计算误差(目前只计算了第一个结果) */
|
|
||||||
|
|
||||||
let t0 = t.get(0, 0).getOrThrow().getValue()
|
|
||||||
//let t1 = t.get(0, 1).getOrThrow().getValue()
|
|
||||||
let r0 = res.get(0, 0).getOrThrow().getValue()
|
|
||||||
//let r1 = res.get(0, 1).getOrThrow().getValue()
|
|
||||||
let loss0 = 0.5 * (t0 - r0) * (t0 - r0)
|
|
||||||
//let loss1 = 0.5 * (t1 - r1) * (t1 - r1)
|
|
||||||
//return (loss0 + loss1) / 2.0
|
|
||||||
return loss0
|
|
||||||
}
|
|
||||||
|
|
||||||
/* -------- (旧版)加载多层神经网络 --------
|
|
||||||
TODO: 也许用类封装
|
|
||||||
INFO(20260517-qing): 新版使用仓颉原生矩阵,然后与MatrixObject类互转 */
|
|
||||||
//public func load_nns(filename: String): ArrayList<Nn> {
|
|
||||||
// let nns = ArrayList<Nn>()
|
|
||||||
// let file = File(filename, OpenMode.Read)
|
|
||||||
// let sr = StringReader<InputStream>(file)
|
|
||||||
// while (true) {
|
|
||||||
// if (let Some(S) <- sr.readUntil(r' ')) {
|
|
||||||
// var S2 = S.trimAscii()
|
|
||||||
// if (S2 == "FullConnectLayer") {
|
|
||||||
// nns.add(Nn.load(sr)) }
|
|
||||||
// else if (S2 == "") { break }
|
|
||||||
// else { throw IllegalArgumentException() }
|
|
||||||
// }
|
|
||||||
// else { break }
|
|
||||||
// }
|
|
||||||
|
|
||||||
// file.close()
|
|
||||||
// return nns
|
|
||||||
//}
|
|
||||||
|
|
||||||
/* -------- (旧版)储存多层神经网络 -------- */
|
|
||||||
//public func save_nns(nns: ArrayList<Nn>, filename: String) {
|
|
||||||
//
|
|
||||||
// let file: File
|
|
||||||
// if (exists(filename)) {
|
|
||||||
// file = File(filename, OpenMode.Write) }
|
|
||||||
// else { file= File.create(filename) }
|
|
||||||
// for (nn in nns) { /* 存储全部的层 */
|
|
||||||
// nn.save(file) }
|
|
||||||
// file.close()
|
|
||||||
//}
|
|
||||||
|
|
||||||
/* -------- 旧版神经网络推理 -------- */
|
|
||||||
//public func forward(nns: ArrayList<Nn>, m: Matrix) {
|
|
||||||
// var out_dof = Matrix(1, 1)
|
|
||||||
// for (idx in 0..nns.size) {
|
|
||||||
// let x = if (idx == 0) { m } else { out_dof }
|
|
||||||
// let res = nns[idx].forward(x)
|
|
||||||
// out_dof = nns[idx].fdo(res)
|
|
||||||
// }
|
|
||||||
// return out_dof
|
|
||||||
//}
|
|
||||||
|
|
||||||
/* -------- 旧版的神经网络训练 -------- */
|
|
||||||
//public func train(nns: ArrayList<Nn>, m: Matrix, t: Matrix, discount: Float32) {
|
|
||||||
//
|
|
||||||
// /* ----------------- */
|
|
||||||
// let outs_dof = ArrayList<Matrix> ()
|
|
||||||
// let outs_dotf = ArrayList<Matrix> ()
|
|
||||||
//
|
|
||||||
// /* -------- 正向传播 -------- */
|
|
||||||
// for (idx in 0..nns.size) {
|
|
||||||
// let x = if (idx == 0) { m } else { outs_dof[idx-1] }
|
|
||||||
// let res = nns[idx].forward(x)
|
|
||||||
// let out_dof = nns[idx].fdo(res)
|
|
||||||
// let out_dotf = nns[idx].fdot(out_dof)
|
|
||||||
// outs_dof.add(out_dof)
|
|
||||||
// outs_dotf.add(out_dotf)
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 计算误差 -------- */
|
|
||||||
// var errs = outs_dof[nns.size-1] - t
|
|
||||||
//
|
|
||||||
// /* -------- 反向传播 -------- */
|
|
||||||
// for (idx in (nns.size-1)..=0:-1) {
|
|
||||||
// let up = if (idx == 0) { m} else { outs_dof[idx-1] }
|
|
||||||
// let bias_grad = errs
|
|
||||||
// //let delta = nns[idx].bak_bsgrad(errs, outs_dotf[idx])
|
|
||||||
// let weights_grad = nns[idx].bak_wsgrad(up, bias_grad)
|
|
||||||
// if (idx > 0) {
|
|
||||||
// errs = nns[idx].bak_uperrs(bias_grad, outs_dotf[idx-1]) }
|
|
||||||
// nns[idx].update(weights_grad, bias_grad, discount)
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
// /* -------- 输出损失 -------- */
|
|
||||||
// unsafe {
|
|
||||||
// let t0 = t.get(0, 0)
|
|
||||||
// let t1 = t.get(0, 1)
|
|
||||||
// let r0 = outs_dof[nns.size-1].get(0, 0)
|
|
||||||
// let r1 = outs_dof[nns.size-1].get(0, 1)
|
|
||||||
// let loss0 = -t0 * log(r0) - (1.0f32-t0) * log(1.0f32-r0);
|
|
||||||
// let loss1 = -t1 * log(r1) - (1.0f32-t1) * log(1.0f32-r1);
|
|
||||||
// return 0.5 * (loss0 * loss0 + loss1 * loss1)
|
|
||||||
// }
|
|
||||||
//
|
|
||||||
//}
|
|
||||||
@@ -1,418 +0,0 @@
|
|||||||
// pooling_layer.cj
|
|
||||||
package bot4cj.math
|
|
||||||
|
|
||||||
import std.math.*
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 池化层类
|
|
||||||
* 设计风格与 Nn1 和 ConvLayer 保持一致
|
|
||||||
*/
|
|
||||||
public class PoolingLayer {
|
|
||||||
|
|
||||||
// ==================== 层配置参数 ====================
|
|
||||||
let nChannels: Int64
|
|
||||||
let poolHeight: Int64
|
|
||||||
let poolWidth: Int64
|
|
||||||
let stride: Int64
|
|
||||||
let padding: Int64
|
|
||||||
let poolType: String
|
|
||||||
let f: Af
|
|
||||||
let learningRateWeight: FP32
|
|
||||||
let learningRateBias: FP32
|
|
||||||
|
|
||||||
// ==================== 可学习参数 ====================
|
|
||||||
var weights: Matrix1
|
|
||||||
var biases: Matrix1
|
|
||||||
|
|
||||||
// ==================== 运行时缓存 ====================
|
|
||||||
var prevInput: Option<Matrix1> = Option.None
|
|
||||||
var prevInputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
|
||||||
var prevCombination: Option<Matrix1> = Option.None
|
|
||||||
var prevOutput: Option<Matrix1> = Option.None
|
|
||||||
var prevOutputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
|
|
||||||
var prevMaxIndices: Option<Matrix1> = Option.None
|
|
||||||
|
|
||||||
var weightsGrad: Option<Matrix1> = Option.None
|
|
||||||
var biasGrad: Option<Matrix1> = Option.None
|
|
||||||
|
|
||||||
// ==================== 构造函数 ====================
|
|
||||||
public init(
|
|
||||||
nChannels: Int64,
|
|
||||||
poolSize: (Int64, Int64),
|
|
||||||
stride: Int64,
|
|
||||||
padding: Int64,
|
|
||||||
poolType: String,
|
|
||||||
activation: Af,
|
|
||||||
lrWeight: FP32,
|
|
||||||
lrBias: FP32
|
|
||||||
) {
|
|
||||||
this.nChannels = nChannels
|
|
||||||
poolHeight = poolSize[0]
|
|
||||||
poolWidth = poolSize[1]
|
|
||||||
this.stride = stride
|
|
||||||
this.padding = padding
|
|
||||||
this.poolType = poolType
|
|
||||||
this.f = activation
|
|
||||||
learningRateWeight = lrWeight
|
|
||||||
learningRateBias = lrBias
|
|
||||||
|
|
||||||
// 初始化权值
|
|
||||||
weights = Matrix1(1, nChannels)
|
|
||||||
if (poolType == "average") {
|
|
||||||
let avgWeight = FP32(1.0) / FP32(Float32(poolHeight * poolWidth))
|
|
||||||
for (c in 0..nChannels) {
|
|
||||||
weights.set(0, c, avgWeight)
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
for (c in 0..nChannels) {
|
|
||||||
weights.set(0, c, FP32(1.0))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 初始化偏置为 0
|
|
||||||
biases = Matrix1(1, nChannels)
|
|
||||||
for (c in 0..nChannels) {
|
|
||||||
biases.set(0, c, FP32(0.0))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
public init(
|
|
||||||
nChannels: Int64,
|
|
||||||
poolSize: (Int64, Int64),
|
|
||||||
poolType: String,
|
|
||||||
activation: Af
|
|
||||||
) {
|
|
||||||
this.nChannels = nChannels
|
|
||||||
poolHeight = poolSize[0]
|
|
||||||
poolWidth = poolSize[1]
|
|
||||||
stride = poolHeight
|
|
||||||
padding = 0
|
|
||||||
this.poolType = poolType
|
|
||||||
this.f = activation
|
|
||||||
learningRateWeight = FP32(0.1)
|
|
||||||
learningRateBias = FP32(0.1)
|
|
||||||
|
|
||||||
weights = Matrix1(1, nChannels)
|
|
||||||
if (poolType == "average") {
|
|
||||||
let avgWeight = FP32(1.0) / FP32(Float32(poolHeight * poolWidth))
|
|
||||||
for (c in 0..nChannels) {
|
|
||||||
weights.set(0, c, avgWeight)
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
for (c in 0..nChannels) {
|
|
||||||
weights.set(0, c, FP32(1.0))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
biases = Matrix1(1, nChannels)
|
|
||||||
for (c in 0..nChannels) {
|
|
||||||
biases.set(0, c, FP32(0.0))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 前向传播 ====================
|
|
||||||
public func forward(x: Matrix1, inputShape: (Int64, Int64, Int64, Int64)): Matrix1 {
|
|
||||||
prevInput = Option.Some(x)
|
|
||||||
prevInputShape = Option.Some(inputShape)
|
|
||||||
|
|
||||||
let batch = inputShape[0]
|
|
||||||
let inH = inputShape[1]
|
|
||||||
let inW = inputShape[2]
|
|
||||||
|
|
||||||
let outH = (inH + 2 * padding - poolHeight) / stride + 1
|
|
||||||
let outW = (inW + 2 * padding - poolWidth) / stride + 1
|
|
||||||
let outSize = batch * outH * outW
|
|
||||||
|
|
||||||
var combination = Matrix1(outSize, nChannels)
|
|
||||||
var output = Matrix1(outSize, nChannels)
|
|
||||||
var maxIndices = Matrix1(outSize, nChannels)
|
|
||||||
|
|
||||||
for (b in 0..batch) {
|
|
||||||
for (oh in 0..outH) {
|
|
||||||
for (ow in 0..outW) {
|
|
||||||
let outRow = b * outH * outW + oh * outW + ow
|
|
||||||
|
|
||||||
for (ch in 0..nChannels) {
|
|
||||||
var pooledVal = FP32(0.0)
|
|
||||||
var maxIdx = 0
|
|
||||||
|
|
||||||
if (poolType == "max") {
|
|
||||||
var maxVal = FP32(-1.0e30)
|
|
||||||
var kernelIdx = 0
|
|
||||||
for (kh in 0..poolHeight) {
|
|
||||||
for (kw in 0..poolWidth) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
let inRow = b * inH * inW + ih * inW + iw
|
|
||||||
let val = x.get(inRow, ch).getOrThrow()
|
|
||||||
if (val.getValue() > maxVal.getValue()) {
|
|
||||||
maxVal = val
|
|
||||||
maxIdx = kernelIdx
|
|
||||||
}
|
|
||||||
}
|
|
||||||
kernelIdx += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
pooledVal = maxVal
|
|
||||||
} else {
|
|
||||||
var sum = FP32(0.0)
|
|
||||||
var count = 0
|
|
||||||
for (kh in 0..poolHeight) {
|
|
||||||
for (kw in 0..poolWidth) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
let inRow = b * inH * inW + ih * inW + iw
|
|
||||||
sum = sum + x.get(inRow, ch).getOrThrow()
|
|
||||||
count += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if (count > 0) {
|
|
||||||
pooledVal = sum / FP32(Float32(count))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
combination.set(outRow, ch, pooledVal)
|
|
||||||
maxIndices.set(outRow, ch, FP32(Float32(maxIdx)))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 逐通道应用权值和偏置
|
|
||||||
for (r in 0..outSize) {
|
|
||||||
for (ch in 0..nChannels) {
|
|
||||||
let val = combination.get(r, ch).getOrThrow()
|
|
||||||
let w = weights.get(0, ch).getOrThrow()
|
|
||||||
let b = biases.get(0, ch).getOrThrow()
|
|
||||||
combination.set(r, ch, val * w + b)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
prevCombination = Option.Some(combination)
|
|
||||||
|
|
||||||
output = f.forward(combination)
|
|
||||||
prevOutput = Option.Some(output)
|
|
||||||
prevOutputShape = Option.Some((batch, outH, outW, nChannels))
|
|
||||||
prevMaxIndices = Option.Some(maxIndices)
|
|
||||||
|
|
||||||
return output
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 反向传播 ====================
|
|
||||||
public func backward(errs: Matrix1): Matrix1 {
|
|
||||||
var x = Matrix1(1, 1)
|
|
||||||
match (prevInput) {
|
|
||||||
case Some(val) => x = val
|
|
||||||
case _ => throw Exception("No input data")
|
|
||||||
}
|
|
||||||
|
|
||||||
var shape = (0, 0, 0, 0)
|
|
||||||
match (prevInputShape) {
|
|
||||||
case Some(val) => shape = val
|
|
||||||
case _ => throw Exception("No input shape")
|
|
||||||
}
|
|
||||||
|
|
||||||
var combination = Matrix1(1, 1)
|
|
||||||
match (prevCombination) {
|
|
||||||
case Some(val) => combination = val
|
|
||||||
case _ => throw Exception("No combination")
|
|
||||||
}
|
|
||||||
|
|
||||||
var outShape = (0, 0, 0, 0)
|
|
||||||
match (prevOutputShape) {
|
|
||||||
case Some(val) => outShape = val
|
|
||||||
case _ => throw Exception("No output shape")
|
|
||||||
}
|
|
||||||
|
|
||||||
var maxIndices = Matrix1(1, 1)
|
|
||||||
let hasMaxIndices = match (prevMaxIndices) {
|
|
||||||
case Some(_) => true
|
|
||||||
case _ => false
|
|
||||||
}
|
|
||||||
if (hasMaxIndices) {
|
|
||||||
maxIndices = match (prevMaxIndices) {
|
|
||||||
case Some(val) => val
|
|
||||||
case _ => Matrix1(1, 1)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let batch = shape[0]
|
|
||||||
let inH = shape[1]
|
|
||||||
let inW = shape[2]
|
|
||||||
let outH = outShape[1]
|
|
||||||
let outW = outShape[2]
|
|
||||||
|
|
||||||
let delta = errs * f.backward(combination)
|
|
||||||
|
|
||||||
calcWeightBiasGradient(delta, combination)
|
|
||||||
|
|
||||||
let errsUp = calcUpsampledError(delta, maxIndices, batch, inH, inW, outH, outW)
|
|
||||||
|
|
||||||
return errsUp
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算权值和偏置梯度
|
|
||||||
*/
|
|
||||||
private func calcWeightBiasGradient(delta: Matrix1, combination: Matrix1): Unit {
|
|
||||||
let outSize = delta.get_row()
|
|
||||||
|
|
||||||
var wGrad = Matrix1(1, nChannels)
|
|
||||||
var bGrad = Matrix1(1, nChannels)
|
|
||||||
|
|
||||||
for (ch in 0..nChannels) {
|
|
||||||
var wSum = FP32(0.0)
|
|
||||||
var bSum = FP32(0.0)
|
|
||||||
for (r in 0..outSize) {
|
|
||||||
let d = delta.get(r, ch).getOrThrow()
|
|
||||||
bSum = bSum + d
|
|
||||||
let comb = combination.get(r, ch).getOrThrow()
|
|
||||||
wSum = wSum + d * comb
|
|
||||||
}
|
|
||||||
wGrad.set(0, ch, wSum)
|
|
||||||
bGrad.set(0, ch, bSum)
|
|
||||||
}
|
|
||||||
|
|
||||||
weightsGrad = Option.Some(wGrad)
|
|
||||||
biasGrad = Option.Some(bGrad)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 上采样误差(池化层的反向传播核心)
|
|
||||||
*/
|
|
||||||
private func calcUpsampledError(
|
|
||||||
delta: Matrix1,
|
|
||||||
maxIndices: Matrix1,
|
|
||||||
batch: Int64,
|
|
||||||
inH: Int64,
|
|
||||||
inW: Int64,
|
|
||||||
outH: Int64,
|
|
||||||
outW: Int64
|
|
||||||
): Matrix1 {
|
|
||||||
let inSize = batch * inH * inW
|
|
||||||
var errsUp = Matrix1(inSize, nChannels)
|
|
||||||
|
|
||||||
for (b in 0..batch) {
|
|
||||||
for (oh in 0..outH) {
|
|
||||||
for (ow in 0..outW) {
|
|
||||||
let deltaRow = b * outH * outW + oh * outW + ow
|
|
||||||
|
|
||||||
for (ch in 0..nChannels) {
|
|
||||||
let deltaVal = delta.get(deltaRow, ch).getOrThrow()
|
|
||||||
let w = weights.get(0, ch).getOrThrow()
|
|
||||||
|
|
||||||
if (poolType == "max") {
|
|
||||||
let maxIdx = Int64(maxIndices.get(deltaRow, ch).getOrThrow().getValue())
|
|
||||||
|
|
||||||
var kernelIdx = 0
|
|
||||||
var found = false
|
|
||||||
for (kh in 0..poolHeight) {
|
|
||||||
for (kw in 0..poolWidth) {
|
|
||||||
if (kernelIdx == maxIdx && !found) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
let inRow = b * inH * inW + ih * inW + iw
|
|
||||||
errsUp.set(inRow, ch, deltaVal * w)
|
|
||||||
found = true
|
|
||||||
}
|
|
||||||
}
|
|
||||||
kernelIdx += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
var count = 0
|
|
||||||
for (kh in 0..poolHeight) {
|
|
||||||
for (kw in 0..poolWidth) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
count += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if (count > 0) {
|
|
||||||
let avgDelta = deltaVal * w / FP32(Float32(count))
|
|
||||||
for (kh in 0..poolHeight) {
|
|
||||||
for (kw in 0..poolWidth) {
|
|
||||||
let ih = oh * stride + kh - padding
|
|
||||||
let iw = ow * stride + kw - padding
|
|
||||||
if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
|
|
||||||
let inRow = b * inH * inW + ih * inW + iw
|
|
||||||
let oldVal = errsUp.get(inRow, ch).getOrThrow()
|
|
||||||
errsUp.set(inRow, ch, oldVal + avgDelta)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return errsUp
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 参数更新 ====================
|
|
||||||
public func update(discount: Float32): Unit {
|
|
||||||
var wGrad = Matrix1(1, 1)
|
|
||||||
match (weightsGrad) {
|
|
||||||
case Some(val) => wGrad = val
|
|
||||||
case _ => throw Exception("No weight gradients")
|
|
||||||
}
|
|
||||||
|
|
||||||
var bGrad = Matrix1(1, 1)
|
|
||||||
match (biasGrad) {
|
|
||||||
case Some(val) => bGrad = val
|
|
||||||
case _ => throw Exception("No bias gradients")
|
|
||||||
}
|
|
||||||
|
|
||||||
for (ch in 0..nChannels) {
|
|
||||||
let oldW = weights.get(0, ch).getOrThrow()
|
|
||||||
let gW = wGrad.get(0, ch).getOrThrow()
|
|
||||||
weights.set(0, ch, oldW + gW * learningRateWeight * FP32(discount))
|
|
||||||
}
|
|
||||||
|
|
||||||
for (ch in 0..nChannels) {
|
|
||||||
let oldB = biases.get(0, ch).getOrThrow()
|
|
||||||
let gB = bGrad.get(0, ch).getOrThrow()
|
|
||||||
biases.set(0, ch, oldB + gB * learningRateBias * FP32(discount))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ==================== 形状访问器 ====================
|
|
||||||
public func getInShape(): (Int64, Int64) {
|
|
||||||
(nChannels, poolHeight * poolWidth)
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getOutShape(): Int64 {
|
|
||||||
nChannels
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getPoolType(): String {
|
|
||||||
poolType
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getLR(): FP32 {
|
|
||||||
learningRateWeight
|
|
||||||
}
|
|
||||||
|
|
||||||
public func getFType(): String {
|
|
||||||
f.get_type()
|
|
||||||
}
|
|
||||||
|
|
||||||
public func calculateOutputSize(inputHeight: Int64, inputWidth: Int64): (Int64, Int64) {
|
|
||||||
let outH = (inputHeight + 2 * padding - poolHeight) / stride + 1
|
|
||||||
let outW = (inputWidth + 2 * padding - poolWidth) / stride + 1
|
|
||||||
(outH, outW)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,310 +0,0 @@
|
|||||||
// test_cnn.cj - 使用 Adam 优化器的版本(所有层)
|
|
||||||
package bot4cj
|
|
||||||
import bot4cj.math.*
|
|
||||||
import std.math.*
|
|
||||||
import std.random.*
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
func generateSolidColorData(numSamples: Int64): (Array<Matrix1>, Array<Int64>) {
|
|
||||||
var inputs = Array(numSamples, { _ => Matrix1(1, 1) })
|
|
||||||
var labels = Array(numSamples, { _ => 0 })
|
|
||||||
let rng = Random()
|
|
||||||
|
|
||||||
for (i in 0..numSamples) {
|
|
||||||
let colorClass = Int64(rng.nextFloat64() * 3.0)
|
|
||||||
labels[i] = colorClass
|
|
||||||
|
|
||||||
let h = 8
|
|
||||||
let w = 8
|
|
||||||
var matrix = Matrix1(h * w, 1)
|
|
||||||
|
|
||||||
let colorValue = match (colorClass) {
|
|
||||||
case 0 => FP32(1.0)
|
|
||||||
case 1 => FP32(0.5)
|
|
||||||
case _ => FP32(0.0)
|
|
||||||
}
|
|
||||||
|
|
||||||
for (r in 0..h * w) {
|
|
||||||
let noise = FP32(Float32(rng.nextFloat64() - 0.5) * 0.1)
|
|
||||||
var val = colorValue + noise
|
|
||||||
if (val.getValue() > 1.0) {
|
|
||||||
val = FP32(1.0)
|
|
||||||
} else if (val.getValue() < 0.0) {
|
|
||||||
val = FP32(0.0)
|
|
||||||
}
|
|
||||||
matrix.set(r, 0, val)
|
|
||||||
}
|
|
||||||
|
|
||||||
inputs[i] = matrix
|
|
||||||
}
|
|
||||||
|
|
||||||
(inputs, labels)
|
|
||||||
}
|
|
||||||
|
|
||||||
func printSeparator(): Unit {
|
|
||||||
println("==================================================")
|
|
||||||
}
|
|
||||||
|
|
||||||
func testCNN(): Unit {
|
|
||||||
printSeparator()
|
|
||||||
println("CNN 纯色矩阵识别测试 (所有层使用 Adam 优化器)")
|
|
||||||
printSeparator()
|
|
||||||
println()
|
|
||||||
|
|
||||||
let nTrainSamples = 200
|
|
||||||
let nTestSamples = 60
|
|
||||||
let nEpochs = 50
|
|
||||||
let batchSize = 10
|
|
||||||
|
|
||||||
println("训练样本数: ${nTrainSamples}")
|
|
||||||
println("测试样本数: ${nTestSamples}")
|
|
||||||
println("训练轮数: ${nEpochs}")
|
|
||||||
println("Batch 大小: ${batchSize}")
|
|
||||||
println()
|
|
||||||
|
|
||||||
println("创建网络结构...")
|
|
||||||
println("Input(8x8x1) -> Conv(3x3,1->4) -> Pool(2x2) -> Conv(3x3,4->8) -> Flatten -> FC(8,3)")
|
|
||||||
println()
|
|
||||||
|
|
||||||
let lr = FP32(0.001)
|
|
||||||
|
|
||||||
var conv1 = ConvLayer(1, 4, (3, 3), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
|
||||||
var pool1 = PoolingLayer(4, (2, 2), 2, 0, "max", Af.ReLU, lr, lr)
|
|
||||||
var conv2 = ConvLayer(4, 8, (3, 3), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
|
||||||
var flatten = FlattenLayer(8, 1, 1)
|
|
||||||
var fc = Nn1(8, 3, Float32(0.4), Af.Softmax)
|
|
||||||
|
|
||||||
// 创建 Adam 优化器(为每个可训练层创建独立的优化器)
|
|
||||||
var adamConv1 = Adam(FP32(0.001)) // 卷积层1
|
|
||||||
var adamConv2 = Adam(FP32(0.001)) // 卷积层2
|
|
||||||
var adamFC = Adam(FP32(0.0005)) // 全连接层(使用更小的学习率)
|
|
||||||
|
|
||||||
println("网络创建完成 (Conv1, Conv2, FC 使用 Adam 优化器)")
|
|
||||||
println()
|
|
||||||
|
|
||||||
println("生成数据...")
|
|
||||||
let (trainInputs, trainLabels) = generateSolidColorData(nTrainSamples)
|
|
||||||
let (testInputs, testLabels) = generateSolidColorData(nTestSamples)
|
|
||||||
println("训练集: ${nTrainSamples} 样本")
|
|
||||||
println("测试集: ${nTestSamples} 样本")
|
|
||||||
println()
|
|
||||||
|
|
||||||
println("开始训练...")
|
|
||||||
println()
|
|
||||||
|
|
||||||
for (epoch in 0..nEpochs) {
|
|
||||||
var epochLoss = Float32(0.0)
|
|
||||||
var epochCorrect = 0
|
|
||||||
|
|
||||||
for (i in 0..nTrainSamples) {
|
|
||||||
let x = trainInputs[i]
|
|
||||||
let label = trainLabels[i]
|
|
||||||
|
|
||||||
// ===== 前向传播 =====
|
|
||||||
let c1Out = conv1.forward(x, (1, 8, 8, 1))
|
|
||||||
let p1Out = pool1.forward(c1Out, (1, 6, 6, 4))
|
|
||||||
let c2Out = conv2.forward(p1Out, (1, 3, 3, 4))
|
|
||||||
let flatOut = flatten.forward(c2Out, (1, 1, 1, 8))
|
|
||||||
|
|
||||||
// 调试:确认 flatten.forward 执行成功
|
|
||||||
if (i == 0 && epoch == 0) {
|
|
||||||
println("DEBUG: flatten.forward done, flatOut rows: ${flatOut.get_row()}, cols: ${flatOut.get_col()}")
|
|
||||||
}
|
|
||||||
|
|
||||||
var fcInput = Matrix1(1, 8)
|
|
||||||
for (j in 0..8) {
|
|
||||||
fcInput.set(0, j, flatOut.get(j, 0).getOrThrow())
|
|
||||||
}
|
|
||||||
let fcOut = fc.forward(fcInput)
|
|
||||||
|
|
||||||
// Softmax
|
|
||||||
var softmax = Matrix1(1, 3)
|
|
||||||
var maxVal = Float32(-1.0e30)
|
|
||||||
for (c in 0..3) {
|
|
||||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
|
||||||
if (raw > maxVal) {
|
|
||||||
maxVal = raw
|
|
||||||
}
|
|
||||||
}
|
|
||||||
var expSum = Float32(0.0)
|
|
||||||
for (c in 0..3) {
|
|
||||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
|
||||||
let ev = Float32(exp(raw - maxVal))
|
|
||||||
softmax.set(0, c, FP32(ev))
|
|
||||||
expSum = expSum + ev
|
|
||||||
}
|
|
||||||
for (c in 0..3) {
|
|
||||||
let v = softmax.get(0, c).getOrThrow().getValue()
|
|
||||||
softmax.set(0, c, FP32(v / expSum))
|
|
||||||
}
|
|
||||||
|
|
||||||
let probVal = softmax.get(0, label).getOrThrow().getValue()
|
|
||||||
let loss = -log(max(probVal, Float32(1e-10)))
|
|
||||||
epochLoss = epochLoss + loss
|
|
||||||
|
|
||||||
var maxProb = Float32(-1.0)
|
|
||||||
var predClass = 0
|
|
||||||
for (c in 0..3) {
|
|
||||||
let p = softmax.get(0, c).getOrThrow().getValue()
|
|
||||||
if (p > maxProb) {
|
|
||||||
maxProb = p
|
|
||||||
predClass = Int64(c)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if (predClass == label) {
|
|
||||||
epochCorrect += 1
|
|
||||||
}
|
|
||||||
|
|
||||||
// ===== 反向传播 =====
|
|
||||||
var delta = Matrix1(1, 3)
|
|
||||||
for (c in 0..3) {
|
|
||||||
let prob = softmax.get(0, c).getOrThrow()
|
|
||||||
if (Int64(c) == label) {
|
|
||||||
delta.set(0, c, prob - FP32(1.0))
|
|
||||||
} else {
|
|
||||||
delta.set(0, c, prob)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let fcDelta = fc.backward(delta)
|
|
||||||
|
|
||||||
// 调试:确认 fcDelta 形状
|
|
||||||
if (i == 0 && epoch == 0) {
|
|
||||||
println("DEBUG: fc.backward done, fcDelta rows: ${fcDelta.get_row()}, cols: ${fcDelta.get_col()}")
|
|
||||||
}
|
|
||||||
|
|
||||||
var flatDelta = Matrix1(8, 1)
|
|
||||||
for (j in 0..8) {
|
|
||||||
flatDelta.set(j, 0, fcDelta.get(0, j).getOrThrow())
|
|
||||||
}
|
|
||||||
|
|
||||||
// 调试:确认 flatDelta 形状
|
|
||||||
if (i == 0 && epoch == 0) {
|
|
||||||
println("DEBUG: flatDelta rows: ${flatDelta.get_row()}, cols: ${flatDelta.get_col()}")
|
|
||||||
}
|
|
||||||
|
|
||||||
let c2Delta = flatten.backward(flatDelta)
|
|
||||||
let p1Delta = conv2.backward(c2Delta)
|
|
||||||
let c1Delta = pool1.backward(p1Delta)
|
|
||||||
let _ = conv1.backward(c1Delta)
|
|
||||||
|
|
||||||
// ===== 参数更新 =====
|
|
||||||
if ((i + 1) % batchSize == 0) {
|
|
||||||
// 所有可训练层使用 Adam 更新
|
|
||||||
adamConv1.updateConvLayer(conv1) // 卷积层1 - Adam
|
|
||||||
adamConv2.updateConvLayer(conv2) // 卷积层2 - Adam
|
|
||||||
adamFC.updateFCLayer(fc) // 全连接层 - Adam
|
|
||||||
|
|
||||||
// 池化层使用 SGD 更新(如果需要的话)
|
|
||||||
pool1.update(Float32(1.0))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Epoch 结束时的更新
|
|
||||||
adamConv1.updateConvLayer(conv1)
|
|
||||||
adamConv2.updateConvLayer(conv2)
|
|
||||||
adamFC.updateFCLayer(fc)
|
|
||||||
pool1.update(Float32(1.0))
|
|
||||||
|
|
||||||
let avgLoss = epochLoss / Float32(nTrainSamples)
|
|
||||||
let accuracy = Float32(epochCorrect) / Float32(nTrainSamples) * Float32(100.0)
|
|
||||||
|
|
||||||
println("Epoch ${epoch + 1}/${nEpochs} - Loss: ${avgLoss}, Accuracy: ${accuracy}%")
|
|
||||||
|
|
||||||
// 打印 Adam 的步数信息
|
|
||||||
if (epoch == 0) {
|
|
||||||
println(" Adam Step: ${adamConv1.getStep()}")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
println("")
|
|
||||||
println("开始测试...")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
var testCorrect = 0
|
|
||||||
var confusion = Array(3, { _ => Array(3, { _ => 0 }) })
|
|
||||||
|
|
||||||
for (i in 0..nTestSamples) {
|
|
||||||
let x = testInputs[i]
|
|
||||||
let label = testLabels[i]
|
|
||||||
|
|
||||||
let c1Out = conv1.forward(x, (1, 8, 8, 1))
|
|
||||||
let p1Out = pool1.forward(c1Out, (1, 6, 6, 4))
|
|
||||||
let c2Out = conv2.forward(p1Out, (1, 3, 3, 4))
|
|
||||||
let flatOut = flatten.forward(c2Out, (1, 1, 1, 8))
|
|
||||||
|
|
||||||
var fcInput = Matrix1(1, 8)
|
|
||||||
for (j in 0..8) {
|
|
||||||
fcInput.set(0, j, flatOut.get(j, 0).getOrThrow())
|
|
||||||
}
|
|
||||||
let fcOut = fc.forward(fcInput)
|
|
||||||
|
|
||||||
var softmax = Matrix1(1, 3)
|
|
||||||
var maxVal = Float32(-1.0e30)
|
|
||||||
for (c in 0..3) {
|
|
||||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
|
||||||
if (raw > maxVal) {
|
|
||||||
maxVal = raw
|
|
||||||
}
|
|
||||||
}
|
|
||||||
var expSum = Float32(0.0)
|
|
||||||
for (c in 0..3) {
|
|
||||||
let raw = fcOut.get(0, c).getOrThrow().getValue()
|
|
||||||
let ev = Float32(exp(raw - maxVal))
|
|
||||||
softmax.set(0, c, FP32(ev))
|
|
||||||
expSum = expSum + ev
|
|
||||||
}
|
|
||||||
for (c in 0..3) {
|
|
||||||
let v = softmax.get(0, c).getOrThrow().getValue()
|
|
||||||
softmax.set(0, c, FP32(v / expSum))
|
|
||||||
}
|
|
||||||
|
|
||||||
var maxProb = Float32(-1.0)
|
|
||||||
var predClass = 0
|
|
||||||
for (c in 0..3) {
|
|
||||||
let p = softmax.get(0, c).getOrThrow().getValue()
|
|
||||||
if (p > maxProb) {
|
|
||||||
maxProb = p
|
|
||||||
predClass = Int64(c)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
confusion[label][predClass] += 1
|
|
||||||
if (predClass == label) {
|
|
||||||
testCorrect += 1
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let testAcc = Float32(testCorrect) / Float32(nTestSamples) * Float32(100.0)
|
|
||||||
println("测试准确率: ${testAcc}% (${testCorrect}/${nTestSamples})")
|
|
||||||
print("")
|
|
||||||
|
|
||||||
let classNames = ["红色(高)", "绿色(中)", "蓝色(低)"]
|
|
||||||
println("各类别准确率:")
|
|
||||||
for (c in 0..3) {
|
|
||||||
var total = 0
|
|
||||||
for (p in 0..3) {
|
|
||||||
total += confusion[c][p]
|
|
||||||
}
|
|
||||||
if (total > 0) {
|
|
||||||
let acc = Float32(confusion[c][c]) / Float32(total) * Float32(100.0)
|
|
||||||
print(" ${classNames[c]}: ${confusion[c][c]}/${total} (${acc}%)")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
println("")
|
|
||||||
|
|
||||||
println("混淆矩阵:")
|
|
||||||
println(" 预测红 预测绿 预测蓝")
|
|
||||||
for (t in 0..3) {
|
|
||||||
var line = "真实${classNames[t]}:"
|
|
||||||
for (p in 0..3) {
|
|
||||||
line = line + " ${confusion[t][p]} "
|
|
||||||
}
|
|
||||||
print(line)
|
|
||||||
}
|
|
||||||
println("")
|
|
||||||
|
|
||||||
printSeparator()
|
|
||||||
println("测试完成!")
|
|
||||||
printSeparator()
|
|
||||||
}
|
|
||||||
@@ -1,517 +0,0 @@
|
|||||||
// test_geometry_cnn.cj - 几何图形分类测试(Adam 优化器版本)
|
|
||||||
package bot4cj
|
|
||||||
import bot4cj.math.*
|
|
||||||
import std.math.*
|
|
||||||
import std.random.*
|
|
||||||
import std.collection.*
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算圆形像素值
|
|
||||||
*/
|
|
||||||
func calcCirclePixel(x: Int64, y: Int64, center: Float64, radius: Float64, offsetX: Float64, offsetY: Float64): FP32 {
|
|
||||||
let dx = Float64(x) - center + offsetX
|
|
||||||
let dy = Float64(y) - center + offsetY
|
|
||||||
let dist = sqrt(dx * dx + dy * dy)
|
|
||||||
|
|
||||||
if (dist <= radius) {
|
|
||||||
let intensity = 1.0 - (dist / radius) * 0.3
|
|
||||||
FP32(Float32(intensity))
|
|
||||||
} else if (dist <= radius + 1.5) {
|
|
||||||
let edgeDist = dist - radius
|
|
||||||
let intensity = max(0.0, 1.0 - edgeDist / 1.5) * 0.7
|
|
||||||
FP32(Float32(intensity))
|
|
||||||
} else {
|
|
||||||
FP32(0.0)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算正方形像素值
|
|
||||||
*/
|
|
||||||
func calcSquarePixel(x: Int64, y: Int64, center: Float64, radius: Float64, offsetX: Float64, offsetY: Float64): FP32 {
|
|
||||||
let halfSize = radius * 0.8
|
|
||||||
let left = center - halfSize + offsetX
|
|
||||||
let right = center + halfSize + offsetX
|
|
||||||
let top = center - halfSize + offsetY
|
|
||||||
let bottom = center + halfSize + offsetY
|
|
||||||
|
|
||||||
let fx = Float64(x)
|
|
||||||
let fy = Float64(y)
|
|
||||||
|
|
||||||
if (fx >= left && fx <= right && fy >= top && fy <= bottom) {
|
|
||||||
let distToEdge = min(
|
|
||||||
min(fx - left, right - fx),
|
|
||||||
min(fy - top, bottom - fy)
|
|
||||||
)
|
|
||||||
let intensity = min(1.0, 0.7 + distToEdge / halfSize * 0.3)
|
|
||||||
FP32(Float32(intensity))
|
|
||||||
} else {
|
|
||||||
let distX = max(0.0, max(left - fx, fx - right))
|
|
||||||
let distY = max(0.0, max(top - fy, fy - bottom))
|
|
||||||
let edgeDist = sqrt(distX * distX + distY * distY)
|
|
||||||
if (edgeDist <= 1.5) {
|
|
||||||
let intensity = max(0.0, 1.0 - edgeDist / 1.5) * 0.6
|
|
||||||
FP32(Float32(intensity))
|
|
||||||
} else {
|
|
||||||
FP32(0.0)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 计算三角形像素值
|
|
||||||
*/
|
|
||||||
func calcTrianglePixel(x: Int64, y: Int64, center: Float64, radius: Float64, offsetX: Float64, offsetY: Float64): FP32 {
|
|
||||||
let height = radius * 2.0
|
|
||||||
let baseWidth = height * 1.2
|
|
||||||
let topY = center - height * 0.5 + offsetY
|
|
||||||
let bottomY = center + height * 0.5 + offsetY
|
|
||||||
|
|
||||||
let fx = Float64(x)
|
|
||||||
let fy = Float64(y)
|
|
||||||
|
|
||||||
if (fy >= topY && fy <= bottomY) {
|
|
||||||
let t = (fy - topY) / (bottomY - topY)
|
|
||||||
let halfWidthAtY = baseWidth * (1.0 - t) * 0.5
|
|
||||||
let centerX = center + offsetX
|
|
||||||
|
|
||||||
if (fx >= centerX - halfWidthAtY && fx <= centerX + halfWidthAtY) {
|
|
||||||
let distToEdge = min(
|
|
||||||
fx - (centerX - halfWidthAtY),
|
|
||||||
(centerX + halfWidthAtY) - fx
|
|
||||||
)
|
|
||||||
let intensity = min(1.0, 0.7 + distToEdge / halfWidthAtY * 0.3)
|
|
||||||
FP32(Float32(intensity))
|
|
||||||
} else {
|
|
||||||
let distToLeft = (centerX - halfWidthAtY) - fx
|
|
||||||
let distToRight = fx - (centerX + halfWidthAtY)
|
|
||||||
let edgeDist = max(distToLeft, distToRight)
|
|
||||||
if (edgeDist <= 1.5) {
|
|
||||||
let intensity = max(0.0, 1.0 - edgeDist / 1.5) * 0.6
|
|
||||||
FP32(Float32(intensity))
|
|
||||||
} else {
|
|
||||||
FP32(0.0)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
FP32(0.0)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 生成几何图形数据
|
|
||||||
* 图形类型:0=圆形, 1=正方形, 2=三角形
|
|
||||||
*/
|
|
||||||
func generateGeometryData(numSamples: Int64): (Array<Matrix1>, Array<Int64>) {
|
|
||||||
var inputs = Array(numSamples, { _ => Matrix1(1, 1) })
|
|
||||||
var labels = Array(numSamples, { _ => 0 })
|
|
||||||
let rng = Random()
|
|
||||||
|
|
||||||
let imgSize = 16
|
|
||||||
|
|
||||||
for (i in 0..numSamples) {
|
|
||||||
let shapeType = Int64(rng.nextFloat64() * 3.0)
|
|
||||||
labels[i] = shapeType
|
|
||||||
|
|
||||||
var matrix = Matrix1(imgSize * imgSize, 1)
|
|
||||||
let center = Float64(imgSize) / 2.0
|
|
||||||
let radius = Float64(imgSize) / 3.0 + rng.nextFloat64() * 2.0 - 1.0
|
|
||||||
|
|
||||||
let offsetX = (rng.nextFloat64() - 0.5) * 2.0
|
|
||||||
let offsetY = (rng.nextFloat64() - 0.5) * 2.0
|
|
||||||
|
|
||||||
for (y in 0..imgSize) {
|
|
||||||
for (x in 0..imgSize) {
|
|
||||||
let idx = y * imgSize + x
|
|
||||||
|
|
||||||
var pixelValue = match (shapeType) {
|
|
||||||
case 0 => calcCirclePixel(x, y, center, radius, offsetX, offsetY)
|
|
||||||
case 1 => calcSquarePixel(x, y, center, radius, offsetX, offsetY)
|
|
||||||
case _ => calcTrianglePixel(x, y, center, radius, offsetX, offsetY)
|
|
||||||
}
|
|
||||||
|
|
||||||
let noise = FP32(Float32(rng.nextFloat64() - 0.5) * 0.05)
|
|
||||||
pixelValue = pixelValue + noise
|
|
||||||
|
|
||||||
if (pixelValue.getValue() > 1.0) {
|
|
||||||
pixelValue = FP32(1.0)
|
|
||||||
} else if (pixelValue.getValue() < 0.0) {
|
|
||||||
pixelValue = FP32(0.0)
|
|
||||||
}
|
|
||||||
|
|
||||||
matrix.set(idx, 0, pixelValue)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
inputs[i] = matrix
|
|
||||||
}
|
|
||||||
|
|
||||||
(inputs, labels)
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 可视化一个样本
|
|
||||||
*/
|
|
||||||
func showSample(matrix: Matrix1, imgSize: Int64, label: Int64): Unit {
|
|
||||||
let shapeNames = ["圆形", "正方形", "三角形"]
|
|
||||||
println("标签: ${shapeNames[label]}")
|
|
||||||
|
|
||||||
for (y in 0..imgSize) {
|
|
||||||
var line = ""
|
|
||||||
for (x in 0..imgSize) {
|
|
||||||
let idx = y * imgSize + x
|
|
||||||
let val = matrix.get(idx, 0).getOrThrow().getValue()
|
|
||||||
if (val > 0.8) {
|
|
||||||
line = line + "██"
|
|
||||||
} else if (val > 0.6) {
|
|
||||||
line = line + "▓▓"
|
|
||||||
} else if (val > 0.4) {
|
|
||||||
line = line + "▒▒"
|
|
||||||
} else if (val > 0.2) {
|
|
||||||
line = line + "░░"
|
|
||||||
} else {
|
|
||||||
line = line + " "
|
|
||||||
}
|
|
||||||
}
|
|
||||||
println(line)
|
|
||||||
}
|
|
||||||
println("")
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Softmax 辅助函数
|
|
||||||
*/
|
|
||||||
func computeSoftmax(input: Matrix1): Matrix1 {
|
|
||||||
let cols = input.get_col()
|
|
||||||
var output = Matrix1(1, cols)
|
|
||||||
|
|
||||||
var maxVal = Float32(-1.0e30)
|
|
||||||
for (c in 0..cols) {
|
|
||||||
let raw = input.get(0, c).getOrThrow().getValue()
|
|
||||||
if (raw > maxVal) {
|
|
||||||
maxVal = raw
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
var expSum = Float32(0.0)
|
|
||||||
for (c in 0..cols) {
|
|
||||||
let raw = input.get(0, c).getOrThrow().getValue()
|
|
||||||
let ev = Float32(exp(raw - maxVal))
|
|
||||||
output.set(0, c, FP32(ev))
|
|
||||||
expSum = expSum + ev
|
|
||||||
}
|
|
||||||
|
|
||||||
for (c in 0..cols) {
|
|
||||||
let v = output.get(0, c).getOrThrow().getValue()
|
|
||||||
output.set(0, c, FP32(v / expSum))
|
|
||||||
}
|
|
||||||
|
|
||||||
output
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 获取预测类别
|
|
||||||
*/
|
|
||||||
func predictClass(probs: Matrix1): Int64 {
|
|
||||||
let cols = probs.get_col()
|
|
||||||
var maxProb = Float32(-1.0)
|
|
||||||
var predClass = 0
|
|
||||||
|
|
||||||
for (c in 0..cols) {
|
|
||||||
let p = probs.get(0, c).getOrThrow().getValue()
|
|
||||||
if (p > maxProb) {
|
|
||||||
maxProb = p
|
|
||||||
predClass = Int64(c)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
predClass
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 打印分隔线
|
|
||||||
*/
|
|
||||||
func printLine(): Unit {
|
|
||||||
println("============================================================")
|
|
||||||
}
|
|
||||||
|
|
||||||
/**
|
|
||||||
* 主测试函数
|
|
||||||
*/
|
|
||||||
func testGeometryCNN(): Unit {
|
|
||||||
printLine()
|
|
||||||
println("CNN 几何图形分类测试 (Adam 优化器版本)")
|
|
||||||
printLine()
|
|
||||||
println("")
|
|
||||||
|
|
||||||
let imgSize = 16
|
|
||||||
let nTrainSamples = 200
|
|
||||||
let nTestSamples = 60
|
|
||||||
let nEpochs = 50
|
|
||||||
let batchSize = 10
|
|
||||||
let nClasses = 3
|
|
||||||
|
|
||||||
println("配置信息:")
|
|
||||||
println(" 图像大小: ${imgSize}x${imgSize}")
|
|
||||||
println(" 训练样本数: ${nTrainSamples}")
|
|
||||||
println(" 测试样本数: ${nTestSamples}")
|
|
||||||
println(" 训练轮数: ${nEpochs}")
|
|
||||||
println(" Batch 大小: ${batchSize}")
|
|
||||||
println(" 类别数: ${nClasses} (圆形, 正方形, 三角形)")
|
|
||||||
println(" 优化器: Adam")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
println("网络结构:")
|
|
||||||
println(" Input(16x16x1)")
|
|
||||||
println(" -> Conv1(5x5, 1->8) -> Pool1(2x2)")
|
|
||||||
println(" -> Conv2(3x3, 8->16) -> Pool2(2x2)")
|
|
||||||
println(" -> Flatten -> FC(64->32) -> FC(32->3)")
|
|
||||||
println(" -> Softmax")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
let lr = FP32(0.01)
|
|
||||||
|
|
||||||
// 构建网络
|
|
||||||
var conv1 = ConvLayer(1, 8, (5, 5), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
|
||||||
var pool1 = PoolingLayer(8, (2, 2), 2, 0, "max", Af.ReLU, lr, lr)
|
|
||||||
var conv2 = ConvLayer(8, 16, (3, 3), 1, 0, lr, Af.Leaky_ReLU(FP32(0.01)), Option.None)
|
|
||||||
var pool2 = PoolingLayer(16, (2, 2), 2, 0, "max", Af.ReLU, lr, lr)
|
|
||||||
var flatten = FlattenLayer(16, 2, 2)
|
|
||||||
var fc1 = Nn1(64, 32, Float32(0.3), Af.Leaky_ReLU(FP32(0.01)))
|
|
||||||
var fc2 = Nn1(32, 3, Float32(0.3), Af.Softmax)
|
|
||||||
|
|
||||||
// 创建 Adam 优化器
|
|
||||||
var adamConv1 = Adam(FP32(0.001)) // 卷积层1
|
|
||||||
var adamConv2 = Adam(FP32(0.001)) // 卷积层2
|
|
||||||
var adamFC1 = Adam(FP32(0.0005)) // 全连接层1
|
|
||||||
var adamFC2 = Adam(FP32(0.0005)) // 全连接层2
|
|
||||||
|
|
||||||
println("网络创建完成 (Adam优化器: Conv1, Conv2, FC1, FC2)")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
// 生成数据
|
|
||||||
println("生成训练数据...")
|
|
||||||
let (trainInputs, trainLabels) = generateGeometryData(nTrainSamples)
|
|
||||||
println("生成测试数据...")
|
|
||||||
let (testInputs, testLabels) = generateGeometryData(nTestSamples)
|
|
||||||
println("数据生成完成")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
// 显示几个样本
|
|
||||||
println("样本预览:")
|
|
||||||
for (i in 0..3) {
|
|
||||||
if (i < nTrainSamples) {
|
|
||||||
println("样本 ${i + 1}:")
|
|
||||||
showSample(trainInputs[i], imgSize, trainLabels[i])
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
println("开始训练...")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
var bestAccuracy = Float32(0.0)
|
|
||||||
|
|
||||||
// 计算卷积输出尺寸
|
|
||||||
let conv1OutH = imgSize - 4 // 16 - 5 + 1 = 12
|
|
||||||
let conv1OutW = imgSize - 4
|
|
||||||
let pool1OutH = conv1OutH / 2 // 6
|
|
||||||
let pool1OutW = conv1OutW / 2
|
|
||||||
let conv2OutH = pool1OutH - 2 // 6 - 3 + 1 = 4
|
|
||||||
let conv2OutW = pool1OutW - 2
|
|
||||||
let pool2OutH = conv2OutH / 2 // 2
|
|
||||||
let pool2OutW = conv2OutW / 2
|
|
||||||
let flatSize = 16 * pool2OutH * pool2OutW // 64
|
|
||||||
|
|
||||||
for (epoch in 0..nEpochs) {
|
|
||||||
var epochLoss = Float32(0.0)
|
|
||||||
var epochCorrect = 0
|
|
||||||
|
|
||||||
// 训练
|
|
||||||
for (i in 0..nTrainSamples) {
|
|
||||||
let x = trainInputs[i]
|
|
||||||
let label = trainLabels[i]
|
|
||||||
|
|
||||||
// 前向传播
|
|
||||||
let c1Out = conv1.forward(x, (1, imgSize, imgSize, 1))
|
|
||||||
let p1Out = pool1.forward(c1Out, (1, conv1OutH, conv1OutW, 8))
|
|
||||||
let c2Out = conv2.forward(p1Out, (1, pool1OutH, pool1OutW, 8))
|
|
||||||
let p2Out = pool2.forward(c2Out, (1, conv2OutH, conv2OutW, 16))
|
|
||||||
|
|
||||||
let flatOut = flatten.forward(p2Out, (1, pool2OutH, pool2OutW, 16))
|
|
||||||
|
|
||||||
// 转换为全连接层输入
|
|
||||||
var fc1Input = Matrix1(1, flatSize)
|
|
||||||
for (j in 0..flatSize) {
|
|
||||||
fc1Input.set(0, j, flatOut.get(j, 0).getOrThrow())
|
|
||||||
}
|
|
||||||
|
|
||||||
let fc1Out = fc1.forward(fc1Input)
|
|
||||||
let fc2Out = fc2.forward(fc1Out)
|
|
||||||
|
|
||||||
// Softmax
|
|
||||||
let probs = computeSoftmax(fc2Out)
|
|
||||||
|
|
||||||
// 计算损失和准确率
|
|
||||||
let probVal = probs.get(0, label).getOrThrow().getValue()
|
|
||||||
let loss = -log(max(probVal, Float32(1e-10)))
|
|
||||||
epochLoss = epochLoss + loss
|
|
||||||
|
|
||||||
let predClass = predictClass(probs)
|
|
||||||
if (predClass == label) {
|
|
||||||
epochCorrect += 1
|
|
||||||
}
|
|
||||||
|
|
||||||
// 反向传播
|
|
||||||
var delta = Matrix1(1, nClasses)
|
|
||||||
for (c in 0..nClasses) {
|
|
||||||
let prob = probs.get(0, c).getOrThrow()
|
|
||||||
if (Int64(c) == label) {
|
|
||||||
delta.set(0, c, prob - FP32(1.0))
|
|
||||||
} else {
|
|
||||||
delta.set(0, c, prob)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let fc2Delta = fc2.backward(delta)
|
|
||||||
let fc1Delta = fc1.backward(fc2Delta)
|
|
||||||
|
|
||||||
// 转换回特征图格式
|
|
||||||
var flatDelta = Matrix1(flatSize, 1)
|
|
||||||
for (j in 0..flatSize) {
|
|
||||||
flatDelta.set(j, 0, fc1Delta.get(0, j).getOrThrow())
|
|
||||||
}
|
|
||||||
|
|
||||||
let p2Delta = flatten.backward(flatDelta)
|
|
||||||
let c2Delta = pool2.backward(p2Delta)
|
|
||||||
let p1Delta = conv2.backward(c2Delta)
|
|
||||||
let c1Delta = pool1.backward(p1Delta)
|
|
||||||
let _ = conv1.backward(c1Delta)
|
|
||||||
|
|
||||||
// 批次更新 - 使用 Adam
|
|
||||||
if ((i + 1) % batchSize == 0) {
|
|
||||||
adamConv1.updateConvLayer(conv1) // Conv1 - Adam
|
|
||||||
adamConv2.updateConvLayer(conv2) // Conv2 - Adam
|
|
||||||
adamFC1.updateFCLayer(fc1) // FC1 - Adam
|
|
||||||
adamFC2.updateFCLayer(fc2) // FC2 - Adam
|
|
||||||
|
|
||||||
// 池化层仍使用 SGD
|
|
||||||
pool1.update(Float32(1.0))
|
|
||||||
pool2.update(Float32(1.0))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 最终更新 - 使用 Adam
|
|
||||||
adamConv1.updateConvLayer(conv1)
|
|
||||||
adamConv2.updateConvLayer(conv2)
|
|
||||||
adamFC1.updateFCLayer(fc1)
|
|
||||||
adamFC2.updateFCLayer(fc2)
|
|
||||||
pool1.update(Float32(1.0))
|
|
||||||
pool2.update(Float32(1.0))
|
|
||||||
|
|
||||||
let avgLoss = epochLoss / Float32(nTrainSamples)
|
|
||||||
let accuracy = Float32(epochCorrect) / Float32(nTrainSamples) * Float32(100.0)
|
|
||||||
|
|
||||||
if (accuracy > bestAccuracy) {
|
|
||||||
bestAccuracy = accuracy
|
|
||||||
}
|
|
||||||
|
|
||||||
// 输出训练信息
|
|
||||||
if (epoch % 1 == 0 || epoch == nEpochs - 1) {
|
|
||||||
println("Epoch ${epoch + 1}/${nEpochs} - Loss: ${avgLoss}, Accuracy: ${accuracy}%, Adam Step: ${adamConv1.getStep()}")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
println("")
|
|
||||||
println("训练完成! 最佳训练准确率: ${bestAccuracy}%")
|
|
||||||
println("总更新步数: ${adamConv1.getStep()}")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
// 测试
|
|
||||||
println("开始测试...")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
var testCorrect = 0
|
|
||||||
var confusion = Array(nClasses, { _ => Array(nClasses, { _ => 0 }) })
|
|
||||||
var testLoss = Float32(0.0)
|
|
||||||
|
|
||||||
for (i in 0..nTestSamples) {
|
|
||||||
let x = testInputs[i]
|
|
||||||
let label = testLabels[i]
|
|
||||||
|
|
||||||
// 前向传播
|
|
||||||
let c1Out = conv1.forward(x, (1, imgSize, imgSize, 1))
|
|
||||||
let p1Out = pool1.forward(c1Out, (1, conv1OutH, conv1OutW, 8))
|
|
||||||
let c2Out = conv2.forward(p1Out, (1, pool1OutH, pool1OutW, 8))
|
|
||||||
let p2Out = pool2.forward(c2Out, (1, conv2OutH, conv2OutW, 16))
|
|
||||||
|
|
||||||
let flatOut = flatten.forward(p2Out, (1, pool2OutH, pool2OutW, 16))
|
|
||||||
|
|
||||||
var fc1Input = Matrix1(1, flatSize)
|
|
||||||
for (j in 0..flatSize) {
|
|
||||||
fc1Input.set(0, j, flatOut.get(j, 0).getOrThrow())
|
|
||||||
}
|
|
||||||
|
|
||||||
let fc1Out = fc1.forward(fc1Input)
|
|
||||||
let fc2Out = fc2.forward(fc1Out)
|
|
||||||
|
|
||||||
let probs = computeSoftmax(fc2Out)
|
|
||||||
|
|
||||||
// 计算损失
|
|
||||||
let probVal = probs.get(0, label).getOrThrow().getValue()
|
|
||||||
testLoss = testLoss + Float32(-log(max(probVal, Float32(1e-10))))
|
|
||||||
|
|
||||||
// 预测
|
|
||||||
let predClass = predictClass(probs)
|
|
||||||
|
|
||||||
confusion[label][predClass] += 1
|
|
||||||
if (predClass == label) {
|
|
||||||
testCorrect += 1
|
|
||||||
}
|
|
||||||
|
|
||||||
// 显示前几个错误分类的样本
|
|
||||||
if (predClass != label && testCorrect < 10) {
|
|
||||||
println("错误分类样本 ${i + 1}:")
|
|
||||||
showSample(x, imgSize, label)
|
|
||||||
let shapeNames = ["圆形", "正方形", "三角形"]
|
|
||||||
println(" 真实标签: ${shapeNames[label]}, 预测标签: ${shapeNames[predClass]}")
|
|
||||||
println("")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let testAcc = Float32(testCorrect) / Float32(nTestSamples) * Float32(100.0)
|
|
||||||
let avgTestLoss = testLoss / Float32(nTestSamples)
|
|
||||||
|
|
||||||
println("测试结果:")
|
|
||||||
println(" 准确率: ${testAcc}% (${testCorrect}/${nTestSamples})")
|
|
||||||
println(" 平均损失: ${avgTestLoss}")
|
|
||||||
println("")
|
|
||||||
|
|
||||||
let classNames = ["圆形 ", "正方形", "三角形"]
|
|
||||||
println("各类别准确率:")
|
|
||||||
for (c in 0..nClasses) {
|
|
||||||
var total = 0
|
|
||||||
for (p in 0..nClasses) {
|
|
||||||
total += confusion[c][p]
|
|
||||||
}
|
|
||||||
if (total > 0) {
|
|
||||||
let acc = Float32(confusion[c][c]) / Float32(total) * Float32(100.0)
|
|
||||||
println(" ${classNames[c]}: ${confusion[c][c]}/${total} (${acc}%)")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
println("")
|
|
||||||
|
|
||||||
println("混淆矩阵:")
|
|
||||||
println(" 预测圆形 预测正方形 预测三角形")
|
|
||||||
for (t in 0..nClasses) {
|
|
||||||
var line = "真实${classNames[t]}:"
|
|
||||||
for (p in 0..nClasses) {
|
|
||||||
line = line + " ${confusion[t][p]} "
|
|
||||||
}
|
|
||||||
println(line)
|
|
||||||
}
|
|
||||||
println("")
|
|
||||||
|
|
||||||
printLine()
|
|
||||||
println("几何图形分类测试完成!")
|
|
||||||
printLine()
|
|
||||||
}
|
|
||||||
+1
-1
@@ -73,7 +73,7 @@ gateway 侧仅做**存在性不变量**检查(显式 `Target-Url` 请求必须
|
|||||||
|
|
||||||
### `POST /register` / `POST /register/send_code` / `POST /reset/send_code` / `POST /reset_password`
|
### `POST /register` / `POST /register/send_code` / `POST /reset/send_code` / `POST /reset_password`
|
||||||
|
|
||||||
与 gateway v0.8.0 契约完全一致(详见 `docs/reports/auth_service_plan.md` §5 与黑盒 tests/blackbox/test_gateway.py)。
|
与 gateway v0.8.0 契约完全一致(详见黑盒 tests/blackbox/test_gateway.py)。
|
||||||
|
|
||||||
### `POST /logout` / `POST /refresh_token` / `POST /whoami` / `POST /admin/verify`
|
### `POST /logout` / `POST /refresh_token` / `POST /whoami` / `POST /admin/verify`
|
||||||
|
|
||||||
|
|||||||
@@ -56,7 +56,6 @@ captcha2 :8009 (RpcServer)
|
|||||||
## 状态
|
## 状态
|
||||||
|
|
||||||
- 生产运行中(services.yaml 拉起),Python captcha 已废弃
|
- 生产运行中(services.yaml 拉起),Python captcha 已废弃
|
||||||
- 开发文档见 [docs/reports/captcha_cpp_rewrite_dev_doc.md](../docs/reports/captcha_cpp_rewrite_dev_doc.md)
|
|
||||||
- 安全审计见 [docs/security/captcha2_SECURITY_AUDIT_20260809.md](../docs/security/captcha2_SECURITY_AUDIT_20260809.md)
|
- 安全审计见 [docs/security/captcha2_SECURITY_AUDIT_20260809.md](../docs/security/captcha2_SECURITY_AUDIT_20260809.md)
|
||||||
|
|
||||||
## 许可证
|
## 许可证
|
||||||
|
|||||||
+20
-19
@@ -223,22 +223,23 @@ services:
|
|||||||
depends_on: []
|
depends_on: []
|
||||||
sleep_before: 0
|
sleep_before: 0
|
||||||
|
|
||||||
# --- 端口重定向(最后启动) ---
|
# --- 端口重定向(最后启动)---
|
||||||
- name: redirect80
|
# redirect80 与系统 casaos-gateway 抢占 80 端口导致绑定失败,暂时停用
|
||||||
cmd: bash /opt/micro_servers/redirect80/_run.sh
|
# - name: redirect80
|
||||||
args:
|
# cmd: bash /opt/micro_servers/redirect80/_run.sh
|
||||||
- "--instance"
|
# args:
|
||||||
- "redirect80"
|
# - "--instance"
|
||||||
- "--log-level"
|
# - "redirect80"
|
||||||
- "info"
|
# - "--log-level"
|
||||||
# --- 配置全部经 CLI 传递(原 confs/configs/redirect80.ini 迁入) ---
|
# - "info"
|
||||||
- "--host"
|
# # --- 配置全部经 CLI 传递(原 confs/configs/redirect80.ini 迁入) ---
|
||||||
- "0.0.0.0"
|
# - "--host"
|
||||||
- "--port"
|
# - "0.0.0.0"
|
||||||
- "80"
|
# - "--port"
|
||||||
- "--redirect-port"
|
# - "80"
|
||||||
- "443"
|
# - "--redirect-port"
|
||||||
- "--redirect-code"
|
# - "443"
|
||||||
- "301"
|
# - "--redirect-code"
|
||||||
depends_on: [waf]
|
# - "301"
|
||||||
sleep_before: 0
|
# depends_on: [waf]
|
||||||
|
# sleep_before: 0
|
||||||
@@ -7,7 +7,8 @@ CREATE TABLE IF NOT EXISTS Users (
|
|||||||
password_hash TEXT NOT NULL,
|
password_hash TEXT NOT NULL,
|
||||||
create_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
create_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||||
last_login DATETIME,
|
last_login DATETIME,
|
||||||
is_active BOOLEAN DEFAULT 1
|
is_active BOOLEAN DEFAULT 1,
|
||||||
|
is_admin BOOLEAN NOT NULL DEFAULT 0
|
||||||
);
|
);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS Notes (
|
CREATE TABLE IF NOT EXISTS Notes (
|
||||||
@@ -18,14 +19,19 @@ CREATE TABLE IF NOT EXISTS Notes (
|
|||||||
content TEXT NOT NULL,
|
content TEXT NOT NULL,
|
||||||
views INTEGER NOT NULL DEFAULT 0,
|
views INTEGER NOT NULL DEFAULT 0,
|
||||||
created_at DATETIME,
|
created_at DATETIME,
|
||||||
updated_at DATETIME
|
updated_at DATETIME,
|
||||||
|
is_public BOOLEAN NOT NULL DEFAULT 1,
|
||||||
|
everyone_editable BOOLEAN NOT NULL DEFAULT 1
|
||||||
);
|
);
|
||||||
|
|
||||||
-- 旧版增量
|
-- 旧版增量(存量库请按编号顺序执行对应 patch 脚本)
|
||||||
-- ALTER TABLE Notes ADD COLUMN owner VARCHAR(128) DEFAULT 'Unknow';
|
-- ALTER TABLE Notes ADD COLUMN owner VARCHAR(128) DEFAULT 'Unknow';
|
||||||
-- ALTER TABLE Notes ADD COLUMN views INTEGER NOT NULL DEFAULT 0; -- patch_001 观看人数列
|
-- ALTER TABLE Notes ADD COLUMN views INTEGER NOT NULL DEFAULT 0; -- patch_001 观看人数列
|
||||||
-- ALTER TABLE Notes ADD COLUMN created_at DATETIME; -- patch_002 创建时间
|
-- ALTER TABLE Notes ADD COLUMN created_at DATETIME; -- patch_002 创建时间
|
||||||
-- ALTER TABLE Notes ADD COLUMN updated_at DATETIME; -- patch_002 更新时间
|
-- ALTER TABLE Notes ADD COLUMN updated_at DATETIME; -- patch_002 更新时间
|
||||||
|
-- ALTER TABLE Users ADD COLUMN is_admin BOOLEAN NOT NULL DEFAULT 0; -- patch_003 管理员标记
|
||||||
|
-- ALTER TABLE Notes ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT 1; -- patch_003 是否公开(默认公开)
|
||||||
|
-- ALTER TABLE Notes ADD COLUMN everyone_editable BOOLEAN NOT NULL DEFAULT 1; -- patch_003 是否人人可编辑(默认人人可编辑)
|
||||||
|
|
||||||
-- Notes 索引:get_notes 按 pid 查子节点、所有权校验按 owner(patch_001 并入)
|
-- Notes 索引:get_notes 按 pid 查子节点、所有权校验按 owner(patch_001 并入)
|
||||||
CREATE INDEX IF NOT EXISTS idx_notes_pid ON Notes(pid);
|
CREATE INDEX IF NOT EXISTS idx_notes_pid ON Notes(pid);
|
||||||
|
|||||||
@@ -0,0 +1,17 @@
|
|||||||
|
-- 补丁 003:Users 管理员标记 + Notes 可见性/可编辑性
|
||||||
|
-- 数据库类型: SQLite
|
||||||
|
-- 依赖: create_sqlite3.sql(或等价 schema)已执行,且 patch_001 / patch_002 已按序执行
|
||||||
|
-- 幂等性: ADD COLUMN 为一次性迁移,重复执行会报 "duplicate column name",
|
||||||
|
-- 请按迁移编号顺序执行一次。
|
||||||
|
--
|
||||||
|
-- 语义(默认值向后兼容现有行为):
|
||||||
|
-- Users.is_admin 是否管理员,默认 0(普通用户)
|
||||||
|
-- Notes.is_public 节点是否公开可读,默认 1(公开,与现状一致)
|
||||||
|
-- Notes.everyone_editable 是否任何人都能编辑,默认 1(人人可编辑)
|
||||||
|
|
||||||
|
-- ── Users 新增管理员标记 ─────────────────────────────────────────────
|
||||||
|
ALTER TABLE Users ADD COLUMN is_admin BOOLEAN NOT NULL DEFAULT 0;
|
||||||
|
|
||||||
|
-- ── Notes 新增可见性 / 可编辑性 ──────────────────────────────────────
|
||||||
|
ALTER TABLE Notes ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT 1;
|
||||||
|
ALTER TABLE Notes ADD COLUMN everyone_editable BOOLEAN NOT NULL DEFAULT 1;
|
||||||
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
> 权威架构参考。本文以 `core/services.yaml`(实际运行清单)与 `git HEAD` 为事实来源。
|
> 权威架构参考。本文以 `core/services.yaml`(实际运行清单)与 `git HEAD` 为事实来源。
|
||||||
> 与根目录 `README.md` 的偏差见文末「§10 与 README 的偏差」。
|
> 与根目录 `README.md` 的偏差见文末「§10 与 README 的偏差」。
|
||||||
> 最后核对于:2026-08-15。
|
> 最后核对于:2026-08-20。
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -85,7 +85,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
|||||||
└─────────────────────────────────────┘
|
└─────────────────────────────────────┘
|
||||||
|
|
||||||
独立 / 外设(不在主链路):
|
独立 / 外设(不在主链路):
|
||||||
aibox :8765 (Python WebSocket, 独立产品)
|
aibox :8765 (Python WebSocket, 独立产品; --once CLI 单次对话 + models 多模型档案含 fallback 失败转移 + model_config.py 交互式接入点管理)
|
||||||
BotOps (aibox 的 TUI 前端)
|
BotOps (aibox 的 TUI 前端)
|
||||||
bot4cj (Cangjie + 树莓派 I2C, 硬件)
|
bot4cj (Cangjie + 树莓派 I2C, 硬件)
|
||||||
pybot (Python, bot4cj 的迁移版: .mt 解析/录制/实时闭环)
|
pybot (Python, bot4cj 的迁移版: .mt 解析/录制/实时闭环)
|
||||||
@@ -148,7 +148,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
|||||||
| **aibox** | Python | 8765 | 127.0.0.1 | 独立 AI 智能体(WebSocket) | 独立产品 |
|
| **aibox** | Python | 8765 | 127.0.0.1 | 独立 AI 智能体(WebSocket) | 独立产品 |
|
||||||
| **BotOps** | Python | — | — | aibox 的 TUI 前端 | 独立 |
|
| **BotOps** | Python | — | — | aibox 的 TUI 前端 | 独立 |
|
||||||
| **bot4cj** | Cangjie | — | 树莓派 I2C | 机器人脚本解析器 | 独立硬件 |
|
| **bot4cj** | Cangjie | — | 树莓派 I2C | 机器人脚本解析器 | 独立硬件 |
|
||||||
| **pybot** | Python | — | 树莓派 | bot4cj 的 Python 迁移版(.mt 解析、actpred episode 录制、--loop 实时预测闭环) | 独立硬件 |
|
| **pybot** | Python | — | 树莓派 | bot4cj 的 Python 迁移版(.mt 解析、actpred episode 录制、--loop 实时预测闭环;HEF 后端自动异步,推理与执行上一步动作并行) | 独立硬件 |
|
||||||
| **mtc** | Python | — | — | 命令式动作脚本(Python DSL)→ .mt 编译器(非阻塞时间轴 IR + 10ms 栅格化),pybot/bot4cj 的脚本 authoring 工具 | 开发期工具 |
|
| **mtc** | Python | — | — | 命令式动作脚本(Python DSL)→ .mt 编译器(非阻塞时间轴 IR + 10ms 栅格化),pybot/bot4cj 的脚本 authoring 工具 | 开发期工具 |
|
||||||
| **lci** | C++ | — | /dev/fb0 | 嵌入式帧缓冲 GUI | 独立硬件 |
|
| **lci** | C++ | — | /dev/fb0 | 嵌入式帧缓冲 GUI | 独立硬件 |
|
||||||
| **actpred** | Python | — | — | 动作预测(图像+动作历史→下一动作),bot4cj/pybot 的「大脑」配套 | 实验 |
|
| **actpred** | Python | — | — | 动作预测(图像+动作历史→下一动作),bot4cj/pybot 的「大脑」配套 | 实验 |
|
||||||
@@ -236,6 +236,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
|||||||
- **smtp3 待上线**:C++ 重写已完成并通过安全审计,但 services.yaml 仍指向 Python `smtp`,建议择机平替以消除最后一块 Python 有状态服务。
|
- **smtp3 待上线**:C++ 重写已完成并通过安全审计,但 services.yaml 仍指向 Python `smtp`,建议择机平替以消除最后一块 Python 有状态服务。
|
||||||
- **可观测性缺口**:缺 tracing/metrics/告警,建议后续引入轻量方案(或复用 db2 Logs + 定时聚合)。
|
- **可观测性缺口**:缺 tracing/metrics/告警,建议后续引入轻量方案(或复用 db2 Logs + 定时聚合)。
|
||||||
- **水平扩展限制**:smtp/captcha 内存态为单点,若需多实例需引入 Redis 等共享存储。
|
- **水平扩展限制**:smtp/captcha 内存态为单点,若需多实例需引入 Redis 等共享存储。
|
||||||
|
- **panel 面板微服务(规划中)**:设计稿见 [docs/panel_service.md](panel_service.md)——背景视图 + 各服务自持控件片段(iframe 沙箱 + BFF 转发),解决整页应用互相改前端代码的耦合问题;落地后需三处同步(本文件 / MODULES.md / 根 README)。
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -251,6 +252,7 @@ Blackbird 是一套 **进程托管式轻量微服务** 系统(自研 UniServe
|
|||||||
6. **工作记忆偏差**:✅ 已同步。`ca:8443 → gateway:8445`、`dnfs2`、confs 内嵌 core 均已写入各文档。
|
6. **工作记忆偏差**:✅ 已同步。`ca:8443 → gateway:8445`、`dnfs2`、confs 内嵌 core 均已写入各文档。
|
||||||
7. **dnfs2 README 图示滞后**:✅ 已修正(2026-08-10)。`dnfs2/README.md` 架构图与验签描述已更新为 ca:8443(CLI 参数 `--gateway-url` 指向 ca)。
|
7. **dnfs2 README 图示滞后**:✅ 已修正(2026-08-10)。`dnfs2/README.md` 架构图与验签描述已更新为 ca:8443(CLI 参数 `--gateway-url` 指向 ca)。
|
||||||
8. **memdb / pybot / mtc 未进根 README**:✅ 已回填(2026-08-15)。`memdb:8011`、`pybot`(bot4cj Python 迁移版)、`mtc`(Python DSL → .mt 编译器,2026-08-15 落地)连同 `actpred` 已入根 `README.md` 模块表与端口一览;本文件与 MODULES.md 已同步。
|
8. **memdb / pybot / mtc 未进根 README**:✅ 已回填(2026-08-15)。`memdb:8011`、`pybot`(bot4cj Python 迁移版)、`mtc`(Python DSL → .mt 编译器,2026-08-15 落地)连同 `actpred` 已入根 `README.md` 模块表与端口一览;本文件与 MODULES.md 已同步。
|
||||||
|
9. **aibox 第五轮审计(Batch K)与 model_config.py**:✅ 已同步(2026-08-20)。根 README aibox 小节审计轮次更新为五轮(E/F/G/H/I/K,单测 490 项全过)并增补 `model_config.py` 交互式接入点管理;MODULES.md 同步 aibox 行与文档计数;aibox/README.md 更新日志补 v3.6 条目(J1~J16 + Batch K 波次 1~5 落地)。
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
+4
-4
@@ -44,7 +44,7 @@
|
|||||||
|
|
||||||
| 模块 | 语言 | 端口 | 职责 |
|
| 模块 | 语言 | 端口 | 职责 |
|
||||||
|------|------|------|------|
|
|------|------|------|------|
|
||||||
| aibox | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度/前端与语音客户端 |
|
| aibox | Python | 8765 | 独立 AI 智能体系统(DeepSeek + WebSocket),自带技能系统/记忆/任务调度/前端与语音客户端;`--once` CLI 单次对话(stdout/stderr 分流),`models` 配置节支持多 OpenAI 兼容端点按档案名切换,`model_config.py` 交互式接入点管理(列表/添加/删除/连通性验证,Key 仅记环境变量名不落盘) |
|
||||||
| BotOps | Python/Textual | — | aibox 的终端 TUI 前端,WebSocket 连接 aibox,Markdown 增量流式渲染、思考模式、任务侧栏、自动重连 |
|
| BotOps | Python/Textual | — | aibox 的终端 TUI 前端,WebSocket 连接 aibox,Markdown 增量流式渲染、思考模式、任务侧栏、自动重连 |
|
||||||
| actpred | Python/PyTorch | — | 动作预测(图像 + 动作历史 → 预测下一个动作),bot4cj 的「大脑」配套;支持从 `.mt` 脚本预训练 |
|
| actpred | Python/PyTorch | — | 动作预测(图像 + 动作历史 → 预测下一个动作),bot4cj 的「大脑」配套;支持从 `.mt` 脚本预训练 |
|
||||||
|
|
||||||
@@ -53,7 +53,7 @@
|
|||||||
| 模块 | 语言 | 端口 | 职责 |
|
| 模块 | 语言 | 端口 | 职责 |
|
||||||
|------|------|------|------|
|
|------|------|------|------|
|
||||||
| bot4cj | 仓颉/Cangjie | — (树莓派 I2C) | 机器人硬件脚本解析器,解析 `.mt` 驱动 PCA9685 舵机 / MPU6050 陀螺仪 / 扬声器,面向 Dashun 蜘蛛机器人 |
|
| bot4cj | 仓颉/Cangjie | — (树莓派 I2C) | 机器人硬件脚本解析器,解析 `.mt` 驱动 PCA9685 舵机 / MPU6050 陀螺仪 / 扬声器,面向 Dashun 蜘蛛机器人 |
|
||||||
| pybot | Python | — (树莓派) | bot4cj 的 Python 迁移版:`.mt` 动作脚本解析器(零第三方依赖,与仓颉版逐语义对齐);支持执行 / `--dry-run` / `--record` 录制 actpred episode(含 imu.npy)/ `--loop` 实时预测闭环 |
|
| pybot | Python | — (树莓派) | bot4cj 的 Python 迁移版:`.mt` 动作脚本解析器(零第三方依赖,与仓颉版逐语义对齐);支持执行 / `--dry-run` / `--record` 录制 actpred episode(含 imu.npy)/ `--loop` 实时预测闭环(HEF 后端自动异步:推理与执行上一步动作并行) |
|
||||||
| mtc | Python | — | 命令式动作脚本(Python DSL)→ `.mt` 编译器:`exec` 注入命名空间 → 时间轴 IR(非阻塞语义、重叠截断)→ 10ms 栅格化;产物可由 pybot/bot4cj 直接执行,零第三方依赖 |
|
| mtc | Python | — | 命令式动作脚本(Python DSL)→ `.mt` 编译器:`exec` 注入命名空间 → 时间轴 IR(非阻塞语义、重叠截断)→ 10ms 栅格化;产物可由 pybot/bot4cj 直接执行,零第三方依赖 |
|
||||||
| lci | C++ | — (/dev/fb0) | 嵌入式帧缓冲图形界面(无需 X11/Wayland),矢量 GB2312 中文字库、图片渲染、实时数据可视化 |
|
| lci | C++ | — (/dev/fb0) | 嵌入式帧缓冲图形界面(无需 X11/Wayland),矢量 GB2312 中文字库、图片渲染、实时数据可视化 |
|
||||||
| 3d | — (资产) | — | 机器人 / 机械臂 CAD 模型(STEP/STL/SLDPRT),bot4cj 配套硬件资产 |
|
| 3d | — (资产) | — | 机器人 / 机械臂 CAD 模型(STEP/STL/SLDPRT),bot4cj 配套硬件资产 |
|
||||||
@@ -77,7 +77,7 @@
|
|||||||
| 模块 | 说明 |
|
| 模块 | 说明 |
|
||||||
|------|------|
|
|------|------|
|
||||||
| tests | 黑盒测试套件(`test_db.py`/`test_gateway.py`/`test_waf.py`/`test_dnfs.py`/`test_captcha.py`/`test_smtp.py`/`test_ai_review.py`/`test_static.py`/`test_core_confs.py`),纯标准库 HTTP/TCP 用户旅程模拟,经 `test_runner.py` 统一调度(详见 [tests/blackbox/README.md](../tests/blackbox/README.md)) |
|
| tests | 黑盒测试套件(`test_db.py`/`test_gateway.py`/`test_waf.py`/`test_dnfs.py`/`test_captcha.py`/`test_smtp.py`/`test_ai_review.py`/`test_static.py`/`test_core_confs.py`),纯标准库 HTTP/TCP 用户旅程模拟,经 `test_runner.py` 统一调度(详见 [tests/blackbox/README.md](../tests/blackbox/README.md)) |
|
||||||
| docs | 文档聚合(security/ 审计、reviews/ 审查、reports/ 分析),70 篇 md |
|
| docs | 文档聚合(security/ 审计、reviews/ 审查、reports/ 分析),64 篇 md + 1 篇 html |
|
||||||
|
|
||||||
## 九、根级构建与部署脚本
|
## 九、根级构建与部署脚本
|
||||||
|
|
||||||
@@ -92,5 +92,5 @@
|
|||||||
- **日志链路**:各服务 `log_client.py`(异步批量)→ `db2:8005` Logs 表(fail-safe,db2 不可用时静默丢弃)
|
- **日志链路**:各服务 `log_client.py`(异步批量)→ `db2:8005` Logs 表(fail-safe,db2 不可用时静默丢弃)
|
||||||
- **配置链路**:`gateway` → `core` 内嵌 confs:8010(内存缓存 + mtime 重载,仅 gateway.ini);其余服务配置纯 CLI 传递
|
- **配置链路**:`gateway` → `core` 内嵌 confs:8010(内存缓存 + mtime 重载,仅 gateway.ini);其余服务配置纯 CLI 传递
|
||||||
- **数据/审核/邮件/缓存**:`dnfs2` → `db2:8005`(pymsg 持久化);`dnfs2` → `ai_review:8007`(审核);注册邮件 `ca` → `smtp:8008`;登录验证码 `ca` → `captcha2:8009`;`memdb:8011`(内存缓存,libmsg RPC,暂无业务调用方,预留热点数据/会话缓存)
|
- **数据/审核/邮件/缓存**:`dnfs2` → `db2:8005`(pymsg 持久化);`dnfs2` → `ai_review:8007`(审核);注册邮件 `ca` → `smtp:8008`;登录验证码 `ca` → `captcha2:8009`;`memdb:8011`(内存缓存,libmsg RPC,暂无业务调用方,预留热点数据/会话缓存)
|
||||||
- **机器人数据链**:`mtc` 编译 Python DSL → `.mt`(或手写 `.mt`)→ `pybot` 执行 / `--record` 录制 episode(images + actions.npy + imu.npy)→ `actpred` 两阶段训练 → ONNX/Hailo HEF 部署 → `pybot --loop` 实时预测闭环(bot4cj 语义对齐)
|
- **机器人数据链**:`mtc` 编译 Python DSL → `.mt`(或手写 `.mt`)→ `pybot` 执行 / `--record` 录制 episode(images + actions.npy + imu.npy)→ `actpred` 两阶段训练 → ONNX/Hailo HEF 部署(`predict_async` 异步流水线)→ `pybot --loop` 实时预测闭环(bot4cj 语义对齐,HEF 后端推理与执行并行)
|
||||||
- **端口一览**:80 redirect80 · 443 waf · 8443 ca · 8445 gateway · 5200/8010 core · 9200 fsys · 8006 dnfs2 · 8005 db2 · 8007 ai_review · 8009 captcha2 · 8008 smtp · 8011 memdb · 8765 aibox
|
- **端口一览**:80 redirect80 · 443 waf · 8443 ca · 8445 gateway · 5200/8010 core · 9200 fsys · 8006 dnfs2 · 8005 db2 · 8007 ai_review · 8009 captcha2 · 8008 smtp · 8011 memdb · 8765 aibox
|
||||||
@@ -1,101 +0,0 @@
|
|||||||
# actpred 动作归一化与逐列权重开发方案
|
|
||||||
|
|
||||||
> 文档类型:开发方案
|
|
||||||
> 模块:actpred(联动 pybot 数据契约)
|
|
||||||
> 作者:自动生成
|
|
||||||
> 日期:2026-08-14
|
|
||||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 1. 问题陈述
|
|
||||||
|
|
||||||
闭环评估(R4)暴露:z-score 归一化把 18 列方差拉平,隐式实现"全列等权",导致**舵机精度 ↔ spk.F 频率跟踪的权衡**——未归一化模型舵机准但 spk.F 崩(21k MSE),归一化模型 spk.F 改善 3× 但简单动作恶化 1.3-25×。
|
|
||||||
|
|
||||||
本质:**缩放方式(数值稳定性)与列重要性(优化方向)是两个正交维度**,当前被隐式耦合。方案目标:固定物理范围归一化(数值稳定、免 stats 文件)+ 显式逐列权重(舵机优先),两者解耦。
|
|
||||||
|
|
||||||
## 2. 方案设计
|
|
||||||
|
|
||||||
### 2.1 固定物理范围 min-max 归一化(actpred 侧)
|
|
||||||
|
|
||||||
位置决策:**归一化保留在 actpred 侧,pybot 继续记录原始量纲**。理由:
|
|
||||||
|
|
||||||
- **raw at rest**:episode 是数据资产,原始角度/频率可复用于逆运动学、碰撞检测等分析,归一化不可逆
|
|
||||||
- **反归一化链条**:若 pybot 侧归一化,`--loop` 执行前需反归一化,min/max 表要同时存在于 recorder 与 loop 两端;现状 `predict()` 返回原始量纲,pybot 无感知,边界更干净
|
|
||||||
- **代码库先例**:IMU 即"pybot 记录原始、actpred `imu_norm.py` 按物理量程归一化",动作走同一条路保持一致
|
|
||||||
- **pybot 零依赖承诺**:归一化逻辑入 recorder 会加深 pybot 对 numpy 的依赖面
|
|
||||||
|
|
||||||
固定范围常量表(新增 `actpred/action_norm.py`,类比 `imu_norm.py`):
|
|
||||||
|
|
||||||
| 列 | 物理范围 | 依据 |
|
|
||||||
|---|---|---|
|
|
||||||
| 0-9 舵机角 | **[0, 180]** | `pybot/hd/servo.py` `MAX_ANGLE=180.0`(校准后全行程 0-180°) |
|
|
||||||
| 10 spk.F | [0, 5000] | `@Speaker max_f=5000` |
|
|
||||||
| 11 spk.delay | [0, 2000] | 实用上限(脚本用 10-500ms) |
|
|
||||||
| 12-17 保留位 | 恒 0 | 不参与归一化,损失权重 0 |
|
|
||||||
|
|
||||||
min-max 映射:`x_norm = (x - lo) / (hi - lo)` → [0,1](或映射到 [-1,1],与 IMU 归一化输出区间一致,待实现时定)。
|
|
||||||
|
|
||||||
优点:免 action_stats.npy(消除 stats 路径类 bug 的整类风险)、零状态部署、数值区间有界。风险:min-max 对离群值敏感——若真实数据某列越界,需 clamp 到 [0,1] 并告警(防御口径与 `_action_bounds` 一致)。
|
|
||||||
|
|
||||||
可选细化(后续):按各舵机 `|max|`(s1=90°、s2=180°…)分列归一化,90° 舵机能用满 [0,1],分辨率更高;需要读取 @Servo 块的 max 列构建表。
|
|
||||||
|
|
||||||
### 2.2 逐列加权损失
|
|
||||||
|
|
||||||
```python
|
|
||||||
COLUMN_WEIGHTS = np.array([
|
|
||||||
1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, # 0-9 舵机(安全关键)
|
|
||||||
0.2, 0.2, # 10-11 spk.F / spk.delay(声音与节奏,不影响运动)
|
|
||||||
0.0, 0.0, 0.0, 0.0, 0.0, 0.0, # 12-17 保留位(恒 0,纯噪声)
|
|
||||||
], dtype=np.float32)
|
|
||||||
# loss = mean(w * (pred - target)^2) / mean(w)
|
|
||||||
```
|
|
||||||
|
|
||||||
实施点:
|
|
||||||
|
|
||||||
- `train.py` / `pretrain_mt.py` 的 `train_one_epoch` / `validate`:MSE 处乘权重向量(权重存 DataConfig,随 checkpoint 序列化,推理端无需感知)
|
|
||||||
- 权重可经 CLI `--column-weights` 覆盖(默认用上表)
|
|
||||||
- 训练集/验证集用同一权重;`evaluate.py` 的逐列误差报告不受影响(报告仍按无权重口径,便于对照)
|
|
||||||
|
|
||||||
### 2.3 归一化与权重的组合矩阵
|
|
||||||
|
|
||||||
| 配置 | 数值稳定性 | 舵机精度 | spk.F 跟踪 | 预期适用 |
|
|
||||||
|---|---|---|---|---|
|
|
||||||
| raw + 等权(模型 A) | 差(loss 数百量级) | 好 | 崩 | 仅演示/对照 |
|
|
||||||
| z-score + 等权(模型 B) | 好 | 差 | 好 | 现状 |
|
|
||||||
| **固定范围 + 加权(本方案)** | 好 | 好 | 中(w=0.2 保留梯度) | **目标** |
|
|
||||||
|
|
||||||
## 3. 配套改进(同批实施)
|
|
||||||
|
|
||||||
### 3.1 R4 评估器修正(指标失效修复)
|
|
||||||
|
|
||||||
1. **oob_rate 恒值列退化**:span<阈值(如 1e-3)的列不参与越界判定;或改用固定物理界(舵机 [0,180]、spk.F [0,5000])。推荐后者——与 §2.1 常量表共用,且"越物理界"比"越经验界"更符合安全语义
|
|
||||||
2. **error_growth 长度混叠**:按 episode 归一化后聚合(每 episode 自身首末段比值取中位数),避免短 episode 尖峰主导曲线
|
|
||||||
|
|
||||||
### 3.2 R1-B rollout + scheduled sampling
|
|
||||||
|
|
||||||
闭环误差累积源于 teacher forcing 训练与自预测推理的不匹配。已实现的 `--rollout-len K --ss-prob p` 直接对标:建议 `K=4, p=0.5`,与加权损失同批训练对照。
|
|
||||||
|
|
||||||
### 3.3 数据扩展
|
|
||||||
|
|
||||||
- **实机录制**(根本瓶颈):dry-run 帧为噪声、IMU 全零,视觉/姿态通路无真实信号。实机 `--record` 后复跑本实验
|
|
||||||
- 验证集 1 个 episode 选模噪声大:val_split 0.15→0.3
|
|
||||||
- 补充站立/保持类长序列(当前 std 类仅 1-81 步)
|
|
||||||
|
|
||||||
## 4. 实施步骤
|
|
||||||
|
|
||||||
| 步骤 | 内容 | 产出 |
|
|
||||||
|---|---|---|
|
|
||||||
| 1 | 新增 `actpred/action_norm.py`:固定范围常量表 + normalize/denormalize + clamp 防御 | 独立纯 numpy 模块 + 单测 |
|
|
||||||
| 2 | DataConfig 增 `normalize_mode`(`zscore`/`fixed_range`)与 `column_weights`;dataset/train 接入 | 训练链路支持 + 单测 |
|
|
||||||
| 3 | `evaluate.py` oob 改物理界、error_growth 按 episode 归一化 | 评估器修正 + 单测 |
|
|
||||||
| 4 | 训练对照:固定范围+加权 vs z-score+等权(同种子、同数据) | 模型 C/D checkpoint |
|
|
||||||
| 5 | 复跑 R4 闭环评估,验证权衡解耦 | 对比报告 |
|
|
||||||
|
|
||||||
步骤 1-3 可并行;步骤 4-5 依赖前置完成。预计步骤 1-3 改动集中在 `actpred/`(`action_norm.py` 新增、`config.py`/`dataset.py`/`train.py`/`evaluate.py` 小改),不影响 pybot 与部署链(归一化在推理端内部完成,`predict()` 输出仍为原始量纲)。
|
|
||||||
|
|
||||||
## 5. 决策记录
|
|
||||||
|
|
||||||
- **归一化位置**:讨论过"移入 pybot 按列 min/max 归一化"。结论:不采纳(raw at rest、反归一化链条变长、IMU 先例不一致),但采纳其"固定物理范围"内核,落在 actpred 侧常量表实现
|
|
||||||
- **舵机范围**:±180° 为误述,实际全行程 0-180°(`servo.py MAX_ANGLE=180.0`),归一化表按 [0,180]
|
|
||||||
- **spk.F 对数尺度**:备选未定。若加权后 spk.F 跟踪仍不足,再考虑 log 变换(频率感知是对数的,10→1000 跨两个数量级)
|
|
||||||
@@ -4,8 +4,7 @@
|
|||||||
> 模块:actpred(hailo_infer.py cv2.LUT 快路径)
|
> 模块:actpred(hailo_infer.py cv2.LUT 快路径)
|
||||||
> 作者:自动生成
|
> 作者:自动生成
|
||||||
> 日期:2026-08-17
|
> 日期:2026-08-17
|
||||||
> 前置:`actpred_hailo_preprocess_optimization_plan_20260817.md`(优化方案,已实施)、
|
> 前置:`actpred_inference_speed_benchmark_20260813.md`(旧基准)
|
||||||
> `actpred_inference_speed_benchmark_20260813.md`(旧基准)
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
@@ -5,7 +5,6 @@
|
|||||||
> 作者:自动生成
|
> 作者:自动生成
|
||||||
> 日期:2026-08-14
|
> 日期:2026-08-14
|
||||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`(模型 A/B)
|
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`(模型 A/B)
|
||||||
> 开发方案:`actpred_action_norm_weighting_plan_20260814.md`(本实验为其步骤 4)
|
|
||||||
> 后续实验:模型 D(fixed_range + 加权 + R1-B rollout),训练中
|
> 后续实验:模型 D(fixed_range + 加权 + R1-B rollout),训练中
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|||||||
@@ -1,154 +0,0 @@
|
|||||||
# actpred Hailo 推理预处理优化:cv2.LUT 查表合并归一化与量化(已实施)
|
|
||||||
|
|
||||||
> 文档类型:开发方案(已实施)
|
|
||||||
> 模块:actpred(部署链 hailo_infer.py)
|
|
||||||
> 作者:自动生成
|
|
||||||
> 日期:2026-08-17(原型实测 + 集成实施)
|
|
||||||
> 前置:`actpred_inference_speed_benchmark_20260813.md`(速度基准)、
|
|
||||||
> 2026-08-17 新 640×480 数据集重测(本方案 §2 数据)
|
|
||||||
|
|
||||||
> 📌 **版本历程**
|
|
||||||
> 1. 初版:numpy 查表法为核心,预期端到端 18-20ms
|
|
||||||
> 2. 原型实测(§3.4):**numpy 高级索引查表在 Pi5 上无加速**(fancy
|
|
||||||
> indexing 的 gather 开销与逐元素 float 运算相当),预期不成立;
|
|
||||||
> 合并系数乘加仅 -5ms;**cv2.LUT 实测 15.8× 加速**(§3.4)
|
|
||||||
> 3. **已实施(§5)**:`hailo_infer.py` 集成 cv2.LUT 快路径(可选依赖,
|
|
||||||
> 无 cv2 自动回退 float 路径),端到端 52.4 → **16.0 ms**,输出与
|
|
||||||
> float 路径逐点一致(max_diff = 0)
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 1. 背景与动机
|
|
||||||
|
|
||||||
2026-08-17 用新录制的 640×480 数据集(免 resize 训练样本)重测 HEF 推理速度:
|
|
||||||
|
|
||||||
| 阶段 | 640×480 输入 | 224×224 输入(旧数据) |
|
|
||||||
|---|---|---|
|
|
||||||
| 端到端 predict | **52.4 ms** | 60.0 ms |
|
|
||||||
| 预处理 | **35.9 ms** | 43.1 ms |
|
|
||||||
| NPU run | 12.5 ms | 12.9 ms |
|
|
||||||
| 后处理 | 0.6 ms | 0.6 ms |
|
|
||||||
|
|
||||||
输入尺寸已与模型一致(resize 变 no-op),**预处理仍是端到端延迟的主导项**(68%)。
|
|
||||||
旧基准报告(2026-08-13)已记录同类问题:1080p 采集的预处理缩放约 45ms 主导端到端延迟,
|
|
||||||
当时通过把默认分辨率改回 640×480 缓解;本次分析证明 640×480 下 Python 端预处理
|
|
||||||
仍有约 36ms 的优化空间——但空间比初版估计小(见 §3.4 实测)。
|
|
||||||
|
|
||||||
## 2. 瓶颈分析(100 次均值分段计时)
|
|
||||||
|
|
||||||
| 子步骤 | 耗时 | 说明 |
|
|
||||||
|---|---|---|
|
|
||||||
| `_frame_float`(3 帧) | **27.2 ms** | PIL fromarray/convert + resize(同尺寸 no-op 仍有开销)+ `(x/255-mean)/std` float 运算(92 万像素×3 通道) |
|
|
||||||
| `create_bindings` + 帧量化绑定 | **9.1 ms** | `_bind_frames`:float 归一化帧 → `round(x/q.scale+q.zp)` 再量化回 uint8 |
|
|
||||||
| history/pred 绑定 | 0.07 ms | 可忽略 |
|
|
||||||
|
|
||||||
**关键冗余**:当前流程先做 float 归一化(`_frame_float`),随后又量化回 uint8(`_bind_frames`)——
|
|
||||||
中间 float 帧是纯中间态,归一化与量化可以**合并为一步**。
|
|
||||||
|
|
||||||
## 3. 优化方案
|
|
||||||
|
|
||||||
### 3.1 主方案:合并系数乘加(纯 numpy,无新依赖)
|
|
||||||
|
|
||||||
HEF 帧输入是 uint8 量化值:`q = round(((x/255 - mean)/std) / q_scale + q_zp)`。
|
|
||||||
把 4 步运算合并为**一次乘加**(系数预计算,按通道):
|
|
||||||
|
|
||||||
```
|
|
||||||
q = round(x * A[c] + B[c])
|
|
||||||
A[c] = 1 / (255 * std[c] * q_scale)
|
|
||||||
B[c] = -mean[c] / (std[c] * q_scale) + q_zp
|
|
||||||
```
|
|
||||||
|
|
||||||
- 现流程 4 次全数组运算(`/255`、`-mean`、`/std`、`/q_scale+q_zp`)→ 1 次乘加
|
|
||||||
- 数值完全等价(实测 max_diff = 0 LSB,见 §3.4)
|
|
||||||
- 直通模式(`_normalize_frames=False`)同理:`A=1/q_scale, B=q_zp`
|
|
||||||
|
|
||||||
### 3.2 附带优化(与主方案正交)
|
|
||||||
|
|
||||||
1. **免 PIL 路径**:输入已是 np.ndarray 时直接运算,跳过 `Image.fromarray/convert`
|
|
||||||
2. **resize no-op 跳过**:`image.size == (W,H)` 时不调用 `resize()`
|
|
||||||
3. **缓冲复用**:`np.multiply/np.add` 带 `out=` 参数,减少临时数组分配
|
|
||||||
|
|
||||||
### 3.3 已实施快路径:cv2.LUT(可选依赖,无 cv2 自动回退)
|
|
||||||
|
|
||||||
`cv2.LUT(img, lut)` 是 C 实现的 uint8→uint8 查表,单次遍历,无 float 中间态——
|
|
||||||
把"归一化+量化"合并为一步。集成要点(`hailo_infer.py`):
|
|
||||||
|
|
||||||
- 惰性导入 cv2(`HAS_CV2`);**无 cv2 或非逐帧输入布局(ncfhw/nhwc_fused)自动回退 float 路径**
|
|
||||||
- 加载时按每帧 `quant_info` 构建 256 项表(归一化/直通两种,`_normalize_frames` 判定)
|
|
||||||
- **OpenCV 5.0 的 cv2.LUT 只接受单通道表**((256,3) 会断言失败),逐通道应用后 `cv2.merge`
|
|
||||||
- 输入尺寸与模型不一致时先 PIL resize(仅此情况走 PIL)
|
|
||||||
- 依赖安装:`opencv-python-headless`(5.0.0,清华镜像;官方源过慢)
|
|
||||||
- 注:本方案实测后发现**合并系数乘加(§3.1)未实施**——cv2.LUT 收益已远超,
|
|
||||||
避免冗余代码;§3.1 保留作为无 cv2 环境下的潜在备选
|
|
||||||
|
|
||||||
### 3.4 原型实测评估(2026-08-17,100-200 次均值,640×480×3 帧)
|
|
||||||
|
|
||||||
| 方案 | 3 帧耗时 | 对比现 float 28.5ms | 数值差异 |
|
|
||||||
|---|---|---|---|
|
|
||||||
| 现 float 路径 | 28.5 ms | 基线 | — |
|
|
||||||
| numpy 高级索引查表 | 26.6 ms | **无加速(1.0x)** | 0 LSB |
|
|
||||||
| 合并系数乘加 | 23.2 ms | 快 19%(-5ms) | 0 LSB |
|
|
||||||
| **cv2.LUT(逐通道)** | **1.68 ms** | **快 15.8×** | **0 LSB** |
|
|
||||||
|
|
||||||
**结论**:
|
|
||||||
- numpy fancy indexing 的 gather 开销与逐元素 float 运算相当,查表法**无收益**
|
|
||||||
- 纯 numpy 可达的最佳是合并乘加(-5ms),受限于 92 万像素 float32 数组的内存吞吐
|
|
||||||
- cv2.LUT 的 C 查表把帧预处理降一个数量级,且数值完全等价(0 LSB)
|
|
||||||
|
|
||||||
### 3.5 实测收益(已实施,集成后 NPU 实测)
|
|
||||||
|
|
||||||
| 指标 | 优化前 | 优化后(cv2.LUT) |
|
|
||||||
|---|---|---|
|
|
||||||
| 帧预处理(3 帧) | 26.5 ms | **3.0 ms(分段实测含绑定)** |
|
|
||||||
| 端到端 predict | 52.4 ms | **16.1 ms(3.3×)**(200 次均值,见测试报告) |
|
|
||||||
| 输出一致性 | — | 与 float 路径 max_diff = **0.0000**(7 类输入 + 8 seed 随机) |
|
|
||||||
| HEF vs torch | 0.50° | 0.30°(4 角复验 PASS) |
|
|
||||||
| vs PyTorch CPU | 2.1× | **6.7×**(16.1 vs 107.9 ms) |
|
|
||||||
|
|
||||||
闭环 10Hz(100ms 预算)占用率 52% → **16%**,20Hz 32%,30Hz(33ms)48%。
|
|
||||||
|
|
||||||
> 完整验收数据见 `actpred_cv2_lut_benchmark_report_20260817.md`(测试报告)。
|
|
||||||
|
|
||||||
## 4. 一致性约束(关键)
|
|
||||||
|
|
||||||
- **训练侧不动**:`dataset.py` 的 `Resize+ToTensor+Normalize` 保持 float 路径——
|
|
||||||
合并乘加只是推理侧对"归一化+量化"的合并实现,数值等价(实测 0 LSB)
|
|
||||||
- **三后端口径**:`predict.py`(PyTorch,float 输入不量化)、`onnx_infer.py`(ORT,float)
|
|
||||||
不受影响;优化仅作用于 `hailo_infer.py` 的 uint8 量化输入路径
|
|
||||||
- **帧输入判定不变**:`_normalize_frames`(量化网格负零偏判定)决定走
|
|
||||||
归一化还是直通,逻辑与现状一致
|
|
||||||
- 旧 5 输入 / 旧版单图像 HEF 布局同样适用(按各自 quant_info 算系数)
|
|
||||||
|
|
||||||
## 5. 实施步骤(已完成 ✅)
|
|
||||||
|
|
||||||
| 步骤 | 内容 | 状态 |
|
|
||||||
|---|---|---|
|
|
||||||
| 1 | `hailo_infer.py`:加载时按每帧 quant_info 构建 256 项 LUT(归一化/直通两种) | ✅ `_build_frame_luts` |
|
|
||||||
| 2 | `predict` 的 nhwc_frames 分支走 cv2.LUT 快路径(尺寸不匹配时 PIL resize);无 cv2 回退 float | ✅ `_frame_lut_quant` + `HAS_CV2` |
|
|
||||||
| 3 | NPU 实测:640×480 数据集 100 次均值 | ✅ 16.0 ms(52.4 → 16.0) |
|
|
||||||
| 4 | `_npu_test.py` 4 角 + HEF vs torch 对比复验 | ✅ PASS,0.30°(≤1° 预算) |
|
|
||||||
| 5 | 全套单测(TestHailoPredictor 等) | ✅ 188 tests OK |
|
|
||||||
| 6 | 依赖安装:opencv-python-headless 5.0.0(清华镜像) | ✅ venv 已装 |
|
|
||||||
|
|
||||||
## 6. 风险与对策
|
|
||||||
|
|
||||||
| 风险 | 对策 |
|
|
||||||
|---|---|
|
|
||||||
| 查表取整与 float 路径存在差异 | 实测 0 LSB(表项 float32 与现路径同精度);NPU 输出 vs float 路径 max_diff=0,vs torch 0.30° |
|
|
||||||
| 无 cv2 环境 | 惰性导入 + 自动回退 float 路径,行为与优化前完全一致 |
|
|
||||||
| OpenCV 5.0 多通道 LUT 断言失败 | 逐通道单表 + `cv2.merge`(实测无性能损失) |
|
|
||||||
| 非 nhwc_frames 布局(旧 HEF) | 仅该布局启用 cv2 路径,其余走原 float 路径 |
|
|
||||||
| 未来输入尺寸变更 | LUT 与尺寸无关(逐像素映射);尺寸不一致时才走 PIL resize |
|
|
||||||
|
|
||||||
## 7. 决策记录
|
|
||||||
|
|
||||||
- **查表法(初版核心)实测证伪**:numpy 高级索引 gather 在 Pi5 上与逐元素 float
|
|
||||||
运算同速(26.6 vs 28.5ms),"查表降为内存索引即快"的假设不成立——教训:
|
|
||||||
性能方案须原型实测,不能凭直觉估计
|
|
||||||
- **最终方案选 cv2.LUT**:C 实现查表把帧预处理从 26.5ms 降到 1.7ms(15.8×),
|
|
||||||
端到端 52.4 → 16.0ms;数值 0 LSB 等价。合并系数乘加(-5ms)收益被碾压,
|
|
||||||
不再实施(§3.1 保留为无 cv2 环境的备选)
|
|
||||||
- **依赖取舍**:opencv-python-headless 作为可选依赖(惰性导入 + 自动回退),
|
|
||||||
有 cv2 环境收益 3.3×,无 cv2 环境行为不变——收益远大于依赖成本
|
|
||||||
- **训练侧不改**:训练用 float 归一化(增强/数值精度),推理侧查表实现——两侧
|
|
||||||
口径在"归一化值"层面一致,量化差异由部署链固有预算吸收(HEF vs torch 0.30° 余量充足)
|
|
||||||
@@ -1,82 +0,0 @@
|
|||||||
# actpred 多头 one-hot 分类输出开发方案
|
|
||||||
|
|
||||||
> 文档类型:开发方案
|
|
||||||
> 模块:actpred(联动 pybot 数据契约)
|
|
||||||
> 作者:自动生成
|
|
||||||
> 日期:2026-08-14
|
|
||||||
> 前置实验:`pybot_actpred_dataset_training_eval_20260814.md`、
|
|
||||||
> `actpred_fixed_range_weighted_experiment_20260814.md`(模型 A-E 系列)
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 1. 背景与动机
|
|
||||||
|
|
||||||
模型 A-E 系列实验定位到两个结构性难题:
|
|
||||||
|
|
||||||
1. **尺度放大**:fixed_range 线性 min-max 下 spk.F 的 1 个归一化单位 = 5000Hz 原始量,闭环反馈中误差滚雪球(模型 C/D 发散根因之一)
|
|
||||||
2. **精度错配**:常用值(10Hz → 0.002)挤在线性尺度的低端,[0,1] 精度预算被浪费(log_range 正在验证的对数尺度解法)
|
|
||||||
|
|
||||||
多头 one-hot 是本问题的**结构性解法**:把 18 维连续回归换成 18 组独立分类头,每组输出一个 one-hot 分布。分类天然消除尺度放大(每箱等距)与物理界问题(箱只覆盖物理范围),且多列 softmax 是 NPU(Hailo)最擅长的计算形态。
|
|
||||||
|
|
||||||
## 2. 方案设计
|
|
||||||
|
|
||||||
### 2.1 分箱设计(每组 one-hot 的维度)
|
|
||||||
|
|
||||||
| 列组 | 分箱方式 | 箱数/列 | 依据 |
|
|
||||||
|---|---|---|---|
|
|
||||||
| 舵机 0-9(10 列) | 均匀 **1°/箱** | 180 | 匹配数据粒度:微动作插值 0.75°/步,5° 箱会把相邻步量化进同一箱、毁掉插值结构 |
|
|
||||||
| spk.F(10) | 对数分箱 | 32 | 频率对数感知,低端聚集;log10(1+x) 均匀切 |
|
|
||||||
| spk.delay(11) | 对数分箱 | 32 | 同左 |
|
|
||||||
| 保留位 12-17 | 无头 | 0 | 恒 0,不参与 |
|
|
||||||
|
|
||||||
**输出总量**:10×180 + 32 + 32 = **1864 logits**(当前 18)。最后层 128→1864 ≈ 238K 参数,模型 409K → ~650K,仍属轻量。
|
|
||||||
|
|
||||||
### 2.2 各层改动
|
|
||||||
|
|
||||||
| 文件 | 改动 |
|
|
||||||
|---|---|
|
|
||||||
| `actpred/bins.py`(新增) | 分箱表 + `encode(x)→one-hot` / `decode(p)→连续值`(含对数箱);纯 numpy,供推理/部署复用 |
|
|
||||||
| `model.py` | PredictionHead 输出 1864 维 → reshape (B, 18, bins) → 逐组 softmax |
|
|
||||||
| `dataset.py` | target 从 (B,18) 连续值 → 每组 one-hot;**软标签**:真值高斯扩散到 ±1-2 箱 |
|
|
||||||
| `train.py` | 损失 = Σ 列权重 × CE(组输出, one-hot 真值);列权重沿用 default_column_weights |
|
|
||||||
| `predict.py` | 输出 decode:**期望值** `Σ pᵢ·箱中心`(推荐)或 argmax 箱中心 → 18 维原始量纲 |
|
|
||||||
| `evaluate.py` | 无感——predict 返回仍是 18 维原始量纲 |
|
|
||||||
| `export_deploy` / Hailo | 多组 softmax 为 NPU 友好形态 |
|
|
||||||
|
|
||||||
### 2.3 解码策略(执行时如何取连续值)
|
|
||||||
|
|
||||||
| 方式 | 行为 | 评估 |
|
|
||||||
|---|---|---|
|
|
||||||
| **期望值**(推荐) | `Σ pᵢ × 箱中心` | 平滑连续,分辨率超越箱宽(1° 箱可出 46.4°);单峰假设 |
|
|
||||||
| argmax | 取最大概率箱中心 | 确定性,量化到 1° |
|
|
||||||
| 按分布采样 | 随机抽箱 | 探索用,后续离线 RL 可复用 |
|
|
||||||
|
|
||||||
## 3. 关键决策点与理由
|
|
||||||
|
|
||||||
1. **舵机 1°/箱而非 5°/箱**:录制数据的微动作插值粒度为 0.75°,箱宽必须匹配数据粒度;5° 箱会把 45.0°/45.75° 压进同一箱,模型学不到特意录制的插值结构
|
|
||||||
2. **软标签**:804 样本对 1864 logits 偏少,多数箱极少被访问。真值 one-hot 高斯扩散到邻近箱,未访问箱也能获得梯度
|
|
||||||
3. **期望值解码补序数信息**:one-hot 丢失"45° 和 46° 相邻"先验,期望值解码在推理端部分恢复;训练端靠软标签缓解
|
|
||||||
4. **pybot 契约零变化**:predict 输出仍是 18 维原始量纲,闭环与录制数据无感知
|
|
||||||
5. **与现有归一化链路的兼容**:one-hot 方案独立于 zscore/fixed_range/log_range——分箱编码替代归一化,不叠加使用
|
|
||||||
|
|
||||||
## 4. 实施步骤
|
|
||||||
|
|
||||||
| 步骤 | 内容 | 产出 |
|
|
||||||
|---|---|---|
|
|
||||||
| 1 | `bins.py`:分箱表 + encode/decode + 对数箱 + 软标签工具 | 独立纯 numpy 模块 + 单测 |
|
|
||||||
| 2 | `dataset.py`:目标 one-hot 化(软标签可开关) | 数据链路 + 单测 |
|
|
||||||
| 3 | `model.py` + `train.py`:多头 softmax + 加权 CE | 模型与训练 + 单测 |
|
|
||||||
| 4 | `predict.py`:期望值/argmax 解码集成 + 冒烟 | 推理链路 + 单测 |
|
|
||||||
| 5 | 训练模型 F,与 A-E 闭环五方对比 | 对比报告 |
|
|
||||||
|
|
||||||
## 5. 与其他方案的定位
|
|
||||||
|
|
||||||
- **模型 E(log_range)正在训练**:若对数尺度已解决精度错配与闭环发散,one-hot 降级为可选优化(分类头天然有界、免归一化常量表的优势仍在)
|
|
||||||
- 若 E 仍失败:one-hot 是下一个强候选(结构性消除尺度放大)
|
|
||||||
- **互补关系**:one-hot 方案与逐列权重、R1-B rollout 正交,可叠加(rollout 治 exposure bias 的结论不变,任何损失形式都需要它)
|
|
||||||
|
|
||||||
## 6. 决策记录
|
|
||||||
|
|
||||||
- **箱宽**:初拟 5°/箱(保守,顾虑数据量),讨论后定为 1°/箱——数据插值粒度 0.75° 决定箱宽下限;spk 两列对数 32 箱
|
|
||||||
- **单头多选 vs 多头单选**:softmax 单头在数学上是单选择分布;"多组 one-hot" = 18 个独立 softmax 头并行,每列一个选择
|
|
||||||
- **多峰分布**:当前数据单峰为主,期望值解码够用;若未来出现双峰(如跨脚本姿态冲突),升级为逐箱 sigmoid 或 MDN
|
|
||||||
@@ -1,116 +0,0 @@
|
|||||||
# actpred 训练时安全图像增强
|
|
||||||
|
|
||||||
> 文档类型:特性报告
|
|
||||||
> 模块:actpred
|
|
||||||
> 日期:2026-08-13
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## 1. 背景
|
|
||||||
|
|
||||||
pybot dry-run 录制的 episode 数据样本量有限(11 个脚本 → 12 个 episode → ~728 样本),训练容易过拟合。直接增大数据集需要更多 `.mt` 脚本或更多轮次录制,成本较高。
|
|
||||||
|
|
||||||
图像增强(data augmentation)是低成本扩充样本的标准手段,但机器人动作预测场景有特殊约束:**空间变换(翻转/旋转/裁剪)会改变图像的语义坐标,而动作标签未随之变换,导致训练对错误**。例如水平翻转后,左右舵机角度应互换,但标签不变,形成错误训练信号。
|
|
||||||
|
|
||||||
## 2. 方案
|
|
||||||
|
|
||||||
### 2.1 增强范围
|
|
||||||
|
|
||||||
仅启用**颜色空间变换**(ColorJitter),不涉及任何空间变换:
|
|
||||||
|
|
||||||
| 变换 | 参数 | 说明 |
|
|
||||||
|------|------|------|
|
|
||||||
| 亮度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同光照条件 |
|
|
||||||
| 对比度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同摄像头增益 |
|
|
||||||
| 饱和度 | `±20%`(factor ∈ [0.8, 1.2]) | 模拟不同色温环境 |
|
|
||||||
| 色调 | `±5%`(factor ∈ [-0.05, 0.05]) | 轻微色偏,避免模型过依赖颜色 |
|
|
||||||
|
|
||||||
**不使用的变换**(因需同步变换动作标签):
|
|
||||||
|
|
||||||
| 变换 | 原因 |
|
|
||||||
|------|------|
|
|
||||||
| 水平/垂直翻转 | 左右舵机角度需互换,否则标签错误 |
|
|
||||||
| 旋转 | 重力方向变化,姿态语义改变 |
|
|
||||||
| 裁剪/平移 | 视野变化导致动作-画面对应关系改变 |
|
|
||||||
| 透视变换 | 改变空间关系 |
|
|
||||||
|
|
||||||
### 2.2 同一样本内帧间参数一致性
|
|
||||||
|
|
||||||
一个训练样本包含 `frame_len` 帧连续图像(默认 3 帧,间隔 10ms)。真实场景中连续帧的光照条件不变,因此**同一样本内所有帧共享同一组增强参数**,而非每帧独立随机:
|
|
||||||
|
|
||||||
```python
|
|
||||||
if self.augment:
|
|
||||||
bf = random.uniform(0.8, 1.2) # 亮度因子,所有帧共用
|
|
||||||
cf = random.uniform(0.8, 1.2) # 对比度因子
|
|
||||||
sf = random.uniform(0.8, 1.2) # 饱和度因子
|
|
||||||
hf = random.uniform(-0.05, 0.05) # 色调因子
|
|
||||||
for p in image_paths:
|
|
||||||
image = Image.open(p).convert("RGB")
|
|
||||||
if self.augment:
|
|
||||||
image = TF.adjust_brightness(image, bf)
|
|
||||||
image = TF.adjust_contrast(image, cf)
|
|
||||||
image = TF.adjust_saturation(image, sf)
|
|
||||||
image = TF.adjust_hue(image, hf)
|
|
||||||
frames.append(self.transform(image))
|
|
||||||
```
|
|
||||||
|
|
||||||
不同样本(不同 `__getitem__` 调用)使用独立随机参数,保证整个数据集的增强多样性。
|
|
||||||
|
|
||||||
### 2.3 自适应 frame_len
|
|
||||||
|
|
||||||
增强逻辑在 `for p in image_paths` 循环前生成参数、循环内逐帧施加。`image_paths` 的长度由 `frame_len` 决定(`image_files[t - frame_len + 1 : t + 1]`),循环自动适配。**`frame_len` 为 1、2、3、5、10 等任意值时均无需修改代码**。
|
|
||||||
|
|
||||||
### 2.4 训练/验证隔离
|
|
||||||
|
|
||||||
| 数据集 | 增强 | 说明 |
|
|
||||||
|--------|------|------|
|
|
||||||
| 训练集 | `augment=cfg.train.augment` | CLI `--augment` 控制 |
|
|
||||||
| 验证集 | `augment=False`(默认) | 验证 loss 反映真实分布表现 |
|
|
||||||
|
|
||||||
推理端(`predict.py` / `onnx_infer.py` / `hailo_infer.py`)不涉及增强,使用原始图像。
|
|
||||||
|
|
||||||
## 3. CLI 接口
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# 训练时启用图像增强
|
|
||||||
bash _run.sh train --data-dir ./data --epochs 100 --augment
|
|
||||||
|
|
||||||
# 显式关闭(默认关闭)
|
|
||||||
bash _run.sh train --data-dir ./data --epochs 100 --no-augment
|
|
||||||
```
|
|
||||||
|
|
||||||
`--augment` 使用 `BooleanOptionalAction`,支持 `--augment` / `--no-augment` 双向开关。
|
|
||||||
|
|
||||||
## 4. 验证
|
|
||||||
|
|
||||||
### 4.1 功能验证
|
|
||||||
|
|
||||||
```
|
|
||||||
dataset size: 728 samples
|
|
||||||
augment diff (should be >0): 0.3767 # 增强前后图像有差异
|
|
||||||
two augment calls diff (should be >0): 0.2626 # 两次增强参数不同(随机性)
|
|
||||||
action_history diff (should be 0): 0.0 # 动作标签不受影响
|
|
||||||
target diff (should be 0): 0.0
|
|
||||||
```
|
|
||||||
|
|
||||||
### 4.2 端到端训练
|
|
||||||
|
|
||||||
`--augment --epochs 3`(64×64,batch=8)训练正常完成,loss 正常下降。
|
|
||||||
|
|
||||||
### 4.3 回归测试
|
|
||||||
|
|
||||||
全量单测 90/90 通过,无回归。
|
|
||||||
|
|
||||||
## 5. 改动文件
|
|
||||||
|
|
||||||
| 文件 | 改动 |
|
|
||||||
|------|------|
|
|
||||||
| `actpred/dataset.py` | `ActionDataset` 新增 `augment` 参数;`__getitem__` 在 Resize 前施加 ColorJitter(`torchvision.transforms.functional`) |
|
|
||||||
| `actpred/config.py` | `TrainConfig` 新增 `augment: bool = False`;`build_argparser` 新增 `--augment`;`apply_overrides` 新增映射 |
|
|
||||||
| `actpred/train.py` | 训练集构造传 `augment=cfg.train.augment`,验证集默认不增强 |
|
|
||||||
|
|
||||||
## 6. 后续方向
|
|
||||||
|
|
||||||
- **空间增强(需同步标签变换)**:水平翻转时自动交换左右舵机角度(s1↔s2, s3↔s4, s5↔s6, s7↔s8),可将训练数据量翻倍。需要了解舵机配对关系后实现。
|
|
||||||
- **动作噪声增强**:`pretrain_mt.py` 已支持 `--noise-std`,`train.py` 可同样引入,对 action_history 和 target 施加高斯扰动。
|
|
||||||
- **更多 `.mt` 脚本**:从根本上扩大动作多样性,是提升模型泛化能力的最有效手段。
|
|
||||||
@@ -4,8 +4,7 @@
|
|||||||
> 模块:actpred
|
> 模块:actpred
|
||||||
> 作者:自动生成
|
> 作者:自动生成
|
||||||
> 日期:2026-08-15
|
> 日期:2026-08-15
|
||||||
> 前置:`actpred_zscore_baseline_20260815.md`(模型 B 基线)、
|
> 前置:`actpred_zscore_baseline_20260815.md`(模型 B 基线)
|
||||||
> `actpred_onehot_multihead_dev_plan_20260814.md`(one-hot 方案)
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
@@ -5,8 +5,7 @@
|
|||||||
> 作者:自动生成
|
> 作者:自动生成
|
||||||
> 日期:2026-08-15
|
> 日期:2026-08-15
|
||||||
> 相关实验:`pybot_actpred_dataset_training_eval_20260814.md`(A/B 对照)、
|
> 相关实验:`pybot_actpred_dataset_training_eval_20260814.md`(A/B 对照)、
|
||||||
> `actpred_fixed_range_weighted_experiment_20260814.md`(C/D/E 失败系列)、
|
> `actpred_fixed_range_weighted_experiment_20260814.md`(C/D/E 失败系列)
|
||||||
> `actpred_onehot_multihead_dev_plan_20260814.md`(F 系列结构性方案)
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|||||||
File diff suppressed because it is too large.
Load diff
Loaded 100 of 160 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user