fix(aibox): FTS5 中文检索 + Token 压缩阈值生效——Compressor 抽离解耦
- db.py: FTS5 改用 trigram 分词(中文可检索),旧 unicode61 表自动迁移回灌;search_memories 增加 LIKE 兜底(%/_ 转义、MATCH 异常防护) - conversation_compressor.py: should_compress 补齐 Token 阈值判断(原为死配置),Token 触发时按比例压缩早期消息 - ConversationCompressor 抽离至 subsystems,注入 llm_client 打破与 DeepSeekChat 循环引用;aibox.py 3517→3018 行 - 新增 21 项测试(中文检索/旧表迁移/LIKE 转义/token 估算/压缩流程/失败回退),全量 232 项通过
This commit is contained in:
1 parent
26de527a09
commit
7f21b65428
7 files changed
+940
-533
No files matched your search
+32
-16
@@ -55,7 +55,7 @@ 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`(当前约 3706 行,集成 WebSocket 服务器、AI 引擎、工具/技能系统、任务调度器于一体;自 v3.0 起工具类与子系统已抽离至 `tools/`、`subsystems/`、`common/`、`tests/` 包,aibox.py 仅保留核心类与编排逻辑)。`client.py` 为示例 WebSocket 客户端,`voice_client.py` 为语音客户端,`llm_client.py` 为多厂商大模型客户端抽象层,`_run.sh` 为被 `core` 拉起的启动脚本。
|
> 注意:本模块入口文件为 `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` 拉起的启动脚本。
|
||||||
|
|
||||||
## 📁 目录结构
|
## 📁 目录结构
|
||||||
|
|
||||||
@@ -73,11 +73,13 @@ aibox/
|
|||||||
│ ├── constants.py # 常量定义
|
│ ├── constants.py # 常量定义
|
||||||
│ ├── i18n.py # 多语言
|
│ ├── i18n.py # 多语言
|
||||||
│ └── logger.py # 日志器
|
│ └── logger.py # 日志器
|
||||||
├── subsystems/ # 子系统包(Wave 1:从 aibox.py 抽离的基础设施)
|
├── subsystems/ # 子系统包(Wave 1/2:从 aibox.py 抽离的基础设施)
|
||||||
│ ├── executors.py # CommandExecutor + SubprocessCommandExecutor
|
│ ├── executors.py # CommandExecutor + SubprocessCommandExecutor
|
||||||
│ ├── prompt_manager.py # PromptManager(单例,提示词按数字前缀排序加载)
|
│ ├── prompt_manager.py # PromptManager(单例,提示词按数字前缀排序加载)
|
||||||
│ ├── db.py # MemoryDatabase + MemoDatabase(SQLite)
|
│ ├── db.py # MemoryDatabase + MemoDatabase(SQLite,FTS5 trigram 中文检索 + LIKE 兜底)
|
||||||
│ └── skill_manager.py # SkillManager(技能发现/加载/热重载)
|
│ ├── skill_manager.py # SkillManager(技能发现/加载/热重载)
|
||||||
|
│ ├── conversation_history.py # ConversationHistory(对话历史管理)
|
||||||
|
│ └── conversation_compressor.py # ConversationCompressor(对话压缩,注入 llm_client 解耦)
|
||||||
├── tools/ # 工具系统包(A 批 + B 批共 17 个工具类)
|
├── tools/ # 工具系统包(A 批 + B 批共 17 个工具类)
|
||||||
│ ├── base.py # Tool 抽象基类
|
│ ├── base.py # Tool 抽象基类
|
||||||
│ ├── registry.py # ToolRegistry 工具注册器
|
│ ├── registry.py # ToolRegistry 工具注册器
|
||||||
@@ -88,14 +90,16 @@ aibox/
|
|||||||
│ ├── 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/ # 单元测试套件(171 项全过)
|
├── tests/ # 单元测试套件(232 项全过)
|
||||||
│ ├── 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 项
|
||||||
│ ├── test_executors.py # 命令执行器
|
│ ├── test_executors.py # 命令执行器
|
||||||
│ ├── test_prompt_manager.py # 提示词管理器
|
│ ├── test_prompt_manager.py # 提示词管理器
|
||||||
│ ├── test_db.py # 记忆/备忘录数据库
|
│ ├── test_db.py # 记忆/备忘录数据库(含中文检索/迁移)
|
||||||
│ └── test_skill_manager.py # 技能管理器
|
│ ├── test_conversation_history.py # 对话历史
|
||||||
|
│ ├── test_conversation_compressor.py # 对话压缩器
|
||||||
|
│ └── test_skill_manager.py # 技能管理器
|
||||||
├── skills/ # 内置扩展技能(*.py 可执行脚本 / *.md Markdown 技能)
|
├── skills/ # 内置扩展技能(*.py 可执行脚本 / *.md Markdown 技能)
|
||||||
│ ├── hello-world.py
|
│ ├── hello-world.py
|
||||||
│ ├── aes-crypto.py
|
│ ├── aes-crypto.py
|
||||||
@@ -130,7 +134,7 @@ common ← subsystems ← tools ← aibox.py
|
|||||||
```bash
|
```bash
|
||||||
# 从仓库根目录运行
|
# 从仓库根目录运行
|
||||||
py -m pytest aibox/tests -q
|
py -m pytest aibox/tests -q
|
||||||
# 当前:171 项全过,约 11s
|
# 当前:232 项全过,约 11s
|
||||||
```
|
```
|
||||||
|
|
||||||
运行期数据存放于 `~/.aibox/`:
|
运行期数据存放于 `~/.aibox/`:
|
||||||
@@ -791,6 +795,18 @@ elif command == "/mycommand":
|
|||||||
|
|
||||||
## 📄 更新日志
|
## 📄 更新日志
|
||||||
|
|
||||||
|
### v3.1
|
||||||
|
- ✅ **修复 FTS5 中文搜索失效**([subsystems/db.py](subsystems/db.py)):
|
||||||
|
- FTS5 虚拟表改用 `tokenize='trigram'` 分词器(需 SQLite >= 3.34,环境不支持时自动退回 unicode61)
|
||||||
|
- 旧 unicode61 表在启动时自动迁移:DROP 重建后从 memories 表回灌全量数据
|
||||||
|
- `search_memories` 增加 LIKE 模糊检索兜底(trigram 对 <3 字符中文短词无法命中;关键词中 `%`/`_` 已转义为字面字符;MATCH 语法错误不再抛异常)
|
||||||
|
- ✅ **修复对话压缩 Token 阈值不生效**:
|
||||||
|
- `should_compress` 补齐 Token 数判断:消息条数或 Token 数任一达到阈值即触发压缩(此前 `compress_token_threshold` 为死配置)
|
||||||
|
- `compress_conversation` 配套:Token 触发(消息数未超阈值)时按比例压缩早期消息,不再因目标消息数计算而跳过
|
||||||
|
- ✅ **ConversationCompressor 抽离至 [subsystems/conversation_compressor.py](subsystems/conversation_compressor.py)**(Wave 3):
|
||||||
|
- 构造注入 `llm_client`(客户端抽象层)替代 AI 引擎实例,打破与 `DeepSeekChat` 的循环引用
|
||||||
|
- aibox.py 3517 → 3018 行;配套 17 项单元测试(估算/阈值判断/压缩流程/LLM 失败回退/孤立 tool 消息清理)
|
||||||
|
|
||||||
### v3.0
|
### v3.0
|
||||||
- ✅ **模块化重构(aibox.py 8450 行 → 3706 行,缩减 56%)**:
|
- ✅ **模块化重构(aibox.py 8450 行 → 3706 行,缩减 56%)**:
|
||||||
- **A 批基础工具抽离**([tools/basic_tools.py](tools/basic_tools.py)):`FileToolBase` / `ReadFileTool` / `WriteFileTool` / `ListFilesTool` / `SimpleWebCrawlerTool` / `GetCurrentTimeTool` / `CalculatorTool` 共 7 类,无外部子系统依赖;配套 43 项单元测试
|
- **A 批基础工具抽离**([tools/basic_tools.py](tools/basic_tools.py)):`FileToolBase` / `ReadFileTool` / `WriteFileTool` / `ListFilesTool` / `SimpleWebCrawlerTool` / `GetCurrentTimeTool` / `CalculatorTool` 共 7 类,无外部子系统依赖;配套 43 项单元测试
|
||||||
@@ -799,7 +815,7 @@ elif command == "/mycommand":
|
|||||||
- **测试套件**:[tests/](tests/) 目录 171 项全过,覆盖 A/B 批工具与 Wave 1 子系统
|
- **测试套件**:[tests/](tests/) 目录 171 项全过,覆盖 A/B 批工具与 Wave 1 子系统
|
||||||
- **修复**:`subsystems/db.py` 的 `get_recent_memories(limit=None)` SQL datatype mismatch bug(`LIMIT ?` 传 None 时移除 LIMIT 子句)
|
- **修复**:`subsystems/db.py` 的 `get_recent_memories(limit=None)` SQL datatype mismatch bug(`LIMIT ?` 传 None 时移除 LIMIT 子句)
|
||||||
- **架构约束**:建立 `common ← subsystems ← tools ← aibox` 单向依赖,禁止反向引用;`tools/__init__.py` 与 `subsystems/__init__.py` 采用"重导出"策略,调用点无需关心子模块位置
|
- **架构约束**:建立 `common ← subsystems ← tools ← aibox` 单向依赖,禁止反向引用;`tools/__init__.py` 与 `subsystems/__init__.py` 采用"重导出"策略,调用点无需关心子模块位置
|
||||||
- **未抽离**:Wave 2(`TaskScheduler` / `ConversationHistory`)、Wave 3(`ConversationCompressor`)、C 批工具(`ExpressEmotionTool` / `SubAgentTool`)、核心类(`DeepSeekChat` / `WebSocketHandler` / `EnhancedAIManager`)仍保留在 `aibox.py` 内
|
- **未抽离**:`TaskScheduler`、C 批工具(`ExpressEmotionTool` / `SubAgentTool`)、核心类(`DeepSeekChat` / `WebSocketHandler` / `EnhancedAIManager`)仍保留在 `aibox.py` 内(`ConversationHistory` 于 Wave 2、`ConversationCompressor` 于 v3.1 抽离)
|
||||||
|
|
||||||
### v2.9
|
### v2.9
|
||||||
- ✅ **编排者-执行者工具分层(工具集 profile)**:
|
- ✅ **编排者-执行者工具分层(工具集 profile)**:
|
||||||
|
|||||||
+4
-502
@@ -56,516 +56,18 @@ from tools import (
|
|||||||
AskUserTool,
|
AskUserTool,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 子系统(Wave 1 + Wave 2 抽离):命令执行器 / 提示词管理器 / 数据库 / 技能管理器 / 对话历史
|
# 子系统(Wave 1 + Wave 2 抽离):命令执行器 / 提示词管理器 / 数据库 / 技能管理器 / 对话历史 / 对话压缩
|
||||||
from subsystems import (
|
from subsystems import (
|
||||||
CommandExecutor, SubprocessCommandExecutor,
|
CommandExecutor, SubprocessCommandExecutor,
|
||||||
PromptManager,
|
PromptManager,
|
||||||
MemoryDatabase, MemoDatabase,
|
MemoryDatabase, MemoDatabase,
|
||||||
SkillManager,
|
SkillManager,
|
||||||
ConversationHistory,
|
ConversationHistory,
|
||||||
|
ConversationCompressor,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
class ConversationCompressor:
|
|
||||||
"""对话压缩管理器 - 优先压缩较早的消息,将tool调用信息一并压缩"""
|
|
||||||
|
|
||||||
def __init__(self, ai: 'DeepSeekChat', logger: Logger = None):
|
|
||||||
self.ai = ai
|
|
||||||
self.logger = logger or Logger("Compressor")
|
|
||||||
self.config = ConfigManager()
|
|
||||||
# 从配置获取最小值
|
|
||||||
self.min_messages_to_compress = constants.COMPRESSOR_MIN_MESSAGES
|
|
||||||
# 权重衰减因子:从配置获取,设为0.6使加权计算更合理
|
|
||||||
self.weight_decay_factor = 0.6 # 使用更高的衰减因子
|
|
||||||
|
|
||||||
def estimate_tokens(self, messages: List[Dict]) -> int:
|
|
||||||
"""估算消息的token数"""
|
|
||||||
total_tokens = 0
|
|
||||||
for msg in messages:
|
|
||||||
content = msg.get("content", "")
|
|
||||||
if content:
|
|
||||||
chinese_chars = len([c for c in content if '\u4e00' <= c <= '\u9fff'])
|
|
||||||
english_words = len(re.findall(r'[a-zA-Z]+', content))
|
|
||||||
other_chars = len(content) - chinese_chars - english_words
|
|
||||||
|
|
||||||
estimated = chinese_chars * 1.5 + english_words * 1.3 + other_chars * 0.8
|
|
||||||
total_tokens += estimated
|
|
||||||
total_tokens += 10
|
|
||||||
|
|
||||||
if msg.get("tool_calls"):
|
|
||||||
# 工具调用也计算token
|
|
||||||
for tc in msg.get("tool_calls", []):
|
|
||||||
func = tc.get("function", {})
|
|
||||||
if func.get("name"):
|
|
||||||
total_tokens += len(func.get("name", "")) * 1.5
|
|
||||||
if func.get("arguments"):
|
|
||||||
total_tokens += len(func.get("arguments", "")) * 0.8
|
|
||||||
total_tokens += 50
|
|
||||||
|
|
||||||
return int(total_tokens)
|
|
||||||
|
|
||||||
def should_compress(self, messages: List[Dict]) -> Tuple[bool, str]:
|
|
||||||
"""检查是否需要压缩对话 - 优先根据消息条数判断"""
|
|
||||||
if not constants.COMPRESS_ENABLED:
|
|
||||||
return False, "压缩功能已禁用"
|
|
||||||
|
|
||||||
conversation_messages = [m for m in messages if m.get("role") != "system"]
|
|
||||||
|
|
||||||
# 优先检查消息条数是否达到阈值
|
|
||||||
message_count = len(conversation_messages)
|
|
||||||
if message_count < constants.COMPRESS_MESSAGE_THRESHOLD:
|
|
||||||
return False, f"对话消息数量({message_count})小于压缩阈值({constants.COMPRESS_MESSAGE_THRESHOLD})"
|
|
||||||
|
|
||||||
# 消息条数达到阈值,进行压缩
|
|
||||||
print(f"⚠️ {i18n.get('compress_threshold_reached', message_count)}")
|
|
||||||
return True, f"消息数({message_count})超过阈值({constants.COMPRESS_MESSAGE_THRESHOLD})"
|
|
||||||
|
|
||||||
def _calculate_target_count(self, current_count: int, threshold: int, keep_ratio: float = 0.7) -> int:
|
|
||||||
"""
|
|
||||||
计算目标消息数量 - 压缩到阈值的 keep_ratio 倍
|
|
||||||
|
|
||||||
Args:
|
|
||||||
current_count: 当前消息数量
|
|
||||||
threshold: 压缩阈值
|
|
||||||
keep_ratio: 压缩后保留的比例(相对于阈值),默认0.7表示压缩到阈值的70%
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
目标消息数量
|
|
||||||
"""
|
|
||||||
# 压缩到阈值以下,保留阈值的 keep_ratio 倍
|
|
||||||
target = int(threshold * keep_ratio)
|
|
||||||
|
|
||||||
# 确保至少有 min_messages_to_compress 条消息
|
|
||||||
target = max(target, self.min_messages_to_compress)
|
|
||||||
|
|
||||||
# 确保不超过当前数量
|
|
||||||
target = min(target, current_count - 1)
|
|
||||||
|
|
||||||
return target
|
|
||||||
|
|
||||||
def _find_compressible_block_by_count(self, messages: List[Dict], target_count: int) -> Tuple[int, int]:
|
|
||||||
"""
|
|
||||||
根据目标消息数量找到可压缩块 - 不保护tool消息,直接压缩早期消息
|
|
||||||
|
|
||||||
Args:
|
|
||||||
messages: 消息列表
|
|
||||||
target_count: 目标消息数量
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
(start, end) 压缩块的起始和结束索引
|
|
||||||
"""
|
|
||||||
if len(messages) <= target_count:
|
|
||||||
return 0, 0
|
|
||||||
|
|
||||||
# 计算需要压缩的消息数量
|
|
||||||
messages_to_remove = len(messages) - target_count
|
|
||||||
|
|
||||||
# 确定要保留的最近消息数量
|
|
||||||
keep_recent = constants.COMPRESS_KEEP_RECENT
|
|
||||||
if keep_recent > 0:
|
|
||||||
keep_recent = max(keep_recent, 10) # 至少保留最近10条
|
|
||||||
|
|
||||||
# 压缩块应该从最早的消息开始,但需要保留前几轮对话
|
|
||||||
# 计算最早可以开始压缩的位置(跳过前2轮对话)
|
|
||||||
start_min = 0
|
|
||||||
|
|
||||||
# 尝试找到前2个对话轮次结束的位置
|
|
||||||
conversation_turns = self._build_conversation_turns(messages)
|
|
||||||
if len(conversation_turns) >= 2:
|
|
||||||
# 前2轮结束的位置
|
|
||||||
first_turn = conversation_turns[0]
|
|
||||||
second_turn = conversation_turns[1]
|
|
||||||
# 取第二轮结束的位置
|
|
||||||
if second_turn["assistant_index"] is not None:
|
|
||||||
start_min = second_turn["assistant_index"] + 1
|
|
||||||
else:
|
|
||||||
start_min = second_turn["user_index"] + 1
|
|
||||||
|
|
||||||
# 压缩结束的位置 = 总消息数 - 保留的最近消息数
|
|
||||||
end_max = len(messages) - keep_recent
|
|
||||||
|
|
||||||
# 确保压缩块有效
|
|
||||||
if end_max <= start_min + 2:
|
|
||||||
# 如果保护区域太大,尝试压缩更多
|
|
||||||
end_max = len(messages) - max(keep_recent // 2, 5)
|
|
||||||
|
|
||||||
# 计算需要压缩的块大小
|
|
||||||
block_size = min(messages_to_remove + 10, end_max - start_min)
|
|
||||||
block_size = max(block_size, 10) # 至少压缩10条
|
|
||||||
|
|
||||||
# 压缩块从 start_min 开始
|
|
||||||
start = start_min
|
|
||||||
end = min(start + block_size, end_max)
|
|
||||||
|
|
||||||
# 确保压缩块至少2条消息
|
|
||||||
if end - start < 2:
|
|
||||||
# 如果块太小,尝试从更早的位置开始
|
|
||||||
start = max(0, end_max - block_size)
|
|
||||||
end = end_max
|
|
||||||
|
|
||||||
self.logger.debug(f"找到压缩块: [{start}:{end}], 移除{end-start}条消息")
|
|
||||||
return start, end
|
|
||||||
|
|
||||||
def _build_conversation_turns(self, messages: List[Dict]) -> List[Dict]:
|
|
||||||
"""构建对话轮次"""
|
|
||||||
conversation_turns = []
|
|
||||||
|
|
||||||
i = 0
|
|
||||||
while i < len(messages):
|
|
||||||
msg = messages[i]
|
|
||||||
role = msg.get("role")
|
|
||||||
|
|
||||||
# 跳过系统消息
|
|
||||||
if role == "system":
|
|
||||||
i += 1
|
|
||||||
continue
|
|
||||||
|
|
||||||
# 遇到用户消息,开始一个新的轮次
|
|
||||||
if role == "user":
|
|
||||||
turn = {"user_index": i, "assistant_index": None, "tool_indices": []}
|
|
||||||
|
|
||||||
# 向后查找对应的助手消息和工具消息
|
|
||||||
j = i + 1
|
|
||||||
while j < len(messages):
|
|
||||||
next_msg = messages[j]
|
|
||||||
next_role = next_msg.get("role")
|
|
||||||
|
|
||||||
# 跳过系统消息
|
|
||||||
if next_role == "system":
|
|
||||||
j += 1
|
|
||||||
continue
|
|
||||||
|
|
||||||
# 找到助手消息
|
|
||||||
if next_role == "assistant":
|
|
||||||
turn["assistant_index"] = j
|
|
||||||
# 检查这个助手消息是否有 tool_calls
|
|
||||||
if next_msg.get("tool_calls"):
|
|
||||||
# 收集这个助手消息的 tool_call_ids
|
|
||||||
tool_call_ids = {tc.get("id") for tc in next_msg.get("tool_calls", []) if tc.get("id")}
|
|
||||||
# 继续查找对应的 tool 消息
|
|
||||||
k = j + 1
|
|
||||||
while k < len(messages) and len(turn["tool_indices"]) < len(tool_call_ids):
|
|
||||||
tool_msg = messages[k]
|
|
||||||
if tool_msg.get("role") == "tool":
|
|
||||||
tool_call_id = tool_msg.get("tool_call_id")
|
|
||||||
if tool_call_id and tool_call_id in tool_call_ids:
|
|
||||||
turn["tool_indices"].append(k)
|
|
||||||
k += 1
|
|
||||||
elif tool_msg.get("role") == "user":
|
|
||||||
break
|
|
||||||
else:
|
|
||||||
k += 1
|
|
||||||
break
|
|
||||||
|
|
||||||
if next_role == "user":
|
|
||||||
break
|
|
||||||
|
|
||||||
j += 1
|
|
||||||
|
|
||||||
conversation_turns.append(turn)
|
|
||||||
i += 1
|
|
||||||
else:
|
|
||||||
i += 1
|
|
||||||
|
|
||||||
return conversation_turns
|
|
||||||
|
|
||||||
def _is_range_safe(self, messages: List[Dict], start: int, end: int, protected: set) -> bool:
|
|
||||||
"""检查消息范围是否安全 - 简化版,不检查tool完整性"""
|
|
||||||
return True # 不保护tool消息,允许压缩
|
|
||||||
|
|
||||||
def _get_protected_indices(self, messages: List[Dict]) -> set:
|
|
||||||
"""
|
|
||||||
找出所有受保护的消息索引 - 简化版,只保护最近的对话轮次
|
|
||||||
"""
|
|
||||||
protected_indices = set()
|
|
||||||
|
|
||||||
# ========== 1. 保护最近的消息(由配置决定) ==========
|
|
||||||
keep_recent = constants.COMPRESS_KEEP_RECENT
|
|
||||||
if keep_recent > 0:
|
|
||||||
recent_start = max(0, len(messages) - keep_recent)
|
|
||||||
for i in range(recent_start, len(messages)):
|
|
||||||
protected_indices.add(i)
|
|
||||||
|
|
||||||
# ========== 2. 保护前2个对话轮次 ==========
|
|
||||||
conversation_turns = self._build_conversation_turns(messages)
|
|
||||||
|
|
||||||
if len(conversation_turns) >= 2:
|
|
||||||
# 保护第1轮
|
|
||||||
first_turn = conversation_turns[0]
|
|
||||||
protected_indices.add(first_turn["user_index"])
|
|
||||||
if first_turn["assistant_index"] is not None:
|
|
||||||
protected_indices.add(first_turn["assistant_index"])
|
|
||||||
for tool_idx in first_turn["tool_indices"]:
|
|
||||||
protected_indices.add(tool_idx)
|
|
||||||
|
|
||||||
# 保护第2轮
|
|
||||||
second_turn = conversation_turns[1]
|
|
||||||
protected_indices.add(second_turn["user_index"])
|
|
||||||
if second_turn["assistant_index"] is not None:
|
|
||||||
protected_indices.add(second_turn["assistant_index"])
|
|
||||||
for tool_idx in second_turn["tool_indices"]:
|
|
||||||
protected_indices.add(tool_idx)
|
|
||||||
|
|
||||||
# 调试日志
|
|
||||||
if self.logger.level >= 3:
|
|
||||||
self.logger.debug(f"受保护的消息索引: {sorted(protected_indices)}")
|
|
||||||
|
|
||||||
return protected_indices
|
|
||||||
|
|
||||||
def _build_compression_prompt(self, messages: List[Dict]) -> str:
|
|
||||||
"""构建压缩提示词 - 包含工具调用信息"""
|
|
||||||
conversation_lines = []
|
|
||||||
|
|
||||||
for msg in messages:
|
|
||||||
role = msg.get("role")
|
|
||||||
content = msg.get("content", "")
|
|
||||||
|
|
||||||
if role == "user":
|
|
||||||
if content:
|
|
||||||
conversation_lines.append(f"用户: {content}")
|
|
||||||
|
|
||||||
elif role == "assistant":
|
|
||||||
if content:
|
|
||||||
conversation_lines.append(f"助手: {content}")
|
|
||||||
|
|
||||||
# 包含工具调用信息
|
|
||||||
tool_calls = msg.get("tool_calls", [])
|
|
||||||
for tc in tool_calls:
|
|
||||||
func = tc.get("function", {})
|
|
||||||
tool_name = func.get("name", "")
|
|
||||||
arguments = func.get("arguments", "")
|
|
||||||
if tool_name:
|
|
||||||
if arguments:
|
|
||||||
conversation_lines.append(f"[调用工具: {tool_name}({arguments[:100]})]")
|
|
||||||
else:
|
|
||||||
conversation_lines.append(f"[调用工具: {tool_name}]")
|
|
||||||
|
|
||||||
elif role == "tool":
|
|
||||||
if content:
|
|
||||||
useful = self._extract_useful_content(content)
|
|
||||||
if useful:
|
|
||||||
conversation_lines.append(f"[工具结果: {useful}]")
|
|
||||||
|
|
||||||
conversation_text = "\n".join(conversation_lines)
|
|
||||||
|
|
||||||
max_input_chars = constants.COMPRESSOR_MAX_INPUT_CHARS
|
|
||||||
if len(conversation_text) > max_input_chars:
|
|
||||||
conversation_text = conversation_text[:max_input_chars] + "\n...(内容过长,已截断)..."
|
|
||||||
|
|
||||||
prompt = f"""请对以下对话内容进行智能压缩和总结。要求:
|
|
||||||
|
|
||||||
1. **提取核心信息**:对话的主要主题、关键问题、重要结论、用户需求
|
|
||||||
2. **保留重要细节**:用户的重要要求、AI的重要回复、达成的共识
|
|
||||||
3. **保留工具调用摘要**:简要说明使用了什么工具,得到了什么关键结果
|
|
||||||
4. **结构化输出**:使用清晰的段落结构
|
|
||||||
|
|
||||||
对话内容:
|
|
||||||
---
|
|
||||||
{conversation_text}
|
|
||||||
---
|
|
||||||
|
|
||||||
请输出压缩后的对话摘要:"""
|
|
||||||
|
|
||||||
return prompt
|
|
||||||
|
|
||||||
def _extract_useful_content(self, content: str) -> str:
|
|
||||||
"""从工具返回内容中提取有用信息"""
|
|
||||||
if not content:
|
|
||||||
return ""
|
|
||||||
|
|
||||||
if "✅" in content:
|
|
||||||
lines = content.split("\n")
|
|
||||||
for line in lines[:3]:
|
|
||||||
if "✅" in line:
|
|
||||||
return line.strip()
|
|
||||||
|
|
||||||
if "❌" in content:
|
|
||||||
lines = content.split("\n")
|
|
||||||
for line in lines[:2]:
|
|
||||||
if "❌" in line:
|
|
||||||
return line.strip()
|
|
||||||
|
|
||||||
if len(content) > 100:
|
|
||||||
return content[:100] + "..."
|
|
||||||
|
|
||||||
return content
|
|
||||||
|
|
||||||
def _compress_via_llm(self, messages: List[Dict]) -> Optional[str]:
|
|
||||||
"""通过语言模型压缩消息(经由客户端抽象层调用,支持多厂商)"""
|
|
||||||
if not messages:
|
|
||||||
return None
|
|
||||||
|
|
||||||
prompt = self._build_compression_prompt(messages)
|
|
||||||
|
|
||||||
print(f"📦 {i18n.get('conversation_compress')}: 正在压缩 {len(messages)} 条消息...")
|
|
||||||
|
|
||||||
try:
|
|
||||||
summary = self.ai.client.chat_completion(
|
|
||||||
[
|
|
||||||
{"role": "system", "content": "你是一个专业的对话摘要助手,擅长提取对话核心信息并进行结构化总结。请保留工具调用的关键结果。"},
|
|
||||||
{"role": "user", "content": prompt}
|
|
||||||
],
|
|
||||||
temperature=0.3,
|
|
||||||
max_tokens=constants.COMPRESSOR_LLM_MAX_TOKENS
|
|
||||||
)
|
|
||||||
return summary
|
|
||||||
except Exception as e:
|
|
||||||
self.logger.error(f"语言模型压缩异常: {e}")
|
|
||||||
return None
|
|
||||||
|
|
||||||
def _validate_message_structure(self, messages: List[Dict]) -> List[Dict]:
|
|
||||||
"""验证并修复消息结构,移除孤立的 tool 消息"""
|
|
||||||
if not messages:
|
|
||||||
return messages
|
|
||||||
|
|
||||||
# 找出所有 assistant 消息的 tool_call_ids
|
|
||||||
assistant_tool_ids = set()
|
|
||||||
for msg in messages:
|
|
||||||
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
|
||||||
for tc in msg.get("tool_calls", []):
|
|
||||||
tc_id = tc.get("id")
|
|
||||||
if tc_id:
|
|
||||||
assistant_tool_ids.add(tc_id)
|
|
||||||
|
|
||||||
# 过滤掉没有对应 assistant 的 tool 消息
|
|
||||||
filtered = []
|
|
||||||
for msg in messages:
|
|
||||||
if msg.get("role") == "tool":
|
|
||||||
tool_call_id = msg.get("tool_call_id")
|
|
||||||
if tool_call_id and tool_call_id not in assistant_tool_ids:
|
|
||||||
print(f" ⚠️ 移除孤立的 tool 消息: {tool_call_id}")
|
|
||||||
continue
|
|
||||||
filtered.append(msg)
|
|
||||||
|
|
||||||
return filtered
|
|
||||||
|
|
||||||
def compress_conversation(self, messages: List[Dict]) -> List[Dict]:
|
|
||||||
"""
|
|
||||||
压缩对话内容 - 优先压缩较早的消息,包含tool调用信息一并压缩
|
|
||||||
当消息条数超过阈值时,压缩到阈值以下
|
|
||||||
"""
|
|
||||||
if not messages:
|
|
||||||
return messages
|
|
||||||
|
|
||||||
# 分离系统消息和对话消息
|
|
||||||
system_messages = []
|
|
||||||
conversation_messages = []
|
|
||||||
|
|
||||||
for msg in messages:
|
|
||||||
if msg.get("role") == "system":
|
|
||||||
system_messages.append(msg)
|
|
||||||
else:
|
|
||||||
conversation_messages.append(msg)
|
|
||||||
|
|
||||||
if not conversation_messages:
|
|
||||||
return system_messages
|
|
||||||
|
|
||||||
original_count = len(conversation_messages)
|
|
||||||
original_tokens = self.estimate_tokens(conversation_messages)
|
|
||||||
threshold = constants.COMPRESS_MESSAGE_THRESHOLD
|
|
||||||
|
|
||||||
print(f"\n📊 对话压缩检查:")
|
|
||||||
print(f" 当前消息数: {original_count} 条 (阈值: {threshold} 条)")
|
|
||||||
print(f" 当前Token数: {original_tokens} 个")
|
|
||||||
|
|
||||||
# 计算目标消息数量 - 压缩到阈值的70%
|
|
||||||
target_count = self._calculate_target_count(original_count, threshold, keep_ratio=0.7)
|
|
||||||
print(f" 目标消息数: {target_count} 条 (压缩后应少于阈值)")
|
|
||||||
|
|
||||||
# 如果目标数量大于等于当前数量,不需要压缩
|
|
||||||
if target_count >= original_count:
|
|
||||||
print(f" 目标消息数大于等于当前,跳过压缩")
|
|
||||||
return messages
|
|
||||||
|
|
||||||
# 计算需要压缩的消息数量
|
|
||||||
messages_to_compress = original_count - target_count
|
|
||||||
print(f" 需要压缩的消息数量: {messages_to_compress} 条")
|
|
||||||
|
|
||||||
# 根据消息数量查找可压缩块
|
|
||||||
compress_start, compress_end = self._find_compressible_block_by_count(
|
|
||||||
conversation_messages, target_count
|
|
||||||
)
|
|
||||||
|
|
||||||
if compress_end - compress_start >= 2:
|
|
||||||
compressible = conversation_messages[compress_start:compress_end]
|
|
||||||
compress_tokens = self.estimate_tokens(compressible)
|
|
||||||
compress_count = len(compressible)
|
|
||||||
|
|
||||||
print(f" 压缩块范围: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
|
||||||
print(f" 压缩块token数: {compress_tokens}")
|
|
||||||
|
|
||||||
# 如果压缩块大小仍然不够,尝试扩大压缩块
|
|
||||||
if compress_count < messages_to_compress:
|
|
||||||
print(f" 压缩块大小不足({compress_count} < {messages_to_compress}),扩大压缩块...")
|
|
||||||
# 扩大压缩块
|
|
||||||
additional = messages_to_compress - compress_count + 5
|
|
||||||
new_start = max(0, compress_start - additional // 2)
|
|
||||||
new_end = min(len(conversation_messages), compress_end + additional // 2)
|
|
||||||
compress_start, compress_end = new_start, new_end
|
|
||||||
compressible = conversation_messages[compress_start:compress_end]
|
|
||||||
compress_count = len(compressible)
|
|
||||||
print(f" 扩大后压缩块: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
|
||||||
else:
|
|
||||||
print(f" 未找到合适的压缩块,使用默认压缩策略")
|
|
||||||
# 使用默认压缩策略:压缩最早的一半消息,保留最近的消息
|
|
||||||
keep_recent = constants.COMPRESS_KEEP_RECENT
|
|
||||||
compress_end = len(conversation_messages) - keep_recent
|
|
||||||
compress_start = max(0, compress_end - messages_to_compress - 5)
|
|
||||||
compress_start = max(compress_start, 0)
|
|
||||||
compressible = conversation_messages[compress_start:compress_end]
|
|
||||||
compress_count = len(compressible)
|
|
||||||
print(f" 默认压缩块: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
|
||||||
|
|
||||||
# 确保压缩块有足够多的消息
|
|
||||||
if compress_count < 5:
|
|
||||||
print(f" 压缩块太小({compress_count}条),跳过压缩")
|
|
||||||
return messages
|
|
||||||
|
|
||||||
# 调用语言模型压缩
|
|
||||||
summary = self._compress_via_llm(compressible)
|
|
||||||
|
|
||||||
if not summary:
|
|
||||||
print(f" ❌ 压缩失败,保留原始对话")
|
|
||||||
return messages
|
|
||||||
|
|
||||||
# 构建压缩后的消息列表
|
|
||||||
compressed_result = []
|
|
||||||
|
|
||||||
# 添加压缩块之前的消息
|
|
||||||
compressed_result.extend(conversation_messages[:compress_start])
|
|
||||||
|
|
||||||
# 插入压缩摘要(作为系统消息)
|
|
||||||
compressed_result.append({
|
|
||||||
"role": "system",
|
|
||||||
"content": i18n.get('compressed_summary', summary)
|
|
||||||
})
|
|
||||||
|
|
||||||
# 添加压缩块之后的消息(包括最近的消息)
|
|
||||||
compressed_result.extend(conversation_messages[compress_end:])
|
|
||||||
|
|
||||||
# 验证并修复消息结构
|
|
||||||
compressed_result = self._validate_message_structure(compressed_result)
|
|
||||||
|
|
||||||
# 合并系统消息
|
|
||||||
final_messages = system_messages + compressed_result
|
|
||||||
|
|
||||||
# 统计结果
|
|
||||||
final_count = len([m for m in final_messages if m.get("role") != "system"])
|
|
||||||
final_tokens = self.estimate_tokens([m for m in final_messages if m.get("role") != "system"])
|
|
||||||
|
|
||||||
print(f"\n✅ {i18n.get('conversation_compressed', original_count, final_count)}")
|
|
||||||
print(f" Token: {original_tokens} -> {final_tokens} (压缩率: {final_tokens/original_tokens:.1%})")
|
|
||||||
|
|
||||||
# 检查是否压缩到阈值以下
|
|
||||||
if final_count >= threshold:
|
|
||||||
print(f" ⚠️ 警告: 压缩后消息数({final_count})仍大于等于阈值({threshold})")
|
|
||||||
|
|
||||||
return final_messages
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# ==================== 任务调度器 ====================
|
# ==================== 任务调度器 ====================
|
||||||
|
|
||||||
class TaskScheduler:
|
class TaskScheduler:
|
||||||
@@ -1404,8 +906,8 @@ class DeepSeekChat:
|
|||||||
# 共享技能管理器:传入时复用(避免每个任务重复扫描技能),否则按需创建
|
# 共享技能管理器:传入时复用(避免每个任务重复扫描技能),否则按需创建
|
||||||
self._shared_skill_manager = skill_manager
|
self._shared_skill_manager = skill_manager
|
||||||
|
|
||||||
# 添加对话压缩器
|
# 添加对话压缩器(注入 llm 客户端而非 AI 实例,避免循环引用)
|
||||||
self.compressor = ConversationCompressor(self, self.logger)
|
self.compressor = ConversationCompressor(self.client, self.logger)
|
||||||
|
|
||||||
self.max_history = max_history or constants.DEFAULT_MAX_HISTORY
|
self.max_history = max_history or constants.DEFAULT_MAX_HISTORY
|
||||||
self.preserve_system = constants.PRESERVE_SYSTEM_PROMPTS
|
self.preserve_system = constants.PRESERVE_SYSTEM_PROMPTS
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ subsystems - aibox 子系统模块(包)
|
|||||||
- skill_manager.py : SkillManager
|
- skill_manager.py : SkillManager
|
||||||
Wave 2:
|
Wave 2:
|
||||||
- conversation_history.py : ConversationHistory(对话历史管理)
|
- conversation_history.py : ConversationHistory(对话历史管理)
|
||||||
|
- conversation_compressor.py : ConversationCompressor(对话压缩,注入 llm_client)
|
||||||
|
|
||||||
依赖方向:common <- subsystems <- aibox
|
依赖方向:common <- subsystems <- aibox
|
||||||
不反向依赖 aibox 任何符号。
|
不反向依赖 aibox 任何符号。
|
||||||
@@ -22,6 +23,7 @@ from .prompt_manager import PromptManager
|
|||||||
from .db import MemoryDatabase, MemoDatabase
|
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
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"CommandExecutor",
|
"CommandExecutor",
|
||||||
@@ -31,4 +33,5 @@ __all__ = [
|
|||||||
"MemoDatabase",
|
"MemoDatabase",
|
||||||
"SkillManager",
|
"SkillManager",
|
||||||
"ConversationHistory",
|
"ConversationHistory",
|
||||||
|
"ConversationCompressor",
|
||||||
]
|
]
|
||||||
@@ -0,0 +1,523 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
subsystems.conversation_compressor - 对话压缩子系统
|
||||||
|
|
||||||
|
抽离自 aibox.py,负责对话上下文压缩:消息条数或 Token 数达到阈值时,
|
||||||
|
将早期对话(含工具调用信息)经 LLM 压缩为摘要,保留最近消息。
|
||||||
|
|
||||||
|
依赖注入:构造时传入 llm_client(llm_client.LLMClient 实例)而非 AI
|
||||||
|
引擎实例,避免与 DeepSeekChat 形成循环引用;依赖方向 common <- subsystems。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
from typing import Dict, List, Optional, Tuple
|
||||||
|
|
||||||
|
from common import constants, i18n, Logger, ConfigManager
|
||||||
|
|
||||||
|
|
||||||
|
class ConversationCompressor:
|
||||||
|
"""对话压缩管理器 - 优先压缩较早的消息,将tool调用信息一并压缩"""
|
||||||
|
|
||||||
|
def __init__(self, llm_client=None, logger: Logger = None):
|
||||||
|
self.llm_client = llm_client
|
||||||
|
self.logger = logger or Logger("Compressor")
|
||||||
|
self.config = ConfigManager()
|
||||||
|
# 从配置获取最小值
|
||||||
|
self.min_messages_to_compress = constants.COMPRESSOR_MIN_MESSAGES
|
||||||
|
# 权重衰减因子:从配置获取,设为0.6使加权计算更合理
|
||||||
|
self.weight_decay_factor = 0.6 # 使用更高的衰减因子
|
||||||
|
|
||||||
|
def estimate_tokens(self, messages: List[Dict]) -> int:
|
||||||
|
"""估算消息的token数"""
|
||||||
|
total_tokens = 0
|
||||||
|
for msg in messages:
|
||||||
|
content = msg.get("content", "")
|
||||||
|
if content:
|
||||||
|
chinese_chars = len([c for c in content if '\u4e00' <= c <= '\u9fff'])
|
||||||
|
english_words = len(re.findall(r'[a-zA-Z]+', content))
|
||||||
|
other_chars = len(content) - chinese_chars - english_words
|
||||||
|
|
||||||
|
estimated = chinese_chars * 1.5 + english_words * 1.3 + other_chars * 0.8
|
||||||
|
total_tokens += estimated
|
||||||
|
total_tokens += 10
|
||||||
|
|
||||||
|
if msg.get("tool_calls"):
|
||||||
|
# 工具调用也计算token
|
||||||
|
for tc in msg.get("tool_calls", []):
|
||||||
|
func = tc.get("function", {})
|
||||||
|
if func.get("name"):
|
||||||
|
total_tokens += len(func.get("name", "")) * 1.5
|
||||||
|
if func.get("arguments"):
|
||||||
|
total_tokens += len(func.get("arguments", "")) * 0.8
|
||||||
|
total_tokens += 50
|
||||||
|
|
||||||
|
return int(total_tokens)
|
||||||
|
|
||||||
|
def should_compress(self, messages: List[Dict]) -> Tuple[bool, str]:
|
||||||
|
"""检查是否需要压缩对话 - 消息条数或 Token 数任一达到阈值即压缩"""
|
||||||
|
if not constants.COMPRESS_ENABLED:
|
||||||
|
return False, "压缩功能已禁用"
|
||||||
|
|
||||||
|
conversation_messages = [m for m in messages if m.get("role") != "system"]
|
||||||
|
|
||||||
|
# 优先检查消息条数是否达到阈值
|
||||||
|
message_count = len(conversation_messages)
|
||||||
|
if message_count >= constants.COMPRESS_MESSAGE_THRESHOLD:
|
||||||
|
print(f"⚠️ {i18n.get('compress_threshold_reached', message_count)}")
|
||||||
|
return True, f"消息数({message_count})超过阈值({constants.COMPRESS_MESSAGE_THRESHOLD})"
|
||||||
|
|
||||||
|
# 再检查 Token 数是否达到阈值(长文档/长工具输出场景)
|
||||||
|
token_count = self.estimate_tokens(conversation_messages)
|
||||||
|
if token_count >= constants.COMPRESS_TOKEN_THRESHOLD:
|
||||||
|
print(f"⚠️ {i18n.get('compress_token_threshold_reached', token_count)}")
|
||||||
|
return True, f"Token数({token_count})超过阈值({constants.COMPRESS_TOKEN_THRESHOLD})"
|
||||||
|
|
||||||
|
return False, f"消息数({message_count})与Token数({token_count})均未达到阈值"
|
||||||
|
|
||||||
|
def _calculate_target_count(self, current_count: int, threshold: int, keep_ratio: float = 0.7) -> int:
|
||||||
|
"""
|
||||||
|
计算目标消息数量 - 压缩到阈值的 keep_ratio 倍
|
||||||
|
|
||||||
|
Args:
|
||||||
|
current_count: 当前消息数量
|
||||||
|
threshold: 压缩阈值
|
||||||
|
keep_ratio: 压缩后保留的比例(相对于阈值),默认0.7表示压缩到阈值的70%
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
目标消息数量
|
||||||
|
"""
|
||||||
|
# 压缩到阈值以下,保留阈值的 keep_ratio 倍
|
||||||
|
target = int(threshold * keep_ratio)
|
||||||
|
|
||||||
|
# 确保至少有 min_messages_to_compress 条消息
|
||||||
|
target = max(target, self.min_messages_to_compress)
|
||||||
|
|
||||||
|
# 确保不超过当前数量
|
||||||
|
target = min(target, current_count - 1)
|
||||||
|
|
||||||
|
return target
|
||||||
|
|
||||||
|
def _find_compressible_block_by_count(self, messages: List[Dict], target_count: int) -> Tuple[int, int]:
|
||||||
|
"""
|
||||||
|
根据目标消息数量找到可压缩块 - 不保护tool消息,直接压缩早期消息
|
||||||
|
|
||||||
|
Args:
|
||||||
|
messages: 消息列表
|
||||||
|
target_count: 目标消息数量
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(start, end) 压缩块的起始和结束索引
|
||||||
|
"""
|
||||||
|
if len(messages) <= target_count:
|
||||||
|
return 0, 0
|
||||||
|
|
||||||
|
# 计算需要压缩的消息数量
|
||||||
|
messages_to_remove = len(messages) - target_count
|
||||||
|
|
||||||
|
# 确定要保留的最近消息数量
|
||||||
|
keep_recent = constants.COMPRESS_KEEP_RECENT
|
||||||
|
if keep_recent > 0:
|
||||||
|
keep_recent = max(keep_recent, 10) # 至少保留最近10条
|
||||||
|
|
||||||
|
# 压缩块应该从最早的消息开始,但需要保留前几轮对话
|
||||||
|
# 计算最早可以开始压缩的位置(跳过前2轮对话)
|
||||||
|
start_min = 0
|
||||||
|
|
||||||
|
# 尝试找到前2个对话轮次结束的位置
|
||||||
|
conversation_turns = self._build_conversation_turns(messages)
|
||||||
|
if len(conversation_turns) >= 2:
|
||||||
|
# 前2轮结束的位置
|
||||||
|
first_turn = conversation_turns[0]
|
||||||
|
second_turn = conversation_turns[1]
|
||||||
|
# 取第二轮结束的位置
|
||||||
|
if second_turn["assistant_index"] is not None:
|
||||||
|
start_min = second_turn["assistant_index"] + 1
|
||||||
|
else:
|
||||||
|
start_min = second_turn["user_index"] + 1
|
||||||
|
|
||||||
|
# 压缩结束的位置 = 总消息数 - 保留的最近消息数
|
||||||
|
end_max = len(messages) - keep_recent
|
||||||
|
|
||||||
|
# 确保压缩块有效
|
||||||
|
if end_max <= start_min + 2:
|
||||||
|
# 如果保护区域太大,尝试压缩更多
|
||||||
|
end_max = len(messages) - max(keep_recent // 2, 5)
|
||||||
|
|
||||||
|
# 计算需要压缩的块大小
|
||||||
|
block_size = min(messages_to_remove + 10, end_max - start_min)
|
||||||
|
block_size = max(block_size, 10) # 至少压缩10条
|
||||||
|
|
||||||
|
# 压缩块从 start_min 开始
|
||||||
|
start = start_min
|
||||||
|
end = min(start + block_size, end_max)
|
||||||
|
|
||||||
|
# 确保压缩块至少2条消息
|
||||||
|
if end - start < 2:
|
||||||
|
# 如果块太小,尝试从更早的位置开始
|
||||||
|
start = max(0, end_max - block_size)
|
||||||
|
end = end_max
|
||||||
|
|
||||||
|
self.logger.debug(f"找到压缩块: [{start}:{end}], 移除{end-start}条消息")
|
||||||
|
return start, end
|
||||||
|
|
||||||
|
def _build_conversation_turns(self, messages: List[Dict]) -> List[Dict]:
|
||||||
|
"""构建对话轮次"""
|
||||||
|
conversation_turns = []
|
||||||
|
|
||||||
|
i = 0
|
||||||
|
while i < len(messages):
|
||||||
|
msg = messages[i]
|
||||||
|
role = msg.get("role")
|
||||||
|
|
||||||
|
# 跳过系统消息
|
||||||
|
if role == "system":
|
||||||
|
i += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 遇到用户消息,开始一个新的轮次
|
||||||
|
if role == "user":
|
||||||
|
turn = {"user_index": i, "assistant_index": None, "tool_indices": []}
|
||||||
|
|
||||||
|
# 向后查找对应的助手消息和工具消息
|
||||||
|
j = i + 1
|
||||||
|
while j < len(messages):
|
||||||
|
next_msg = messages[j]
|
||||||
|
next_role = next_msg.get("role")
|
||||||
|
|
||||||
|
# 跳过系统消息
|
||||||
|
if next_role == "system":
|
||||||
|
j += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 找到助手消息
|
||||||
|
if next_role == "assistant":
|
||||||
|
turn["assistant_index"] = j
|
||||||
|
# 检查这个助手消息是否有 tool_calls
|
||||||
|
if next_msg.get("tool_calls"):
|
||||||
|
# 收集这个助手消息的 tool_call_ids
|
||||||
|
tool_call_ids = {tc.get("id") for tc in next_msg.get("tool_calls", []) if tc.get("id")}
|
||||||
|
# 继续查找对应的 tool 消息
|
||||||
|
k = j + 1
|
||||||
|
while k < len(messages) and len(turn["tool_indices"]) < len(tool_call_ids):
|
||||||
|
tool_msg = messages[k]
|
||||||
|
if tool_msg.get("role") == "tool":
|
||||||
|
tool_call_id = tool_msg.get("tool_call_id")
|
||||||
|
if tool_call_id and tool_call_id in tool_call_ids:
|
||||||
|
turn["tool_indices"].append(k)
|
||||||
|
k += 1
|
||||||
|
elif tool_msg.get("role") == "user":
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
k += 1
|
||||||
|
break
|
||||||
|
|
||||||
|
if next_role == "user":
|
||||||
|
break
|
||||||
|
|
||||||
|
j += 1
|
||||||
|
|
||||||
|
conversation_turns.append(turn)
|
||||||
|
i += 1
|
||||||
|
else:
|
||||||
|
i += 1
|
||||||
|
|
||||||
|
return conversation_turns
|
||||||
|
|
||||||
|
def _is_range_safe(self, messages: List[Dict], start: int, end: int, protected: set) -> bool:
|
||||||
|
"""检查消息范围是否安全 - 简化版,不检查tool完整性"""
|
||||||
|
return True # 不保护tool消息,允许压缩
|
||||||
|
|
||||||
|
def _get_protected_indices(self, messages: List[Dict]) -> set:
|
||||||
|
"""
|
||||||
|
找出所有受保护的消息索引 - 简化版,只保护最近的对话轮次
|
||||||
|
"""
|
||||||
|
protected_indices = set()
|
||||||
|
|
||||||
|
# ========== 1. 保护最近的消息(由配置决定) ==========
|
||||||
|
keep_recent = constants.COMPRESS_KEEP_RECENT
|
||||||
|
if keep_recent > 0:
|
||||||
|
recent_start = max(0, len(messages) - keep_recent)
|
||||||
|
for i in range(recent_start, len(messages)):
|
||||||
|
protected_indices.add(i)
|
||||||
|
|
||||||
|
# ========== 2. 保护前2个对话轮次 ==========
|
||||||
|
conversation_turns = self._build_conversation_turns(messages)
|
||||||
|
|
||||||
|
if len(conversation_turns) >= 2:
|
||||||
|
# 保护第1轮
|
||||||
|
first_turn = conversation_turns[0]
|
||||||
|
protected_indices.add(first_turn["user_index"])
|
||||||
|
if first_turn["assistant_index"] is not None:
|
||||||
|
protected_indices.add(first_turn["assistant_index"])
|
||||||
|
for tool_idx in first_turn["tool_indices"]:
|
||||||
|
protected_indices.add(tool_idx)
|
||||||
|
|
||||||
|
# 保护第2轮
|
||||||
|
second_turn = conversation_turns[1]
|
||||||
|
protected_indices.add(second_turn["user_index"])
|
||||||
|
if second_turn["assistant_index"] is not None:
|
||||||
|
protected_indices.add(second_turn["assistant_index"])
|
||||||
|
for tool_idx in second_turn["tool_indices"]:
|
||||||
|
protected_indices.add(tool_idx)
|
||||||
|
|
||||||
|
# 调试日志
|
||||||
|
if self.logger.level >= 3:
|
||||||
|
self.logger.debug(f"受保护的消息索引: {sorted(protected_indices)}")
|
||||||
|
|
||||||
|
return protected_indices
|
||||||
|
|
||||||
|
def _build_compression_prompt(self, messages: List[Dict]) -> str:
|
||||||
|
"""构建压缩提示词 - 包含工具调用信息"""
|
||||||
|
conversation_lines = []
|
||||||
|
|
||||||
|
for msg in messages:
|
||||||
|
role = msg.get("role")
|
||||||
|
content = msg.get("content", "")
|
||||||
|
|
||||||
|
if role == "user":
|
||||||
|
if content:
|
||||||
|
conversation_lines.append(f"用户: {content}")
|
||||||
|
|
||||||
|
elif role == "assistant":
|
||||||
|
if content:
|
||||||
|
conversation_lines.append(f"助手: {content}")
|
||||||
|
|
||||||
|
# 包含工具调用信息
|
||||||
|
tool_calls = msg.get("tool_calls", [])
|
||||||
|
for tc in tool_calls:
|
||||||
|
func = tc.get("function", {})
|
||||||
|
tool_name = func.get("name", "")
|
||||||
|
arguments = func.get("arguments", "")
|
||||||
|
if tool_name:
|
||||||
|
if arguments:
|
||||||
|
conversation_lines.append(f"[调用工具: {tool_name}({arguments[:100]})]")
|
||||||
|
else:
|
||||||
|
conversation_lines.append(f"[调用工具: {tool_name}]")
|
||||||
|
|
||||||
|
elif role == "tool":
|
||||||
|
if content:
|
||||||
|
useful = self._extract_useful_content(content)
|
||||||
|
if useful:
|
||||||
|
conversation_lines.append(f"[工具结果: {useful}]")
|
||||||
|
|
||||||
|
conversation_text = "\n".join(conversation_lines)
|
||||||
|
|
||||||
|
max_input_chars = constants.COMPRESSOR_MAX_INPUT_CHARS
|
||||||
|
if len(conversation_text) > max_input_chars:
|
||||||
|
conversation_text = conversation_text[:max_input_chars] + "\n...(内容过长,已截断)..."
|
||||||
|
|
||||||
|
prompt = f"""请对以下对话内容进行智能压缩和总结。要求:
|
||||||
|
|
||||||
|
1. **提取核心信息**:对话的主要主题、关键问题、重要结论、用户需求
|
||||||
|
2. **保留重要细节**:用户的重要要求、AI的重要回复、达成的共识
|
||||||
|
3. **保留工具调用摘要**:简要说明使用了什么工具,得到了什么关键结果
|
||||||
|
4. **结构化输出**:使用清晰的段落结构
|
||||||
|
|
||||||
|
对话内容:
|
||||||
|
---
|
||||||
|
{conversation_text}
|
||||||
|
---
|
||||||
|
|
||||||
|
请输出压缩后的对话摘要:"""
|
||||||
|
|
||||||
|
return prompt
|
||||||
|
|
||||||
|
def _extract_useful_content(self, content: str) -> str:
|
||||||
|
"""从工具返回内容中提取有用信息"""
|
||||||
|
if not content:
|
||||||
|
return ""
|
||||||
|
|
||||||
|
if "✅" in content:
|
||||||
|
lines = content.split("\n")
|
||||||
|
for line in lines[:3]:
|
||||||
|
if "✅" in line:
|
||||||
|
return line.strip()
|
||||||
|
|
||||||
|
if "❌" in content:
|
||||||
|
lines = content.split("\n")
|
||||||
|
for line in lines[:2]:
|
||||||
|
if "❌" in line:
|
||||||
|
return line.strip()
|
||||||
|
|
||||||
|
if len(content) > 100:
|
||||||
|
return content[:100] + "..."
|
||||||
|
|
||||||
|
return content
|
||||||
|
|
||||||
|
def _compress_via_llm(self, messages: List[Dict]) -> Optional[str]:
|
||||||
|
"""通过语言模型压缩消息(经由注入的客户端抽象层调用,支持多厂商)"""
|
||||||
|
if not messages or not self.llm_client:
|
||||||
|
return None
|
||||||
|
|
||||||
|
prompt = self._build_compression_prompt(messages)
|
||||||
|
|
||||||
|
print(f"📦 {i18n.get('conversation_compress')}: 正在压缩 {len(messages)} 条消息...")
|
||||||
|
|
||||||
|
try:
|
||||||
|
summary = self.llm_client.chat_completion(
|
||||||
|
[
|
||||||
|
{"role": "system", "content": "你是一个专业的对话摘要助手,擅长提取对话核心信息并进行结构化总结。请保留工具调用的关键结果。"},
|
||||||
|
{"role": "user", "content": prompt}
|
||||||
|
],
|
||||||
|
temperature=0.3,
|
||||||
|
max_tokens=constants.COMPRESSOR_LLM_MAX_TOKENS
|
||||||
|
)
|
||||||
|
return summary
|
||||||
|
except Exception as e:
|
||||||
|
self.logger.error(f"语言模型压缩异常: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
def _validate_message_structure(self, messages: List[Dict]) -> List[Dict]:
|
||||||
|
"""验证并修复消息结构,移除孤立的 tool 消息"""
|
||||||
|
if not messages:
|
||||||
|
return messages
|
||||||
|
|
||||||
|
# 找出所有 assistant 消息的 tool_call_ids
|
||||||
|
assistant_tool_ids = set()
|
||||||
|
for msg in messages:
|
||||||
|
if msg.get("role") == "assistant" and msg.get("tool_calls"):
|
||||||
|
for tc in msg.get("tool_calls", []):
|
||||||
|
tc_id = tc.get("id")
|
||||||
|
if tc_id:
|
||||||
|
assistant_tool_ids.add(tc_id)
|
||||||
|
|
||||||
|
# 过滤掉没有对应 assistant 的 tool 消息
|
||||||
|
filtered = []
|
||||||
|
for msg in messages:
|
||||||
|
if msg.get("role") == "tool":
|
||||||
|
tool_call_id = msg.get("tool_call_id")
|
||||||
|
if tool_call_id and tool_call_id not in assistant_tool_ids:
|
||||||
|
print(f" ⚠️ 移除孤立的 tool 消息: {tool_call_id}")
|
||||||
|
continue
|
||||||
|
filtered.append(msg)
|
||||||
|
|
||||||
|
return filtered
|
||||||
|
|
||||||
|
def compress_conversation(self, messages: List[Dict]) -> List[Dict]:
|
||||||
|
"""
|
||||||
|
压缩对话内容 - 优先压缩较早的消息,包含tool调用信息一并压缩
|
||||||
|
当消息条数或 Token 数超过阈值时,压缩到阈值以下
|
||||||
|
"""
|
||||||
|
if not messages:
|
||||||
|
return messages
|
||||||
|
|
||||||
|
# 分离系统消息和对话消息
|
||||||
|
system_messages = []
|
||||||
|
conversation_messages = []
|
||||||
|
|
||||||
|
for msg in messages:
|
||||||
|
if msg.get("role") == "system":
|
||||||
|
system_messages.append(msg)
|
||||||
|
else:
|
||||||
|
conversation_messages.append(msg)
|
||||||
|
|
||||||
|
if not conversation_messages:
|
||||||
|
return system_messages
|
||||||
|
|
||||||
|
original_count = len(conversation_messages)
|
||||||
|
original_tokens = self.estimate_tokens(conversation_messages)
|
||||||
|
threshold = constants.COMPRESS_MESSAGE_THRESHOLD
|
||||||
|
|
||||||
|
print(f"\n📊 对话压缩检查:")
|
||||||
|
print(f" 当前消息数: {original_count} 条 (阈值: {threshold} 条)")
|
||||||
|
print(f" 当前Token数: {original_tokens} 个")
|
||||||
|
|
||||||
|
# 计算目标消息数量 - 压缩到阈值的70%
|
||||||
|
target_count = self._calculate_target_count(original_count, threshold, keep_ratio=0.7)
|
||||||
|
|
||||||
|
if original_count < threshold:
|
||||||
|
# 消息数未超阈值(由 Token 阈值触发):按比例压缩早期消息
|
||||||
|
target_count = max(self.min_messages_to_compress, original_count // 2)
|
||||||
|
print(f" Token 触发压缩:目标消息数改为 {target_count} 条 (当前的一半)")
|
||||||
|
else:
|
||||||
|
print(f" 目标消息数: {target_count} 条 (压缩后应少于阈值)")
|
||||||
|
|
||||||
|
# 如果目标数量大于等于当前数量,不需要压缩
|
||||||
|
if target_count >= original_count:
|
||||||
|
print(f" 目标消息数大于等于当前,跳过压缩")
|
||||||
|
return messages
|
||||||
|
|
||||||
|
# 计算需要压缩的消息数量
|
||||||
|
messages_to_compress = original_count - target_count
|
||||||
|
print(f" 需要压缩的消息数量: {messages_to_compress} 条")
|
||||||
|
|
||||||
|
# 根据消息数量查找可压缩块
|
||||||
|
compress_start, compress_end = self._find_compressible_block_by_count(
|
||||||
|
conversation_messages, target_count
|
||||||
|
)
|
||||||
|
|
||||||
|
if compress_end - compress_start >= 2:
|
||||||
|
compressible = conversation_messages[compress_start:compress_end]
|
||||||
|
compress_tokens = self.estimate_tokens(compressible)
|
||||||
|
compress_count = len(compressible)
|
||||||
|
|
||||||
|
print(f" 压缩块范围: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
||||||
|
print(f" 压缩块token数: {compress_tokens}")
|
||||||
|
|
||||||
|
# 如果压缩块大小仍然不够,尝试扩大压缩块
|
||||||
|
if compress_count < messages_to_compress:
|
||||||
|
print(f" 压缩块大小不足({compress_count} < {messages_to_compress}),扩大压缩块...")
|
||||||
|
# 扩大压缩块
|
||||||
|
additional = messages_to_compress - compress_count + 5
|
||||||
|
new_start = max(0, compress_start - additional // 2)
|
||||||
|
new_end = min(len(conversation_messages), compress_end + additional // 2)
|
||||||
|
compress_start, compress_end = new_start, new_end
|
||||||
|
compressible = conversation_messages[compress_start:compress_end]
|
||||||
|
compress_count = len(compressible)
|
||||||
|
print(f" 扩大后压缩块: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
||||||
|
else:
|
||||||
|
print(f" 未找到合适的压缩块,使用默认压缩策略")
|
||||||
|
# 使用默认压缩策略:压缩最早的一半消息,保留最近的消息
|
||||||
|
keep_recent = constants.COMPRESS_KEEP_RECENT
|
||||||
|
compress_end = len(conversation_messages) - keep_recent
|
||||||
|
compress_start = max(0, compress_end - messages_to_compress - 5)
|
||||||
|
compress_start = max(compress_start, 0)
|
||||||
|
compressible = conversation_messages[compress_start:compress_end]
|
||||||
|
compress_count = len(compressible)
|
||||||
|
print(f" 默认压缩块: [{compress_start}:{compress_end}], 共 {compress_count} 条消息")
|
||||||
|
|
||||||
|
# 确保压缩块有足够多的消息
|
||||||
|
if compress_count < 5:
|
||||||
|
print(f" 压缩块太小({compress_count}条),跳过压缩")
|
||||||
|
return messages
|
||||||
|
|
||||||
|
# 调用语言模型压缩
|
||||||
|
summary = self._compress_via_llm(compressible)
|
||||||
|
|
||||||
|
if not summary:
|
||||||
|
print(f" ❌ 压缩失败,保留原始对话")
|
||||||
|
return messages
|
||||||
|
|
||||||
|
# 构建压缩后的消息列表
|
||||||
|
compressed_result = []
|
||||||
|
|
||||||
|
# 添加压缩块之前的消息
|
||||||
|
compressed_result.extend(conversation_messages[:compress_start])
|
||||||
|
|
||||||
|
# 插入压缩摘要(作为系统消息)
|
||||||
|
compressed_result.append({
|
||||||
|
"role": "system",
|
||||||
|
"content": i18n.get('compressed_summary', summary)
|
||||||
|
})
|
||||||
|
|
||||||
|
# 添加压缩块之后的消息(包括最近的消息)
|
||||||
|
compressed_result.extend(conversation_messages[compress_end:])
|
||||||
|
|
||||||
|
# 验证并修复消息结构
|
||||||
|
compressed_result = self._validate_message_structure(compressed_result)
|
||||||
|
|
||||||
|
# 合并系统消息
|
||||||
|
final_messages = system_messages + compressed_result
|
||||||
|
|
||||||
|
# 统计结果
|
||||||
|
final_count = len([m for m in final_messages if m.get("role") != "system"])
|
||||||
|
final_tokens = self.estimate_tokens([m for m in final_messages if m.get("role") != "system"])
|
||||||
|
|
||||||
|
print(f"\n✅ {i18n.get('conversation_compressed', original_count, final_count)}")
|
||||||
|
print(f" Token: {original_tokens} -> {final_tokens} (压缩率: {final_tokens/original_tokens:.1%})")
|
||||||
|
|
||||||
|
# 检查是否压缩到阈值以下
|
||||||
|
if final_count >= threshold:
|
||||||
|
print(f" ⚠️ 警告: 压缩后消息数({final_count})仍大于等于阈值({threshold})")
|
||||||
|
|
||||||
|
return final_messages
|
||||||
+101
-15
@@ -51,13 +51,57 @@ class MemoryDatabase:
|
|||||||
)
|
)
|
||||||
''')
|
''')
|
||||||
|
|
||||||
cursor.execute('''
|
self._init_fts_table(conn)
|
||||||
CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts
|
|
||||||
USING fts5(topic, summary, content, tags, content=memories)
|
|
||||||
''')
|
|
||||||
|
|
||||||
conn.commit()
|
conn.commit()
|
||||||
|
|
||||||
|
def _fts_trigram_supported(self, conn) -> bool:
|
||||||
|
"""检测当前 SQLite 是否支持 trigram 分词器(需 SQLite >= 3.34)"""
|
||||||
|
try:
|
||||||
|
conn.execute("CREATE VIRTUAL TABLE temp.__trigram_probe USING fts5(x, tokenize='trigram')")
|
||||||
|
conn.execute("DROP TABLE temp.__trigram_probe")
|
||||||
|
return True
|
||||||
|
except sqlite3.OperationalError:
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _init_fts_table(self, conn):
|
||||||
|
"""初始化 FTS5 全文索引表。
|
||||||
|
|
||||||
|
优先使用 trigram 分词器(中文可检索,unicode61 对中文不分词);
|
||||||
|
检测到旧 unicode61 表时自动迁移:DROP 重建后从 memories 表回灌全量数据。
|
||||||
|
环境不支持 trigram 时退回 unicode61 建表,中文搜索由 LIKE 兜底。
|
||||||
|
"""
|
||||||
|
cursor = conn.cursor()
|
||||||
|
trigram_ok = self._fts_trigram_supported(conn)
|
||||||
|
tokenize_opt = ", tokenize='trigram'" if trigram_ok else ""
|
||||||
|
create_sql = (
|
||||||
|
"CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts "
|
||||||
|
f"USING fts5(topic, summary, content, tags, content=memories{tokenize_opt})"
|
||||||
|
)
|
||||||
|
|
||||||
|
row = cursor.execute(
|
||||||
|
"SELECT sql FROM sqlite_master WHERE type='table' AND name='memories_fts'"
|
||||||
|
).fetchone()
|
||||||
|
|
||||||
|
need_backfill = False
|
||||||
|
if row is None:
|
||||||
|
need_backfill = True # 新库(memories 可能已有数据,回灌无害)
|
||||||
|
elif trigram_ok and "trigram" not in (row[0] or ""):
|
||||||
|
# 旧 unicode61 表 → 迁移为 trigram
|
||||||
|
cursor.execute("DROP TABLE memories_fts")
|
||||||
|
need_backfill = True
|
||||||
|
|
||||||
|
cursor.execute(create_sql)
|
||||||
|
|
||||||
|
if need_backfill:
|
||||||
|
count = cursor.execute('''
|
||||||
|
INSERT INTO memories_fts (rowid, topic, summary, content, tags)
|
||||||
|
SELECT id, topic, summary, COALESCE(content, ''), COALESCE(tags, '')
|
||||||
|
FROM memories
|
||||||
|
''').rowcount
|
||||||
|
if count:
|
||||||
|
self.logger.info(f"FTS5 索引已使用 trigram 分词重建并回灌 {count} 条记忆")
|
||||||
|
|
||||||
def add_memory(self, topic: str, summary: str, content: str = None,
|
def add_memory(self, topic: str, summary: str, content: str = None,
|
||||||
start_time: str = None, end_time: str = None,
|
start_time: str = None, end_time: str = None,
|
||||||
tags: List[str] = None, importance: int = None) -> int:
|
tags: List[str] = None, importance: int = None) -> int:
|
||||||
@@ -106,7 +150,11 @@ class MemoryDatabase:
|
|||||||
return deleted
|
return deleted
|
||||||
|
|
||||||
def search_memories(self, keywords: List[str], match_all: bool = False, limit: int = None) -> List[Dict]:
|
def search_memories(self, keywords: List[str], match_all: bool = False, limit: int = None) -> List[Dict]:
|
||||||
"""多关键词模糊搜索记忆"""
|
"""多关键词模糊搜索记忆。
|
||||||
|
|
||||||
|
优先走 FTS5 全文检索;FTS 无结果时(trigram 对 <3 字符的中文短词
|
||||||
|
无法命中、旧表 unicode61 对中文不分词等场景)用 LIKE 兜底。
|
||||||
|
"""
|
||||||
if not keywords:
|
if not keywords:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
@@ -117,11 +165,18 @@ class MemoryDatabase:
|
|||||||
conn.row_factory = sqlite3.Row
|
conn.row_factory = sqlite3.Row
|
||||||
cursor = conn.cursor()
|
cursor = conn.cursor()
|
||||||
|
|
||||||
if match_all:
|
results = self._search_via_fts(cursor, keywords, match_all, limit)
|
||||||
query = " AND ".join([f'"{kw}"' for kw in keywords])
|
if not results:
|
||||||
else:
|
results = self._search_via_like(cursor, keywords, match_all, limit)
|
||||||
query = " OR ".join([f'"{kw}"' for kw in keywords])
|
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
|
def _search_via_fts(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]:
|
||||||
|
"""FTS5 全文检索(关键词含特殊字符导致 MATCH 语法错误时返回空,走 LIKE 兜底)"""
|
||||||
|
joiner = " AND " if match_all else " OR "
|
||||||
|
query = joiner.join([f'"{kw}"' for kw in keywords])
|
||||||
|
|
||||||
|
try:
|
||||||
cursor.execute('''
|
cursor.execute('''
|
||||||
SELECT m.*, rank
|
SELECT m.*, rank
|
||||||
FROM memories m
|
FROM memories m
|
||||||
@@ -130,14 +185,45 @@ class MemoryDatabase:
|
|||||||
ORDER BY m.end_time DESC, rank, m.importance DESC
|
ORDER BY m.end_time DESC, rank, m.importance DESC
|
||||||
LIMIT ?
|
LIMIT ?
|
||||||
''', (query, limit))
|
''', (query, limit))
|
||||||
|
except sqlite3.OperationalError as e:
|
||||||
|
self.logger.warning(f"FTS5 检索失败(改用 LIKE 兜底): {e}")
|
||||||
|
return []
|
||||||
|
|
||||||
results = []
|
return self._rows_to_memories(cursor.fetchall())
|
||||||
for row in cursor.fetchall():
|
|
||||||
memory = dict(row)
|
|
||||||
if memory.get('tags'):
|
|
||||||
memory['tags'] = memory['tags'].split(',')
|
|
||||||
results.append(memory)
|
|
||||||
|
|
||||||
|
def _search_via_like(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]:
|
||||||
|
"""LIKE 模糊检索兜底(转义 %/_/\\ 防止关键词被当作通配符)"""
|
||||||
|
fields = ["topic", "summary", "content", "tags"]
|
||||||
|
|
||||||
|
def escape(kw: str) -> str:
|
||||||
|
return kw.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||||
|
|
||||||
|
conditions = []
|
||||||
|
params = []
|
||||||
|
for kw in keywords:
|
||||||
|
pattern = f"%{escape(kw)}%"
|
||||||
|
per_field = " OR ".join([f"{f} LIKE ? ESCAPE '\\'" for f in fields])
|
||||||
|
conditions.append(f"({per_field})")
|
||||||
|
params.extend([pattern] * len(fields))
|
||||||
|
|
||||||
|
joiner = " AND " if match_all else " OR "
|
||||||
|
cursor.execute(
|
||||||
|
f"SELECT * FROM memories WHERE {joiner.join(conditions)} "
|
||||||
|
f"ORDER BY end_time DESC, importance DESC LIMIT ?",
|
||||||
|
(*params, limit),
|
||||||
|
)
|
||||||
|
|
||||||
|
return self._rows_to_memories(cursor.fetchall())
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _rows_to_memories(rows) -> List[Dict]:
|
||||||
|
"""sqlite3.Row 列表转记忆字典列表(tags 逗号串拆分为列表)"""
|
||||||
|
results = []
|
||||||
|
for row in rows:
|
||||||
|
memory = dict(row)
|
||||||
|
if memory.get('tags'):
|
||||||
|
memory['tags'] = memory['tags'].split(',')
|
||||||
|
results.append(memory)
|
||||||
return results
|
return results
|
||||||
|
|
||||||
def get_recent_memories(self, limit: int = 10) -> List[Dict]:
|
def get_recent_memories(self, limit: int = 10) -> List[Dict]:
|
||||||
|
|||||||
@@ -0,0 +1,196 @@
|
|||||||
|
"""subsystems.conversation_compressor 单元测试
|
||||||
|
|
||||||
|
覆盖:
|
||||||
|
- estimate_tokens: 中英文/工具调用 token 估算
|
||||||
|
- should_compress: 消息数阈值触发 / Token 阈值触发 / 均未触发 / 功能禁用
|
||||||
|
- compress_conversation: 消息数触发压缩 / Token 触发压缩 / LLM 失败保留原文
|
||||||
|
- _validate_message_structure: 移除孤立 tool 消息
|
||||||
|
通过注入 FakeLLMClient 模拟 LLM 调用,不依赖真实 API。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from subsystems import ConversationCompressor
|
||||||
|
from common import constants, Logger
|
||||||
|
|
||||||
|
|
||||||
|
class FakeLLMClient:
|
||||||
|
"""模拟 llm_client 抽象层的 chat_completion"""
|
||||||
|
|
||||||
|
def __init__(self, reply="这是压缩摘要", error=None):
|
||||||
|
self.reply = reply
|
||||||
|
self.error = error
|
||||||
|
self.calls = []
|
||||||
|
|
||||||
|
def chat_completion(self, messages, temperature=0.7, max_tokens=None, **kwargs):
|
||||||
|
self.calls.append({"max_tokens": max_tokens, "n_msgs": len(messages)})
|
||||||
|
if self.error:
|
||||||
|
raise self.error
|
||||||
|
return self.reply
|
||||||
|
|
||||||
|
|
||||||
|
def make_messages(n_pairs=15, long_content=False):
|
||||||
|
"""构造 n_pairs 轮 user/assistant 对话(含 1 条前导 system)"""
|
||||||
|
msgs = [{"role": "system", "content": "SYS_PROMPT"}]
|
||||||
|
for i in range(n_pairs):
|
||||||
|
content = f"用户消息内容{i}"
|
||||||
|
if long_content:
|
||||||
|
content += "很长的文档内容" * 200
|
||||||
|
msgs.append({"role": "user", "content": content})
|
||||||
|
msgs.append({"role": "assistant", "content": f"助手回复{i}"})
|
||||||
|
return msgs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def compressor():
|
||||||
|
return ConversationCompressor(llm_client=FakeLLMClient(), logger=Logger("TestCompressor", 0))
|
||||||
|
|
||||||
|
|
||||||
|
def patch_const(monkeypatch, name, value):
|
||||||
|
"""把 Constants 类属性 property 替换为固定值(参照 test_db.py 既有做法)"""
|
||||||
|
monkeypatch.setattr(type(constants), name, property(lambda self: value))
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== estimate_tokens ====================
|
||||||
|
|
||||||
|
def test_estimate_tokens_chinese(compressor):
|
||||||
|
assert compressor.estimate_tokens([{"role": "user", "content": "你好"}]) == 13 # 2*1.5 + 10
|
||||||
|
|
||||||
|
|
||||||
|
def test_estimate_tokens_english(compressor):
|
||||||
|
# 2 词 * 1.3 + (11 字符 - 2 词) * 0.8 + 10 = 19.8 → int 截断为 19
|
||||||
|
# (既有算法:other_chars 按词数抵扣,字母被计入 other,保持原行为)
|
||||||
|
assert compressor.estimate_tokens([{"role": "user", "content": "hello world"}]) == 19
|
||||||
|
|
||||||
|
|
||||||
|
def test_estimate_tokens_tool_calls(compressor):
|
||||||
|
msgs = [{
|
||||||
|
"role": "assistant", "content": "",
|
||||||
|
"tool_calls": [{"function": {"name": "run_command", "arguments": "{\"a\":1}"}}],
|
||||||
|
}]
|
||||||
|
# name 11 字符*1.5=16.5 + arguments 8 字符*0.8=6.4 + 50 = 72.9 → 72
|
||||||
|
assert compressor.estimate_tokens(msgs) == 72
|
||||||
|
|
||||||
|
|
||||||
|
def test_estimate_tokens_empty(compressor):
|
||||||
|
assert compressor.estimate_tokens([]) == 0
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== should_compress ====================
|
||||||
|
|
||||||
|
def test_should_compress_below_thresholds(compressor, monkeypatch):
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 1000)
|
||||||
|
patch_const(monkeypatch, "COMPRESS_TOKEN_THRESHOLD", 10 ** 9)
|
||||||
|
should, reason = compressor.should_compress(make_messages(5))
|
||||||
|
assert should is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_compress_by_message_count(compressor, monkeypatch):
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 5)
|
||||||
|
should, reason = compressor.should_compress(make_messages(5)) # 10 条对话消息
|
||||||
|
assert should is True
|
||||||
|
assert "消息数" in reason
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_compress_by_token_count(compressor, monkeypatch):
|
||||||
|
"""消息数未超阈值但 Token 数超阈值时也应触发(v3.1 修复项)。"""
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 1000)
|
||||||
|
patch_const(monkeypatch, "COMPRESS_TOKEN_THRESHOLD", 1)
|
||||||
|
should, reason = compressor.should_compress(make_messages(5))
|
||||||
|
assert should is True
|
||||||
|
assert "Token" in reason
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_compress_disabled(compressor, monkeypatch):
|
||||||
|
patch_const(monkeypatch, "COMPRESS_ENABLED", False)
|
||||||
|
should, _ = compressor.should_compress(make_messages(100))
|
||||||
|
assert should is False
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== compress_conversation ====================
|
||||||
|
|
||||||
|
def test_compress_by_message_count(compressor, monkeypatch):
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||||
|
messages = make_messages(15) # 30 条对话消息 + 1 system
|
||||||
|
result = compressor.compress_conversation(messages)
|
||||||
|
|
||||||
|
non_system = [m for m in result if m["role"] != "system"]
|
||||||
|
assert len(non_system) < 30 # 消息数下降
|
||||||
|
# 原始 system 提示词保留
|
||||||
|
assert any(m["content"] == "SYS_PROMPT" for m in result)
|
||||||
|
# 压缩摘要已注入
|
||||||
|
assert any(m["role"] == "system" and "这是压缩摘要" in m["content"] for m in result)
|
||||||
|
# 最近消息保留
|
||||||
|
assert any(m.get("content") == "助手回复14" for m in non_system)
|
||||||
|
# LLM 被调用了一次
|
||||||
|
assert len(compressor.llm_client.calls) == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_compress_by_token_count(compressor, monkeypatch):
|
||||||
|
"""Token 触发(消息数远低于阈值)时也应实际执行压缩(v3.1 修复项)。"""
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 1000)
|
||||||
|
patch_const(monkeypatch, "COMPRESS_TOKEN_THRESHOLD", 1)
|
||||||
|
messages = make_messages(15, long_content=True) # 30 条长消息
|
||||||
|
result = compressor.compress_conversation(messages)
|
||||||
|
|
||||||
|
non_system = [m for m in result if m["role"] != "system"]
|
||||||
|
assert len(non_system) < 30
|
||||||
|
assert any(m["role"] == "system" and "这是压缩摘要" in m["content"] for m in result)
|
||||||
|
assert len(compressor.llm_client.calls) == 1
|
||||||
|
# Token 数下降
|
||||||
|
assert compressor.estimate_tokens(non_system) < compressor.estimate_tokens(
|
||||||
|
[m for m in messages if m["role"] != "system"])
|
||||||
|
|
||||||
|
|
||||||
|
def test_compress_llm_failure_keeps_original(monkeypatch):
|
||||||
|
comp = ConversationCompressor(
|
||||||
|
llm_client=FakeLLMClient(error=RuntimeError("api down")),
|
||||||
|
logger=Logger("TestCompressor", 0),
|
||||||
|
)
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||||
|
messages = make_messages(15)
|
||||||
|
result = comp.compress_conversation(messages)
|
||||||
|
assert result == messages # 失败时原样返回
|
||||||
|
|
||||||
|
|
||||||
|
def test_compress_without_llm_client_keeps_original(monkeypatch):
|
||||||
|
comp = ConversationCompressor(llm_client=None)
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||||
|
messages = make_messages(15)
|
||||||
|
assert comp.compress_conversation(messages) == messages
|
||||||
|
|
||||||
|
|
||||||
|
def test_compress_empty_messages(compressor):
|
||||||
|
assert compressor.compress_conversation([]) == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_compress_too_few_messages_skips(compressor, monkeypatch):
|
||||||
|
patch_const(monkeypatch, "COMPRESS_MESSAGE_THRESHOLD", 10)
|
||||||
|
messages = make_messages(3) # 6 条对话消息,压缩块必然 < 5
|
||||||
|
assert compressor.compress_conversation(messages) == messages
|
||||||
|
assert compressor.llm_client.calls == []
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== _validate_message_structure ====================
|
||||||
|
|
||||||
|
def test_validate_removes_orphan_tool_messages(compressor):
|
||||||
|
msgs = [
|
||||||
|
{"role": "user", "content": "hi"},
|
||||||
|
{"role": "assistant", "content": "",
|
||||||
|
"tool_calls": [{"id": "t1", "function": {"name": "f", "arguments": "{}"}}]},
|
||||||
|
{"role": "tool", "tool_call_id": "t1", "content": "ok"},
|
||||||
|
{"role": "tool", "tool_call_id": "orphan", "content": "孤儿消息"},
|
||||||
|
]
|
||||||
|
result = compressor._validate_message_structure(msgs)
|
||||||
|
assert len(result) == 3
|
||||||
|
assert all(m.get("tool_call_id") != "orphan" for m in result)
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_keeps_paired_tool_messages(compressor):
|
||||||
|
msgs = [
|
||||||
|
{"role": "user", "content": "hi"},
|
||||||
|
{"role": "assistant", "content": "",
|
||||||
|
"tool_calls": [{"id": "t1", "function": {"name": "f", "arguments": "{}"}}]},
|
||||||
|
{"role": "tool", "tool_call_id": "t1", "content": "ok"},
|
||||||
|
]
|
||||||
|
assert len(compressor._validate_message_structure(msgs)) == 3
|
||||||
@@ -70,6 +70,87 @@ def test_memory_search_empty_keywords_returns_empty(mem_db):
|
|||||||
assert mem_db.search_memories([]) == []
|
assert mem_db.search_memories([]) == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_memory_search_chinese_keyword(mem_db):
|
||||||
|
"""中文关键词应可检索(trigram 分词或 LIKE 兜底)。"""
|
||||||
|
mem_db.add_memory(topic="天气查询", summary="今天天气晴朗", content="气温25度")
|
||||||
|
mem_db.add_memory(topic="会议安排", summary="下午三点开会", content="会议室A")
|
||||||
|
|
||||||
|
results = mem_db.search_memories(["天气"])
|
||||||
|
assert len(results) == 1
|
||||||
|
assert results[0]["topic"] == "天气查询"
|
||||||
|
|
||||||
|
|
||||||
|
def test_memory_search_chinese_match_all(mem_db):
|
||||||
|
"""中文多关键词 match_all:全部命中才返回。"""
|
||||||
|
mem_db.add_memory(topic="天气查询", summary="今天天气晴朗", content="气温25度")
|
||||||
|
mem_db.add_memory(topic="天气与会议", summary="开会前先查天气", content="会议室A")
|
||||||
|
|
||||||
|
results = mem_db.search_memories(["天气", "开会"], match_all=True)
|
||||||
|
assert len(results) == 1
|
||||||
|
assert results[0]["topic"] == "天气与会议"
|
||||||
|
|
||||||
|
|
||||||
|
def test_memory_search_like_escape_wildcards(mem_db):
|
||||||
|
"""LIKE 兜底路径中 %/_ 应被转义为字面字符,不当作通配符。"""
|
||||||
|
import sqlite3
|
||||||
|
mem_db.add_memory(topic="progress", summary="100% done")
|
||||||
|
mem_db.add_memory(topic="task_a", summary="plain text")
|
||||||
|
mem_db.add_memory(topic="taskXa", summary="wildcard trap")
|
||||||
|
|
||||||
|
def like_search(kw):
|
||||||
|
with sqlite3.connect(mem_db.db_path) as conn:
|
||||||
|
conn.row_factory = sqlite3.Row
|
||||||
|
return mem_db._search_via_like(conn.cursor(), [kw], match_all=False, limit=100)
|
||||||
|
|
||||||
|
# "%" 为字面字符:"100%" 只命中 progress,不当作任意串通配
|
||||||
|
results = like_search("100%")
|
||||||
|
assert [r["topic"] for r in results] == ["progress"]
|
||||||
|
|
||||||
|
# "_" 为字面字符:"ask_" 只命中 task_a;若未转义,taskXa 也会被 %ask_% 误匹配
|
||||||
|
results = like_search("ask_")
|
||||||
|
assert [r["topic"] for r in results] == ["task_a"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_memory_fts_migration_from_unicode61(tmp_path):
|
||||||
|
"""旧 unicode61 FTS 表在重新打开时应自动迁移为 trigram 并回灌数据。"""
|
||||||
|
import sqlite3
|
||||||
|
db_path = str(tmp_path / "legacy.db")
|
||||||
|
|
||||||
|
# 模拟旧库:unicode61 建表 + 写入数据
|
||||||
|
conn = sqlite3.connect(db_path)
|
||||||
|
conn.execute('''CREATE TABLE memories (
|
||||||
|
id INTEGER PRIMARY KEY AUTOINCREMENT, topic TEXT NOT NULL, summary TEXT NOT NULL,
|
||||||
|
content TEXT, start_time TEXT NOT NULL, end_time TEXT NOT NULL,
|
||||||
|
created_at TEXT NOT NULL, tags TEXT, importance INTEGER DEFAULT 1)''')
|
||||||
|
conn.execute('''CREATE VIRTUAL TABLE memories_fts
|
||||||
|
USING fts5(topic, summary, content, tags, content=memories)''')
|
||||||
|
conn.execute("INSERT INTO memories (topic, summary, start_time, end_time, created_at) "
|
||||||
|
"VALUES ('天气查询', '今天天气晴朗', 't', 't', 't')")
|
||||||
|
conn.execute("INSERT INTO memories_fts (rowid, topic, summary) VALUES (1, '天气查询', '今天天气晴朗')")
|
||||||
|
conn.commit()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
# 重新打开 → 自动迁移 → 中文可搜
|
||||||
|
db = MemoryDatabase(db_path=db_path)
|
||||||
|
results = db.search_memories(["天气"])
|
||||||
|
assert len(results) == 1
|
||||||
|
assert results[0]["topic"] == "天气查询"
|
||||||
|
|
||||||
|
# 迁移后 FTS 表应为 trigram 分词(环境支持时)
|
||||||
|
conn = sqlite3.connect(db_path)
|
||||||
|
sql = conn.execute("SELECT sql FROM sqlite_master WHERE name='memories_fts'").fetchone()[0]
|
||||||
|
fts_count = conn.execute("SELECT COUNT(*) FROM memories_fts").fetchone()[0]
|
||||||
|
conn.close()
|
||||||
|
assert fts_count == 1 # 数据已回灌
|
||||||
|
try:
|
||||||
|
probe = sqlite3.connect(":memory:")
|
||||||
|
probe.execute("CREATE VIRTUAL TABLE t USING fts5(x, tokenize='trigram')")
|
||||||
|
probe.close()
|
||||||
|
assert "trigram" in sql
|
||||||
|
except sqlite3.OperationalError:
|
||||||
|
pass # 环境不支持 trigram,保持 unicode61 属预期行为
|
||||||
|
|
||||||
|
|
||||||
def test_memory_delete(mem_db):
|
def test_memory_delete(mem_db):
|
||||||
mid = mem_db.add_memory(topic="待删除", summary="摘要")
|
mid = mem_db.add_memory(topic="待删除", summary="摘要")
|
||||||
assert mem_db.delete_memory(mid) is True
|
assert mem_db.delete_memory(mid) is True
|
||||||
|
|||||||
Reference in new issue
Block a user