教程：在 asyncio 中使用 AsyncMilvusClient

AsyncMilvusClient是一种异步 MilvusClient，它提供了基于例程的 API，可通过asyncio 以非阻塞方式访问 Milvus。在本文中，您将了解调用 AsyncMilvusClient 提供的 API 的过程以及需要注意的方面。

概述

Asyncio 是一个使用async/await语法编写并发代码的库，也是 Milvus 高性能异步客户端的基础，它将适合您在 asyncio 上运行的代码库。

AsyncMilvusClient 提供的方法具有与 MilvusClient 相同的参数集和行为。唯一的区别在于调用它们的方式。下表列出了 AsyncMilvusClient 中可用的方法。

客户端
close()
Collection & Partition
create_collection()	drop_collection()	create_partition()
drop_partition()
索引
create_index()	drop_index()	load_collection()
release_collection()	load_partitions()	release_partitions()
向量
insert()	upsert()	delete()
search()	query()	hybrid_search()
get()

如果您仍然需要其他 MilvusClient 方法的异步版本，可以向pymilvus软件仓库提交功能请求。我们也欢迎你贡献代码。

创建事件循环

使用 asyncio 的应用程序通常将事件循环作为管理异步任务和 I/O 操作的协调器。在本教程中，我们将从 asyncio 中获取一个事件循环，并将其用作协调器。

import asyncio
import numpy as np
from scipy.sparse import csr_matrix
from pymilvus import MilvusClient, AsyncMilvusClient, DataType, RRFRanker, AnnSearchRequest

loop = asyncio.get_event_loop()

连接 AsyncMilvusClient

下面的示例演示了如何以异步方式连接 Milvus。

# Connect to Milvus server using AsyncMilvusClient
async_client = AsyncMilvusClient(
    uri="http://localhost:19530",
    token="root:Milvus"
)

创建 Schema

目前，AsyncMilvusClient 中没有create_schema() 。相反，我们将使用 MilvusClient 为 Collections 创建 Schema。

schema = async_client.create_schema(
    auto_id=False,
    description="This is a sample schema",
)

schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("dense_vector", DataType.FLOAT_VECTOR, dim=5)
schema.add_field("sparse_vector", DataType.SPARSE_FLOAT_VECTOR)
schema.add_field("text", DataType.VARCHAR, max_length=512)

创建 Collections

现在，我们将使用 Schema 创建一个 Collection。请注意，您需要在调用AsyncMilvusClient 方法时加上await 关键字，并将调用放在async 函数中，如下所示。

async def create_my_collection(collection_name, schema):
    if (client.has_collection(collection_name)):
        await async_client.drop_collection(collection_name)

    await async_client.create_collection(
        collection_name=collection_name,
        schema=schema
    )

    if (client.has_collection(collection_name)):
        print("Collection created successfully")
    else:
        print("Failed to create collection")
        
# Call the above function asynchronously 
loop.run_until_complete(create_my_collection("my_collection", schema))

# Output
#
# Collection created successfully

创建索引

您还需要为所有向量字段和可选标量字段创建索引。根据上文定义的 Schema，Collection 中有两个向量字段，您将为它们创建索引，如下所示。

async def create_indexes(collection_name):
    index_params = client.prepare_index_params()

    index_params.add_index(field_name="dense_vector", index_type="AUTOINDEX", metric_type="IP")
    index_params.add_index(field_name="sparse_vector", index_type="AUTOINDEX", metric_type="IP")
    index_params.add_index(field_name="text", index_type="AUTOINDEX")

    await async_client.create_index(collection_name, index_params)

# Call the above function asynchronously 
loop.run_until_complete(create_indexes("my_collection"))

加载 Collections

在为必要的字段建立索引后，就可以加载 Collections 了。下面的代码演示了如何异步加载 Collections。

async def load_my_collection(collection_name):
    await async_client.load_collection(collection_name)
    print(client.get_load_state(collection_name))
    
# Call the above function asynchronously 
loop.run_until_complete(load_my_collection("my_collection"))

# Output
#
# {'state': <LoadState: Loaded>}

插入数据

您可以使用 pymilvus 中提供的嵌入模型为文本生成向量嵌入。详情请参阅嵌入概述。在本节中，我们将向 Collections 插入随机生成的数据。

async def insert_sample_data(collection_name):
    # Randomly generated data will be used here
    rng = np.random.default_rng(42)

    def generate_random_text(length):
        seed = "this is a seed paragraph to generate random text, which is used for testing purposes. Specifically, a random text is generated by randomly selecting words from this sentence."
        words = seed.split()
        return " ".join(rng.choice(words, length))
    
    data = [{
        'id': i, 
        'dense_vector': rng.random(5).tolist(), 
        'sparse_vector': csr_matrix(rng.random(5)), 
        'text': generate_random_text(10)
    } for i in range(10000)]

    res = await async_client.insert(collection_name, data)

    print(res)

# Call the above function asynchronously 
loop.run_until_complete(insert_sample_data("my_collection"))

# Output
#
# {'insert_count': 10000, 'ids': [0, 1, 2, 3, ..., 9999]}

查询

加载并填充数据后，就可以在 Collection 中进行搜索和查询了。在本节中，您将在名为my_collection 的 Collections 中查找text 字段中以random 开头的实体数量。

async def query_my_collection(collection_name):
    # Find the number of entities with the `text` fields starting with the word "random" in the `my_collection` collection.

    res = await async_client.query(
        collection_name="my_collection",
        filter='text like "%random%"',
        output_fields=["count(*)"]
    )

    print(res) 
    
# Call the above function asynchronously   
loop.run_until_complete(query_my_collection("my_collection"))

# Output
#
# data: ["{'count(*)': 6802}"] 

搜索

在本节中，您将对目标 Collections 的密集和稀疏向量场进行向量搜索。

async def conduct_vector_search(collection_name, type, field):
    # Generate a set of three random query vectors
    query_vectors = []
    if type == "dense":
        query_vectors = [ rng.random(5) for _ in range(3) ]
    
    if type == "sparse":
        query_vectors = [ csr_matrix(rng.random(5)) for _ in range(3) ]

    print(query_vectors)

    res = await async_client.search(
        collection_name="my_collection",
        data=query_vectors,
        anns_field=field,
        output_fields=["text", field]
    )

    print(res)
    
# To search against the dense vector field asynchronously 
loop.run_until_complete(conduct_vector_search("my_collection", "dense", "dense_vector"))

# To search against the sparse vector field asynchronously 
loop.run_until_complete(conduct_vector_search("my_collection", "sparse", "sparse_vector"))

搜索输出应列出与指定查询向量相对应的三组结果。

混合搜索

混合搜索结合了多个搜索结果，并对其进行重新排序，以获得更好的召回效果。本节将使用密集向量场和稀疏向量场进行混合搜索。

async def conduct_hybrid_search(collection_name):
    req_dense = AnnSearchRequest(
        data=[ rng.random(5) for _ in range(3) ],
        anns_field="dense_vector",
        param={"metric_type": "IP"},
        limit=10
    )

    req_sparse = AnnSearchRequest(
        data=[ csr_matrix(rng.random(5)) for _ in range(3) ],
        anns_field="sparse_vector",
        param={"metric_type": "IP"},
        limit=10
    )

    reqs = [req_dense, req_sparse]

    ranker = RRFRanker()

    res = await async_client.hybrid_search(
        collection_name="my_collection",
        reqs=reqs,
        ranker=ranker,
        output_fields=["text", "dense_vector", "sparse_vector"]
    )

    print(res)
    
# Call the above function asynchronously  
loop.run_until_complete(conduct_hybrid_search("my_collection"))