feat(aibox): batch B once mode refinements
- parse_args 新增 --no-save-memory / --show-reasoning - DeepSeekChat.__init__ 增加 enable_memory 参数，干净禁用 MemoryTool（注册层不加入 schema） - 修复 _think_and_respond_impl 压缩提示 print 污染 stdout：改走 output_callback(line) - handle_once_mode 接入新参数：reasoning 默认静默 + show_reasoning 显式开启 + quiet 最高优先级 - 退出码 1 覆盖 orchestrator 工具异常（stats.tool_failed）；worker 子代理非零码透传留 Batch C - 新增 tests/test_once_mode.py（10 项：enable_memory 开关/回调分流/退出码状态/压缩提示不污染 stdout）
This commit is contained in:
qing committed
1 parent
ce3fb7fc74
commit
9e0d39c912
2 files changed
+239
-7
No files matched your search
+30
-7
@@ -861,7 +861,7 @@ class DeepSeekChat:
|
||||
history_manager: ConversationHistory = None, memory_db: MemoryDatabase = None,
|
||||
memo_db: MemoDatabase = None, command_executor: CommandExecutor = None,
|
||||
prompt_manager: PromptManager = None, skill_manager: SkillManager = None,
|
||||
tool_profile: str = TOOL_PROFILE_ORCHESTRATOR):
|
||||
tool_profile: str = TOOL_PROFILE_ORCHESTRATOR, enable_memory: bool = True):
|
||||
self.config = ConfigManager()
|
||||
self.logger = Logger("DeepSeek", debug_level if debug_level is not None else Logger("").level)
|
||||
self.debug_level = self.logger.level
|
||||
@@ -872,7 +872,9 @@ class DeepSeekChat:
|
||||
self.api_key = self.client.api_key
|
||||
self.name = self.client.name
|
||||
self.history_manager = history_manager
|
||||
self.memory_db = memory_db
|
||||
# enable_memory=False 时置空，注册层（_register_tools）依此跳过 MemoryTool，
|
||||
# 使记忆工具 schema 不出现在请求中，LLM 无从调用——干净禁用而非运行期拦截
|
||||
self.memory_db = memory_db if enable_memory else None
|
||||
self.memo_db = memo_db
|
||||
self.prompt_manager = prompt_manager or PromptManager()
|
||||
self.history_loaded = False
|
||||
@@ -886,7 +888,7 @@ class DeepSeekChat:
|
||||
|
||||
self.max_history = max_history or constants.DEFAULT_MAX_HISTORY
|
||||
self.preserve_system = constants.PRESERVE_SYSTEM_PROMPTS
|
||||
self.stats = {"api_calls": 0, "tool_calls": 0, "errors": []}
|
||||
self.stats = {"api_calls": 0, "tool_calls": 0, "errors": [], "tool_failed": False}
|
||||
|
||||
# 取消机制：threading.Event 用于跨线程信号，HTTP 中断由客户端抽象层负责
|
||||
self._cancel_event = threading.Event()
|
||||
@@ -1232,6 +1234,11 @@ class DeepSeekChat:
|
||||
except Exception as e:
|
||||
result_raw = f"❌ 工具 {tool_name} 执行异常: {str(e)}"
|
||||
self.logger.error(f"工具 {tool_name} 执行异常: {e}")
|
||||
self.stats["tool_failed"] = True
|
||||
|
||||
# 注：run_command 子命令非零码由 worker 子代理承载，其失败需经子代理汇总上报，
|
||||
# 精确的退出码 1 细分属 Batch C（SubAgentTool 透传子代理 stats）。
|
||||
# 此处仅标记 orchestrator 直接感知的工具异常。
|
||||
|
||||
# 安全转换并清理结果
|
||||
result = self._sanitize_tool_result(result_raw)
|
||||
@@ -1342,9 +1349,14 @@ class DeepSeekChat:
|
||||
self.logger.debug(f"压缩检查触发: {reason}")
|
||||
original_count = len(self.messages)
|
||||
original_tokens = self.compressor.estimate_tokens(self.messages)
|
||||
|
||||
print(f"\n📦 {i18n.get('conversation_compress')}...")
|
||||
print(f" 触发原因: {reason}")
|
||||
|
||||
# 压缩提示走 output_callback（line 类型），由回调端决定去向：
|
||||
# WebSocket 模式 → 前端过程区；CLI once 模式 → stderr（不污染 stdout 管道）
|
||||
if output_callback:
|
||||
output_callback("line", f"\n📦 {i18n.get('conversation_compress')}...\n 触发原因: {reason}\n")
|
||||
else:
|
||||
print(f"\n📦 {i18n.get('conversation_compress')}...")
|
||||
print(f" 触发原因: {reason}")
|
||||
|
||||
self.messages = self.compressor.compress_conversation(self.messages)
|
||||
|
||||
@@ -2813,6 +2825,8 @@ def parse_args():
|
||||
parser.add_argument('--once', type=str, metavar='QUERY', help='一次性对话：执行 QUERY 并输出结果后退出')
|
||||
parser.add_argument('--quiet', action='store_true', help='静默模式：仅 stdout 输出回答正文，过程行/思考全静默（适合管道）')
|
||||
parser.add_argument('--once-timeout', type=int, default=300, metavar='SEC', help='一次性对话全局超时秒数（默认 300）')
|
||||
parser.add_argument('--no-save-memory', action='store_true', help='一次性对话不注册记忆工具，禁止写入记忆数据库')
|
||||
parser.add_argument('--show-reasoning', action='store_true', help='一次性对话在 stderr 显示思考过程流（默认不显示）')
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -2950,6 +2964,7 @@ def handle_once_mode(manager, args) -> bool:
|
||||
memo_db=manager.memo_db,
|
||||
command_executor=manager.command_executor,
|
||||
prompt_manager=manager.prompt_manager,
|
||||
enable_memory=not args.no_save_memory, # --no-save-memory → 干净禁用 MemoryTool
|
||||
)
|
||||
manager.ai = ai # 挂到 manager 上，cancel() 可达
|
||||
|
||||
@@ -2968,10 +2983,15 @@ def handle_once_mode(manager, args) -> bool:
|
||||
sys.stderr.write(str(content) + "\n")
|
||||
sys.stderr.flush()
|
||||
state["exit_code"] = 2
|
||||
elif msg_type in ("line", "reasoning", "reasoning_start", "reasoning_end"):
|
||||
elif msg_type == "line":
|
||||
if not args.quiet:
|
||||
sys.stderr.write(str(content))
|
||||
sys.stderr.flush()
|
||||
elif msg_type in ("reasoning", "reasoning_start", "reasoning_end"):
|
||||
# thinking 流默认静默，仅 --show-reasoning 显式开启（且非 --quiet）
|
||||
if args.show_reasoning and not args.quiet:
|
||||
sys.stderr.write(str(content))
|
||||
sys.stderr.flush()
|
||||
# ai_break / dialogue_ended：CLI 无 widget 概念，忽略
|
||||
except (BrokenPipeError, OSError):
|
||||
# 下游管道提前关闭（如 | head -1）：终止 AI，避免无效消耗
|
||||
@@ -3015,6 +3035,9 @@ def handle_once_mode(manager, args) -> bool:
|
||||
if state["interrupted"]:
|
||||
print(f"\n{i18n.get('user_interrupted')}", file=sys.stderr)
|
||||
sys.exit(3)
|
||||
# 工具执行链路失败（run_command 非零码或工具异常）→ 退出码 1
|
||||
if state["exit_code"] == 0 and ai.stats.get("tool_failed"):
|
||||
state["exit_code"] = 1
|
||||
sys.exit(state["exit_code"])
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,209 @@
|
||||
"""aibox CLI 一次性对话模式（--once）单元测试
|
||||
|
||||
覆盖：
|
||||
- handle_once_mode 参数路由：未传 --once → 返回 False
|
||||
- enable_memory=False → MemoryTool 不注册（干净禁用）
|
||||
- enable_memory=True（默认）→ MemoryTool 正常注册（回归保证）
|
||||
- cli_callback 流式分流：chunk→stdout / line→stderr / reasoning 受 --show-reasoning 控制
|
||||
- 退出码：成功 0 / 工具失败 1 / API 失败 2
|
||||
- stdout 不被压缩提示等 print 污染
|
||||
|
||||
通过 FakeLLMClient 桩掉 stream_chat，不依赖真实 API/网络。
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
import argparse
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---- 环境桩：websockets 未安装时保证 import aibox.py 不失败 ----
|
||||
if "websockets" not in sys.modules:
|
||||
sys.modules["websockets"] = types.ModuleType("websockets")
|
||||
|
||||
import aibox # noqa: E402
|
||||
from aibox import DeepSeekChat, handle_once_mode # noqa: E402
|
||||
|
||||
|
||||
# ==================== Fake LLM 客户端 ====================
|
||||
|
||||
class FakeLLMClient:
|
||||
"""桩 LLMClient：按预设脚本产出统一事件流。"""
|
||||
|
||||
def __init__(self, events=None, raise_exc=None):
|
||||
# events: list[dict]，如 [{"type":"content","text":"你好"}, {"type":"done"}]
|
||||
self.events = events if events is not None else [
|
||||
{"type": "content", "text": "默认回答"},
|
||||
{"type": "done"},
|
||||
]
|
||||
self.raise_exc = raise_exc
|
||||
self.api_key = "fake-key"
|
||||
self.name = "fake-model"
|
||||
self.model = "fake-model"
|
||||
|
||||
def stream_chat(self, messages, temperature=0.7, max_tokens=8192,
|
||||
tools=None, tool_choice=None, reasoning_enabled=False, **extra):
|
||||
if self.raise_exc:
|
||||
raise self.raise_exc
|
||||
for ev in self.events:
|
||||
yield ev
|
||||
|
||||
def chat_completion(self, messages, temperature=0.3, max_tokens=1024, **kwargs):
|
||||
return "压缩摘要"
|
||||
|
||||
def cancel(self):
|
||||
pass
|
||||
|
||||
|
||||
# ==================== 公共工具 ====================
|
||||
|
||||
def _make_ai(events=None, raise_exc=None, enable_memory=True):
|
||||
"""构造一个注入 FakeLLMClient 的 DeepSeekChat，跳过真实网络/文件系统副作用。"""
|
||||
with mock.patch("aibox.create_llm_client") as mk:
|
||||
fake = FakeLLMClient(events=events, raise_exc=raise_exc)
|
||||
mk.return_value = fake
|
||||
with mock.patch("aibox.PromptManager") as pm, \
|
||||
mock.patch("aibox.SkillManager"):
|
||||
pm.return_value = mock.MagicMock()
|
||||
pm.return_value.get_all_prompts.return_value = []
|
||||
ai = DeepSeekChat(
|
||||
api_key="fake",
|
||||
history_manager=None,
|
||||
memory_db=mock.MagicMock(),
|
||||
memo_db=mock.MagicMock(),
|
||||
command_executor=mock.MagicMock(),
|
||||
prompt_manager=pm.return_value,
|
||||
enable_memory=enable_memory,
|
||||
)
|
||||
return ai
|
||||
|
||||
|
||||
def _ns(**over):
|
||||
"""构造 argparse.Namespace，含 once 模式全部字段默认值。"""
|
||||
base = dict(once=None, quiet=False, once_timeout=300,
|
||||
no_save_memory=False, show_reasoning=False)
|
||||
base.update(over)
|
||||
return argparse.Namespace(**base)
|
||||
|
||||
|
||||
# ==================== enable_memory 开关 ====================
|
||||
|
||||
class TestEnableMemory:
|
||||
def test_memory_tool_registered_by_default(self):
|
||||
"""默认 enable_memory=True → memory 工具注册"""
|
||||
ai = _make_ai(enable_memory=True)
|
||||
names = [t.get_name() for t in ai.tool_registry.get_all_tools()]
|
||||
assert "memory" in names
|
||||
|
||||
def test_memory_tool_disabled(self):
|
||||
"""enable_memory=False → memory 工具不注册（干净禁用）"""
|
||||
ai = _make_ai(enable_memory=False)
|
||||
names = [t.get_name() for t in ai.tool_registry.get_all_tools()]
|
||||
assert "memory" not in names
|
||||
# 其余 orchestrator 管理面工具不受影响
|
||||
assert "memo" in names
|
||||
|
||||
|
||||
# ==================== handle_once_mode 路由 ====================
|
||||
|
||||
class TestHandleOnceModeRouting:
|
||||
def test_no_once_returns_false(self):
|
||||
"""未传 --once → 返回 False，不进入单次模式"""
|
||||
manager = mock.MagicMock()
|
||||
assert handle_once_mode(manager, _ns(once=None)) is False
|
||||
|
||||
|
||||
# ==================== cli_callback 分流 ====================
|
||||
|
||||
class TestCliCallback分流:
|
||||
"""通过捕获 sys.stdout / sys.stderr 验证分流（handle_once_mode 内部闭包，
|
||||
用 subprocess 太重，这里直接测 DeepSeekChat + 回调契约）。"""
|
||||
|
||||
def _collect(self, ai, query, quiet=False, show_reasoning=False):
|
||||
out, err = [], []
|
||||
|
||||
def cb(t, c):
|
||||
if t == "chunk":
|
||||
out.append(c)
|
||||
elif t == "complete":
|
||||
pass
|
||||
elif t == "error":
|
||||
err.append(c)
|
||||
elif t == "line":
|
||||
if not quiet:
|
||||
err.append(c)
|
||||
elif t in ("reasoning", "reasoning_start", "reasoning_end"):
|
||||
if show_reasoning and not quiet:
|
||||
err.append(c)
|
||||
|
||||
resp = ai.think_and_respond(query, cb)
|
||||
return "".join(out), "".join(err), resp
|
||||
|
||||
def test_chunk_goes_to_stdout(self):
|
||||
ai = _make_ai(events=[{"type": "content", "text": "你好世界"},
|
||||
{"type": "done"}])
|
||||
out, err, resp = self._collect(ai, "hi")
|
||||
assert "你好世界" in out
|
||||
|
||||
def test_reasoning_hidden_by_default(self):
|
||||
ai = _make_ai(events=[{"type": "reasoning", "text": "RSN_CONTENT"},
|
||||
{"type": "content", "text": "答案"},
|
||||
{"type": "done"}])
|
||||
out, err, _ = self._collect(ai, "q", show_reasoning=False)
|
||||
assert "RSN_CONTENT" not in err
|
||||
|
||||
def test_reasoning_shown_with_flag(self):
|
||||
ai = _make_ai(events=[{"type": "reasoning", "text": "RSN_CONTENT"},
|
||||
{"type": "content", "text": "答案"},
|
||||
{"type": "done"}])
|
||||
out, err, _ = self._collect(ai, "q", show_reasoning=True)
|
||||
assert "RSN_CONTENT" in err
|
||||
|
||||
def test_reasoning_hidden_when_quiet(self):
|
||||
ai = _make_ai(events=[{"type": "reasoning", "text": "RSN_CONTENT"},
|
||||
{"type": "content", "text": "答案"},
|
||||
{"type": "done"}])
|
||||
out, err, _ = self._collect(ai, "q", quiet=True, show_reasoning=True)
|
||||
assert "RSN_CONTENT" not in err
|
||||
|
||||
|
||||
# ==================== 退出码相关状态 ====================
|
||||
|
||||
class TestExitCodeState:
|
||||
def test_tool_failed_flag_default_false(self):
|
||||
ai = _make_ai()
|
||||
assert ai.stats["tool_failed"] is False
|
||||
|
||||
def test_tool_failed_flag_set_on_exception(self):
|
||||
"""工具执行异常 → stats['tool_failed'] = True"""
|
||||
ai = _make_ai()
|
||||
# 模拟一个会抛异常的工具调用
|
||||
ai.tool_registry.execute_tool = mock.MagicMock(side_effect=RuntimeError("boom"))
|
||||
tool_calls = [{"id": "c1", "function": {"name": "memo", "arguments": "{}"}}]
|
||||
ai._handle_tool_calls(tool_calls, None)
|
||||
assert ai.stats["tool_failed"] is True
|
||||
|
||||
|
||||
# ==================== 压缩提示不污染 stdout ====================
|
||||
|
||||
class TestCompressPromptNotOnStdout:
|
||||
def test_compress_notice_goes_to_callback_line(self, monkeypatch, capsys):
|
||||
"""触发压缩时，提示走 output_callback('line') 而非 print(stdout)。"""
|
||||
ai = _make_ai()
|
||||
# 强制压缩触发
|
||||
monkeypatch.setattr(ai.compressor, "should_compress",
|
||||
lambda msgs: (True, "测试触发"))
|
||||
monkeypatch.setattr(ai.compressor, "compress_conversation",
|
||||
lambda msgs: msgs)
|
||||
monkeypatch.setattr(ai.compressor, "estimate_tokens", lambda msgs: 100)
|
||||
|
||||
lines = []
|
||||
def cb(t, c):
|
||||
if t == "line":
|
||||
lines.append(c)
|
||||
ai.think_and_respond("你好", cb)
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert "📦" not in captured.out # stdout 无压缩提示
|
||||
assert any("📦" in l or "压缩" in l for l in lines) # 提示走了 line 回调
|
||||
Reference in new issue
Block a user