Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import uuid
+import logging
 from datetime import datetime
 from fastapi import FastAPI, HTTPException, Depends, Request
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
@@ -9,36 +10,63 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 from typing import List, Optional
 
-# ------------------- 1. 基础配置（缓存 + 环境变量） -------------------
+# ------------------- 1. 日志配置 -------------------
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S"
+)
+logger = logging.getLogger("cross-encoder-api")
+
+# ------------------- 2. 基础配置（缓存 + 环境变量） -------------------
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"
 os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/huggingface_cache"
 
 # 从环境变量获取 API Key（OpenAI 风格）
 API_KEY = os.getenv("OPENAI_API_KEY")
 if not API_KEY:
+    logger.error("环境变量 OPENAI_API_KEY 未设置")
     raise ValueError("请设置环境变量 OPENAI_API_KEY")
+logger.info("API Key 加载成功")
 
-# ------------------- 2. 初始化 FastAPI 应用 -------------------
+# ------------------- 3. 初始化 FastAPI 应用 -------------------
 app = FastAPI(
     title="OpenAI 兼容的 Cross-Encoder 重排序 API",
     description="基于 cross-encoder/ms-marco-MiniLM-L-6-v2 的文本相关性排序接口",
     version="1.0.0"
 )
 
-# ------------------- 3. OpenAI 风格认证（Bearer Token） -------------------
+# ------------------- 4. OpenAI 风格认证（Bearer Token） -------------------
 oauth2_scheme = HTTPBearer(auto_error=False)
 
 def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(oauth2_scheme)):
     """验证 API Key：必须通过 Authorization: Bearer YOUR_API_KEY 传递"""
-    if not credentials or credentials.scheme != "Bearer" or credentials.credentials != API_KEY:
+    request_id = str(uuid.uuid4())[:8]  # 生成短请求ID用于日志追踪
+    if not credentials:
+        logger.warning(f"请求 {request_id}：缺少认证信息")
+        raise HTTPException(
+            status_code=401,
+            detail="缺少认证信息（请使用 'Authorization: Bearer YOUR_API_KEY'）",
+            headers={"WWW-Authenticate": "Bearer"}
+        )
+    if credentials.scheme != "Bearer":
+        logger.warning(f"请求 {request_id}：认证方案错误，应为 Bearer，实际为 {credentials.scheme}")
+        raise HTTPException(
+            status_code=401,
+            detail="认证方案错误（请使用 'Bearer' 方案）",
+            headers={"WWW-Authenticate": "Bearer"}
+        )
+    if credentials.credentials != API_KEY:
+        logger.warning(f"请求 {request_id}：无效的 API Key")
         raise HTTPException(
             status_code=401,
-            detail="无效的 API Key（请使用 'Authorization: Bearer YOUR_API_KEY'）",
+            detail="无效的 API Key",
             headers={"WWW-Authenticate": "Bearer"}
         )
-    return credentials.credentials
+    logger.info(f"请求 {request_id}：API Key 验证通过")
+    return (credentials.credentials, request_id)  # 返回API Key和请求ID
 
-# ------------------- 4. 数据模型定义 -------------------
+# ------------------- 5. 数据模型定义 -------------------
 class RerankRequest(BaseModel):
     query: str
     documents: List[str]
@@ -81,10 +109,12 @@ class GPTResponse(BaseModel):
     choices: List[Choice]
     usage: dict = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
 
-# ------------------- 5. 加载 Cross-Encoder 模型 -------------------
+# ------------------- 6. 加载 Cross-Encoder 模型 -------------------
 class CrossEncoderModel:
     def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
         self.model_name = model_name
+        logger.info(f"开始加载模型：{model_name}")
+        
         # 验证缓存目录可写
         cache_dir = os.environ.get("TRANSFORMERS_CACHE", "/tmp/huggingface_cache")
         try:
@@ -92,50 +122,94 @@ class CrossEncoderModel:
             with open(test_file, "w") as f:
                 f.write("test")
             os.remove(test_file)
-            print(f"缓存目录可写：{cache_dir}")
+            logger.info(f"缓存目录可写：{cache_dir}")
         except Exception as e:
+            logger.error(f"缓存目录不可写：{str(e)}")
             raise RuntimeError(f"缓存目录不可写：{str(e)}")
+        
         # 加载模型
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
-        self.model = AutoModelForSequenceClassification.from_pretrained(model_name, cache_dir=cache_dir)
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(self.device)
-        self.model.eval()
-        print(f"模型加载完成，设备：{self.device}")
+        try:
+            logger.info("开始加载分词器...")
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+            logger.info("分词器加载完成")
+            
+            logger.info("开始加载模型权重...")
+            self.model = AutoModelForSequenceClassification.from_pretrained(model_name, cache_dir=cache_dir)
+            logger.info("模型权重加载完成")
+            
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.model.to(self.device)
+            self.model.eval()
+            logger.info(f"模型加载完成，使用设备：{self.device}")
+        except Exception as e:
+            logger.error(f"模型加载失败：{str(e)}")
+            raise
 
-    def rerank(self, query: str, documents: List[str], top_k: int, truncation: bool) -> List[DocumentScore]:
+    def rerank(self, query: str, documents: List[str], top_k: int, truncation: bool, request_id: str) -> List[DocumentScore]:
+        """核心重排序逻辑，增加详细日志"""
+        logger.info(f"请求 {request_id}：开始重排序处理，查询长度: {len(query)}, 文档数量: {len(documents)}, top_k: {top_k}")
+        
+        # 参数校验
         if not documents:
+            logger.warning(f"请求 {request_id}：候选文档列表为空")
             raise ValueError("候选文档不能为空")
         if top_k <= 0:
+            logger.warning(f"请求 {request_id}：无效的 top_k 值: {top_k}")
             raise ValueError("top_k 必须为正整数")
+        
         # 自动将 top_k 限制为文档数量（避免超出）
-        top_k = min(top_k, len(documents))
+        adjusted_top_k = min(top_k, len(documents))
+        if adjusted_top_k != top_k:
+            logger.info(f"请求 {request_id}：top_k 从 {top_k} 调整为 {adjusted_top_k}（文档数量限制）")
+        
+        # 计算每篇文档的相关性分数
         doc_scores = []
-        for doc in documents:
-            inputs = self.tokenizer(
-                f"{query} {self.tokenizer.sep_token} {doc}",
-                return_tensors="pt",
-                padding="max_length",
-                truncation=truncation,
-                max_length=512
-            ).to(self.device)
-            with torch.no_grad():
-                outputs = self.model(**inputs)
-            score = outputs.logits.item()
-            doc_scores.append((doc, score))
-        sorted_docs = sorted(doc_scores, key=lambda x: x[1], reverse=True)[:top_k]
-        return [
-            DocumentScore(document=doc, score=round(score, 4), rank=i+1)
-            for i, (doc, score) in enumerate(sorted_docs)
-        ]
+        try:
+            for i, doc in enumerate(documents):
+                if i % 5 == 0:  # 每处理5个文档输出一次日志
+                    logger.info(f"请求 {request_id}：正在处理第 {i+1}/{len(documents)} 个文档")
+                
+                inputs = self.tokenizer(
+                    f"{query} {self.tokenizer.sep_token} {doc}",
+                    return_tensors="pt",
+                    padding="max_length",
+                    truncation=truncation,
+                    max_length=512
+                ).to(self.device)
+                
+                with torch.no_grad():
+                    outputs = self.model(**inputs)
+                
+                score = outputs.logits.item()
+                doc_scores.append((doc, score))
+                logger.debug(f"请求 {request_id}：文档 {i+1} 分数: {score:.4f}")
+            
+            # 排序并返回结果
+            sorted_docs = sorted(doc_scores, key=lambda x: x[1], reverse=True)[:adjusted_top_k]
+            logger.info(f"请求 {request_id}：重排序完成，返回 {len(sorted_docs)} 个结果")
+            
+            return [
+                DocumentScore(document=doc, score=round(score, 4), rank=i+1)
+                for i, (doc, score) in enumerate(sorted_docs)
+            ]
+        except Exception as e:
+            logger.error(f"请求 {request_id}：重排序过程出错: {str(e)}")
+            raise
 
-reranker = CrossEncoderModel()
+# 初始化模型（全局唯一）
+try:
+    reranker = CrossEncoderModel()
+except Exception as e:
+    logger.critical(f"模型初始化失败，服务无法启动: {str(e)}")
+    raise
 
-# ------------------- 6. API 端点（OpenAI 风格路径） -------------------
-# 6.1 根路径首页
+# ------------------- 7. API 端点（OpenAI 风格路径） -------------------
+# 7.1 根路径首页
 @app.get("/", response_class=HTMLResponse)
-async def home_page():
+async def home_page(request: Request):
+    client_ip = request.client.host
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    logger.info(f"首页访问来自 {client_ip}")
     return f"""
 <!DOCTYPE html>
 <html lang="zh-CN">
@@ -208,55 +282,94 @@ print(response.choices[0].message.content)</code></pre>
 </html>
 """
 
-# 6.2 基础重排序接口（/v1/rerank）
+# 7.2 基础重排序接口（/v1/rerank）
 @app.post("/v1/rerank", response_model=RerankResponse)
 async def base_rerank(
     request: RerankRequest,
-    api_key: str = Depends(verify_api_key)
+    auth_result: tuple = Depends(verify_api_key)
 ):
+    api_key, request_id = auth_result
     try:
-        print(f"接收到的请求：{request.dict()}")  # 打印请求内容
+        logger.info(f"请求 {request_id}：收到 /v1/rerank 请求，query: {request.query[:50]}...（截断显示）")
+        
+        # 执行重排序
         results = reranker.rerank(
             query=request.query,
             documents=request.documents,
             top_k=request.top_k,
-            truncation=request.truncation
+            truncation=request.truncation,
+            request_id=request_id
         )
-        return RerankResponse(
-            request_id=str(uuid.uuid4()),
+        
+        # 构建响应
+        response = RerankResponse(
+            request_id=request_id,
             query=request.query,
-            top_k=request.top_k,
+            top_k=min(request.top_k, len(request.documents)),
             results=results
         )
+        
+        logger.info(f"请求 {request_id}：处理完成，返回 {len(results)} 个结果")
+        return response
+        
     except ValueError as e:
+        logger.warning(f"请求 {request_id}：参数错误 - {str(e)}")
         raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
+        logger.error(f"请求 {request_id}：服务器错误 - {str(e)}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"服务器错误：{str(e)}")
 
-# 6.3 GPT 兼容接口（/v1/chat/completions）
+# 7.3 GPT 兼容接口（/v1/chat/completions）
 @app.post("/v1/chat/completions", response_model=GPTResponse)
 async def gpt_compatible_rerank(
     request: GPTRequest,
-    api_key: str = Depends(verify_api_key)
+    auth_result: tuple = Depends(verify_api_key)
 ):
+    api_key, request_id = auth_result
     try:
+        logger.info(f"请求 {request_id}：收到 /v1/chat/completions 请求，模型: {request.model}")
+        
+        # 验证模型名
         if request.model != reranker.model_name:
-            raise ValueError(f"仅支持模型：{reranker.model_name}")
-        if not request.messages or request.messages[-1].role != "user":
-            raise ValueError("最后一条消息必须是 'user' 角色")
+            error_msg = f"仅支持模型：{reranker.model_name}，实际请求：{request.model}"
+            logger.warning(f"请求 {request_id}：{error_msg}")
+            raise ValueError(error_msg)
+        
+        # 验证消息格式
+        if not request.messages:
+            logger.warning(f"请求 {request_id}：消息列表为空")
+            raise ValueError("消息列表不能为空")
+        if request.messages[-1].role != "user":
+            error_msg = f"最后一条消息必须是 'user' 角色，实际为：{request.messages[-1].role}"
+            logger.warning(f"请求 {request_id}：{error_msg}")
+            raise ValueError(error_msg)
+        
+        # 解析输入内容
         content = request.messages[-1].content
+        logger.info(f"请求 {request_id}：用户输入: {content[:100]}...（截断显示）")
+        
         if "; documents: " not in content:
-            raise ValueError("输入格式需为 'query: [查询]; documents: [文档1]; [文档2]; ...'")
+            error_msg = "输入格式需为 'query: [查询]; documents: [文档1]; [文档2]; ...'"
+            logger.warning(f"请求 {request_id}：{error_msg}")
+            raise ValueError(error_msg)
+        
         query_part, docs_part = content.split("; documents: ")
         query = query_part.replace("query: ", "").strip()
         documents = [doc.strip() for doc in docs_part.split(";") if doc.strip()]
+        
+        logger.info(f"请求 {request_id}：解析完成，query: {query[:50]}..., 文档数量: {len(documents)}")
+        
+        # 执行重排序
         results = reranker.rerank(
             query=query,
             documents=documents,
             top_k=request.top_k,
-            truncation=True
+            truncation=True,
+            request_id=request_id
         )
-        return GPTResponse(
+        
+        # 构建 GPT 风格响应
+        response = GPTResponse(
             model=request.model,
             choices=[
                 Choice(
@@ -268,22 +381,38 @@ async def gpt_compatible_rerank(
                 )
             ]
         )
+        
+        logger.info(f"请求 {request_id}：处理完成，返回 {len(results)} 个结果")
+        return response
+        
     except ValueError as e:
+        logger.warning(f"请求 {request_id}：参数错误 - {str(e)}")
         raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
+        logger.error(f"请求 {request_id}：服务器错误 - {str(e)}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"服务器错误：{str(e)}")
 
-# 6.4 健康检查接口（/v1/health）
+# 7.4 健康检查接口（/v1/health）
 @app.get("/v1/health")
-async def health_check():
-    return {
+async def health_check(request: Request):
+    client_ip = request.client.host
+    status = {
         "status": "healthy",
         "model": reranker.model_name,
         "device": reranker.device,
-        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "uptime": datetime.now().strftime("%Y-%m-%d %H:%M:%S")  # 简化版uptime
     }
+    logger.info(f"健康检查来自 {client_ip}：{status['status']}")
+    return status
 
-# ------------------- 7. 本地运行入口 -------------------
+# ------------------- 8. 本地运行入口 -------------------
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    logger.info("启动本地开发服务器...")
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860,
+        log_config=None  # 使用自定义日志配置
+    )