e2hang
2a844e83a8
- EmotionEngine: 5状态马尔可夫情绪机 + 蒙特卡洛转移 - VectorMemory: TF-IDF向量记忆 + SQLite持久化 + RAG检索 - AgentBrain: Ollama/OpenAI/Dummy三后端LLM - BehaviorScheduler: 优先级/冷却/活跃度调度 - FastAPI服务器 + WebSocket实时推送 - perception: 键鼠监控 + 屏幕截图 - ui/pet_window: PySide6桌宠窗口 + 像素动画 - assets/pet: 5情绪各2帧像素艺术资源
456 lines
14 KiB
Python
456 lines
14 KiB
Python
#!/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)
|