from typing import List, Dict, Any import asyncio from langchain_openai import ChatOpenAI from langchain.prompts import PromptTemplate from services.vector_store import AsyncVectorStore import os import time class AsyncRAGService: """异步 RAG 服务主类""" def __init__(self): self.vector_store = AsyncVectorStore() 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"), ) self.prompt_template = PromptTemplate( input_variables=["context", "question"], template=""" 基于以下上下文回答问题。如果上下文中没有相关信息,请说明无法从提供的文档中找到答案。 上下文: {context} 问题:{question} 答案:""", ) async def add_document_async(self, content: str, filename: str) -> str: """异步添加文档""" return await self.vector_store.add_document_async(content, filename) async def chat_async( self, question: str, top_k: int = 3, temperature: float = 0.7 ) -> Dict[str, Any]: """异步聊天问答""" start_time = time.time() # 异步检索相关文档 search_results = await self.vector_store.search_async(question, top_k) if not search_results: return { "answer": "抱歉,我无法在现有文档中找到相关信息来回答您的问题。", "sources": [], "processing_time": time.time() - start_time, } # 并行执行上下文构建和 LLM 调用准备 context_task = asyncio.create_task(self._build_context_async(search_results)) sources_task = asyncio.create_task(self._format_sources_async(search_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 return { "answer": response.content, "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 _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)