This commit is contained in:
2025-09-07 08:08:16 +08:00
parent ee76c69ac8
commit dd5ed7fc53
2 changed files with 31 additions and 25 deletions

View File

@@ -1,3 +1,4 @@
import logging
import os
import time
@@ -14,11 +15,16 @@ AI 就能用刚写入的记忆来回答。
import json
import time
from VikingDBMemoryService import initialize_services, ensure_collection_exists, search_relevant_memories
# 控制日志输出
logger = logging.getLogger('CollectionMemory')
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(handler)
def handle_conversation_turn(memory_service, llm_client, collection_name, user_id, user_message, conversation_history):
"""处理一轮对话包括记忆搜索和LLM响应。"""
print("\n" + "=" * 60)
print(f"用户: {user_message}")
logger.info("\n" + "=" * 60)
logger.info(f"用户: {user_message}")
relevant_memories = search_relevant_memories(memory_service, collection_name, user_id, user_message)
@@ -28,7 +34,7 @@ def handle_conversation_turn(memory_service, llm_client, collection_name, user_i
[f"- {json.dumps(mem['memory_info'], ensure_ascii=False)}" for mem in relevant_memories])
system_prompt += f"\n\n这是我们过去的一些对话记忆,请参考:\n{memory_context}\n\n请利用这些信息来更好地理解和回应用户。"
print("AI正在思考...")
logger.info("AI正在思考...")
try:
messages = [{"role": "system", "content": system_prompt}] + conversation_history + [
@@ -39,10 +45,10 @@ def handle_conversation_turn(memory_service, llm_client, collection_name, user_i
)
assistant_reply = completion.choices[0].message.content
except Exception as e:
print(f"LLM调用失败: {e}")
logger.info(f"LLM调用失败: {e}")
assistant_reply = "抱歉,我现在有点混乱,无法回应。我们可以稍后再聊吗?"
print(f"伙伴: {assistant_reply}")
logger.info(f"伙伴: {assistant_reply}")
conversation_history.extend([
{"role": "user", "content": user_message},
@@ -53,10 +59,10 @@ def handle_conversation_turn(memory_service, llm_client, collection_name, user_i
def archive_conversation(memory_service, collection_name, user_id, assistant_id, conversation_history, topic_name):
"""将对话历史归档到记忆数据库。"""
if not conversation_history:
print("没有对话可以归档。")
logger.info("没有对话可以归档。")
return False
print(f"\n正在归档关于 '{topic_name}' 的对话...")
logger.info(f"\n正在归档关于 '{topic_name}' 的对话...")
session_id = f"{topic_name}_{int(time.time())}"
metadata = {
"default_user_id": user_id,
@@ -71,11 +77,11 @@ def archive_conversation(memory_service, collection_name, user_id, assistant_id,
messages=conversation_history,
metadata=metadata
)
print(f"对话已成功归档会话ID: {session_id}")
print("正在等待记忆索引更新...")
logger.info(f"对话已成功归档会话ID: {session_id}")
logger.info("正在等待记忆索引更新...")
return True
except Exception as e:
print(f"归档对话失败: {e}")
logger.info(f"归档对话失败: {e}")
return False
def setup_memory_collection(collection_name="emotional_support"):
@@ -83,18 +89,18 @@ def setup_memory_collection(collection_name="emotional_support"):
try:
memory_service, _ = initialize_services()
ensure_collection_exists(memory_service, collection_name)
print(f"记忆体 '{collection_name}' 创建/验证成功")
logger.info(f"记忆体 '{collection_name}' 创建/验证成功")
# 添加集合就绪等待
print("等待集合准备就绪...")
logger.info("等待集合准备就绪...")
if wait_for_collection_ready(memory_service, collection_name):
print(f"集合 '{collection_name}' 已就绪")
logger.info(f"集合 '{collection_name}' 已就绪")
return memory_service
else:
print(f"集合 '{collection_name}' 未能就绪")
logger.info(f"集合 '{collection_name}' 未能就绪")
return None
except Exception as e:
print(f"记忆体创建失败: {e}")
logger.info(f"记忆体创建失败: {e}")
return None
def wait_for_collection_ready(memory_service, collection_name, timeout=300, interval=10):
@@ -112,18 +118,18 @@ def wait_for_collection_ready(memory_service, collection_name, timeout=300, inte
collection_info = memory_service.get_collection(collection_name)
# 根据Volcengine API文档状态可能在Status字段中值可能为"READY"、"CREATING"等
status = collection_info.get("Status", "UNKNOWN")
print(f"集合 '{collection_name}' 当前状态: {status}")
logger.info(f"集合 '{collection_name}' 当前状态: {status}")
if status == "READY":
return True
time.sleep(interval)
except Exception as e:
print(f"检查集合状态失败: {e}")
logger.info(f"检查集合状态失败: {e}")
time.sleep(interval)
print(f"集合 '{collection_name}'{timeout}秒内未就绪")
logger.info(f"集合 '{collection_name}'{timeout}秒内未就绪")
return False
def main():
print("开始端到端记忆测试...")
logger.info("开始端到端记忆测试...")
collection_name="emotional_support"
try:
@@ -138,10 +144,10 @@ def main():
user_id = "xiaoming" # 用户ID:小明
assistant_id = "assistant1" # 助手ID:助手1
except Exception as e:
print(f"初始化失败: {e}")
logger.info(f"初始化失败: {e}")
return
print("\n--- 阶段 1: 初始对话 ---")
logger.info("\n--- 阶段 1: 初始对话 ---")
initial_conversation_history = []
handle_conversation_turn(
memory_service, llm_client, collection_name, user_id,
@@ -154,13 +160,13 @@ def main():
initial_conversation_history
)
print("\n--- 阶段 2: 归档记忆 ---")
logger.info("\n--- 阶段 2: 归档记忆 ---")
archive_conversation(
memory_service, collection_name, user_id, assistant_id,
initial_conversation_history, "study_stress_discussion"
)
print("\n--- 阶段 3: 验证记忆 ---")
logger.info("\n--- 阶段 3: 验证记忆 ---")
verification_conversation_history = []
handle_conversation_turn(
memory_service, llm_client, collection_name, user_id,
@@ -168,7 +174,7 @@ def main():
verification_conversation_history
)
print("\n端到端记忆测试完成!")
logger.info("\n端到端记忆测试完成!")
if __name__ == "__main__":
setup_memory_collection()