From c8b1754e73fb8f4cd24bb42b0848e66e4d73f032 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Wed, 9 Jul 2025 23:59:12 +0800 Subject: [PATCH] feat: add rerank --- services/rag_service.py | 75 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 69 insertions(+), 6 deletions(-) diff --git a/services/rag_service.py b/services/rag_service.py index 783d18b..0dc9bae 100644 --- a/services/rag_service.py +++ b/services/rag_service.py @@ -21,9 +21,16 @@ class AsyncRAGService: self.llm = ChatOpenAI( model="deepseek-r1:8b", temperature=0.7, - openai_api_key=os.getenv("OPENAI_API_KEY"), - openai_api_base=os.getenv("OPENAI_BASE_URL"), + openai_api_key=self.openai_api_key, + openai_api_base=self.openai_api_base, ) + self.rerank_llm = ChatOpenAI( + model="deepseek-r1:8b", + temperature=0.7, + openai_api_key=self.openai_api_key, + openai_api_base=self.openai_api_base, + ) + self.prompt_template = PromptTemplate( input_variables=["context", "question"], template=""" @@ -76,12 +83,15 @@ class AsyncRAGService: "processing_time": time.time() - start_time, } + # rerank + reranked_results = await self._rerank_results(question, search_results) + # 并行执行上下文构建和 LLM 调用准备 context_task = asyncio.create_task( - self._build_context_async(search_results) + self._build_context_async(reranked_results) ) sources_task = asyncio.create_task( - self._format_sources_async(search_results) + self._format_sources_async(reranked_results) ) # 等待上下文构建完成 @@ -128,9 +138,12 @@ class AsyncRAGService: } return + # rerank + reranked_results = await self._rerank_results(question, search_results) + # 构建上下文和源信息 - context_task = self._build_context_async(search_results) - sources_task = self._format_sources_async(search_results) + context_task = self._build_context_async(reranked_results) + sources_task = self._format_sources_async(reranked_results) context = await context_task @@ -177,6 +190,56 @@ class AsyncRAGService: """异步删除文档""" return await self.vector_store.delete_document_async(doc_id) + async def _format_sources_async( + self, search_results: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + def _format_sources(): + return [ + { + "filename": r["metadata"]["filename"], + "content": ( + (r["content"][:200] + "...") + if len(r["content"]) > 200 + else r["content"] + ), + "similarity": 1 - r["distance"], + "rerank_score": r.get("rerank_score", None), + } + for r in search_results + ] + + return await asyncio.to_thread(_format_sources) + + async def _rerank_results( + self, question: str, search_results: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """使用 rerank LLM 对搜索结果重新排序""" + + async def score_result(result: Dict[str, Any]) -> float: + prompt = f""" +你是一个智能评分助手,请判断以下“文档片段”与“用户问题”的相关程度。 +请只输出一个介于 0 到 1 之间的分数,数值越高表示相关性越强。 + +用户问题: +{question} + +文档片段: +{result['content'][:1000]} +""" + try: + response = await asyncio.to_thread(self.rerank_llm.invoke, prompt) + score = float(response.content.strip()) + return max(0.0, min(score, 1.0)) + except Exception as e: + self.logger.warning(f"rerank 评分失败,fallback 使用向量相似度: {e}") + return 1 - result["distance"] + + scores = await asyncio.gather(*[score_result(r) for r in search_results]) + for r, score in zip(search_results, scores): + r["rerank_score"] = score + + return sorted(search_results, key=lambda r: r["rerank_score"], reverse=True) + async def _build_context_async(self, search_results: List[Dict[str, Any]]) -> str: """异步构建上下文"""