107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
import json
|
||
import logging
|
||
import time
|
||
from Volcengine.VikingDBMemoryService import initialize_services, MEMORY_COLLECTION_NAME
|
||
from volcenginesdkarkruntime import Ark
|
||
from Config.Config import VOLC_API_KEY
|
||
|
||
# 配置日志
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger('MemoryChatSimulator')
|
||
|
||
def simulate_chat_session(memory_service, llm_client, user_id, session_topic):
|
||
"""模拟完整对话会话：创建对话→存储记忆→搜索记忆"""
|
||
# 初始化对话历史
|
||
conversation_history = []
|
||
logger.info(f"\n=== 开始新对话会话: {session_topic} ===")
|
||
|
||
# 模拟两轮对话
|
||
user_messages = [
|
||
f"你好，我是{user_id}，{session_topic}",
|
||
"我最近感觉压力很大，不知道该怎么办"
|
||
]
|
||
|
||
for message in user_messages:
|
||
# 处理对话轮次
|
||
assistant_reply = memory_service.handle_conversation_turn(
|
||
llm_client=llm_client,
|
||
user_id=user_id,
|
||
user_message=message,
|
||
conversation_history=conversation_history
|
||
)
|
||
time.sleep(1) # 模拟思考延迟
|
||
|
||
# 归档对话记忆
|
||
logger.info("\n=== 归档对话记忆 ===")
|
||
archive_success = memory_service.archive_conversation(
|
||
user_id=user_id,
|
||
assistant_id="simulated_assistant",
|
||
conversation_history=conversation_history,
|
||
topic_name=session_topic
|
||
)
|
||
|
||
# 搜索相关记忆
|
||
if archive_success:
|
||
logger.info("\n=== 搜索相关记忆 ===")
|
||
# 等待索引更新
|
||
logger.info("等待记忆库索引更新...")
|
||
time.sleep(5) # 实际环境可能需要更长时间
|
||
|
||
# 搜索与当前主题相关的记忆
|
||
search_query = f"{session_topic}相关的问题"
|
||
relevant_memories = memory_service.search_relevant_memories(
|
||
collection_name=MEMORY_COLLECTION_NAME,
|
||
user_id=user_id,
|
||
query=search_query,
|
||
limit=3
|
||
)
|
||
|
||
logger.info(f"\n搜索到 '{search_query}' 的相关记忆 ({len(relevant_memories)}条):")
|
||
for i, memory in enumerate(relevant_memories, 1):
|
||
logger.info(f"\n记忆 {i} (相关度: {memory['score']:.3f}):")
|
||
logger.info(f"内容: {memory['memory_info']}")
|
||
|
||
return archive_success
|
||
|
||
def main():
|
||
"""主函数：初始化服务并执行模拟测试"""
|
||
logger.info("===== 记忆库存储与搜索模拟测试 ====")
|
||
|
||
try:
|
||
# 初始化服务
|
||
logger.info("初始化记忆库服务和LLM客户端...")
|
||
memory_service, llm_client = initialize_services()
|
||
|
||
# 确保集合就绪
|
||
logger.info(f"检查集合 '{MEMORY_COLLECTION_NAME}' 是否就绪...")
|
||
|
||
# 先验证集合是否存在
|
||
try:
|
||
collection_info = memory_service.get_collection(MEMORY_COLLECTION_NAME)
|
||
logger.info(f"集合 '{MEMORY_COLLECTION_NAME}' 存在，详细信息: {json.dumps(collection_info, ensure_ascii=False)}...")
|
||
except Exception as e:
|
||
logger.error(f"集合 '{MEMORY_COLLECTION_NAME}' 不存在: {str(e)}")
|
||
logger.error("请先运行T2_CreateIndex.py创建集合")
|
||
return
|
||
|
||
if not memory_service.wait_for_collection_ready(timeout=120):
|
||
logger.error("集合未就绪，无法继续测试")
|
||
return
|
||
|
||
# 模拟两个不同用户的对话会话
|
||
user_sessions = [
|
||
("student_001", "学习压力问题"),
|
||
("teacher_001", "教学方法讨论")
|
||
]
|
||
|
||
for user_id, topic in user_sessions:
|
||
simulate_chat_session(memory_service, llm_client, user_id, topic)
|
||
time.sleep(2) # 会话间隔
|
||
|
||
logger.info("\n===== 模拟测试完成 ====")
|
||
|
||
except Exception as e:
|
||
logger.error(f"测试过程中发生错误: {str(e)}", exc_info=True)
|
||
|
||
if __name__ == "__main__":
|
||
main() |