from WxMini.Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
from WxMini.Milvus.Utils.MilvusConnectionPool import *
from WxMini.Milvus.Config.MulvusConfig 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. 查询所有数据
try:
    # 使用 Milvus 的 query 方法查询所有数据
    results = collection_manager.collection.query(
        expr="",  # 空表达式表示查询所有数据
        output_fields=["id", "text", "embedding"],  # 指定返回的字段
        limit=1000  # 设置最大返回记录数
    )
    print("查询结果：")
    if results:
        for result in results:
            try:
                text = result["text"]  # 获取 text 字段
                embedding = result["embedding"]  # 获取 embedding 字段
                print(f"ID: {result['id']}, Text: {text}, Embedding: {embedding[:5]}...")  # 只打印前 5 维向量
            except Exception as e:
                print(f"查询失败: {e}")
    else:
        print("未找到相关数据，请检查查询参数或数据。")
except Exception as e:
    print(f"查询失败: {e}")

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

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