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

@ -25,3 +25,7 @@ LOG_DIR=./logs
# Tokenizers 配置 # 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 配置
TOKENIZERS_PARALLELISM = os.getenv("TOKENIZERS_PARALLELISM", "false") 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 @classmethod
def validate(cls): def validate(cls):
"""验证配置""" """验证配置"""

47
main.py
View File

@ -1,6 +1,7 @@
from fastapi import FastAPI, File, UploadFile, HTTPException, Depends from fastapi import FastAPI, File, UploadFile, HTTPException, Depends
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import uvicorn import uvicorn
import os import os
from typing import List from typing import List
@ -63,12 +64,37 @@ logger.info(f"正在启动 {config.APP_NAME} v{config.APP_VERSION}")
rag_service = AsyncRAGService() rag_service = AsyncRAGService()
logger.info("RAG服务实例创建完成") logger.info("RAG服务实例创建完成")
# 创建认证方案
security = HTTPBearer(auto_error=False)
def get_rag_service() -> AsyncRAGService: def get_rag_service() -> AsyncRAGService:
"""依赖注入获取RAG服务实例""" """依赖注入获取RAG服务实例"""
return rag_service 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) @app.get("/", response_model=dict)
async def root(): async def root():
"""根路径 - 服务健康检查""" """根路径 - 服务健康检查"""
@ -89,7 +115,9 @@ 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(...), service: AsyncRAGService = Depends(get_rag_service) file: UploadFile = File(...),
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token)
): ):
"""上传文档接口""" """上传文档接口"""
start_time = time.time() start_time = time.time()
@ -153,7 +181,9 @@ async def upload_document(
@app.post("/chat", response_model=ChatResponse) @app.post("/chat", response_model=ChatResponse)
async def chat( 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() start_time = time.time()
@ -184,7 +214,9 @@ async def chat(
@app.post("/chat/stream") @app.post("/chat/stream")
async def 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]}...") logger.info(f"开始处理流式问答: {request.question[:50]}...")
@ -224,7 +256,10 @@ async def chat_stream(
@app.get("/documents", response_model=List[DocumentInfo]) @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: try:
logger.info("获取文档列表") logger.info("获取文档列表")
@ -247,7 +282,9 @@ async def get_documents(service: AsyncRAGService = Depends(get_rag_service)):
@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, service: AsyncRAGService = Depends(get_rag_service) doc_id: str,
service: AsyncRAGService = Depends(get_rag_service),
_: bool = Depends(verify_token)
): ):
"""删除文档接口""" """删除文档接口"""
try: try: