87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
启用bbav5持久化记忆
|
||
|
|
"""
|
||
|
|
import asyncio
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from gfa_core.mas.memory import MemoryManager
|
||
|
|
from gfa_core.mas.memory.persistent import PersistentMemory
|
||
|
|
from gfa_core.mas.memory.session import SessionMemory, SessionData
|
||
|
|
from gfa_core.mas.memory.case import CaseMemory, AnalysisCase
|
||
|
|
from gfa_core.mas.memory.knowledge import KnowledgeMemory, KnowledgeEntry
|
||
|
|
|
||
|
|
|
||
|
|
async def create_persistent_memory_manager(storage_dir: Path = None) -> MemoryManager:
|
||
|
|
"""创建带持久化的MemoryManager"""
|
||
|
|
|
||
|
|
if storage_dir is None:
|
||
|
|
storage_dir = Path("/home/xia-lab/wutx/00test/bbav5/memory_storage")
|
||
|
|
|
||
|
|
storage_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
# 创建持久化存储
|
||
|
|
session_storage = PersistentMemory(storage_dir / "sessions")
|
||
|
|
case_storage = PersistentMemory(storage_dir / "cases")
|
||
|
|
knowledge_storage = PersistentMemory(storage_dir / "knowledge")
|
||
|
|
|
||
|
|
# 创建MemoryManager
|
||
|
|
manager = MemoryManager(
|
||
|
|
session_memory=SessionMemory(), # 内存缓存 + 持久化备份
|
||
|
|
case_memory=CaseMemory(),
|
||
|
|
knowledge_memory=KnowledgeMemory(),
|
||
|
|
)
|
||
|
|
|
||
|
|
# 可选：添加持久化包装器
|
||
|
|
# 或修改SessionMemory等使用PersistentMemory作为后端
|
||
|
|
|
||
|
|
print(f"✅ MemoryManager创建成功")
|
||
|
|
print(f" 存储目录: {storage_dir}")
|
||
|
|
print(f" Session: {storage_dir / 'sessions'}")
|
||
|
|
print(f" Case: {storage_dir / 'cases'}")
|
||
|
|
print(f" Knowledge: {storage_dir / 'knowledge'}")
|
||
|
|
|
||
|
|
return manager
|
||
|
|
|
||
|
|
|
||
|
|
async def test_persistent_memory():
|
||
|
|
"""测试持久化记忆"""
|
||
|
|
manager = await create_persistent_memory_manager()
|
||
|
|
|
||
|
|
# 测试Session
|
||
|
|
session_id = "test_session_001"
|
||
|
|
session = await manager.create_session(session_id, {"test": "data"})
|
||
|
|
print(f"\n✅ Session创建: {session_id}")
|
||
|
|
|
||
|
|
# 添加消息
|
||
|
|
await manager.add_message_to_session(session_id, "user", "Hello")
|
||
|
|
await manager.add_message_to_session(session_id, "assistant", "Hi there")
|
||
|
|
print(f"✅ 消息添加完成")
|
||
|
|
|
||
|
|
# 测试Case
|
||
|
|
case = await manager.create_case(
|
||
|
|
case_id="test_case_001",
|
||
|
|
user_id="user_001",
|
||
|
|
gene_family="SPL",
|
||
|
|
species="Arabidopsis",
|
||
|
|
)
|
||
|
|
print(f"✅ Case创建: {case.case_id}")
|
||
|
|
|
||
|
|
# 更新结果
|
||
|
|
await manager.update_case_results("test_case_001", {"status": "completed"})
|
||
|
|
print(f"✅ Case结果更新")
|
||
|
|
|
||
|
|
# 检索验证
|
||
|
|
retrieved_session = await manager.get_session(session_id)
|
||
|
|
retrieved_case = await manager.get_case("test_case_001")
|
||
|
|
|
||
|
|
print(f"\n✅ 检索验证:")
|
||
|
|
print(f" Session: {retrieved_session.session_id if retrieved_session else 'None'}")
|
||
|
|
print(f" Case: {retrieved_case.case_id if retrieved_case else 'None'}")
|
||
|
|
|
||
|
|
print("\n🎉 持久化记忆测试完成！")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
asyncio.run(test_persistent_memory())
|