142 lines
5.3 KiB
Python
142 lines
5.3 KiB
Python
import os
|
||
import time
|
||
|
||
from volcenginesdkarkruntime import Ark
|
||
|
||
from Config.Config import VOLC_API_KEY
|
||
|
||
"""
|
||
在记忆库准备好后，我们先模拟一段包含两轮的完整对话。
|
||
对话结束后，把这段对话历史消息写入记忆库。然后再开启一个新话题，提出和刚才相关的问题，
|
||
AI 就能用刚写入的记忆来回答。
|
||
注意：首次写入需要 3–5 分钟建立索引，这段时间内检索会报错。
|
||
"""
|
||
import json
|
||
import time
|
||
from VikingDBMemoryService import initialize_services, ensure_collection_exists, search_relevant_memories
|
||
|
||
def handle_conversation_turn(memory_service, llm_client, collection_name, user_id, user_message, conversation_history):
|
||
"""处理一轮对话，包括记忆搜索和LLM响应。"""
|
||
print("\n" + "=" * 60)
|
||
print(f"用户: {user_message}")
|
||
|
||
relevant_memories = search_relevant_memories(memory_service, collection_name, user_id, user_message)
|
||
|
||
system_prompt = "你是一个富有同情心、善于倾听的AI伙伴，拥有长期记忆能力。你的目标是为用户提供情感支持和温暖的陪伴。"
|
||
if relevant_memories:
|
||
memory_context = "\n".join(
|
||
[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正在思考...")
|
||
|
||
try:
|
||
messages = [{"role": "system", "content": system_prompt}] + conversation_history + [
|
||
{"role": "user", "content": user_message}]
|
||
completion = llm_client.chat.completions.create(
|
||
model="doubao-seed-1-6-flash-250715",
|
||
messages=messages
|
||
)
|
||
assistant_reply = completion.choices[0].message.content
|
||
except Exception as e:
|
||
print(f"LLM调用失败: {e}")
|
||
assistant_reply = "抱歉，我现在有点混乱，无法回应。我们可以稍后再聊吗？"
|
||
|
||
print(f"伙伴: {assistant_reply}")
|
||
|
||
conversation_history.extend([
|
||
{"role": "user", "content": user_message},
|
||
{"role": "assistant", "content": assistant_reply}
|
||
])
|
||
return assistant_reply
|
||
|
||
def archive_conversation(memory_service, collection_name, user_id, assistant_id, conversation_history, topic_name):
|
||
"""将对话历史归档到记忆数据库。"""
|
||
if not conversation_history:
|
||
print("没有对话可以归档。")
|
||
return False
|
||
|
||
print(f"\n正在归档关于 '{topic_name}' 的对话...")
|
||
session_id = f"{topic_name}_{int(time.time())}"
|
||
metadata = {
|
||
"default_user_id": user_id,
|
||
"default_assistant_id": assistant_id,
|
||
"time": int(time.time() * 1000)
|
||
}
|
||
|
||
try:
|
||
memory_service.add_session(
|
||
collection_name=collection_name,
|
||
session_id=session_id,
|
||
messages=conversation_history,
|
||
metadata=metadata
|
||
)
|
||
print(f"对话已成功归档，会话ID: {session_id}")
|
||
print("正在等待记忆索引更新...")
|
||
return True
|
||
except Exception as e:
|
||
print(f"归档对话失败: {e}")
|
||
return False
|
||
|
||
def setup_memory_collection(collection_name="emotional_support"):
|
||
"""独立封装记忆体创建逻辑，返回memory_service供测试使用"""
|
||
try:
|
||
memory_service, _ = initialize_services()
|
||
ensure_collection_exists(memory_service, collection_name)
|
||
print(f"记忆体 '{collection_name}' 创建/验证成功")
|
||
return memory_service
|
||
except Exception as e:
|
||
print(f"记忆体创建失败: {e}")
|
||
return None
|
||
|
||
def main():
|
||
print("开始端到端记忆测试...")
|
||
collection_name="emotional_support"
|
||
|
||
try:
|
||
# 调用封装的记忆体创建函数
|
||
memory_service = setup_memory_collection()
|
||
if not memory_service:
|
||
return
|
||
llm_client = Ark(
|
||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||
api_key=VOLC_API_KEY
|
||
)
|
||
user_id = "xiaoming" # 用户ID:小明
|
||
assistant_id = "assistant1" # 助手ID:助手1
|
||
except Exception as e:
|
||
print(f"初始化失败: {e}")
|
||
return
|
||
|
||
print("\n--- 阶段 1: 初始对话 ---")
|
||
initial_conversation_history = []
|
||
handle_conversation_turn(
|
||
memory_service, llm_client, collection_name, user_id,
|
||
"你好，我是小明，今年18岁，但压力好大。",
|
||
initial_conversation_history
|
||
)
|
||
handle_conversation_turn(
|
||
memory_service, llm_client, collection_name, user_id,
|
||
"马上就要高考了，家里人的期待好高。",
|
||
initial_conversation_history
|
||
)
|
||
|
||
print("\n--- 阶段 2: 归档记忆 ---")
|
||
archive_conversation(
|
||
memory_service, collection_name, user_id, assistant_id,
|
||
initial_conversation_history, "study_stress_discussion"
|
||
)
|
||
|
||
print("\n--- 阶段 3: 验证记忆 ---")
|
||
verification_conversation_history = []
|
||
handle_conversation_turn(
|
||
memory_service, llm_client, collection_name, user_id,
|
||
"我最近很焦虑，不知道该怎么办。",
|
||
verification_conversation_history
|
||
)
|
||
|
||
print("\n端到端记忆测试完成！")
|
||
|
||
if __name__ == "__main__":
|
||
setup_memory_collection()
|
||
# main() |