rag_server/applications/utils/milvus/functions.py

import asyncio
from typing import Dict, List

import pymilvus


async def async_insert_chunk(collection: pymilvus.Collection, data: Dict) -> List[int]:
    """
    :param collection:
    :param data: insert data
    :return:
    """
    result = await asyncio.to_thread(collection.insert, [data])
    return result.primary_keys


async def async_delete_chunk(
    collection: pymilvus.Collection, ids: List[int]
) -> List[int]:
    """
    Delete entities by ids from a Milvus collection asynchronously.

    :param collection: Milvus Collection object
    :param ids: List of primary key ids to delete
    :return: List of successfully deleted ids
    """
    if not ids:
        return []
    expr = f"id in {ids}"
    result = await asyncio.to_thread(collection.delete, expr)
    await asyncio.to_thread(collection.flush)
    success_count = result.succ_count
    if success_count == len(ids):
        return ids
    else:
        return ids[:success_count]