from pymilvus import (
    connections,
    utility,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    MilvusClient,
    db,
    Function,
    FunctionType,
    AnnSearchRequest,
    WeightedRanker
)
import os

class Milvus_Database:
    def __init__(self, user=os.getenv("Milvus_USER", "root"), password=os.getenv("Milvus_PASSWORD", "Milvus"), uri=os.getenv("Milvus_URI")):
        self.client = MilvusClient(
            uri="http://172.18.30.165:19530",
            token=f"{user}:{password}"
        )
    
    def search_resource(self, collection_name, database_name,resource):
        self.client.using_database(db_name=database_name)
        iterator = self.client.query_iterator(
            collection_name=collection_name,
            batch_size=5,
            filter=f"resource like \"{resource}\"",
            output_fields=["text","resource"],
        )

        results = []

        while True:
            result = iterator.next()
            if not result:
                iterator.close()
                break

            print(result)
            results += result
        return results
        # res = self.client.query(
        #     collection_name=collection_name,
        #     filter=f"resource like \"{resource}\"",
        #     output_fields=["resource"],
        #     limit=None  # 先限制数量测试
        # )
        # return res

if __name__ == "__main__":
    db = Milvus_Database()
    res = db.search_resource("DockerTest", database_name="XIAN",resource="测试.pdf")
    for i,re in enumerate(res):
        print(i)
    # print(res)