feat: API 鉴权

This commit is contained in:
李如威 2025-07-09 00:36:00 +08:00
parent abe584f232
commit 2563d4c3c8
3 changed files with 51 additions and 6 deletions

View File

@ -24,4 +24,8 @@ LOG_LEVEL=INFO
LOG_DIR=./logs
# Tokenizers 配置
TOKENIZERS_PARALLELISM=false
TOKENIZERS_PARALLELISM=false
# 认证配置
AUTH_ENABLED=false
API_TOKEN=695ee365-efbe-4a97-8ff0-73195c23d31f

View File

@ -35,6 +35,10 @@ class Config:
# Tokenizers 配置
TOKENIZERS_PARALLELISM = os.getenv("TOKENIZERS_PARALLELISM", "false")
# 认证配置
AUTH_ENABLED = os.getenv("AUTH_ENABLED", "true").lower() == "true"
API_TOKEN = os.getenv("API_TOKEN", "easy-rag-token-2025")
@classmethod
def validate(cls):
"""验证配置"""

47
main.py
View File

@ -1,6 +1,7 @@
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
@ -63,12 +64,37 @@ 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
@app.get("/", response_model=dict)
async def root():
"""根路径 - 服务健康检查"""
@ -89,7 +115,9 @@ async def health_check():
@app.post("/upload", response_model=SuccessResponse)
async def upload_document(
file: UploadFile = File(...), service: AsyncRAGService = Depends(get_rag_service)
file: UploadFile = File(...),
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token)
):
"""上传文档接口"""
start_time = time.time()
@ -153,7 +181,9 @@ async def upload_document(
@app.post("/chat", response_model=ChatResponse)
async def chat(
request: ChatRequest, service: AsyncRAGService = Depends(get_rag_service)
request: ChatRequest,
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token)
):
"""聊天问答接口"""
start_time = time.time()
@ -184,7 +214,9 @@ async def chat(
@app.post("/chat/stream")
async def chat_stream(
request: StreamChatRequest, service: AsyncRAGService = Depends(get_rag_service)
request: StreamChatRequest,
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token)
):
"""流式聊天问答接口"""
logger.info(f"开始处理流式问答: {request.question[:50]}...")
@ -224,7 +256,10 @@ async def chat_stream(
@app.get("/documents", response_model=List[DocumentInfo])
async def get_documents(service: AsyncRAGService = Depends(get_rag_service)):
async def get_documents(
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token)
):
"""获取文档列表接口"""
try:
logger.info("获取文档列表")
@ -247,7 +282,9 @@ async def get_documents(service: AsyncRAGService = Depends(get_rag_service)):
@app.delete("/documents/{doc_id}", response_model=SuccessResponse)
async def delete_document(
doc_id: str, service: AsyncRAGService = Depends(get_rag_service)
doc_id: str,
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token)
):
"""删除文档接口"""
try: