feat: API 鉴权
This commit is contained in:
parent
abe584f232
commit
2563d4c3c8
|
@ -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
|
|
@ -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
47
main.py
|
@ -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:
|
||||
|
|
Loading…
Reference in New Issue