'commit'
This commit is contained in:
Binary file not shown.
@@ -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()
|
||||
|
Reference in New Issue
Block a user