62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
import logging
|
||
|
||
from Volcengine.VikingDBMemoryService import MEMORY_COLLECTION_NAME, initialize_services, ensure_collection_exists
|
||
from Volcengine.chat import handle_conversation_turn, archive_conversation
|
||
|
||
# 控制日志输出
|
||
logger = logging.getLogger('CollectionMemory')
|
||
logger.setLevel(logging.INFO)
|
||
|
||
# 只添加一次处理器，避免重复日志
|
||
if not logger.handlers:
|
||
handler = logging.StreamHandler()
|
||
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
|
||
logger.addHandler(handler)
|
||
|
||
|
||
|
||
def main():
|
||
logger.info("开始端到端记忆测试...")
|
||
|
||
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)
|
||
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端到端记忆测试完成！")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |