From 2563d4c3c8d56362958d9681dffbb3d1eb0367ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Wed, 9 Jul 2025 00:36:00 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20API=20=E9=89=B4=E6=9D=83?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 6 +++++- config.py | 4 ++++ main.py | 47 ++++++++++++++++++++++++++++++++++++++++++----- 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/.env.example b/.env.example index c50a162..768382a 100644 --- a/.env.example +++ b/.env.example @@ -24,4 +24,8 @@ LOG_LEVEL=INFO LOG_DIR=./logs # Tokenizers 配置 -TOKENIZERS_PARALLELISM=false \ No newline at end of file +TOKENIZERS_PARALLELISM=false + +# 认证配置 +AUTH_ENABLED=false +API_TOKEN=695ee365-efbe-4a97-8ff0-73195c23d31f \ No newline at end of file diff --git a/config.py b/config.py index 6671221..257b1d6 100644 --- a/config.py +++ b/config.py @@ -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): """验证配置""" diff --git a/main.py b/main.py index 422df5b..df9ede4 100644 --- a/main.py +++ b/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: