#!/usr/bin/env python3 """ brain.py 单元测试 运行: python -m agent.test_brain """ import asyncio import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).parent.parent)) from agent.brain import ( AgentBrain, OllamaBackend, OpenAIBackend, DummyLLMBackend, DEFAULT_SYSTEM_PROMPT, ) # ================================================================ # 辅助 # ================================================================ def _dummy_emotion(): """返回一个假的 emotion 引擎（只有 get_state / get_display_name）。""" class FakeEmotion: def get_state(self): return "happy" def get_display_name(self): return "开心 (happy)" return FakeEmotion() # ================================================================ # 测试 LLM 后端工厂 # ================================================================ def test_make_backend_ollama(): backend = AgentBrain._make_backend("ollama", {"model": "llama3"}) assert isinstance(backend, OllamaBackend) assert backend.model == "llama3" assert backend.base_url == OllamaBackend.DEFAULT_URL print("[PASS] test_make_backend_ollama") def test_make_backend_openai(): backend = AgentBrain._make_backend("openai", {"model": "gpt-4"}) assert isinstance(backend, OpenAIBackend) assert backend.model == "gpt-4" print("[PASS] test_make_backend_openai") def test_make_backend_dummy(): backend = AgentBrain._make_backend("dummy", {}) assert isinstance(backend, DummyLLMBackend) print("[PASS] test_make_backend_dummy") def test_make_backend_unknown(): try: AgentBrain._make_backend("unknown_backend", {}) assert False, "应抛出 ValueError" except ValueError as e: assert "unknown" in str(e).lower() print("[PASS] test_make_backend_unknown") # ================================================================ # 测试 Dummy LLM # ================================================================ def test_dummy_llm_sequential(): """测试 Dummy LLM 顺序返回不同回复。""" backend = DummyLLMBackend() async def run(): r1 = await backend.generate("hello") r2 = await backend.generate("world") r3 = await backend.generate("test") assert r1 in DummyLLMBackend.RESPONSES assert r2 in DummyLLMBackend.RESPONSES assert r1 != r2 # 顺序轮换 asyncio.run(run()) print("[PASS] test_dummy_llm_sequential") # ================================================================ # 测试 AgentBrain 初始化 # ================================================================ def test_brain_init_defaults(): """测试默认参数初始化。""" brain = AgentBrain( llm_backend="dummy", emotion_engine=_dummy_emotion(), ) assert brain._backend_type == "dummy" assert brain._session_history_limit == 10 assert brain._activity_threshold == 0.3 assert "emotion_display" in brain._system_prompt print("[PASS] test_brain_init_defaults") def test_brain_init_custom_system_prompt(): """测试自定义系统提示词。""" custom = "你是一个严肃的机器人。" brain = AgentBrain( llm_backend="dummy", system_prompt=custom, ) assert brain._system_prompt == custom print("[PASS] test_brain_init_custom_system_prompt") def test_brain_init_custom_llm_config(): """测试自定义 LLM 配置。""" brain = AgentBrain( llm_backend="ollama", llm_config={ "model": "deepseek-r1", "base_url": "http://localhost:8000", "timeout": 30.0, }, ) assert brain._llm_config["model"] == "deepseek-r1" assert brain._llm_config["base_url"] == "http://localhost:8000" print("[PASS] test_brain_init_custom_llm_config") # ================================================================ # 测试 think() 核心流程 # ================================================================ def test_think_returns_structure(): """测试 think() 返回结构完整性。""" async def run(): brain = AgentBrain( llm_backend="dummy", emotion_engine=_dummy_emotion(), ) result = await brain.think("你好！") assert isinstance(result, dict) assert "text" in result assert "action" in result assert "emotion_state" in result assert "memory_id" in result assert isinstance(result["text"], str) assert result["emotion_state"] == "happy" asyncio.run(run()) print("[PASS] test_think_returns_structure") def test_think_without_emotion_engine(): """测试无 emotion_engine 时仍可正常运行。""" async def run(): brain = AgentBrain(llm_backend="dummy") result = await brain.think("测试") assert result["emotion_state"] == "idle" assert isinstance(result["text"], str) asyncio.run(run()) print("[PASS] test_think_without_emotion_engine") def test_think_history_tracking(): """测试对话历史正确追踪。""" async def run(): brain = AgentBrain(llm_backend="dummy") await brain.think("你好") await brain.think("今天天气") await brain.think("你在干嘛") history = brain.get_history() assert len(history) == 6 # 3 轮 = 6 条消息 # 验证顺序 roles = [m["role"] for m in history] assert roles == ["user", "assistant", "user", "assistant", "user", "assistant"] asyncio.run(run()) print("[PASS] test_think_history_tracking") def test_think_history_limit(): """测试历史超过限制后自动截断。""" async def run(): brain = AgentBrain(llm_backend="dummy", session_history=3) for i in range(10): await brain.think(f"消息 {i}") history = brain.get_history() # 最多 3 轮 = 6 条消息 assert len(history) <= 6 asyncio.run(run()) print("[PASS] test_think_history_limit") def test_think_activity_context(): """测试 activity_level 正确注入到 prompt 并传给 LLM（不抛异常）。""" async def run(): brain = AgentBrain(llm_backend="dummy") # 高活跃度 result = await brain.think( "hello", context={"activity_level": 0.9}, ) assert isinstance(result["text"], str) # 低活跃度 result2 = await brain.think( "world", context={"activity_level": 0.1}, ) assert isinstance(result2["text"], str) asyncio.run(run()) print("[PASS] test_think_activity_context") # ================================================================ # 测试主动行为决策 # ================================================================ def test_decide_action_idle(): """测试空闲状态下无行为。""" brain = AgentBrain(llm_backend="dummy") # 无历史，刚初始化，cooldown 全部可触发 # 但 probability < threshold，所以返回 None result = brain.decide_action(emotion="idle", user_context={"activity_level": 0.9}) # 大概率是 None（随机概率触发） # 不做严格断言，只验证不抛异常 assert result is None or isinstance(result, dict) print("[PASS] test_decide_action_idle") def test_decide_action_high_priority_health(): """测试高优先级健康提醒（用户极度专注时强制提醒）。""" brain = AgentBrain(llm_backend="dummy") # 极低活跃度 + 不烦躁 → 应触发伸展提醒 result = brain.decide_action( emotion="idle", user_context={"activity_level": 0.05}, ) assert result is not None assert result["type"] in ("remind_stretch", "remind_water") assert result["priority"] == 0 print("[PASS] test_decide_action_high_priority_health") def test_decide_action_activity_threshold(): """测试用户高活跃度时不会打扰。""" brain = AgentBrain(llm_backend="dummy", activity_threshold=0.5) # 刚触发过一次，cooldown 中，返回 None result = brain.decide_action( emotion="happy", user_context={"activity_level": 0.8}, ) # 可能返回 None（冷却中或概率未触发） assert result is None or isinstance(result, dict) print("[PASS] test_decide_action_activity_threshold") def test_decide_action_annoyed_blocks_health(): """测试烦躁状态不会强制健康提醒（避免激怒用户）。""" brain = AgentBrain(llm_backend="dummy") result = brain.decide_action( emotion="annoyed", user_context={"activity_level": 0.05}, ) # 烦躁时不强制提醒（但可能在随机概率下触发 nudge） assert result is None or result["type"] != "remind_stretch" print("[PASS] test_decide_action_annoyed_blocks_health") def test_emotion_trigger_prob(): """测试各情绪对应的触发概率。""" brain = AgentBrain(llm_backend="dummy") probs = { "happy": brain._get_emotion_trigger_prob("happy"), "idle": brain._get_emotion_trigger_prob("idle"), "focused": brain._get_emotion_trigger_prob("focused"), "annoyed": brain._get_emotion_trigger_prob("annoyed"), "sleepy": brain._get_emotion_trigger_prob("sleepy"), } # 开心触发率最高，专注最低 assert probs["happy"] > probs["focused"] assert probs["happy"] > probs["sleepy"] assert probs["focused"] == 0.05 # 专注几乎不打扰 print(f"[PASS] test_emotion_trigger_prob: {probs}") def test_cooldown_mechanism(): """测试冷却机制。""" import time brain = AgentBrain(llm_backend="dummy") # 第一次应返回 True（无冷却记录） assert brain._check_cooldown("test_action", cooldown=60.0) is True # 立即再次调用应返回 False（冷却中） assert brain._check_cooldown("test_action", cooldown=60.0) is False # 用极短冷却测试（0.01s） brain._action_cooldown["test_fast"] = time.time() - 0.1 assert brain._check_cooldown("test_fast", cooldown=0.05) is True print("[PASS] test_cooldown_mechanism") def test_action_parse_basic(): """测试 [ACTION: type:message] 标签解析。""" brain = AgentBrain(llm_backend="dummy") result = brain._parse_action( "好的，我去提醒你喝水。[ACTION: remind_water:该喝水了！]" ) assert result is not None assert result["type"] == "remind_water" assert result["message"] == "该喝水了！" assert result["priority"] == 3 print("[PASS] test_action_parse_basic") def test_action_parse_no_tag(): """测试无 ACTION 标签时返回 None。""" brain = AgentBrain(llm_backend="dummy") result = brain._parse_action("你好呀！今天过得怎么样？") assert result is None print("[PASS] test_action_parse_no_tag") def test_action_parse_whitespace(): """测试带多余空格的 ACTION 标签。""" brain = AgentBrain(llm_backend="dummy") result = brain._parse_action( "[ACTION: remind_water : 喝点水吧！ ]" ) assert result is not None assert result["type"] == "remind_water" print("[PASS] test_action_parse_whitespace") # ================================================================ # 测试辅助方法 # ================================================================ def test_get_status(): """测试 get_status() 返回完整状态。""" brain = AgentBrain( llm_backend="dummy", emotion_engine=_dummy_emotion(), ) status = brain.get_status() assert status["backend"] == "dummy" assert status["emotion"] == "happy" assert "history_turns" in status assert "cooldowns" in status print(f"[PASS] test_get_status: {status}") def test_clear_history(): """测试清空历史。""" async def run(): brain = AgentBrain(llm_backend="dummy") await brain.think("hello") await brain.think("world") assert len(brain.get_history()) > 0 brain.clear_history() assert len(brain.get_history()) == 0 asyncio.run(run()) print("[PASS] test_clear_history") # ================================================================ # 运行 # ================================================================ def main(): print("=" * 60) print("brain.py 测试套件") print("=" * 60) tests = [ # 后端工厂 test_make_backend_ollama, test_make_backend_openai, test_make_backend_dummy, test_make_backend_unknown, # Dummy LLM test_dummy_llm_sequential, # 初始化 test_brain_init_defaults, test_brain_init_custom_system_prompt, test_brain_init_custom_llm_config, # think() test_think_returns_structure, test_think_without_emotion_engine, test_think_history_tracking, test_think_history_limit, test_think_activity_context, # 主动行为 test_decide_action_idle, test_decide_action_high_priority_health, test_decide_action_activity_threshold, test_decide_action_annoyed_blocks_health, test_emotion_trigger_prob, test_cooldown_mechanism, # Action 解析 test_action_parse_basic, test_action_parse_no_tag, test_action_parse_whitespace, # 辅助 test_get_status, test_clear_history, ] passed = failed = 0 for test in tests: try: test() passed += 1 except Exception as exc: print(f"[FAIL] {test.__name__}: {exc}") import traceback traceback.print_exc() failed += 1 print("=" * 60) print(f"测试结果: {passed}/{passed+failed} 通过", end="") if failed: print(f", {failed} 失败", end="" ) print() print("=" * 60) return failed == 0 if __name__ == "__main__": success = main() sys.exit(0 if success else 1)