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_update_embedding(collection: pymilvus.Collection, data: Dict, flush: bool = True) -> List[int]:

    """
    Update entities by ids from a Milvus collection asynchronously.
    :param collection: Milvus Collection object
    :param data: Dictionary of fields to update, with keys as field names and values as new values
    :param flush: Whether to flush the collection after update
    :return: List of primary key ids that were updated
    """
    result = await asyncio.to_thread(
        collection.upsert,
        data=data
    )

    if flush:
        await asyncio.to_thread(collection.flush)

    # Prefer server-returned PKs; fallback to input pk
    try:
        print("update successful")
        return list(result.primary_keys)
    except Exception:
        return [data['id']]


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]