feat: add rerank
This commit is contained in:
parent
d7a4d671f0
commit
c8b1754e73
|
@ -21,9 +21,16 @@ class AsyncRAGService:
|
||||||
self.llm = ChatOpenAI(
|
self.llm = ChatOpenAI(
|
||||||
model="deepseek-r1:8b",
|
model="deepseek-r1:8b",
|
||||||
temperature=0.7,
|
temperature=0.7,
|
||||||
openai_api_key=os.getenv("OPENAI_API_KEY"),
|
openai_api_key=self.openai_api_key,
|
||||||
openai_api_base=os.getenv("OPENAI_BASE_URL"),
|
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(
|
self.prompt_template = PromptTemplate(
|
||||||
input_variables=["context", "question"],
|
input_variables=["context", "question"],
|
||||||
template="""
|
template="""
|
||||||
|
@ -76,12 +83,15 @@ class AsyncRAGService:
|
||||||
"processing_time": time.time() - start_time,
|
"processing_time": time.time() - start_time,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# rerank
|
||||||
|
reranked_results = await self._rerank_results(question, search_results)
|
||||||
|
|
||||||
# 并行执行上下文构建和 LLM 调用准备
|
# 并行执行上下文构建和 LLM 调用准备
|
||||||
context_task = asyncio.create_task(
|
context_task = asyncio.create_task(
|
||||||
self._build_context_async(search_results)
|
self._build_context_async(reranked_results)
|
||||||
)
|
)
|
||||||
sources_task = asyncio.create_task(
|
sources_task = asyncio.create_task(
|
||||||
self._format_sources_async(search_results)
|
self._format_sources_async(reranked_results)
|
||||||
)
|
)
|
||||||
|
|
||||||
# 等待上下文构建完成
|
# 等待上下文构建完成
|
||||||
|
@ -128,9 +138,12 @@ class AsyncRAGService:
|
||||||
}
|
}
|
||||||
return
|
return
|
||||||
|
|
||||||
|
# rerank
|
||||||
|
reranked_results = await self._rerank_results(question, search_results)
|
||||||
|
|
||||||
# 构建上下文和源信息
|
# 构建上下文和源信息
|
||||||
context_task = self._build_context_async(search_results)
|
context_task = self._build_context_async(reranked_results)
|
||||||
sources_task = self._format_sources_async(search_results)
|
sources_task = self._format_sources_async(reranked_results)
|
||||||
|
|
||||||
context = await context_task
|
context = await context_task
|
||||||
|
|
||||||
|
@ -177,6 +190,56 @@ class AsyncRAGService:
|
||||||
"""异步删除文档"""
|
"""异步删除文档"""
|
||||||
return await self.vector_store.delete_document_async(doc_id)
|
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:
|
async def _build_context_async(self, search_results: List[Dict[str, Any]]) -> str:
|
||||||
"""异步构建上下文"""
|
"""异步构建上下文"""
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue