feat: API 鉴权
This commit is contained in:
parent
abe584f232
commit
2563d4c3c8
|
@ -24,4 +24,8 @@ LOG_LEVEL=INFO
|
||||||
LOG_DIR=./logs
|
LOG_DIR=./logs
|
||||||
|
|
||||||
# Tokenizers 配置
|
# 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 配置
|
||||||
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
47
main.py
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue