113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
import logging
|
||
import sys
|
||
import time
|
||
|
||
from Volcengine.Kit.VikingDBMemoryService import MEMORY_COLLECTION_NAME, initialize_services, \
|
||
handle_conversation_turn, archive_conversation
|
||
|
||
# 控制日志输出
|
||
logger = logging.getLogger('ChatWithMemory')
|
||
logger.setLevel(logging.INFO)
|
||
|
||
# 只添加一次处理器，避免重复日志
|
||
if not logger.handlers:
|
||
handler = logging.StreamHandler()
|
||
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
|
||
logger.addHandler(handler)
|
||
|
||
|
||
def main():
|
||
logger.info("开始测试大模型记忆功能...")
|
||
|
||
try:
|
||
# 使用initialize_services函数初始化服务和LLM客户端
|
||
memory_service, llm_client = initialize_services()
|
||
|
||
collection_name = MEMORY_COLLECTION_NAME
|
||
user_id = "liming" # 用户李明
|
||
assistant_id = "assistant" # 助手ID:助手
|
||
|
||
# 告知大模型用户信息
|
||
logger.info("告知大模型用户信息...")
|
||
user_info = "李明，男生，15岁，家住长春"
|
||
|
||
# 记录信息到记忆体
|
||
logger.info("记录用户信息到记忆体...")
|
||
|
||
# 准备对话历史
|
||
conversation_history = []
|
||
|
||
# 使用正确的handle_conversation_turn方法参数
|
||
response = handle_conversation_turn(
|
||
memory_service=memory_service,
|
||
llm_client=llm_client,
|
||
collection_name=collection_name,
|
||
user_id=user_id,
|
||
user_message=f"请记住以下用户信息：{user_info}",
|
||
conversation_history=conversation_history
|
||
)
|
||
|
||
logger.info(f"模型回复: {response}")
|
||
|
||
# 归档对话
|
||
archive_conversation(
|
||
memory_service=memory_service,
|
||
collection_name=collection_name,
|
||
user_id=user_id,
|
||
assistant_id=assistant_id,
|
||
conversation_history=conversation_history,
|
||
topic_name="user_info"
|
||
)
|
||
|
||
# 等待一段时间确保索引更新
|
||
logger.info("等待索引更新...")
|
||
time.sleep(5)
|
||
|
||
# 验证大模型是否记住个人信息
|
||
logger.info("验证大模型是否记住个人信息...")
|
||
test_conversation_history = []
|
||
|
||
test_response = handle_conversation_turn(
|
||
memory_service=memory_service,
|
||
llm_client=llm_client,
|
||
collection_name=collection_name,
|
||
user_id=user_id,
|
||
user_message="请告诉我李明的个人信息",
|
||
conversation_history=test_conversation_history
|
||
)
|
||
|
||
logger.info(f"测试回复: {test_response}")
|
||
|
||
# 检查回复中是否包含关键信息
|
||
keywords = ["李明", "男", "15", "长春"]
|
||
found_keywords = [kw for kw in keywords if kw in test_response]
|
||
|
||
if len(found_keywords) == len(keywords):
|
||
logger.info("✅ 大模型成功记住了用户信息！")
|
||
else:
|
||
logger.info(f"❌ 大模型可能没有完全记住用户信息。找到的关键词: {found_keywords}")
|
||
|
||
# 尝试直接搜索记忆
|
||
logger.info("尝试直接搜索记忆...")
|
||
filter_params = {
|
||
"user_id": [user_id],
|
||
"assistant_id": assistant_id, # 添加这一行
|
||
"memory_type": ["sys_event_v1", "sys_profile_v1"]
|
||
}
|
||
search_result = memory_service.search_memory(
|
||
collection_name=collection_name,
|
||
query="李明 15岁 长春 男生",
|
||
filter=filter_params,
|
||
limit=5
|
||
)
|
||
|
||
logger.info(f"搜索结果: {search_result}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"操作失败: {e}")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|