# 通过pymilvus sdk来操作数据库

from pymilvus import (
    connections,
    db,
    connections,
    Collection,
    utility,
)
import numpy as np
import random


def test_connection(
    db_name: str = "blog",
    collection_name: str = "articles",
) -> None:
    try:
        # 建立与milvus的连接
        connections.connect(
            "default", host="49.233.56.55", port="19530", timeout=10, db_name=db_name
        )
        # 使用或者切换到db_name
        db.using_database(db_name)
        if not utility.has_collection(collection_name):
            raise ValueError(f"集合{collection_name}不存在，无法执行操作")
        # 获取集合对象
        collection = Collection(collection_name)
        print(collection)
        # # 插入数据，分三列 id，title，title_vector
        # title_list = ["中国", "必胜"]
        # title_vector = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
        # title_vector = np.array(title_vector, dtype=np.float16)
        data = [
            {
                "tile_vector": np.array(
                    [random.random() for _ in range(3)], dtype=np.float16
                ),
                "title": "title" + str(random.random()),
            }
            for _ in range(1000)
        ]
        insert_result = collection.insert(data)
        print("插入数据库成功", insert_result.insert_count)

        # id_list = [462281535447972304]
        # title_list = ["伟大的中国必胜"]
        # title_vector = [
        #     [0.1, 0.2, 0.3],
        # ]
        # title_vector = np.array(title_vector, dtype=np.float16)
        # upsert_result = collection.upsert([id_list, title_list, title_vector])
        # print("Upsert 数据成功:", upsert_result.insert_count)

        # delete_result = collection.delete(
        #     expr="id in [462281535447972245,462281535447972246,462281535447972249]"
        # )
        # print("删除数据成功:", delete_result.delete_count)

    except Exception as exc:
        print("操作失败:", exc)
        raise
    finally:
        connections.disconnect("default")


if __name__ == "__main__":
    if not test_connection():
        raise SystemExit(1)

{
    "id": "100124410301",
    "confirm_rebate_record": {"actual_rebate_amount": "10", "comment": "我是备注"},
}