feat: 格式化

This commit is contained in:
李如威 2025-07-09 00:39:54 +08:00
parent 2563d4c3c8
commit 9a676f601c
1 changed files with 26 additions and 22 deletions

48
main.py
View File

@ -34,7 +34,7 @@ setup_logger(
name=config.APP_NAME, name=config.APP_NAME,
level=config.LOG_LEVEL, level=config.LOG_LEVEL,
log_dir=config.LOG_DIR, log_dir=config.LOG_DIR,
use_async=True # 使用异步日志,避免阻塞 use_async=True, # 使用异步日志,避免阻塞
) )
logger = get_logger(__name__) logger = get_logger(__name__)
@ -77,21 +77,21 @@ async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(secur
"""验证Bearer token""" """验证Bearer token"""
if not config.AUTH_ENABLED: if not config.AUTH_ENABLED:
return True return True
if credentials is None: if credentials is None:
raise HTTPException( raise HTTPException(
status_code=401, status_code=401,
detail="认证失败缺少Authorization header", detail="认证失败缺少Authorization header",
headers={"WWW-Authenticate": "Bearer"} headers={"WWW-Authenticate": "Bearer"},
) )
if credentials.credentials != config.API_TOKEN: if credentials.credentials != config.API_TOKEN:
raise HTTPException( raise HTTPException(
status_code=401, status_code=401,
detail="认证失败无效的token", detail="认证失败无效的token",
headers={"WWW-Authenticate": "Bearer"} headers={"WWW-Authenticate": "Bearer"},
) )
return True return True
@ -115,16 +115,16 @@ async def health_check():
@app.post("/upload", response_model=SuccessResponse) @app.post("/upload", response_model=SuccessResponse)
async def upload_document( async def upload_document(
file: UploadFile = File(...), file: UploadFile = File(...),
service: AsyncRAGService = Depends(get_rag_service), service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token) _: bool = Depends(verify_token),
): ):
"""上传文档接口""" """上传文档接口"""
start_time = time.time() start_time = time.time()
try: try:
logger.info(f"开始上传文档: {file.filename}") logger.info(f"开始上传文档: {file.filename}")
# 验证文件类型 # 验证文件类型
if not is_supported_file_type(file.filename): if not is_supported_file_type(file.filename):
logger.warning(f"不支持的文件类型: {file.filename}") logger.warning(f"不支持的文件类型: {file.filename}")
@ -160,7 +160,9 @@ async def upload_document(
f.write(content) f.write(content)
duration = time.time() - start_time duration = time.time() - start_time
logger.info(f"文档上传成功: {file.filename}, 文档ID: {doc_id}, 耗时: {duration:.2f}s") logger.info(
f"文档上传成功: {file.filename}, 文档ID: {doc_id}, 耗时: {duration:.2f}s"
)
return SuccessResponse( return SuccessResponse(
message="文档上传成功", message="文档上传成功",
@ -175,22 +177,25 @@ async def upload_document(
raise raise
except Exception as e: except Exception as e:
duration = time.time() - start_time duration = time.time() - start_time
logger.error(f"文档上传失败: {file.filename}, 错误: {str(e)}, 耗时: {duration:.2f}s", exc_info=True) logger.error(
f"文档上传失败: {file.filename}, 错误: {str(e)}, 耗时: {duration:.2f}s",
exc_info=True,
)
raise HTTPException(status_code=500, detail=f"文档处理失败: {str(e)}") raise HTTPException(status_code=500, detail=f"文档处理失败: {str(e)}")
@app.post("/chat", response_model=ChatResponse) @app.post("/chat", response_model=ChatResponse)
async def chat( async def chat(
request: ChatRequest, request: ChatRequest,
service: AsyncRAGService = Depends(get_rag_service), service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token) _: bool = Depends(verify_token),
): ):
"""聊天问答接口""" """聊天问答接口"""
start_time = time.time() start_time = time.time()
try: try:
logger.info(f"开始处理问答: {request.question[:50]}...") logger.info(f"开始处理问答: {request.question[:50]}...")
result = await service.chat_async( result = await service.chat_async(
question=request.question, question=request.question,
top_k=request.top_k, top_k=request.top_k,
@ -214,9 +219,9 @@ async def chat(
@app.post("/chat/stream") @app.post("/chat/stream")
async def chat_stream( async def chat_stream(
request: StreamChatRequest, request: StreamChatRequest,
service: AsyncRAGService = Depends(get_rag_service), service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token) _: bool = Depends(verify_token),
): ):
"""流式聊天问答接口""" """流式聊天问答接口"""
logger.info(f"开始处理流式问答: {request.question[:50]}...") logger.info(f"开始处理流式问答: {request.question[:50]}...")
@ -257,8 +262,7 @@ async def chat_stream(
@app.get("/documents", response_model=List[DocumentInfo]) @app.get("/documents", response_model=List[DocumentInfo])
async def get_documents( async def get_documents(
service: AsyncRAGService = Depends(get_rag_service), service: AsyncRAGService = Depends(get_rag_service), _: bool = Depends(verify_token)
_: bool = Depends(verify_token)
): ):
"""获取文档列表接口""" """获取文档列表接口"""
try: try:
@ -282,9 +286,9 @@ async def get_documents(
@app.delete("/documents/{doc_id}", response_model=SuccessResponse) @app.delete("/documents/{doc_id}", response_model=SuccessResponse)
async def delete_document( async def delete_document(
doc_id: str, doc_id: str,
service: AsyncRAGService = Depends(get_rag_service), service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token) _: bool = Depends(verify_token),
): ):
"""删除文档接口""" """删除文档接口"""
try: try: