dsProject/dsLightRag/Volcengine/T2_CreateIndex.py

import logging
import logging
import threading

from volcengine.ApiInfo import ApiInfo
from volcengine.Credentials import Credentials
from volcengine.ServiceInfo import ServiceInfo
from volcengine.auth.SignerV4 import SignerV4
from volcengine.base.Service import Service

import json
import time
from dotenv import load_dotenv
from volcenginesdkarkruntime import Ark

from Config.Config import VOLC_SECRETKEY, VOLC_API_KEY, VOLC_ACCESSKEY
from Volcengine.VikingDBMemoryService import MEMORY_COLLECTION_NAME, VikingDBMemoryException, VikingDBMemoryService, \
    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()