update

app.py

@@ -1,80 +1,86 @@
-from fastapi import FastAPI, HTTPException, Depends, Header
+from fastapi import FastAPI, HTTPException, Depends, Header, Request
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel, Field
-from sentence_transformers import CrossEncoder  # 关键修改：使用 CrossEncoder 而非 SentenceTransformer
+from sentence_transformers import CrossEncoder
 import logging
 import os
-from typing import List
+from typing import List, Dict
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# 依赖项：校验 Authorization Header
+# 鉴权
 async def verify_auth(authorization: str = Header(..., alias="Authorization")):
     if not authorization.startswith("Bearer "):
-        raise HTTPException(status_code=401, detail="Invalid token format. Use 'Bearer YOUR_TOKEN'")
+        raise HTTPException(401, detail="Invalid token format")
     token = authorization[len("Bearer "):]
     if token != os.getenv("AUTHORIZATION"):
-        raise HTTPException(status_code=401, detail="Invalid token")
+        raise HTTPException(401, detail="Invalid token")
     return token
 
 app = FastAPI()
 
-# 加载重排序模型（初始化时加载，避免每次请求重复加载）
+# 模型配置
+MODEL_NAME = "BAAI/bge-reranker-large"  # 确保名称正确
+
 try:
     model = CrossEncoder(
-        "BAAI/bge-reranker-large",
-        tokenizer_args={"truncation": True},  # 在初始化时设置截断
-        max_length=512  # 可选：限制最大长度
+        MODEL_NAME,
+        tokenizer_args={"truncation": True},
+        max_length=512
     )
-    logger.info("Model loaded with truncation support")
+    # 健康检查
+    test_score = model.predict([("test", "test")])[0]
+    logger.info(f"Model loaded. Test score: {test_score}")
 except Exception as e:
     logger.critical(f"Model load failed: {str(e)}")
-    raise RuntimeError("Model initialization error")
+    raise RuntimeError("Model init failed")
 
-# 请求体模型
+# 请求/响应模型
 class RerankRequest(BaseModel):
     query: str = Field(..., min_length=1, max_length=8192)
     documents: List[str] = Field(..., min_items=1)
     top_k: int = Field(None, ge=1, le=100)
 
-# 响应模型
 class RerankResult(BaseModel):
     index: int
     score: float
     document: str
 
-@app.post("/rerank", response_model=List[RerankResult])
+# 统一错误响应
+@app.exception_handler(HTTPException)
+async def handle_errors(request: Request, exc: HTTPException):
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"error": {"message": exc.detail, "type": "api_error"}}
+    )
+
+@app.post("/rerank")
 async def rerank(
     request: RerankRequest,
-    token: str = Depends(verify_auth)  # 依赖权限校验
-):
+    token: str = Depends(verify_auth)
+) -> Dict:
     try:
-        # 创建 (query, document) 对
-        pairs = [(request.query, doc) for doc in request.documents]
+        logger.info(f"Processing query: {request.query[:50]}...")
         
-        # 获取预测分数（添加 truncation 处理长文本）
-        scores = model.predict(pairs) 
+        pairs = [(request.query, doc) for doc in request.documents]
+        scores = model.predict(pairs)
         
-        # 构建结果列表
         results = [
             {"index": idx, "score": float(score), "document": doc}
             for idx, (doc, score) in enumerate(zip(request.documents, scores))
         ]
-        
-        # 按分数排序
         sorted_results = sorted(results, key=lambda x: x["score"], reverse=True)
         
-        # 返回 top_k 结果
-        if request.top_k is not None:
-            return sorted_results[:request.top_k]
-        # 返回符合 OpenAI 风格的格式
+        if request.top_k:
+            sorted_results = sorted_results[:request.top_k]
+        
         return {
             "object": "list",
-            "data": sorted_results,  # 结果包装在 data 字段中
-            "model": "BAAI/bge-reranker-large"
+            "data": sorted_results,
+            "model": MODEL_NAME
         }
         
     except Exception as e:
-        logger.error(f"API Error: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Error: {str(e)}", exc_info=True)
+        raise HTTPException(500, detail="Internal server error")