'commit'
This commit is contained in:
@@ -1,5 +1,10 @@
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import os
|
||||
|
||||
# 添加当前目录到系统路径
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME
|
||||
from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
|
||||
|
@@ -1,7 +1,12 @@
|
||||
import logging
|
||||
import sys
|
||||
import os
|
||||
import time
|
||||
import json
|
||||
|
||||
from Volcengine.VikingDBMemoryService import MEMORY_COLLECTION_NAME, initialize_services, ensure_collection_exists
|
||||
from Volcengine.chat import handle_conversation_turn, archive_conversation
|
||||
from volcenginesdkarkruntime import Ark
|
||||
from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY, VOLC_API_KEY
|
||||
from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME
|
||||
|
||||
# 控制日志输出
|
||||
logger = logging.getLogger('CollectionMemory')
|
||||
@@ -14,48 +19,97 @@ if not logger.handlers:
|
||||
logger.addHandler(handler)
|
||||
|
||||
|
||||
|
||||
def main():
|
||||
logger.info("开始端到端记忆测试...")
|
||||
|
||||
logger.info("开始创建索引...")
|
||||
|
||||
# 初始化记忆库服务
|
||||
memory_service = VikingDBMemoryService(
|
||||
ak=VOLC_ACCESSKEY,
|
||||
sk=VOLC_SECRETKEY,
|
||||
host="api-knowledgebase.mlp.cn-beijing.volces.com",
|
||||
region="cn-beijing"
|
||||
)
|
||||
|
||||
# 初始化LLM客户端
|
||||
llm_client = Ark(
|
||||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||||
api_key=VOLC_API_KEY
|
||||
)
|
||||
|
||||
try:
|
||||
memory_service, llm_client = initialize_services()
|
||||
collection_name = MEMORY_COLLECTION_NAME
|
||||
user_id = "system"
|
||||
assistant_id = "assistant"
|
||||
ensure_collection_exists(memory_service, collection_name)
|
||||
|
||||
# 确保集合存在
|
||||
logger.info("检查/创建集合...")
|
||||
memory_service.ensure_collection_exists(collection_name)
|
||||
|
||||
# 添加测试数据以触发索引构建
|
||||
logger.info("添加测试数据...")
|
||||
test_messages = [
|
||||
{"role": "user", "content": "你好,我是测试用户"},
|
||||
{"role": "assistant", "content": "你好,我是测试助手"}
|
||||
]
|
||||
|
||||
test_metadata = {
|
||||
"default_user_id": user_id,
|
||||
"default_assistant_id": assistant_id,
|
||||
"time": int(time.time() * 1000)
|
||||
}
|
||||
|
||||
session_id = f"test_session_{int(time.time())}"
|
||||
memory_service.add_session(
|
||||
collection_name=collection_name,
|
||||
session_id=session_id,
|
||||
messages=test_messages,
|
||||
metadata=test_metadata
|
||||
)
|
||||
|
||||
logger.info("测试数据添加成功,等待索引构建...")
|
||||
|
||||
# 使用与chat.py.backup相同的等待索引就绪的逻辑
|
||||
max_retries = 30
|
||||
retry_interval = 60 # 秒,与chat.py.backup一致
|
||||
|
||||
for retry in range(max_retries):
|
||||
try:
|
||||
# 尝试搜索以验证索引是否就绪
|
||||
filter_params = {
|
||||
"user_id": user_id,
|
||||
"memory_type": ["sys_event_v1", "sys_profile_v1"]
|
||||
}
|
||||
|
||||
response = memory_service.search_memory(
|
||||
collection_name=collection_name,
|
||||
query="测试查询",
|
||||
filter=filter_params,
|
||||
limit=1
|
||||
)
|
||||
|
||||
# 如果搜索成功,说明索引已就绪
|
||||
logger.info(f"索引已就绪,找到 {response.get('data', {}).get('count', 0)} 条记录")
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
if "1000023" in error_message: # 与chat.py.backup中的错误码一致
|
||||
retry_attempt = retry + 1
|
||||
logger.info(f"记忆索引正在构建中。将在{retry_interval}秒后重试... (尝试次数 {retry_attempt})")
|
||||
time.sleep(retry_interval)
|
||||
else:
|
||||
logger.error(f"搜索时发生错误: {error_message}")
|
||||
raise
|
||||
else:
|
||||
# 如果循环正常结束(未break),说明超时
|
||||
logger.error(f"索引构建超时,已尝试 {max_retries} 次")
|
||||
sys.exit(1)
|
||||
|
||||
logger.info("索引创建和测试完成!")
|
||||
|
||||
except Exception as e:
|
||||
logger.info(f"初始化失败: {e}")
|
||||
return
|
||||
|
||||
logger.info("\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
|
||||
)
|
||||
|
||||
logger.info("\n--- 阶段 2: 归档记忆 ---")
|
||||
archive_conversation(
|
||||
memory_service, collection_name, user_id, assistant_id,
|
||||
initial_conversation_history, "study_stress_discussion"
|
||||
)
|
||||
|
||||
logger.info("\n--- 阶段 3: 验证记忆 ---")
|
||||
verification_conversation_history = []
|
||||
handle_conversation_turn(
|
||||
memory_service, llm_client, collection_name, user_id,
|
||||
"我最近很焦虑,不知道该怎么办。",
|
||||
verification_conversation_history
|
||||
)
|
||||
|
||||
logger.info("\n端到端记忆测试完成!")
|
||||
logger.error(f"操作失败: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
@@ -4,9 +4,10 @@ pip install --upgrade "volcengine-python-sdk[ark]"
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import sys
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from volcengine.ApiInfo import ApiInfo
|
||||
@@ -21,107 +22,16 @@ from Config.Config import VOLC_SECRETKEY, VOLC_ACCESSKEY, VOLC_API_KEY
|
||||
# 配置日志
|
||||
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)
|
||||
|
||||
# 只添加一次处理器,避免重复日志
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler()
|
||||
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
|
||||
logger.addHandler(handler)
|
||||
|
||||
# 记忆体集合名称
|
||||
MEMORY_COLLECTION_NAME="dsideal_collection"
|
||||
|
||||
def initialize_services(ak=None, sk=None, ark_api_key=None):
|
||||
"""初始化记忆数据库服务和LLM客户端"""
|
||||
load_dotenv()
|
||||
# 如果参数未提供,尝试从环境变量获取
|
||||
if not ak:
|
||||
ak = VOLC_ACCESSKEY
|
||||
if not sk:
|
||||
sk = VOLC_SECRETKEY
|
||||
if not ark_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。")
|
||||
|
||||
memory_service = VikingDBMemoryService(ak=ak, sk=sk)
|
||||
llm_client = Ark(
|
||||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||||
api_key=ark_api_key,
|
||||
)
|
||||
return memory_service, llm_client
|
||||
|
||||
def ensure_collection_exists(memory_service, collection_name, description="",
|
||||
builtin_event_types=["sys_event_v1", "sys_profile_collect_v1"],
|
||||
builtin_entity_types=["sys_profile_v1"]):
|
||||
"""检查记忆集合是否存在,如果不存在则创建。"""
|
||||
try:
|
||||
memory_service.get_collection(collection_name)
|
||||
print(f"记忆集合 '{collection_name}' 已存在。")
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
if "collection not exist" in error_message:
|
||||
print(f"记忆集合 '{collection_name}' 未找到,正在创建...")
|
||||
try:
|
||||
memory_service.create_collection(
|
||||
collection_name=collection_name,
|
||||
description=description,
|
||||
builtin_event_types=builtin_event_types,
|
||||
builtin_entity_types=builtin_entity_types
|
||||
)
|
||||
print(f"记忆集合 '{collection_name}' 创建成功。")
|
||||
print("等待集合准备就绪...")
|
||||
except Exception as create_e:
|
||||
print(f"创建集合失败: {create_e}")
|
||||
raise
|
||||
else:
|
||||
print(f"检查集合时出错: {e}")
|
||||
raise
|
||||
|
||||
def search_relevant_memories(memory_service, collection_name, user_id, query, limit=3):
|
||||
"""搜索与用户查询相关的记忆,并在索引构建中时重试。"""
|
||||
print(f"正在搜索与 '{query}' 相关的记忆...")
|
||||
retry_attempt = 0
|
||||
while True:
|
||||
try:
|
||||
filter_params = {
|
||||
"user_id": user_id, # 修正为字符串类型
|
||||
"memory_type": ["sys_event_v1", "sys_profile_v1"]
|
||||
}
|
||||
response = memory_service.search_memory(
|
||||
collection_name=collection_name,
|
||||
query=query,
|
||||
filter=filter_params,
|
||||
limit=limit
|
||||
)
|
||||
|
||||
memories = []
|
||||
if response.get('data', {}).get('count', 0) > 0:
|
||||
for result in response['data']['result_list']:
|
||||
if 'memory_info' in result and result['memory_info']:
|
||||
memories.append({
|
||||
'memory_info': result['memory_info'],
|
||||
'score': result['score']
|
||||
})
|
||||
|
||||
if memories:
|
||||
if retry_attempt > 0:
|
||||
print("重试后搜索成功。")
|
||||
print(f"找到 {len(memories)} 条相关记忆:")
|
||||
for i, memory in enumerate(memories, 1):
|
||||
print(
|
||||
f" {i}. (相关度: {memory['score']:.3f}): {json.dumps(memory['memory_info'], ensure_ascii=False, indent=2)}")
|
||||
else:
|
||||
print("未找到相关记忆。")
|
||||
return memories
|
||||
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
if "1000023" in error_message:
|
||||
retry_attempt += 1
|
||||
print(f"记忆索引正在构建中。将在60秒后重试... (尝试次数 {retry_attempt})")
|
||||
time.sleep(60)
|
||||
else:
|
||||
print(f"搜索记忆时出错 (不可重试): {e}")
|
||||
return []
|
||||
|
||||
class VikingDBMemoryException(Exception):
|
||||
def __init__(self, code, request_id, message=None):
|
||||
@@ -309,7 +219,8 @@ class VikingDBMemoryService(Service):
|
||||
logger.info("\n" + "=" * 60)
|
||||
logger.info(f"用户: {user_message}")
|
||||
|
||||
relevant_memories = self.search_relevant_memories(user_id, user_message)
|
||||
# 修复:调用正确的search_relevant_memories方法
|
||||
relevant_memories = self.search_relevant_memories(MEMORY_COLLECTION_NAME, user_id, user_message)
|
||||
|
||||
system_prompt = "你是一个富有同情心、善于倾听的AI伙伴,拥有长期记忆能力。你的目标是为用户提供情感支持和温暖的陪伴。"
|
||||
if relevant_memories:
|
||||
@@ -395,10 +306,105 @@ class VikingDBMemoryService(Service):
|
||||
logger.error(f"集合 '{MEMORY_COLLECTION_NAME}' 在{timeout}秒内未就绪")
|
||||
return False
|
||||
|
||||
def initialize_services(self, ak=None, sk=None, ark_api_key=None):
|
||||
"""初始化记忆数据库服务和LLM客户端"""
|
||||
load_dotenv()
|
||||
# 如果参数未提供,尝试从环境变量获取
|
||||
if not ak:
|
||||
ak = VOLC_ACCESSKEY
|
||||
if not sk:
|
||||
sk = VOLC_SECRETKEY
|
||||
if not ark_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。")
|
||||
|
||||
memory_service = VikingDBMemoryService(ak=ak, sk=sk)
|
||||
llm_client = Ark(
|
||||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||||
api_key=ark_api_key,
|
||||
)
|
||||
return memory_service, llm_client
|
||||
|
||||
def ensure_collection_exists(self, collection_name, description="",
|
||||
builtin_event_types=["sys_event_v1", "sys_profile_collect_v1"],
|
||||
builtin_entity_types=["sys_profile_v1"]):
|
||||
"""检查记忆集合是否存在,如果不存在则创建。"""
|
||||
try:
|
||||
self.get_collection(collection_name)
|
||||
logger.info(f"记忆集合 '{collection_name}' 已存在。")
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
if "collection not exist" in error_message:
|
||||
logger.info(f"记忆集合 '{collection_name}' 未找到,正在创建...")
|
||||
try:
|
||||
self.create_collection(
|
||||
collection_name=collection_name,
|
||||
description=description,
|
||||
builtin_event_types=builtin_event_types,
|
||||
builtin_entity_types=builtin_entity_types
|
||||
)
|
||||
logger.info(f"记忆集合 '{collection_name}' 创建成功。")
|
||||
logger.info("等待集合准备就绪...")
|
||||
except Exception as create_e:
|
||||
logger.info(f"创建集合失败: {create_e}")
|
||||
raise
|
||||
else:
|
||||
logger.info(f"检查集合时出错: {e}")
|
||||
raise
|
||||
|
||||
def search_relevant_memories(self, collection_name, user_id, query, limit=3):
|
||||
"""搜索与用户查询相关的记忆,并在索引构建中时重试。"""
|
||||
logger.info(f"正在搜索与 '{query}' 相关的记忆...")
|
||||
retry_attempt = 0
|
||||
while True:
|
||||
try:
|
||||
filter_params = {
|
||||
"user_id": user_id, # 修正为字符串类型
|
||||
"memory_type": ["sys_event_v1", "sys_profile_v1"]
|
||||
}
|
||||
response = self.search_memory(
|
||||
collection_name=collection_name,
|
||||
query=query,
|
||||
filter=filter_params,
|
||||
limit=limit
|
||||
)
|
||||
|
||||
memories = []
|
||||
if response.get('data', {}).get('count', 0) > 0:
|
||||
for result in response['data']['result_list']:
|
||||
if 'memory_info' in result and result['memory_info']:
|
||||
memories.append({
|
||||
'memory_info': result['memory_info'],
|
||||
'score': result['score']
|
||||
})
|
||||
|
||||
if memories:
|
||||
if retry_attempt > 0:
|
||||
logger.info("重试后搜索成功。")
|
||||
logger.info(f"找到 {len(memories)} 条相关记忆:")
|
||||
for i, memory in enumerate(memories, 1):
|
||||
logger.info(
|
||||
f" {i}. (相关度: {memory['score']:.3f}): {json.dumps(memory['memory_info'], ensure_ascii=False, indent=2)}")
|
||||
else:
|
||||
logger.info("未找到相关记忆。")
|
||||
return memories
|
||||
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
if "1000023" in error_message:
|
||||
retry_attempt += 1
|
||||
logger.info(f"记忆索引正在构建中。将在60秒后重试... (尝试次数 {retry_attempt})")
|
||||
time.sleep(60)
|
||||
else:
|
||||
logger.info(f"搜索记忆时出错 (不可重试): {e}")
|
||||
return []
|
||||
|
||||
def setup_memory_collection(self):
|
||||
"""独立封装记忆体创建逻辑,返回memory_service供测试使用"""
|
||||
try:
|
||||
ensure_collection_exists(self, MEMORY_COLLECTION_NAME)
|
||||
self.ensure_collection_exists(MEMORY_COLLECTION_NAME)
|
||||
logger.info(f"记忆体 '{MEMORY_COLLECTION_NAME}' 创建/验证成功")
|
||||
|
||||
# 添加集合就绪等待
|
||||
@@ -411,58 +417,4 @@ class VikingDBMemoryService(Service):
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.info(f"记忆体创建失败: {e}")
|
||||
return None
|
||||
|
||||
def run_end_to_end_test(self):
|
||||
"""端到端记忆测试的主函数"""
|
||||
logger.info("开始端到端记忆测试...")
|
||||
|
||||
try:
|
||||
# 调用封装的记忆体创建函数
|
||||
memory_service = self.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:
|
||||
logger.info(f"初始化失败: {e}")
|
||||
return
|
||||
|
||||
logger.info("\n--- 阶段 1: 初始对话 ---")
|
||||
initial_conversation_history = []
|
||||
self.handle_conversation_turn(
|
||||
llm_client, user_id,
|
||||
"你好,我是小明,今年18岁,但压力好大。",
|
||||
initial_conversation_history
|
||||
)
|
||||
self.handle_conversation_turn(
|
||||
llm_client, user_id,
|
||||
"马上就要高考了,家里人的期待好高。",
|
||||
initial_conversation_history
|
||||
)
|
||||
|
||||
logger.info("\n--- 阶段 2: 归档记忆 ---")
|
||||
self.archive_conversation(
|
||||
user_id, assistant_id,
|
||||
initial_conversation_history, "study_stress_discussion"
|
||||
)
|
||||
|
||||
logger.info("\n--- 阶段 3: 验证记忆 ---")
|
||||
verification_conversation_history = []
|
||||
self.handle_conversation_turn(
|
||||
llm_client, user_id,
|
||||
"我最近很焦虑,不知道该怎么办。",
|
||||
verification_conversation_history
|
||||
)
|
||||
|
||||
logger.info("\n端到端记忆测试完成!")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 初始化服务
|
||||
memory_service, _ = initialize_services()
|
||||
# 运行端到端测试
|
||||
memory_service.run_end_to_end_test()
|
||||
return None
|
Binary file not shown.
Binary file not shown.
@@ -1,12 +1,13 @@
|
||||
import json
|
||||
import threading
|
||||
|
||||
import time
|
||||
from dotenv import load_dotenv
|
||||
from volcengine.ApiInfo import ApiInfo
|
||||
from volcengine.Credentials import Credentials
|
||||
from volcengine.base.Service import Service
|
||||
from volcengine.ServiceInfo import ServiceInfo
|
||||
from volcengine.auth.SignerV4 import SignerV4
|
||||
from volcengine.base.Request import Request
|
||||
from volcengine.base.Service import Service
|
||||
from volcenginesdkarkruntime import Ark
|
||||
|
||||
from Config.Config import VOLC_SECRETKEY, VOLC_API_KEY, VOLC_ACCESSKEY
|
||||
|
||||
@@ -193,11 +194,6 @@ class VikingDBMemoryService(Service):
|
||||
return json.loads(res)
|
||||
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from dotenv import load_dotenv
|
||||
from volcenginesdkarkruntime import Ark
|
||||
|
||||
|
||||
def initialize_services():
|
Reference in New Issue
Block a user