353 lines
11 KiB
Python
353 lines
11 KiB
Python
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()
|