77 lines
2.7 KiB
Python
77 lines
2.7 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("验证大模型是否记住个人信息...")
|
||
|
|
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()
|