feat: add rerank

This commit is contained in:
李如威 2025-07-09 23:59:12 +08:00
parent d7a4d671f0
commit c8b1754e73
1 changed files with 69 additions and 6 deletions

View File

@ -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:
"""异步构建上下文"""