|
@@ -0,0 +1,65 @@
|
|
|
|
|
+"""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]
|