easy-rag/main.py

353 lines
11 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 fastapi import FastAPI, File, UploadFile, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import uvicorn
import os
from typing import List
from io import BytesIO
import time
from config import config
from models import (
ChatRequest,
ChatResponse,
StreamChatRequest,
StreamChatChunk,
DocumentInfo,
ErrorResponse,
SuccessResponse,
)
from services import AsyncRAGService
from utils import (
extract_text_from_pdf_async,
validate_file_size,
ensure_directory_exists,
is_supported_file_type,
)
from utils.logger import setup_logger, get_logger, cleanup_logger
# 初始化日志系统
setup_logger(
name=config.APP_NAME,
level=config.LOG_LEVEL,
log_dir=config.LOG_DIR,
use_async=True, # 使用异步日志,避免阻塞
)
logger = get_logger(__name__)
# 创建FastAPI应用
app = FastAPI(
title=config.APP_NAME,
version=config.APP_VERSION,
description="高效简洁的RAG服务API",
docs_url="/api/docs",
redoc_url="/api/redoc",
)
# 创建API路由器
from fastapi import APIRouter
api_router = APIRouter(prefix="/api")
# 添加CORS中间件
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 确保上传目录存在
ensure_directory_exists(config.UPLOAD_DIR)
# 创建RAG服务实例
logger.info(f"正在启动 {config.APP_NAME} v{config.APP_VERSION}")
rag_service = AsyncRAGService()
logger.info("RAG服务实例创建完成")
# 创建认证方案
security = HTTPBearer(auto_error=False)
def get_rag_service() -> AsyncRAGService:
"""依赖注入获取RAG服务实例"""
return rag_service
async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
"""验证Bearer token"""
if not config.AUTH_ENABLED:
return True
if credentials is None:
raise HTTPException(
status_code=401,
detail="认证失败缺少Authorization header",
headers={"WWW-Authenticate": "Bearer"},
)
if credentials.credentials != config.API_TOKEN:
raise HTTPException(
status_code=401,
detail="认证失败无效的token",
headers={"WWW-Authenticate": "Bearer"},
)
return True
@api_router.get("/", response_model=dict)
async def root():
"""根路径 - 服务健康检查"""
logger.info("根路径访问")
return {
"message": f"欢迎使用 {config.APP_NAME}",
"version": config.APP_VERSION,
"status": "running",
}
@api_router.get("/health")
async def health_check():
"""健康检查接口"""
logger.debug("健康检查请求")
return {"status": "healthy", "service": config.APP_NAME}
@api_router.post("/upload", response_model=SuccessResponse)
async def upload_document(
file: UploadFile = File(...),
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token),
):
"""上传文档接口"""
start_time = time.time()
try:
logger.info(f"开始上传文档: {file.filename}")
# 验证文件类型
if not is_supported_file_type(file.filename):
logger.warning(f"不支持的文件类型: {file.filename}")
raise HTTPException(
status_code=400, detail="不支持的文件类型。目前支持PDF, TXT"
)
# 验证文件大小
content = await file.read()
if not validate_file_size(len(content), config.MAX_FILE_SIZE):
logger.warning(f"文件过大: {file.filename}, 大小: {len(content)} bytes")
raise HTTPException(
status_code=400,
detail=f"文件过大。最大支持 {config.MAX_FILE_SIZE // 1024 // 1024}MB",
)
# 提取文本内容
if file.filename.lower().endswith(".pdf"):
text_content = await extract_text_from_pdf_async(BytesIO(content))
else: # txt文件
text_content = content.decode("utf-8")
if not text_content.strip():
logger.warning(f"文件内容为空: {file.filename}")
raise HTTPException(status_code=400, detail="文件内容为空或无法提取文本")
# 添加到向量库
doc_id = await service.add_document_async(text_content, file.filename)
# 保存文件到本地(可选)
file_path = os.path.join(config.UPLOAD_DIR, f"{doc_id}_{file.filename}")
with open(file_path, "wb") as f:
f.write(content)
duration = time.time() - start_time
logger.info(
f"文档上传成功: {file.filename}, 文档ID: {doc_id}, 耗时: {duration:.2f}s"
)
return SuccessResponse(
message="文档上传成功",
data={
"document_id": doc_id,
"filename": file.filename,
"size": len(content),
},
)
except HTTPException:
raise
except Exception as e:
duration = time.time() - start_time
logger.error(
f"文档上传失败: {file.filename}, 错误: {str(e)}, 耗时: {duration:.2f}s",
exc_info=True,
)
raise HTTPException(status_code=500, detail=f"文档处理失败: {str(e)}")
@api_router.post("/chat", response_model=ChatResponse)
async def chat(
request: ChatRequest,
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token),
):
"""聊天问答接口"""
start_time = time.time()
try:
logger.info(f"开始处理问答: {request.question[:50]}...")
result = await service.chat_async(
question=request.question,
top_k=request.top_k,
temperature=request.temperature,
)
duration = time.time() - start_time
logger.info(f"问答处理完成, 耗时: {duration:.2f}s")
return ChatResponse(
answer=result["answer"],
sources=result["sources"],
processing_time=result["processing_time"],
)
except Exception as e:
duration = time.time() - start_time
logger.error(f"问答处理失败: {str(e)}, 耗时: {duration:.2f}s", exc_info=True)
raise HTTPException(status_code=500, detail=f"问答处理失败: {str(e)}")
@api_router.post("/chat/stream")
async def chat_stream(
request: StreamChatRequest,
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token),
):
"""流式聊天问答接口"""
logger.info(f"开始处理流式问答: {request.question[:50]}...")
async def generate_stream():
try:
async for chunk_data in service.chat_stream_async(
question=request.question,
top_k=request.top_k,
temperature=request.temperature,
):
# 将数据转换为 JSON 格式并添加换行符
chunk = StreamChatChunk(**chunk_data)
yield f"data: {chunk.model_dump_json()}\n\n"
except Exception as e:
logger.error(f"流式问答处理失败: {str(e)}", exc_info=True)
# 发生错误时发送错误信息
error_chunk = StreamChatChunk(
content=f"生成回答时发生错误: {str(e)}",
is_final=True,
sources=[],
processing_time=0.0,
)
yield f"data: {error_chunk.model_dump_json()}\n\n"
return StreamingResponse(
generate_stream(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Headers": "*",
},
)
@api_router.get("/documents", response_model=List[DocumentInfo])
async def get_documents(
service: AsyncRAGService = Depends(get_rag_service), _: bool = Depends(verify_token)
):
"""获取文档列表接口"""
try:
logger.info("获取文档列表")
docs = await service.get_documents_async()
logger.info(f"获取到 {len(docs)} 个文档")
return [
DocumentInfo(
id=doc["id"],
filename=doc["filename"],
upload_time=doc["upload_time"],
size=0, # 可以后续添加文件大小信息
chunks_count=doc["chunks_count"],
)
for doc in docs
]
except Exception as e:
logger.error(f"获取文档列表失败: {str(e)}", exc_info=True)
raise HTTPException(status_code=500, detail=f"获取文档列表失败: {str(e)}")
@api_router.delete("/documents/{doc_id}", response_model=SuccessResponse)
async def delete_document(
doc_id: str,
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token),
):
"""删除文档接口"""
try:
logger.info(f"删除文档: {doc_id}")
success = await service.delete_document_async(doc_id)
if not success:
logger.warning(f"文档不存在: {doc_id}")
raise HTTPException(status_code=404, detail="文档不存在")
logger.info(f"文档删除成功: {doc_id}")
return SuccessResponse(message="文档删除成功")
except HTTPException:
raise
except Exception as e:
logger.error(f"删除文档失败: {doc_id}, 错误: {str(e)}", exc_info=True)
raise HTTPException(status_code=500, detail=f"删除文档失败: {str(e)}")
@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
"""全局异常处理器"""
logger.error(f"全局异常: {str(exc)}", exc_info=True)
return JSONResponse(
status_code=500,
content=ErrorResponse(
error="内部服务器错误", detail=str(exc) if config.DEBUG else "请联系管理员"
).dict(),
)
# 包含API路由器
app.include_router(api_router)
if __name__ == "__main__":
# 验证配置
try:
config.validate()
logger.info("配置验证通过")
except ValueError as e:
logger.error(f"配置错误: {e}")
print(f"配置错误: {e}")
exit(1)
# 启动服务
logger.info(f"启动服务在 {config.HOST}:{config.PORT}")
try:
uvicorn.run(
"main:app",
host=config.HOST,
port=config.PORT,
reload=config.DEBUG,
log_level="info",
)
finally:
# 应用退出时清理日志资源
cleanup_logger()