feat: 格式化
This commit is contained in:
parent
2563d4c3c8
commit
9a676f601c
48
main.py
48
main.py
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue