Update main.py

main.py

@@ -1,5 +1,5 @@
 # ===================================================================
-# main.py (已修改以适配 Hugging Face Secrets)
+# main.py (已修复 Hugging Face 健康检查问题)
 # ===================================================================
 
 import json
@@ -15,7 +15,7 @@ from fastapi.responses import StreamingResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel, Field
 
-# --- 类型定义和全局变量 (与原始文件一致) ---
+# --- 类型定义和全局变量 ---
 class CodeGeeXToken(TypedDict):
     token: str
     is_valid: bool
@@ -30,12 +30,12 @@ MAX_ERROR_COUNT = 3
 ERROR_COOLDOWN = 300
 DEBUG_MODE = os.environ.get("DEBUG_MODE", "false").lower() == "true"
 
-# --- Pydantic 模型 (与原始文件一致) ---
+# --- Pydantic 模型 ---
 class ChatMessage(BaseModel):
     role: str
     content: Union[str, List[Dict[str, Any]]]
     reasoning_content: Optional[str] = None
-
+# ... (其他模型保持不变)
 class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[ChatMessage]
@@ -53,11 +53,12 @@ class ModelInfo(BaseModel):
 class ModelList(BaseModel):
     object: str = "list"
     data: List[ModelInfo]
-# ... (其他 Pydantic 模型与原始文件一致)
+
 class ChatCompletionChoice(BaseModel):
     message: ChatMessage
     index: int = 0
     finish_reason: str = "stop"
+
 class ChatCompletionResponse(BaseModel):
     id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")
     object: str = "chat.completion"
@@ -65,10 +66,12 @@ class ChatCompletionResponse(BaseModel):
     model: str
     choices: List[ChatCompletionChoice]
     usage: Dict[str, int] = Field(default_factory=lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
+
 class StreamChoice(BaseModel):
     delta: Dict[str, Any] = Field(default_factory=dict)
     index: int = 0
     finish_reason: Optional[str] = None
+
 class StreamResponse(BaseModel):
     id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")
     object: str = "chat.completion.chunk"
@@ -84,43 +87,58 @@ def log_debug(message: str):
     if DEBUG_MODE:
         print(f"[DEBUG] {message}")
 
-# --- 配置加载函数 (已修改为从 Secrets 读取) ---
+# --- 配置加载函数 (从 Secrets 读取) ---
 def load_client_api_keys_from_secrets():
-    """从环境变量加载客户端 API Keys"""
     global VALID_CLIENT_KEYS
     try:
         keys_str = os.environ.get("CLIENT_API_KEYS")
-        if not keys_str:
-            raise ValueError("Secret 'CLIENT_API_KEYS' not found.")
+        if not keys_str: raise ValueError("Secret 'CLIENT_API_KEYS' not found.")
         keys = json.loads(keys_str)
         VALID_CLIENT_KEYS = set(keys) if isinstance(keys, list) else set()
-        print(f"Successfully loaded {len(VALID_CLIENT_KEYS)} client API keys from secrets.")
+        print(f"Successfully loaded {len(VALID_CLIENT_KEYS)} client API keys.")
     except Exception as e:
-        print(f"FATAL: Error loading client API keys from secrets: {e}")
+        print(f"FATAL: Error loading client API keys: {e}")
         VALID_CLIENT_KEYS = set()
 
 def load_codegeex_tokens_from_secrets():
-    """从环境变量加载 CodeGeeX Tokens"""
     global CODEGEEX_TOKENS
     CODEGEEX_TOKENS = []
     try:
         tokens_str = os.environ.get("CODEGEEX_TOKENS")
-        if not tokens_str:
-            raise ValueError("Secret 'CODEGEEX_TOKENS' not found.")
-        tokens = json.loads(tokens_str) # 假设Secret是一个JSON数组
-        if not isinstance(tokens, list):
-             raise TypeError("Secret 'CODEGEEX_TOKENS' must be a JSON list of strings.")
-
+        if not tokens_str: raise ValueError("Secret 'CODEGEEX_TOKENS' not found.")
+        tokens = json.loads(tokens_str)
+        if not isinstance(tokens, list): raise TypeError("Secret 'CODEGEEX_TOKENS' must be a JSON list.")
         for token in tokens:
             if isinstance(token, str) and token:
-                CODEGEEX_TOKENS.append({
-                    "token": token, "is_valid": True, "last_used": 0, "error_count": 0
-                })
-        print(f"Successfully loaded {len(CODEGEEX_TOKENS)} CodeGeeX tokens from secrets.")
+                CODEGEEX_TOKENS.append({"token": token, "is_valid": True, "last_used": 0, "error_count": 0})
+        print(f"Successfully loaded {len(CODEGEEX_TOKENS)} CodeGeeX tokens.")
     except Exception as e:
-        print(f"FATAL: Error loading CodeGeeX tokens from secrets: {e}")
+        print(f"FATAL: Error loading CodeGeeX tokens: {e}")
+
+@app.on_event("startup")
+async def startup():
+    print("Starting CodeGeeX OpenAI API Adapter server...")
+    load_client_api_keys_from_secrets()
+    load_codegeex_tokens_from_secrets()
+    print("Server initialization completed.")
+
+# --- THE FIX: Add a root endpoint for the health check ---
+@app.get("/")
+def health_check():
+    """Health check endpoint for Hugging Face."""
+    return {"status": "ok", "message": "CodeGeeX API Adapter is running."}
+# ---------------------------------------------------------
 
-# --- 核心逻辑 (与原始文件保持一致) ---
+# --- 原始路由 ---
+def get_models_list_response() -> ModelList:
+    return ModelList(data=[ModelInfo(id=model, created=int(time.time()), owned_by="anthropic") for model in CODEGEEX_MODELS])
+
+@app.get("/v1/models", response_model=ModelList)
+async def list_v1_models(_: None = Depends(authenticate_client)):
+    return get_models_list_response()
+
+# ... (所有其他函数和路由都与之前的文件一致，无需改动)
+# The rest of the original code is unchanged.
 def get_best_codegeex_token() -> Optional[CodeGeeXToken]:
     with token_rotation_lock:
         now = time.time()
@@ -155,28 +173,10 @@ async def authenticate_client(auth: Optional[HTTPAuthorizationCredentials] = Dep
     if not auth or not auth.credentials: raise HTTPException(status_code=401, detail="API key required.", headers={"WWW-Authenticate": "Bearer"})
     if auth.credentials not in VALID_CLIENT_KEYS: raise HTTPException(status_code=403, detail="Invalid client API key.")
 
-@app.on_event("startup")
-async def startup():
-    print("Starting CodeGeeX OpenAI API Adapter server...")
-    load_client_api_keys_from_secrets()
-    load_codegeex_tokens_from_secrets()
-    print("Server initialization completed.")
-
-def get_models_list_response() -> ModelList:
-    return ModelList(data=[ModelInfo(id=model, created=int(time.time()), owned_by="anthropic") for model in CODEGEEX_MODELS])
-
-@app.get("/v1/models", response_model=ModelList)
-async def list_v1_models(_: None = Depends(authenticate_client)):
-    return get_models_list_response()
-
 @app.get("/models", response_model=ModelList)
 async def list_models_no_auth():
     return get_models_list_response()
 
-# ... (所有路由和核心函数都与原始文件一致, 此处省略以保持简洁, 但它们都在上面的完整代码块中)
-# --- The rest of the original code follows ---
-# This includes _codegeex_stream_generator, _build_codegeex_non_stream_response, chat_completions, etc.
-# They are included in the full code block above.
 def _codegeex_stream_generator(response, model: str):
     stream_id = f"chatcmpl-{uuid.uuid4().hex}"
     created_time = int(time.time())