chat-python/rag/vector_db/milvus_vector.py

import time
import numpy as np
from pymilvus import (
    MilvusClient,
    DataType,
    Function,
    FunctionType,
    AnnSearchRequest,
    RRFRanker,
)
# from pymilvus.model.hybrid import BGEM3EmbeddingFunction
from pymilvus import model
from rag.load_model import sentence_transformer_ef
from utils.get_logger import setup_logger
import torch
#device = "cpu" if torch.cuda.is_available() else "cuda"
device = "cuda" if torch.cuda.is_available() else "cpu"
logger = setup_logger(__name__)

# embedding_path = r"G:/work/code/models/multilingual-e5-large-instruct/"
# sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(model_name=embedding_path,device=device)


embedding_mapping = {
    "e5": sentence_transformer_ef,
    "multilingual-e5-large-instruct": sentence_transformer_ef,
}

class HybridRetriever:
    def __init__(self, uri, embedding_name:str="e5", collection_name:str ="hybrid"):
        self.uri = uri
        self.collection_name = collection_name
        # self.embedding_function = sentence_transformer_ef
        self.embedding_function = embedding_mapping.get(embedding_name, "e5")
        self.use_reranker = True
        self.use_sparse = True
        self.client = MilvusClient(uri=uri)
    
    def has_collection(self):
        try:
            collection_flag = self.client.has_collection(self.collection_name)
            logger.info(f"查询向量库的结果:{collection_flag}")
        except Exception as e:
            logger.info(f"查询向量库是否存在出错：{e}")
            collection_flag = False
        return collection_flag

    def build_collection(self):
        if isinstance(self.embedding_function.dim, dict):
            dense_dim = self.embedding_function.dim["dense"]
        else:
            dense_dim = self.embedding_function.dim
        logger.info(f"创建数据库的向量维度：{dense_dim}")
        analyzer_params={
            "type": "chinese"
        }

        schema = MilvusClient.create_schema()
        schema.add_field(
            field_name="pk",
            datatype=DataType.VARCHAR,
            is_primary=True,
            auto_id=True,
            max_length=100,
        )
        schema.add_field(
            field_name="content",
            datatype=DataType.VARCHAR,
            max_length=65535,
            analyzer_params=analyzer_params,
            enable_match=True,
            enable_analyzer=True,
        )
        schema.add_field(
            field_name="sparse_vector", datatype=DataType.SPARSE_FLOAT_VECTOR
        )
        schema.add_field(
            field_name="dense_vector", datatype=DataType.FLOAT_VECTOR, dim=dense_dim
        )
        schema.add_field(field_name="doc_id", datatype=DataType.VARCHAR, max_length=64)
        schema.add_field(
            field_name="chunk_id", datatype=DataType.VARCHAR, max_length=64
        )
        schema.add_field(field_name="metadata", datatype=DataType.JSON)

        functions = Function(
            name="bm25",
            function_type=FunctionType.BM25,
            input_field_names=["content"],
            output_field_names="sparse_vector",
        )

        schema.add_function(functions)

        index_params = MilvusClient.prepare_index_params()
        index_params.add_index(
            field_name="sparse_vector",
            index_type="SPARSE_INVERTED_INDEX",
            metric_type="BM25",
        )
        index_params.add_index(
            field_name="dense_vector", index_type="FLAT", metric_type="IP"
        )
        try:
            self.client.create_collection(
                collection_name=self.collection_name,
                schema=schema,
                index_params=index_params,
            )
            return "create_collection_success"
        except Exception as e:
            logger.error(f"创建{self.collection_name}数据库失败：{e}")
            return "create_collection_error"

    def insert_data(self, chunk, metadata):
        logger.info("准备插入数据")
        with torch.no_grad():
            embedding = self.embedding_function([chunk])
        logger.info("获取文本的向量信息。")
        if isinstance(embedding, dict) and "dense" in embedding:
            # bge embedding 获取embedding的方式
            dense_vec = embedding["dense"][0]
        else:
            dense_vec = embedding[0]
        
        try:
            self.client.insert(
                self.collection_name, {"dense_vector": dense_vec, **metadata}
            )
            logger.info("插入一条数据成功。")
            return True, "success"
        except Exception as e:
            doc_id = metadata.get("doc_id")
            logger.error(f"处理文档：{doc_id}，插入数据出错：{e}")
            self.delete_by_doc_id(doc_id=doc_id)
            return False, str(e)
        
    def batch_insert_data(self, chunks, metadatas):
        logger.info("准备插入数据")
        embedding_lists = self.embedding_function.encode_documents(chunks)
        logger.info("获取文本的向量信息。")
        record_lists = []
        for embedding, metadata in zip(embedding_lists, metadatas):
            if isinstance(embedding, dict) and "dense" in embedding:
                # bge embedding 获取embedding的方式
                dense_vec = embedding["dense"][0]
            else:
                dense_vec = embedding.tolist()

            # if hasattr(dense_vec, 'tolist'):
            #     dense_vec = dense_vec.tolist()

            # logger.info(f"向量维度：{dense_vec}")
            
            # if isinstance(dense_vec, (float, int)):
            #     dense_vec = [dense_vec]
            
            # if isinstance(dense_vec, np.float32):
            #     dense_vec = [float(dense_vec)]
            
            record = {"dense_vector": dense_vec}
            record.update(metadata)
            record_lists.append(record)
        
        try:
            self.client.insert(
                self.collection_name, record_lists
            )
            logger.info("插入数据成功。")
            return True, "success"
        except Exception as e:
            doc_id = metadata.get("doc_id")
            logger.error(f"处理文档：{doc_id}，插入数据出错：{e}")
            self.delete_by_doc_id(doc_id=doc_id)
            return False, str(e)
                

    def search(self, query: str, k: int = 20, mode="hybrid"):

        output_fields = [
            "content",
            "doc_id",
            "chunk_id",
            "metadata",
        ]
        if mode in ["dense", "hybrid"]:
            with torch.no_grad():
                embedding = self.embedding_function([query])
            if isinstance(embedding, dict) and "dense" in embedding:
                dense_vec = embedding["dense"][0]
            else:
                dense_vec = embedding[0]

        if mode == "sparse":
            results = self.client.search(
                collection_name=self.collection_name,
                data=[query],
                anns_field="sparse_vector",
                limit=k,
                output_fields=output_fields,
            )
        elif mode == "dense":
            results = self.client.search(
                collection_name=self.collection_name,
                data=[dense_vec],
                anns_field="dense_vector",
                limit=k,
                output_fields=output_fields,
            )
        elif mode == "hybrid":
            full_text_search_params = {"metric_type": "BM25"}
            full_text_search_req = AnnSearchRequest(
                [query], "sparse_vector", full_text_search_params, limit=k
            )

            dense_search_params = {"metric_type": "IP"}
            dense_req = AnnSearchRequest(
                [dense_vec], "dense_vector", dense_search_params, limit=k
            )

            results = self.client.hybrid_search(
                self.collection_name,
                [full_text_search_req, dense_req],
                ranker=RRFRanker(),
                limit=k,
                output_fields=output_fields,
            )
        else:
            raise ValueError("Invalid mode")
        return [
            {
                "doc_id": doc["entity"]["doc_id"],
                "chunk_id": doc["entity"]["chunk_id"],
                "content": doc["entity"]["content"],
                "metadata": doc["entity"]["metadata"],
                "score": doc["distance"],
            }
            for doc in results[0]
        ]
    
    def query_filter(self, doc_id, filter_field):
        # doc id 文档id，content中包含 filter_field 字段的
        query_output_field = [
            "content",
            "chunk_id",
            "doc_id",
            "metadata"
        ]
        # query_expr = f"doc_id in {doc_id} && content like '%{filter_field}%'"
        # 根据doc_id查询如果有关键词，根据关键词查询，如果没有关键词，只根据doc_id查询
        if filter_field:
            query_expr = f"doc_id == '{doc_id}' && content like '%{filter_field}%'"
        else:
            query_expr = f"doc_id == '{doc_id}'"

        try:
            query_filter_results = self.client.query(collection_name=self.collection_name, filter=query_expr, output_fields=query_output_field)
        except Exception as e:
            logger.error(f"根据关键词查询数据失败：{e}")
            query_filter_results = [{"code": 500}]
        return query_filter_results
        # for result in query_filter_results:
        #     print(f"根据doc id 和 field 过滤结果: {result}\n\n")
    
    def query_chunk_id(self, chunk_id):
        # chunk id，查询切片
        query_output_field = [
            "content",
            "doc_id",
            "chunk_id",
            # "metadata"
        ]
        query_expr = f"chunk_id == '{chunk_id}'"
        try:
            query_filter_results = self.client.query(collection_name=self.collection_name, filter=query_expr, output_fields=query_output_field)
        except Exception as e:
            logger.info(f"根据chunk id 查询出错：{e}")
            query_filter_results = [{"code": 500}]
        return query_filter_results
    
    def query_chunk_id_list(self, chunk_id_list):
        # chunk id，查询切片
        query_output_field = [
            "content",
            "doc_id",
            "chunk_id",
            # "metadata"
        ]
        query_expr = f"chunk_id in {chunk_id_list}"
        try:
            query_filter_results = self.client.query(collection_name=self.collection_name, filter=query_expr, output_fields=query_output_field)
        except Exception as e:
            logger.info(f"根据chunk id 查询出错：{e}")
            query_filter_results = [{"code": 500}]
        return query_filter_results

    def update_data(self, chunk_id, chunk):
        # 根据chunk id查询对应的信息，
        chunk_expr = f"chunk_id == '{chunk_id}'"
        chunk_output_fields = [
            "pk",
            "doc_id",
            "chunk_id",
            "metadata"
        ]
        try:
            chunk_results = self.client.query(collection_name=self.collection_name, filter=chunk_expr, output_fields=chunk_output_fields)
            # logger.info(f"{chunk_id}更新切片的信息：{chunk_results}")
        except Exception as e:
            logger.error(f"更新切片数据时查询失败：{e}")
            return "update_query_error", ""
        if not chunk_results:
            logger.info(f"根据{chunk_id}未在向量库中查询到对应数据，无法更新数据")
            return "update_query_no_result", ""
        with torch.no_grad():
            embedding = self.embedding_function([chunk])
        if isinstance(embedding, dict) and "dense" in embedding:
            # bge embedding 获取embedding的方式
            dense_vec = embedding["dense"][0]
        else:
            dense_vec = embedding[0]
        chunk_dict = chunk_results[0]
        metadata = chunk_dict.get("metadata")
        old_chunk_len = metadata.get("chunk_len")
        chunk_len = len(chunk)
        metadata["chunk_len"] = chunk_len
        chunk_dict["content"] = chunk
        chunk_dict["dense_vector"] = dense_vec
        chunk_dict["metadata"] = metadata
        try:
            update_res = self.client.upsert(collection_name=self.collection_name, data=[chunk_dict])
            logger.info(f"更新返回的数据：{update_res}")
            return "update_success", chunk_len - old_chunk_len
        except Exception as e:
            logger.error(f"更新数据时出错：{e}")
            return "update_error", ""
            
    def update_data(self, chunk_id, chunk):
        # 根据chunk id查询对应的信息，
        chunk_expr = f"chunk_id == '{chunk_id}'"
        chunk_output_fields = [
            "pk",
            "doc_id",
            "chunk_id",
            "metadata"
        ]
        try:
            chunk_results = self.client.query(collection_name=self.collection_name, filter=chunk_expr, output_fields=chunk_output_fields)
            # logger.info(f"{chunk_id}更新切片的信息：{chunk_results}")
        except Exception as e:
            logger.error(f"更新切片数据时查询失败：{e}")
            return "update_query_error", ""
        if not chunk_results:
            logger.info(f"根据{chunk_id}未查询到对应数据，无法更新数据")
            return "update_query_no_result", ""
        with torch.no_grad():
            embedding = self.embedding_function([chunk])
        if isinstance(embedding, dict) and "dense" in embedding:
            # bge embedding 获取embedding的方式
            dense_vec = embedding["dense"][0]
        else:
            dense_vec = embedding[0]
        chunk_dict = chunk_results[0]
        metadata = chunk_dict.get("metadata")
        old_chunk_len = metadata.get("chunk_len")
        chunk_len = len(chunk)
        metadata["chunk_len"] = chunk_len
        chunk_dict["content"] = chunk
        chunk_dict["dense_vector"] = dense_vec
        chunk_dict["metadata"] = metadata
        try:
            update_res = self.client.upsert(collection_name=self.collection_name, data=[chunk_dict])
            logger.info(f"更新返回的数据：{update_res}")
            return "update_success", chunk_len - old_chunk_len
        except Exception as e:
            logger.error(f"更新数据时出错：{e}")
            return "update_error", ""
        
    def delete_collection(self, collection):
        try:
            self.client.drop_collection(collection_name=collection)
            return "delete_collection_success"
        except Exception as e:
            logger.error(f"删除{collection}失败，出错原因：{e}")
            return "delete_collection_error"

    def delete_by_chunk_id(self, chunk_id:str = None):
        # 根据文档id查询主键值 milvus只支持主键删除
        expr = f"chunk_id == '{chunk_id}'"
        try:
            results = self.client.query(collection_name=self.collection_name, filter=expr, output_fields=["pk","metadata"])  # 获取主键 id
            logger.info(f"根据切片id：{chunk_id}，查询的数据：{results}")
        except Exception as e:
            logger.error(f"根据切片id查询主键失败：{e}")
            return "delete_query_error", []
        if not results:
            print(f"No data found for chunk id: {chunk_id}")
            # return "delete_no_result", []
            return "delete_success", [0]
        
        # 提取主键值
        primary_keys = [result["pk"] for result in results]
        chunk_len = [result["metadata"]["chunk_len"] for result in results]
        logger.info(f"获取到的主键信息：{primary_keys}")
        
        # 执行删除操作
        expr_delete = f"pk in {primary_keys}"  # 构造删除表达式
        try:
            delete_res = self.client.delete(collection_name=self.collection_name, filter=expr_delete)
            self.client.flush(collection_name=self.collection_name)
            self.client.compact(collection_name=self.collection_name)
            logger.info(f"Deleted data with chunk_id: {delete_res}")
            return "delete_success", chunk_len
        except Exception as e:
            logger.error(f"删除数据失败：{e}")
            return "delete_error", []
        
    def delete_by_doc_id(self, doc_id:str =None):
        # 根据文档id查询主键值 milvus只支持主键删除
        expr = f"doc_id == '{doc_id}'"
        try:
            results = self.client.query(collection_name=self.collection_name, filter=expr, output_fields=["pk"])  # 获取主键 id
        except Exception as e:
            logger.error(f"根据切片id查询主键失败：{e}")
            return "delete_query_error"
        if not results:
            print(f"No data found for doc_id: {doc_id}")
            return "delete_no_result"
        
        # 提取主键值
        primary_keys = [result["pk"] for result in results]
        logger.info(f"获取到的主键信息：{primary_keys}")
        
        # 执行删除操作
        try:
            delete_res = self.client.delete(collection_name=self.collection_name, ids=primary_keys)
            self.client.flush(collection_name=self.collection_name)
            self.client.compact(collection_name=self.collection_name)
            logger.info(f"Deleted data with doc_id: {delete_res}")
            return "delete_success"
        except Exception as e:
            logger.error(f"删除数据失败：{e}")
            return "delete_error"
        # # 执行删除操作
        # expr_delete = f"pk in {primary_keys}"  # 构造删除表达式
        # try:
        #     delete_res = self.client.delete(collection_name=self.collection_name, filter=expr_delete)
        #     logger.info(f"Deleted data with doc_id: {delete_res}")
        #     return "delete_success"
        # except Exception as e:
        #     logger.error(f"删除数据失败：{e}")
        #     return "delete_error"
        
    
# 测试
# def parse_json_to_schema_data_format():
#     sql_text_list = load_sql_query_ddl_info()
#     docs = []
#     for sql_info in sql_text_list:
#         sql_text = sql_info.get("sql_text")
#         source = ",".join(sql_info.get("table_list", []))
#         source = sql_info.get("source") if not source else source
#         sql = sql_info.get("sql")
#         ddl = sql_info.get("table_ddl")
#         metadata = {"source": source, "sql": sql, "ddl": ddl}
#         text_list = sql_info.get("sim_sql_text_list", [])
#         text_list.append(sql_text)
#         doc_id = str(uuid4())
#         for text in text_list:
#             chunk_id = str(uuid4())
#             insert_dict = {
#                 "content": text,
#                 "doc_id": doc_id,
#                 "chunk_id": chunk_id,
#                 "metadata": metadata
#             }
#             docs.append(insert_dict)
    
#     return docs

# def insert_data_to_milvus(standard_retriever):
#     sql_dataset = parse_json_to_schema_data_format()
#     standard_retriever.build_collection()
#     for sql_dict in sql_dataset:
#         text = sql_dict["content"]
#         standard_retriever.insert_data(text, sql_dict)

def main():
    # dense_ef = BGEM3EmbeddingFunction()
    embedding_path = r"G:/work/code/models/multilingual-e5-large-instruct/"
    sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(model_name=embedding_path,device=device)
    standard_retriever = HybridRetriever(
        uri="http://localhost:19530",
        collection_name="milvus_hybrid",
        dense_embedding_function=sentence_transformer_ef,
    )
    # 插入数据
    # insert_data_to_milvus(standard_retriever)

    # 查询 混合检索：hybrid ，稀疏检索sparse，向量检索：dense
    # results = standard_retriever.search("查一下加班情况", mode="hybrid", k=3)
    # model=sparse 稀疏检索
    # print(f"稀疏检索结果：{results}")

    # model=dense 向量检索
    # print(f"向量检索结果：{results}")

    # model=hybrid
    # print(f"向量检索结果：{results}")

    # 根据doc id删除数据
    # delete_start_time = time.time()
    # doc_id = "e72ebf78-6c0d-410b-8fbb-9a2057673064"
    # standard_retriever.delete_data(doc_id=doc_id)
    # delete_end_time = time.time()
    # print(f"删除耗时：{delete_end_time-delete_start_time}")

    # 根据chunk_id 更新数据
    # update_start_time = time.time()
    # chunk = "查询一下员工加班的情况"
    # chunk_id = "a5e8dded-f5a7-4a1f-92cd-82fa8113b418"
    # standard_retriever.update_data(chunk_id, chunk)
    # update_end_time = time.time()
    # print(f"更新数据的时间：{update_end_time-update_start_time}")

    # 根据doc id 和关键字查询
    query_start_time = time.time()
    filter_field = "加班"
    doc_id = ["7b73ae0b-db97-4315-ba71-783fe7a69c61", "96bbe5a8-5fcf-4769-8343-938acb8735bd"]
    standard_retriever.query_filter(doc_id, filter_field)
    query_end_time = time.time()
    print(f"关键字搜索数据的时间：{query_end_time-query_start_time}")


if __name__=="__main__":
    main()