diff --git a/aibox/README.md b/aibox/README.md index 2ee5284..4450ff4 100644 --- a/aibox/README.md +++ b/aibox/README.md @@ -55,7 +55,7 @@ python aibox.py 启动后，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 # 常量定义 │ ├── i18n.py # 多语言 │ └── logger.py # 日志器 -├── subsystems/ # 子系统包（Wave 1：从 aibox.py 抽离的基础设施） -│ ├── executors.py # CommandExecutor + SubprocessCommandExecutor -│ ├── prompt_manager.py # PromptManager（单例，提示词按数字前缀排序加载） -│ ├── db.py # MemoryDatabase + MemoDatabase（SQLite） -│ └── skill_manager.py # SkillManager（技能发现/加载/热重载） +├── subsystems/ # 子系统包（Wave 1/2：从 aibox.py 抽离的基础设施） +│ ├── executors.py # CommandExecutor + SubprocessCommandExecutor +│ ├── prompt_manager.py # PromptManager（单例，提示词按数字前缀排序加载） +│ ├── db.py # MemoryDatabase + MemoDatabase（SQLite，FTS5 trigram 中文检索 + LIKE 兜底） +│ ├── skill_manager.py # SkillManager（技能发现/加载/热重载） +│ ├── conversation_history.py # ConversationHistory（对话历史管理） +│ └── conversation_compressor.py # ConversationCompressor（对话压缩，注入 llm_client 解耦） ├── tools/ # 工具系统包（A 批 + B 批共 17 个工具类） │ ├── base.py # Tool 抽象基类 │ ├── registry.py # ToolRegistry 工具注册器 @@ -88,14 +90,16 @@ aibox/ │ ├── prompt_tools.py # B 批：UpdateIdentity/UpdateSoul/SaveMemoryAndEndConversation │ ├── proxy_tools.py # B 批：SetProxy/GetProxy │ └── misc_tools.py # B 批：AskUserTool -├── tests/ # 单元测试套件（171 项全过） -│ ├── conftest.py # sys.path 配置 + websockets 桩 -│ ├── test_basic_tools.py # A 批工具 43 项 -│ ├── test_b_tools.py # B 批工具 58 项 -│ ├── test_executors.py # 命令执行器 -│ ├── test_prompt_manager.py # 提示词管理器 -│ ├── test_db.py # 记忆/备忘录数据库 -│ └── test_skill_manager.py # 技能管理器 +├── tests/ # 单元测试套件（232 项全过） +│ ├── conftest.py # sys.path 配置 + websockets 桩 +│ ├── test_basic_tools.py # A 批工具 43 项 +│ ├── test_b_tools.py # B 批工具 58 项 +│ ├── test_executors.py # 命令执行器 +│ ├── test_prompt_manager.py # 提示词管理器 +│ ├── test_db.py # 记忆/备忘录数据库（含中文检索/迁移） +│ ├── test_conversation_history.py # 对话历史 +│ ├── test_conversation_compressor.py # 对话压缩器 +│ └── test_skill_manager.py # 技能管理器 ├── skills/ # 内置扩展技能（*.py 可执行脚本 / *.md Markdown 技能） │ ├── hello-world.py │ ├── aes-crypto.py @@ -130,7 +134,7 @@ common ← subsystems ← tools ← aibox.py ```bash # 从仓库根目录运行 py -m pytest aibox/tests -q -# 当前：171 项全过，约 11s +# 当前：232 项全过，约 11s ``` 运行期数据存放于 `~/.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 - ✅ **模块化重构（aibox.py 8450 行 → 3706 行，缩减 56%）**： - **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 子系统 - **修复**：`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` 采用"重导出"策略，调用点无需关心子模块位置 - - **未抽离**：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 - ✅ **编排者-执行者工具分层（工具集 profile）**： diff --git a/aibox/aibox.py b/aibox/aibox.py index 280f8f5..d3cdca7 100644 --- a/aibox/aibox.py +++ b/aibox/aibox.py @@ -56,516 +56,18 @@ from tools import ( AskUserTool, ) -# 子系统（Wave 1 + Wave 2 抽离）：命令执行器 / 提示词管理器 / 数据库 / 技能管理器 / 对话历史 +# 子系统（Wave 1 + Wave 2 抽离）：命令执行器 / 提示词管理器 / 数据库 / 技能管理器 / 对话历史 / 对话压缩 from subsystems import ( CommandExecutor, SubprocessCommandExecutor, PromptManager, MemoryDatabase, MemoDatabase, SkillManager, 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: @@ -1404,8 +906,8 @@ class DeepSeekChat: # 共享技能管理器：传入时复用（避免每个任务重复扫描技能），否则按需创建 self._shared_skill_manager = skill_manager - # 添加对话压缩器 - self.compressor = ConversationCompressor(self, self.logger) + # 添加对话压缩器（注入 llm 客户端而非 AI 实例，避免循环引用） + self.compressor = ConversationCompressor(self.client, self.logger) self.max_history = max_history or constants.DEFAULT_MAX_HISTORY self.preserve_system = constants.PRESERVE_SYSTEM_PROMPTS diff --git a/aibox/subsystems/__init__.py b/aibox/subsystems/__init__.py index d6cefe4..70a8293 100644 --- a/aibox/subsystems/__init__.py +++ b/aibox/subsystems/__init__.py @@ -9,6 +9,7 @@ subsystems - aibox 子系统模块（包） - skill_manager.py : SkillManager Wave 2： - conversation_history.py : ConversationHistory（对话历史管理） + - conversation_compressor.py : ConversationCompressor（对话压缩，注入 llm_client） 依赖方向：common <- subsystems <- aibox 不反向依赖 aibox 任何符号。 @@ -22,6 +23,7 @@ from .prompt_manager import PromptManager from .db import MemoryDatabase, MemoDatabase from .skill_manager import SkillManager from .conversation_history import ConversationHistory +from .conversation_compressor import ConversationCompressor __all__ = [ "CommandExecutor", @@ -31,4 +33,5 @@ __all__ = [ "MemoDatabase", "SkillManager", "ConversationHistory", + "ConversationCompressor", ] diff --git a/aibox/subsystems/conversation_compressor.py b/aibox/subsystems/conversation_compressor.py new file mode 100644 index 0000000..a5ebc20 --- /dev/null +++ b/aibox/subsystems/conversation_compressor.py @@ -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 diff --git a/aibox/subsystems/db.py b/aibox/subsystems/db.py index 630f465..7151088 100644 --- a/aibox/subsystems/db.py +++ b/aibox/subsystems/db.py @@ -51,13 +51,57 @@ class MemoryDatabase: ) ''') - cursor.execute(''' - CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts - USING fts5(topic, summary, content, tags, content=memories) - ''') + self._init_fts_table(conn) conn.commit() + def _fts_trigram_supported(self, conn) -> bool: + """检测当前 SQLite 是否支持 trigram 分词器（需 SQLite >= 3.34）""" + try: + conn.execute("CREATE VIRTUAL TABLE temp.__trigram_probe USING fts5(x, tokenize='trigram')") + conn.execute("DROP TABLE temp.__trigram_probe") + return True + except sqlite3.OperationalError: + return False + + def _init_fts_table(self, conn): + """初始化 FTS5 全文索引表。 + + 优先使用 trigram 分词器（中文可检索，unicode61 对中文不分词）； + 检测到旧 unicode61 表时自动迁移：DROP 重建后从 memories 表回灌全量数据。 + 环境不支持 trigram 时退回 unicode61 建表，中文搜索由 LIKE 兜底。 + """ + cursor = conn.cursor() + trigram_ok = self._fts_trigram_supported(conn) + tokenize_opt = ", tokenize='trigram'" if trigram_ok else "" + create_sql = ( + "CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts " + f"USING fts5(topic, summary, content, tags, content=memories{tokenize_opt})" + ) + + row = cursor.execute( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='memories_fts'" + ).fetchone() + + need_backfill = False + if row is None: + need_backfill = True # 新库（memories 可能已有数据，回灌无害） + elif trigram_ok and "trigram" not in (row[0] or ""): + # 旧 unicode61 表 → 迁移为 trigram + cursor.execute("DROP TABLE memories_fts") + need_backfill = True + + cursor.execute(create_sql) + + if need_backfill: + count = cursor.execute(''' + INSERT INTO memories_fts (rowid, topic, summary, content, tags) + SELECT id, topic, summary, COALESCE(content, ''), COALESCE(tags, '') + FROM memories + ''').rowcount + if count: + self.logger.info(f"FTS5 索引已使用 trigram 分词重建并回灌 {count} 条记忆") + def add_memory(self, topic: str, summary: str, content: str = None, start_time: str = None, end_time: str = None, tags: List[str] = None, importance: int = None) -> int: @@ -106,7 +150,11 @@ class MemoryDatabase: return deleted def search_memories(self, keywords: List[str], match_all: bool = False, limit: int = None) -> List[Dict]: - """多关键词模糊搜索记忆""" + """多关键词模糊搜索记忆。 + + 优先走 FTS5 全文检索；FTS 无结果时（trigram 对 <3 字符的中文短词 + 无法命中、旧表 unicode61 对中文不分词等场景）用 LIKE 兜底。 + """ if not keywords: return [] @@ -117,11 +165,18 @@ class MemoryDatabase: conn.row_factory = sqlite3.Row cursor = conn.cursor() - if match_all: - query = " AND ".join([f'"{kw}"' for kw in keywords]) - else: - query = " OR ".join([f'"{kw}"' for kw in keywords]) + results = self._search_via_fts(cursor, keywords, match_all, limit) + if not results: + results = self._search_via_like(cursor, keywords, match_all, limit) + return results + + def _search_via_fts(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]: + """FTS5 全文检索（关键词含特殊字符导致 MATCH 语法错误时返回空，走 LIKE 兜底）""" + joiner = " AND " if match_all else " OR " + query = joiner.join([f'"{kw}"' for kw in keywords]) + + try: cursor.execute(''' SELECT m.*, rank FROM memories m @@ -130,14 +185,45 @@ class MemoryDatabase: ORDER BY m.end_time DESC, rank, m.importance DESC LIMIT ? ''', (query, limit)) + except sqlite3.OperationalError as e: + self.logger.warning(f"FTS5 检索失败（改用 LIKE 兜底）: {e}") + return [] - results = [] - for row in cursor.fetchall(): - memory = dict(row) - if memory.get('tags'): - memory['tags'] = memory['tags'].split(',') - results.append(memory) + return self._rows_to_memories(cursor.fetchall()) + def _search_via_like(self, cursor, keywords: List[str], match_all: bool, limit: int) -> List[Dict]: + """LIKE 模糊检索兜底（转义 %/_/\\ 防止关键词被当作通配符）""" + fields = ["topic", "summary", "content", "tags"] + + def escape(kw: str) -> str: + return kw.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + + conditions = [] + params = [] + for kw in keywords: + pattern = f"%{escape(kw)}%" + per_field = " OR ".join([f"{f} LIKE ? ESCAPE '\\'" for f in fields]) + conditions.append(f"({per_field})") + params.extend([pattern] * len(fields)) + + joiner = " AND " if match_all else " OR " + cursor.execute( + f"SELECT * FROM memories WHERE {joiner.join(conditions)} " + f"ORDER BY end_time DESC, importance DESC LIMIT ?", + (*params, limit), + ) + + return self._rows_to_memories(cursor.fetchall()) + + @staticmethod + def _rows_to_memories(rows) -> List[Dict]: + """sqlite3.Row 列表转记忆字典列表（tags 逗号串拆分为列表）""" + results = [] + for row in rows: + memory = dict(row) + if memory.get('tags'): + memory['tags'] = memory['tags'].split(',') + results.append(memory) return results def get_recent_memories(self, limit: int = 10) -> List[Dict]: diff --git a/aibox/tests/test_conversation_compressor.py b/aibox/tests/test_conversation_compressor.py new file mode 100644 index 0000000..a6fd685 --- /dev/null +++ b/aibox/tests/test_conversation_compressor.py @@ -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 diff --git a/aibox/tests/test_db.py b/aibox/tests/test_db.py index 6a920ff..0f7079b 100644 --- a/aibox/tests/test_db.py +++ b/aibox/tests/test_db.py @@ -70,6 +70,87 @@ def test_memory_search_empty_keywords_returns_empty(mem_db): assert mem_db.search_memories([]) == [] +def test_memory_search_chinese_keyword(mem_db): + """中文关键词应可检索（trigram 分词或 LIKE 兜底）。""" + mem_db.add_memory(topic="天气查询", summary="今天天气晴朗", content="气温25度") + mem_db.add_memory(topic="会议安排", summary="下午三点开会", content="会议室A") + + results = mem_db.search_memories(["天气"]) + assert len(results) == 1 + assert results[0]["topic"] == "天气查询" + + +def test_memory_search_chinese_match_all(mem_db): + """中文多关键词 match_all：全部命中才返回。""" + mem_db.add_memory(topic="天气查询", summary="今天天气晴朗", content="气温25度") + mem_db.add_memory(topic="天气与会议", summary="开会前先查天气", content="会议室A") + + results = mem_db.search_memories(["天气", "开会"], match_all=True) + assert len(results) == 1 + assert results[0]["topic"] == "天气与会议" + + +def test_memory_search_like_escape_wildcards(mem_db): + """LIKE 兜底路径中 %/_ 应被转义为字面字符，不当作通配符。""" + import sqlite3 + mem_db.add_memory(topic="progress", summary="100% done") + mem_db.add_memory(topic="task_a", summary="plain text") + mem_db.add_memory(topic="taskXa", summary="wildcard trap") + + def like_search(kw): + with sqlite3.connect(mem_db.db_path) as conn: + conn.row_factory = sqlite3.Row + return mem_db._search_via_like(conn.cursor(), [kw], match_all=False, limit=100) + + # "%" 为字面字符："100%" 只命中 progress，不当作任意串通配 + results = like_search("100%") + assert [r["topic"] for r in results] == ["progress"] + + # "_" 为字面字符："ask_" 只命中 task_a；若未转义，taskXa 也会被 %ask_% 误匹配 + results = like_search("ask_") + assert [r["topic"] for r in results] == ["task_a"] + + +def test_memory_fts_migration_from_unicode61(tmp_path): + """旧 unicode61 FTS 表在重新打开时应自动迁移为 trigram 并回灌数据。""" + import sqlite3 + db_path = str(tmp_path / "legacy.db") + + # 模拟旧库：unicode61 建表 + 写入数据 + conn = sqlite3.connect(db_path) + conn.execute('''CREATE TABLE memories ( + id INTEGER PRIMARY KEY AUTOINCREMENT, topic TEXT NOT NULL, summary TEXT NOT NULL, + content TEXT, start_time TEXT NOT NULL, end_time TEXT NOT NULL, + created_at TEXT NOT NULL, tags TEXT, importance INTEGER DEFAULT 1)''') + conn.execute('''CREATE VIRTUAL TABLE memories_fts + USING fts5(topic, summary, content, tags, content=memories)''') + conn.execute("INSERT INTO memories (topic, summary, start_time, end_time, created_at) " + "VALUES ('天气查询', '今天天气晴朗', 't', 't', 't')") + conn.execute("INSERT INTO memories_fts (rowid, topic, summary) VALUES (1, '天气查询', '今天天气晴朗')") + conn.commit() + conn.close() + + # 重新打开 → 自动迁移 → 中文可搜 + db = MemoryDatabase(db_path=db_path) + results = db.search_memories(["天气"]) + assert len(results) == 1 + assert results[0]["topic"] == "天气查询" + + # 迁移后 FTS 表应为 trigram 分词（环境支持时） + conn = sqlite3.connect(db_path) + sql = conn.execute("SELECT sql FROM sqlite_master WHERE name='memories_fts'").fetchone()[0] + fts_count = conn.execute("SELECT COUNT(*) FROM memories_fts").fetchone()[0] + conn.close() + assert fts_count == 1 # 数据已回灌 + try: + probe = sqlite3.connect(":memory:") + probe.execute("CREATE VIRTUAL TABLE t USING fts5(x, tokenize='trigram')") + probe.close() + assert "trigram" in sql + except sqlite3.OperationalError: + pass # 环境不支持 trigram，保持 unicode61 属预期行为 + + def test_memory_delete(mem_db): mid = mem_db.add_memory(topic="待删除", summary="摘要") assert mem_db.delete_memory(mid) is True