ZhouYI/03_RAG优化/agent/rag/multi_retriever.py

"""BM25、Milvus 向量和混合检索。"""

from langchain_classic.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_core.documents import Document
from langchain_milvus import Milvus


class MultiRetrieverSystem:
    """多路召回检索系统，默认使用混合检索。"""

    def __init__(self, documents, embedding_model, llm):
        self.documents = documents
        self.embedding_model = embedding_model
        self.llm = llm
        self.setup_retrievers()

    def setup_retrievers(self):
        """初始化 BM25、向量和混合检索器。"""
        self.bm25 = BM25Retriever.from_documents(self.documents)
        self.bm25.k = 10

        self.vectorstore = Milvus(
            embedding_function=self.embedding_model,
            connection_args={"uri": "http://localhost:19530"},
            collection_name="car_info_collection",
            primary_field="id",
            text_field="content",
            vector_field="embedding",
            auto_id=True,
            search_params={"metric_type": "COSINE", "params": {}},
        )
        self._index_documents_if_empty()
        self.vector = self.vectorstore.as_retriever(search_kwargs={"k": 10})

        # LangChain 当前版本使用 RRF 融合排序，不需要 normalize_scores 参数。
        self.ensemble = EnsembleRetriever(
            retrievers=[self.bm25, self.vector],
            weights=[0.4, 0.6],
        )

    def _index_documents_if_empty(self):
        """Milvus 集合为空时才写入分块，避免重复入库。"""
        rows = self.vectorstore.client.query(
            collection_name="car_info_collection",
            filter="id >= 0",
            output_fields=["id"],
            limit=1,
            consistency_level="Strong",
        )
        if not rows:
            self.vectorstore.add_documents(self.documents)

    def search(self, query, mode="ensemble") -> list[Document]:
        """执行检索，mode 可选 bm25、vector、ensemble。"""
        retriever_map = {
            "bm25": self.bm25,
            "vector": self.vector,
            "ensemble": self.ensemble,
        }
        if mode not in retriever_map:
            choices = ", ".join(retriever_map)
            raise ValueError(f"不支持的检索模式: {mode}，可选: {choices}")

        return retriever_map[mode].invoke(query)[:10]