easy-rag/services/rag_service.py

276 lines
9.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from typing import List, Dict, Any
import asyncio
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.callbacks import AsyncIteratorCallbackHandler
from services.vector_store import AsyncVectorStore
from utils.logger import get_logger
import os
import time
class AsyncRAGService:
"""异步 RAG 服务主类"""
def __init__(self):
self.logger = get_logger(__name__)
self.vector_store = AsyncVectorStore()
self.openai_api_base = os.getenv("OPENAI_BASE_URL")
self.openai_api_key = os.getenv("OPENAI_API_KEY")
self.llm = ChatOpenAI(
model="deepseek-r1:8b",
temperature=0.7,
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="""
基于以下上下文回答问题。如果上下文中没有相关信息,请说明无法从提供的文档中找到答案。
上下文:
{context}
问题:{question}
答案:""",
)
self.logger.info("RAG服务初始化完成")
async def add_document_async(self, content: str, filename: str) -> str:
"""异步添加文档"""
start_time = time.time()
try:
self.logger.info(f"开始添加文档: {filename}")
result = await self.vector_store.add_document_async(content, filename)
duration = time.time() - start_time
self.logger.info(f"文档添加成功: {filename}, 耗时: {duration:.2f}s")
return result
except Exception as e:
duration = time.time() - start_time
self.logger.error(
f"文档添加失败: {filename}, 错误: {str(e)}, 耗时: {duration:.2f}s"
)
raise
async def chat_async(
self, question: str, top_k: int = 3, temperature: float = 0.7
) -> Dict[str, Any]:
"""异步聊天问答"""
start_time = time.time()
try:
self.logger.info(f"开始处理问答: {question[:50]}...")
# 异步检索相关文档
search_results = await self.vector_store.search_async(question, top_k)
self.logger.debug(f"检索到 {len(search_results)} 个相关文档")
if not search_results:
self.logger.warning("未找到相关文档")
return {
"answer": "抱歉,我无法在现有文档中找到相关信息来回答您的问题。",
"sources": [],
"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(reranked_results)
)
sources_task = asyncio.create_task(
self._format_sources_async(reranked_results)
)
# 等待上下文构建完成
context = await context_task
# 异步生成回答
self.llm.temperature = temperature
prompt = self.prompt_template.format(context=context, question=question)
response = await asyncio.to_thread(self.llm.invoke, prompt)
# 等待源信息格式化完成
sources = await sources_task
duration = time.time() - start_time
self.logger.info(f"问答处理完成, 耗时: {duration:.2f}s")
return {
"answer": response.content,
"sources": sources,
"processing_time": duration,
}
except Exception as e:
duration = time.time() - start_time
self.logger.error(f"问答处理失败: {str(e)}, 耗时: {duration:.2f}s")
raise
async def chat_stream_async(
self, question: str, top_k: int = 3, temperature: float = 0.7
):
"""异步流式聊天问答"""
start_time = time.time()
# 异步检索相关文档
search_results = await self.vector_store.search_async(question, top_k)
if not search_results:
yield {
"content": "抱歉,我无法在现有文档中找到相关信息来回答您的问题。",
"is_final": True,
"sources": [],
"processing_time": time.time() - start_time,
}
return
# rerank
reranked_results = await self._rerank_results(question, search_results)
# 构建上下文和源信息
context_task = self._build_context_async(reranked_results)
sources_task = self._format_sources_async(reranked_results)
context = await context_task
# 设置 LLM 参数
prompt = self.prompt_template.format(context=context, question=question)
callback_handler = AsyncIteratorCallbackHandler()
stream_llm = ChatOpenAI(
model="deepseek-r1:8b",
streaming=True,
callbacks=[callback_handler],
openai_api_key=self.openai_api_key,
openai_api_base=self.openai_api_base,
)
self.logger.info("启动 LLM 流式生成任务...")
task = asyncio.create_task(stream_llm.ainvoke(prompt))
self.logger.info("LLM 流式生成任务已启动")
async for token in callback_handler.aiter():
yield {
"content": token,
"is_final": False,
"sources": None,
"processing_time": None,
}
await task
# 最后一个数据块包含完整信息
sources = await sources_task
yield {
"content": "",
"is_final": True,
"sources": sources,
"processing_time": time.time() - start_time,
}
async def get_documents_async(self) -> List[Dict[str, Any]]:
"""异步获取文档列表"""
return await self.vector_store.get_documents_async()
async def delete_document_async(self, doc_id: str) -> bool:
"""异步删除文档"""
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:
"""异步构建上下文"""
def _build_context():
return "\n\n".join(
[
f"文档片段 {i+1} (来源: {result['metadata']['filename']}):\n{result['content']}"
for i, result in enumerate(search_results)
]
)
return await asyncio.to_thread(_build_context)
async def _format_sources_async(
self, search_results: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""异步格式化源信息"""
def _format_sources():
return [
{
"filename": result["metadata"]["filename"],
"content": (
result["content"][:200] + "..."
if len(result["content"]) > 200
else result["content"]
),
"similarity": 1 - result["distance"],
}
for result in search_results
]
return await asyncio.to_thread(_format_sources)