update

app.py

@@ -1,12 +1,23 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends, Header
 from pydantic import BaseModel, Field
 from sentence_transformers import SentenceTransformer
-import logging
+import logging, os
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# 定义依赖项来校验 Authorization
+async def check_authorization(authorization: str = Header(..., alias="Authorization")):
+    # 去掉 Bearer 和后面的空格
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid Authorization header format")
+    
+    token = authorization[len("Bearer "):]
+    if token != os.environ.get("AUTHORIZATION"):
+        raise HTTPException(status_code=401, detail="Unauthorized access")
+    return token
+
 app = FastAPI()
 
 try:
@@ -20,9 +31,10 @@ except Exception as e:
 class RerankerRequest(BaseModel):
     query: str = Field(..., min_length=1, max_length=1000, description="The query text.")
     documents: list[str] = Field(..., min_items=2, description="A list of documents to rerank.")
+    truncate: bool = Field(False, description="Whether to truncate the documents.")
 
 @app.post("/rerank")
-async def rerank(request: RerankerRequest):
+async def rerank(request: RerankerRequest, authorization: str = Depends(check_authorization)):
     query = request.query
     documents = request.documents