feat: add rerank
This commit is contained in:
parent
d7a4d671f0
commit
c8b1754e73
|
@ -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:
|
||||
"""异步构建上下文"""
|
||||
|
||||
|
|
Loading…
Reference in New Issue