from Backup.Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
from Config.Config import *

# 1. 使用连接池管理 Milvus 连接
milvus_pool = MilvusConnectionPool(host=MS_HOST, port=MS_PORT, max_connections=MS_MAX_CONNECTIONS)

# 2. 从连接池中获取一个连接
connection = milvus_pool.get_connection()

# 3. 初始化集合管理器
collection_name = MS_COLLECTION_NAME
collection_manager = MilvusCollectionManager(collection_name)

# 4. 加载集合到内存
collection_manager.load_collection()
print(f"集合 '{collection_name}' 已加载到内存。")

# 5. 直接在代码中指定要查询的标签
#query_tag = "MATH_DATA_1"  # 可以修改为MATH_DATA_2或其他需要的标签
query_tag = "MATH_DATA_2"  # 可以修改为MATH_DATA_2或其他需要的标签
expr = f"array_contains(tags['tags'], '{query_tag}')"
print(f"查询表达式: {expr}")

# 6. 查询数据
try:
    results = collection_manager.collection.query(
        expr=expr,
        output_fields=["id", "tags", "user_input", "timestamp", "embedding"],
        limit=1000
    )
    print(f"查询标签 '{query_tag}' 结果：")
    if results:
        for result in results:
            try:
                print(f"ID: {result['id']}")
                print(f"标签: {result['tags']}")
                print(f"用户问题: {result['user_input']}")
                print(f"时间: {result['timestamp']}")
                print(f"向量: {result['embedding'][:5]}...")
                print("-" * 40)
            except Exception as e:
                print(f"处理结果失败: {e}")
    else:
        print(f"未找到标签为 '{query_tag}' 的数据。")
except Exception as e:
    print(f"查询失败: {e}")

# 7. 释放连接
milvus_pool.release_connection(connection)

# 8. 关闭连接池
milvus_pool.close()