diff --git a/dsLightRag/Volcengine/VikingDBMemoryService.py b/dsLightRag/Volcengine/VikingDBMemoryService.py index 9bb15294..619233fa 100644 --- a/dsLightRag/Volcengine/VikingDBMemoryService.py +++ b/dsLightRag/Volcengine/VikingDBMemoryService.py @@ -15,17 +15,19 @@ from volcengine.auth.SignerV4 import SignerV4 from volcengine.base.Service import Service from volcenginesdkarkruntime import Ark +from Config.Config import VOLC_SECRETKEY, VOLC_ACCESSKEY, VOLC_API_KEY + def initialize_services(ak=None, sk=None, ark_api_key=None): """初始化记忆数据库服务和LLM客户端""" load_dotenv() # 如果参数未提供,尝试从环境变量获取 if not ak: - ak = os.getenv("VOLC_ACCESSKEY") + ak = VOLC_ACCESSKEY if not sk: - sk = os.getenv("VOLC_SECRETKEY") + sk = VOLC_SECRETKEY if not ark_api_key: - ark_api_key = os.getenv("VOLC_API_KEY") + ark_api_key = VOLC_API_KEY if not all([ak, sk, ark_api_key]): raise ValueError("必须提供 VOLC_ACCESSKEY, VOLC_SECRETKEY, 和 VOLC_API_KEY。") diff --git a/dsLightRag/Volcengine/__pycache__/VikingDBMemoryService.cpython-310.pyc b/dsLightRag/Volcengine/__pycache__/VikingDBMemoryService.cpython-310.pyc index 92e57fa8..17662bf7 100644 Binary files a/dsLightRag/Volcengine/__pycache__/VikingDBMemoryService.cpython-310.pyc and b/dsLightRag/Volcengine/__pycache__/VikingDBMemoryService.cpython-310.pyc differ diff --git a/dsLightRag/Volcengine/chat.py b/dsLightRag/Volcengine/chat.py index 5e9eb541..e30f992c 100644 --- a/dsLightRag/Volcengine/chat.py +++ b/dsLightRag/Volcengine/chat.py @@ -1,5 +1,10 @@ +import os import time +from volcenginesdkarkruntime import Ark + +from Config.Config import VOLC_API_KEY + """ 在记忆库准备好后,我们先模拟一段包含两轮的完整对话。 对话结束后,把这段对话历史消息写入记忆库。然后再开启一个新话题,提出和刚才相关的问题, @@ -7,7 +12,9 @@ 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) @@ -71,15 +78,32 @@ def archive_conversation(memory_service, collection_name, user_id, assistant_id, 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, llm_client = initialize_services() - collection_name = "emotional_support" - user_id = "xiaoming" # 用户ID:小明 - assistant_id = "assistant1" # 助手ID:助手1 - ensure_collection_exists(memory_service, collection_name) + # 调用封装的记忆体创建函数 + 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 @@ -114,4 +138,5 @@ def main(): print("\n端到端记忆测试完成!") if __name__ == "__main__": - main() \ No newline at end of file + setup_memory_collection() +# main() \ No newline at end of file