Milvus 集合客户端

返回《Milvus 入门与部署》

适用于 Python 3.13，依赖为 pymilvus==3.0.2。MilvusCollectionClient 是示例定义的类，内部使用官方 MilvusClient。一个实例绑定一个集合，集中保存向量字段、度量、搜索参数、返回字段和超时。

接口	作用	调用时机
`create_collection()`	按传入的 Schema 创建集合、建立索引并加载	初始化新集合
`load()`	加载已有集合	连接已有数据时
`insert()` / `upsert()`	分批插入或更新，汇总写入数	文档入库、记录更新
`search()`	单向量搜索、标量过滤、可选阈值筛选	相似度检索
`query()` / `delete()`	按条件查询、按主键删除	记录管理
`close()`	关闭客户端	所有操作完成后

常驻服务在每个工作进程启动时创建客户端并加载集合，请求间复用。停止接收请求、进行中的操作全部完成后关闭客户端。

milvus_client.py 末尾的调用示例用 try/finally 在操作结束后关闭客户端；常驻服务不需要在每次请求中重新连接或加载集合。

集合的 Schema 和文档向量由调用方提供。 类中没有模型下载、业务字段或自动建集合的查询逻辑。search_params 填索引的搜索参数，例如 IVF 的 {"nprobe": 10}；metric_type 必须与已有索引一致。

search() 接收一条向量，返回命中列表。每项保留原有主键、distance 和 entity，本文主键字段名为 id。默认不筛分数；传入 threshold 后，COSINE、IP 保留大于等于阈值的结果，L2 保留小于等于阈值的结果。

分数筛选发生在 Top-K 返回后，结果可能少于 limit。L2 的阈值对应平方距离；IP、COSINE 的阈值对应各自的相似度。阈值应结合所选度量和实际数据确定。

分批写入不是一个整体事务，每批是一次独立请求。后批失败不会撤销已成功的批次；此封装不额外重试，最终异常交给调用方。普通 insert() 不负责主键去重。

connect_timeout 控制连接阶段；timeout 统一传给加载、查询、写入等操作，每批分别计时，不是整个批量任务的总时限。创建集合包含多个操作，也不保证在这个时间内全部完成。

将代码作为 milvus_client.py 放在调用脚本同目录或 Python 可导入路径中。文件末尾提供独立的写入、检索、更新和删除示例，使用手工构造的四维向量，无需下载模型：

代码块 Shell

1
2
export MILVUS_URI="http://localhost:19530"
export MILVUS_COLLECTION="client_example"

开启认证时设置 MILVUS_TOKEN="用户名:密码"；MILVUS_DATABASE 默认是已存在的 default 数据库。这些环境变量名是示例约定。

直接运行会创建 client_example 集合并写入三条记录，最后删除其中一条。 示例不会删除整个集合；重复运行前换一个未使用的集合名。已有集合的查询方式见正文第 8 节。

milvus_client.py Python

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
"""绑定一个集合的 Milvus 客户端，适用于 Python 3.13、pymilvus 3.0.2。

安装依赖：python -m pip install "pymilvus==3.0.2"
直接运行需设置 MILVUS_URI 和 MILVUS_COLLECTION，后者必须是一个尚未使用的集合名。
演示会创建集合并写入三条记录，完成检索、更新和删除后保留集合，不自动删除整个集合。
"""

import os
from itertools import batched

from pymilvus import CollectionSchema, DataType, MilvusClient


class MilvusCollectionClient:
    """固定集合和检索配置；模型编码、Schema 和业务字段由调用方提供。"""

    def __init__(self, uri: str, collection_name: str, vector_field="vector", metric_type="COSINE",
                 search_params=None, output_fields=None, token="", db_name="default", connect_timeout=10, timeout=30):
        if metric_type not in {"COSINE", "IP", "L2"}:
            raise ValueError("此示例仅支持 COSINE、IP、L2 三种浮点向量度量")
        self.collection_name = collection_name
        self.vector_field = vector_field
        self.metric_type = metric_type
        self.search_params = {} if search_params is None else dict(search_params)
        self.output_fields = [] if output_fields is None else list(output_fields)
        self.timeout = timeout
        self.client = MilvusClient(uri=uri, token=token, db_name=db_name, timeout=connect_timeout)

    def create_collection(self, schema: CollectionSchema, index_type="AUTOINDEX", index_params=None):
        """初始化新集合、建立向量索引并加载；已有集合使用 load()。"""
        if self.client.has_collection(collection_name=self.collection_name, timeout=self.timeout):
            raise ValueError(f"集合 {self.collection_name!r} 已存在，请使用 load() 或更换集合名")
        indexes = self.client.prepare_index_params()
        indexes.add_index(field_name=self.vector_field, index_type=index_type, metric_type=self.metric_type,
                          params={} if index_params is None else index_params)
        self.client.create_collection(collection_name=self.collection_name, schema=schema, index_params=indexes,
                                      consistency_level="Strong", timeout=self.timeout)

    def load(self):
        """加载已有集合，在服务启动或首次查询前调用。"""
        self.client.load_collection(collection_name=self.collection_name, timeout=self.timeout)

    def _write_batches(self, records, batch_size: int, upsert: bool) -> int:
        method = self.client.upsert if upsert else self.client.insert
        count_key = "upsert_count" if upsert else "insert_count"
        count = 0
        for batch in batched(records, batch_size):
            result = method(collection_name=self.collection_name, data=list(batch), timeout=self.timeout)
            count += result[count_key]
        return count

    def insert(self, records, batch_size=500) -> int:
        """分批插入并返回成功写入数；此封装不额外重试、去重或逐批 flush。"""
        return self._write_batches(records, batch_size, upsert=False)

    def upsert(self, records, batch_size=500) -> int:
        """按主键分批更新或插入，records 中每条记录包含完整字段。"""
        return self._write_batches(records, batch_size, upsert=True)

    def search(self, vector: list[float], limit=10, filter="", threshold=None) -> list[dict]:
        """搜索一条向量；L2 保留不高于阈值的结果，IP/COSINE 保留不低于阈值的结果。"""
        results = self.client.search(
            collection_name=self.collection_name, data=[vector], anns_field=self.vector_field,
            search_params={"metric_type": self.metric_type, "params": self.search_params},
            filter=filter, limit=limit, output_fields=self.output_fields, timeout=self.timeout,
        )
        hits = results[0]  # 单条查询去掉最外层，保留主键、distance、entity。
        if threshold is None:
            return hits
        if self.metric_type == "L2":
            return [hit for hit in hits if hit["distance"] <= threshold]
        return [hit for hit in hits if hit["distance"] >= threshold]

    def query(self, filter: str, limit=100) -> list[dict]:
        """按标量字段查询，使用当前集合、返回字段和操作超时。"""
        return self.client.query(collection_name=self.collection_name, filter=filter, limit=limit,
                                 output_fields=self.output_fields, timeout=self.timeout)

    def delete(self, ids: list[int | str]):
        """按主键删除，原样返回 SDK 的删除结果。"""
        return self.client.delete(collection_name=self.collection_name, ids=ids, timeout=self.timeout)

    def close(self):
        """所有使用者完成操作后关闭客户端。"""
        self.client.close()


def main():
    store = MilvusCollectionClient(
        uri=os.environ["MILVUS_URI"], collection_name=os.environ["MILVUS_COLLECTION"],
        token=os.environ.get("MILVUS_TOKEN", ""), db_name=os.environ.get("MILVUS_DATABASE", "default"),
        output_fields=["text", "category"],
    )
    try:
        schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
        schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
        schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=4)
        schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=256)
        schema.add_field(field_name="category", datatype=DataType.VARCHAR, max_length=64)
        store.create_collection(schema, index_type="FLAT")

        # 手工向量仅演示接口调用，不表示这些文本的语义，也不需要下载模型。
        records = [
            {"id": 1, "vector": [1.0, 0.0, 0.0, 0.0], "text": "向量检索", "category": "检索"},
            {"id": 2, "vector": [0.8, 0.6, 0.0, 0.0], "text": "字段过滤", "category": "检索"},
            {"id": 3, "vector": [0.0, 0.0, 1.0, 0.0], "text": "日志采集", "category": "运维"},
        ]
        print("插入数：", store.insert(records, batch_size=2))
        hits = store.search([1.0, 0.0, 0.0, 0.0], limit=3, filter='category == "检索"', threshold=0.75)
        for hit in hits:
            print(hit["id"], round(hit["distance"], 4), hit["entity"])

        print("更新数：", store.upsert([{**records[0], "category": "示例"}]))
        print("查询结果：", store.query("id == 1"))
        print("删除结果：", store.delete([3]))
        print("删除后查询：", store.query("id == 3"))
    finally:
        store.close()


if __name__ == "__main__":
    main()