file_rag/milvus_process.py

from pymilvus import  DataType, FieldSchema, CollectionSchema, AnnSearchRequest, WeightedRanker, Function, FunctionType
from config import milvus_url, logger_search
from embedding_search import model_embedding_bge, bge_rerank
import time, asyncio
import traceback
from pymilvus import MilvusClient, DataType, RRFRanker
from util import encode_to_base64
from typing_extensions import Annotated

async def create_collection(collection_name):
    client = MilvusClient(uri=milvus_url)
    if client.has_collection(collection_name=collection_name):
        # client.drop_collection(collection_name=collection_name)
        # print('ok')
        return False
    fields = [
            FieldSchema(name="client_id", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="file_name", dtype=DataType.VARCHAR, max_length=65535),
            FieldSchema(name="filename_embedding", dtype=DataType.FLOAT_VECTOR, dim=1024),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535, is_primary=True, enable_analyzer=True),
            FieldSchema(name="content_embedding", dtype=DataType.FLOAT_VECTOR, dim=1024),
            FieldSchema(name="content_sparse", dtype=DataType.SPARSE_FLOAT_VECTOR)
        ]
    bm25_function = Function(
        name="text_bm25_emb", 
        input_field_names=["content"],
        output_field_names=["content_sparse"], 
        function_type=FunctionType.BM25,
    )
    schema = CollectionSchema(fields=fields, enable_dynamic_field=True, auto_id=False)
    schema.add_function(bm25_function)
    
    index_params = client.prepare_index_params()
    index_params.add_index(
        field_name="client_id",
        index_type="INVERTED"
    )
    index_params.add_index(
        field_name="file_name",
        index_type="INVERTED"
    )
    index_params.add_index(
        field_name="content_embedding", 
        index_type="IVF_FLAT",
        metric_type="COSINE",
        params={ "nlist": 128 }
    )
    index_params.add_index(
        field_name="filename_embedding", 
        index_type="IVF_FLAT",
        metric_type="COSINE",
        params={ "nlist": 128 }
    )
    index_params.add_index(
        field_name="content_sparse",
        index_type="AUTOINDEX", 
        metric_type="BM25"
    )
    client.create_collection(
    collection_name=collection_name, 
    schema=schema, 
    index_params=index_params
)
    return True

async def update_mulvus_file(client_id, file_name, chunks, collection_name='gloria_files'):
    try:
        start_time = time.time()
        # 从图中拿数据
        client = MilvusClient(uri=milvus_url)
        await create_collection(collection_name=collection_name)
        client.load_collection(collection_name=collection_name,)    
        try:
            
            res = client.delete(collection_name=collection_name, filter=f'file_name == "{file_name}"')
        except Exception as e:
            
            print(e)

        final_data = []
        filename_embedding = await asyncio.to_thread(model_embedding_bge, file_name)
        for chunk in chunks:
            chunk_embedding = await asyncio.to_thread(model_embedding_bge, chunk)
            final_data.append({
                'client_id': client_id, 
                'file_name': file_name, 
                "filename_embedding":filename_embedding,
                'content': chunk, 
                'content_embedding': chunk_embedding, 
            })

        insert_dict = await asyncio.to_thread(client.insert, collection_name=collection_name, data=final_data)
        
        status = {"result": "succeed"}
        end_time = time.time()
        logger_search.info(f"{file_name} 嵌入milvus单次耗时: {end_time - start_time} s")
        client.release_collection(collection_name=collection_name)
        client.close()
        return status

    except Exception as e:
        logger_search.error(f"处理 {client_id} 的文件 {file_name} 嵌入milvus报错如下:{e}")
        traceback.print_exc()
        status = {"result": "failed"}
        return status


async def get_filedata_from_milvus(query, collection_name='gloria_files', file_name=None, thread=0.5):
    try:
        client = MilvusClient(uri=milvus_url)
        query_embedding = await asyncio.to_thread(model_embedding_bge, query)
        time0 = time.time()
        client.load_collection(collection_name=collection_name)
        time1 = time.time()
        logger_search.info(f"load 时间: {time1 - time0} s")
        
        search_filter = f'file_name == "{file_name}"' if file_name else ''

        # 创建针对不同字段的ANN搜索请求
        search_param_1 = {
            "data": [query_embedding],          # 确保输入为二维向量
            "anns_field": "filename_embedding",# 文件名向量字段
            "param": {
                "metric_type": "COSINE",
                "params": {"nprobe": 10}
            },
            "limit": 300,
        }
        request_1 = AnnSearchRequest(**search_param_1)

        search_param_2 = {
            "data": [query_embedding],
            "anns_field": "content_embedding",# 段落内容向量字段
            "param": {
                "metric_type": "COSINE",
                "params": {"nprobe": 10}
            },
            "limit": 300,
            
        }
        request_2 = AnnSearchRequest(**search_param_2)

        
        req_temp = AnnSearchRequest(
                data=[query],
                anns_field='content_sparse',
                param={
                    "metric_type": 'BM25',
                    "params": {'drop_ratio_search': 0.2}
                },
                limit=50,
            )
        # 执行混合搜索（文件名权重0.4，内容权重0.6）
        rerank = WeightedRanker(0.2, 0.5, 0.5)
        res = await asyncio.to_thread(
            client.hybrid_search,
            collection_name=collection_name,
            reqs=[request_1, request_2, req_temp],
            ranker=rerank,
            output_fields=['content', 'file_name'],
            filter = search_filter
        )

        client.release_collection(collection_name=collection_name)
        client.close()
        final_answer = []
        for hits in res:
            for hit in hits:
                if hit.get('distance', 0) > thread:
                    entity = hit.get('entity')
                    final_answer.append({'file_name': entity.get("file_name"), 'content': entity.get("content")})
    except Exception as e:
        logger_search.error(e)
        final_answer = []
    logger_search.info(f'从milvus文件知识库搜索到{len(final_answer)}条信息')
    return final_answer

async def rerank_file_data(query, data):
    lookup = {}  # 使用单个字典存储映射关系
    chunks = []
    final_chunks = {}
    for item in data:
        content = item.get('content', '')
        # 直接存储原始数据项的引用，避免后续查找
        lookup[content] = item
        chunks.append(content)
    # 执行重排序
    logger_search.info(chunks)
    reranked = bge_rerank(query=query, indicates=chunks, socre=0.3)

    
    send_data = [lookup[chunk] for chunk in reranked if chunk in lookup]
    
    reranked_content = [i.get('content') for i in send_data]
    

    send_data = [dict(t) for t in {tuple(d.items()) for d in send_data}]
    logger_search.info(send_data)
    for sed in send_data:
        if not sed.get('file_name') in final_chunks:
            final_chunks[sed.get('file_name')] = sed.get('content')
        final_chunks[sed.get('file_name')] += sed.get('content')
    logger_search.info(final_chunks)
    search_results_str = "\n".join([f"[{i+1}]: \n {k}: \n{s} \n" for i, (k,s) in enumerate(final_chunks.items())])

    return final_chunks, search_results_str

async def get_search_results(query: Annotated[str, "需要搜索的问题"], summary=True):
    data = await get_filedata_from_milvus(query=query)
    final_data, search_str = await rerank_file_data(query=query,data=data)
    if not final_data:
        final_chunks = {}
        for sed in data:
            if not sed.get('file_name') in final_chunks:
                final_chunks[sed.get('file_name')] = sed.get('content')
            final_chunks[sed.get('file_name')] += sed.get('content')
        search_str = "\n".join([f"[{i+1}]: \n {k}: \n{s} \n" for i, (k,s) in enumerate(final_chunks.items())])
        final_data = final_chunks

    logger_search.info(search_str)
    return final_data, search_str


if __name__ == '__main__':
    from agent import get_content_summary
    final_data, result = asyncio.run(get_search_results(query='每月盘点时间有要求吗?'))
    data, res = asyncio.run(get_content_summary(question='每月盘点时间有要求吗?',res_info=result, final_data=final_data))
    logger_search.info(data)
    logger_search.info(res)
    # x = asyncio.run(create_collection(collection_name='gloria_files'))
    # print(x)
    pass