ReMe/tests/test_reme.py

118 lines
4.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Test module for ReMe memory system functionality."""
import asyncio
from reme import ReMe
from reme.core.schema import VectorNode, MemoryNode
async def test_reme():
"""Tests ReMe memory system with personal information storage and retrieval."""
# 构建一段包含个人信息的对话
reme = ReMe(vector_store={"collection_name": "reme"})
await reme.start()
# reme = await ReMe.create(vector_store={"collection_name": "reme"})
await reme.vector_store.delete_all()
messages = [
{
"role": "user",
"content": "你好我是张伟今年28岁是一名软件工程师。",
"time_created": "2026-01-29 10:00:00",
},
{
"role": "assistant",
"content": "你好张伟!很高兴认识你。作为一名软件工程师,你主要从事什么方向的开发工作呢?",
"time_created": "2026-01-29 10:00:05",
},
{
"role": "user",
"content": "我主要做后端开发擅长Python和Go语言。最近在研究AI Agent相关的技术。",
"time_created": "2026-01-29 10:00:30",
},
{
"role": "assistant",
"content": "很棒Python和Go都是非常实用的语言。AI Agent是当前很热门的方向你在这方面有什么具体的研究重点吗",
"time_created": "2026-01-29 10:00:35",
},
{
"role": "user",
"content": "我特别关注记忆系统的设计希望能让AI Agent具有长期记忆能力。我的工作地点在北京平时喜欢看技术博客和参加技术分享会。",
"time_created": "2026-01-29 10:01:00",
},
{
"role": "assistant",
"content": "记忆系统确实是AI Agent的核心能力之一。北京有很多优秀的技术社区和活动相信你能找到很多志同道合的朋友。",
"time_created": "2026-01-29 10:01:05",
},
{
"role": "user",
"content": "是的我每周末都会去参加一些技术沙龙。对了我的邮箱是zhangwei@example.com如果有好的技术资料可以发给我。",
"time_created": "2026-01-29 10:01:30",
},
{
"role": "assistant",
"content": "好的我记下了。保持学习的热情很重要祝你在AI Agent领域的研究越来越深入",
"time_created": "2026-01-29 10:01:35",
},
]
print("=" * 60)
print("步骤1: 开始总结对话并生成记忆")
print("=" * 60)
# 对对话进行总结,生成记忆
await reme.summary_memory(
messages=messages,
user_name="zhangwei",
description="用户自我介绍和技术兴趣分享",
)
print("\n✓ 记忆总结完成")
print("\n" + "=" * 60)
print("步骤2: 查看已存储的记忆节点")
print("=" * 60)
# 列出所有存储的记忆节点
nodes: list[VectorNode] = await reme.vector_store.list()
for i, node in enumerate(nodes, 1):
memory_node = MemoryNode.from_vector_node(node)
print(f"{i} {memory_node.model_dump_json()}")
print("\n" + "=" * 60)
print("步骤3: 测试记忆检索 - 验证个人信息")
print("=" * 60)
# 测试问题1: 检索用户姓名
query1 = "用户叫什么名字?"
print(f"\n问题1: {query1}")
result1 = await reme.retrieve_memory(query=query1, user_name="zhangwei")
print(f"检索结果:\n{result1}")
# 测试问题2: 检索技术背景
query2 = "用户擅长什么编程语言和技术方向?"
print(f"\n问题2: {query2}")
result2 = await reme.retrieve_memory(query=query2, user_name="zhangwei")
print(f"检索结果:\n{result2}")
# 测试问题3: 检索个人信息
query3 = "用户的工作地点和联系方式是什么?"
print(f"\n问题3: {query3}")
result3 = await reme.retrieve_memory(query=query3, user_name="zhangwei")
print(f"检索结果:\n{result3}")
# 测试问题4: 检索兴趣爱好
query4 = "用户平时有什么爱好或活动?"
print(f"\n问题4: {query4}")
result4 = await reme.retrieve_memory(query=query4, user_name="zhangwei")
print(f"检索结果:\n{result4}")
print("\n" + "=" * 60)
print("测试完成!")
print("=" * 60)
if __name__ == "__main__":
asyncio.run(test_reme())