Update app.py

app.py

@@ -1,7 +1,5 @@
-from dataclasses import dataclass
-from typing import Any, Dict, Optional
-from fastapi import FastAPI, Request, HTTPException, status
-from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 import httpx
 import json
@@ -9,16 +7,22 @@ import os
 import random
 from datetime import datetime
 import pytz
-import logging
+from typing import Dict, List, Optional, Any, Callable
 
-# 配置日志
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# --------------------------- 1. 环境变量和配置加载 ---------------------------
+# 提取环境变量，并设置默认值(如果不存在)
+SI_KEYS = os.environ.get("SI_KEY", "").split(",")
+MODEL_MAP_JSON = os.environ.get("MODEL_MAP")
+API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.siliconflow.cn/v1")
 
-@dataclass
-class ApiConfig:
-    """API 配置数据类"""
-    base_url: str = os.environ.get("BASE_URL", "https://api.siliconflow.cn/v1")
-    model_map: Dict[str, str] = {
+# 定义模型映射的类型
+ModelMap = Dict[str, str]
+# 定义 Key 余额的类型
+KeyBalance = Dict[str, float]
+
+
+# 加载模型映射，如果环境变量 MODEL_MAP 存在，则使用它，否则使用默认值
+DEFAULT_MODEL_MAP: ModelMap = {
     "qwen-72b": "Qwen/Qwen2.5-72B-Instruct",
     "qwen-32b": "Qwen/Qwen2.5-32B-Instruct",
     "qwen-14b": "Qwen/Qwen2.5-14B-Instruct",
@@ -33,25 +37,31 @@ class ApiConfig:
     "bce": "netease-youdao/bce-embedding-base_v1",
     "bge-m3": "BAAI/bge-m3",
     "bge-zh": "BAAI/bge-large-zh-v1.5",
-    "sd":"stabilityai/stable-diffusion-3-5-large",
-    "sd-turbo":"stabilityai/stable-diffusion-3-5-large-turbo",
-    "flux-s":"black-forest-labs/FLUX.1-schnell",
-    "flux-d":"black-forest-labs/FLUX.1-dev"
+    "sd": "stabilityai/stable-diffusion-3-5-large",
+    "sd-turbo": "stabilityai/stable-diffusion-3-5-large-turbo",
+    "flux-s": "black-forest-labs/FLUX.1-schnell",
+    "flux-d": "black-forest-labs/FLUX.1-dev",
 }
-    api_key_str: Optional[str] = os.environ.get("SI_KEY")
-    timezone: str = "Asia/Shanghai"
-    
-    def __post_init__(self):
-        if not self.api_key_str :
-            logging.error("SI_KEY not found in env")
-            raise EnvironmentError("SI_KEY not found in env")
-        self.api_keys:list[str] = self.api_key_str.split(",")
-        if os.environ.get("MODEL_MAP"):
-          self.model_map = json.loads(os.environ.get("MODEL_MAP"))
 
-config = ApiConfig()
+model_map: ModelMap = json.loads(MODEL_MAP_JSON) if MODEL_MAP_JSON else DEFAULT_MODEL_MAP
+
+# --------------------------- 2. 全局变量初始化 ---------------------------
+keys: List[str] = [key.strip() for key in SI_KEYS if key.strip()]  # 删除空白字符、空值
+key_balance: KeyBalance = {}
+key_balance_notes: str = ""
+tz = pytz.timezone("Asia/Shanghai")
+last_updated_time: str = ""
+
+# --------------------------- 3. 密钥选择策略 ---------------------------
+def get_api_key() -> str:
+    """随机返回一个API密钥."""
+    random.shuffle(keys)
+    return keys[0] if keys else ""   # 添加空列表保护
+
 
+# --------------------------- 4. FastAPI 应用初始化 ---------------------------
 app = FastAPI()
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -60,171 +70,115 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# API Key 管理类
-class ApiKeyManager:
-    def __init__(self, keys:list[str]):
-        self.keys = keys
-        self.key_balance :Dict[str,float] = {}
-        self.key_balance_notes = ""
-        self.last_update_time = ""
-
-    def get_key(self) -> str:
-        """随机获取一个可用 API Key"""
-        if not self.keys:
-            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="No available API keys")
-        random.shuffle(self.keys)
-        return self.keys[0]
+
+# --------------------------- 5. 辅助函数 ---------------------------
+def format_key_balance_note(key: str, balance: float) -> str:
+    """将 key 和 balance 信息格式化为 HTML 片段."""
+    safe_key = f"{key[0:4]}****{key[-4:]}"
+    return f"<h2>{safe_key}————{balance}</h2>"
+
+async def check_key(client: httpx.AsyncClient, key:str) -> Optional[float]:
+    """检查单个密钥是否有效，并返回余额."""
+    url = f"{API_BASE_URL}/user/info"
+    headers = {"Authorization": f"Bearer {key}"}
+    try:
+        res = await client.get(url, headers=headers)
+        res.raise_for_status() # 确保状态码为200
+        balance = res.json()["data"]["balance"]
+        return float(balance)
+    except httpx.HTTPError as e:
+        print(f"Error checking key {key}: {e}")
+        return None
+
+async def forward_request(
+    request: Request,
+    url_path: str,
+    is_stream: bool = False,
+) -> Any:
+    """通用的请求转发函数."""
+    body = await request.json()
+    key = get_api_key()
+    if not key:
+        raise HTTPException(status_code=400, detail="No valid API key available.")
+    headers = {"Authorization": f"Bearer {key}"}
     
-    def remove_key(self,key:str):
-          """移除不可用 API Key"""
-          if key in self.keys:
-              self.keys.remove(key)
-          else:
-              logging.warning("try remove a not exists key from key_pool")
-
-    async def check_keys_balance(self):
-        """检查所有 API Key 的余额，并移除余额不足的 API Key"""
-        self.key_balance_notes = ""
-        key_to_remove = []
-        
-        for key in self.keys:
-            try:
-                balance = await self._fetch_balance(key)
-                if balance < 0.1:
-                    key_to_remove.append(key)
-                else:
-                    balance_info = f"<h2>{key.strip()[0:4]}****{key.strip()[-4:]}————{balance}</h2>"
-                    self.key_balance[key.strip()] = balance
-                    self.key_balance_notes += balance_info
-            except HTTPException as e:
-                    logging.error(f"Key {key} check balance failed, detail:{e.detail}")
-                    key_to_remove.append(key)
-            except Exception as e:
-                logging.error(f"Key {key} check balance failed, unexcept error:{e}")
-                key_to_remove.append(key)
-        for remove_key in key_to_remove:
-            self.remove_key(remove_key)
-        self.last_update_time = datetime.now(pytz.timezone(config.timezone))
-
-    async def _fetch_balance(self, key: str) -> float:
-        """发送 API 请求，获取 API Key 的余额"""
-        url = "https://api.siliconflow.cn/v1/user/info"
-        headers = {"Authorization": f"Bearer {key.strip()}"}
+    if "model" in body and body["model"] in model_map:
+        body["model"] = model_map[body["model"]]
+    if not "stream" in body or not body['stream']
+        is_stream=False
+    if is_stream:
+        async def generate_response():
+            async with httpx.AsyncClient() as client:
+                async with client.stream(
+                    "POST", f"{API_BASE_URL}{url_path}", headers=headers, json=body
+                ) as response:
+                    response.raise_for_status()
+                    async for chunk in response.aiter_bytes():
+                        if chunk:
+                            yield chunk
+        return StreamingResponse(generate_response(), media_type="text/event-stream")
+    else:
         async with httpx.AsyncClient() as client:
-            try:
-                res = await client.get(url, headers=headers)
-                res.raise_for_status()
-                balance = res.json()["data"]["balance"]
-                return float(balance)
-            except httpx.HTTPError as exc:
-                logging.error("httpx request error, detail:" + str(exc))
-                raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Check balance failed with status:{exc.response.status_code},url:{exc.request.url}")
-key_manager = ApiKeyManager(config.api_keys)
-
-
-async def get_model_info() -> str:
-    models = ""
-    for key, value in config.model_map.items():
-        models += f"<h2>{key}————{value}</h2>"
-    return models
+            response = await client.post(
+                f"{API_BASE_URL}{url_path}", headers=headers, json=body
+            )
+            response.raise_for_status()
+            return response.json()
 
+# --------------------------- 6. API 路由处理 ---------------------------
 @app.get("/", response_class=HTMLResponse)
-async def root() -> HTMLResponse:
-    """根路由，返回 HTML 页面，展示模型信息和更新时间"""
-    models_info = await get_model_info()
-    return HTMLResponse(f"""
+async def root():
+    """返回 HTML 格式的页面，显示模型和密钥信息."""
+    models_info = ""
+    for key, value in model_map.items():
+        models_info += f"<h2>{key}————{value}</h2>"
+
+    return f"""
     <html>
         <head>
-            <title>富文本示例</title>
+            <title>API 状态</title>
         </head>
         <body>
-            <h1>有效key数量:{len(key_manager.keys)}</h1>
+            <h1>有效Key数量: {len(keys)}</h1>
             {models_info}
-            <h1>最后更新时间：{key_manager.last_update_time}</h1>
-            {key_manager.key_balance_notes}
+            <h1>最后更新时间：{last_updated_time}</h1>
+            {key_balance_notes}
         </body>
     </html>
-    """)
+    """
 
 
 @app.get("/check")
-async def check() -> str:
-    """手动触发检查 API Key 余额的路由"""
-    await key_manager.check_keys_balance()
-    return f"更新成功：{key_manager.last_update_time}"
-
-def is_chat_stream_request(request:Request) -> bool:
-      headers = request.headers
-      if headers.get("content-type") == 'application/json':
-        try:
-             request_body=json.loads(request.headers._list[4][1])
-             return  request_body.get('model') is not None and  request_body.get("stream") is True
-        except Exception as e:
-            logging.error("parse request body error detail:" + str(e))
-            return False
-      return False
-
-async def _forward_request(request: Request, api_endpoint: str, is_stream:bool = False ) -> Any:
-    """转发请求到硅流 API"""
-    try:
-        key = key_manager.get_key()
-    except HTTPException as e:
-        return e
-    logging.info(f"using key {key[0:4]}***{key[-4:]} to {api_endpoint}")
-    headers = {"Authorization": f"Bearer {key}"}
-    url=f"{config.base_url}{api_endpoint}"
-    if api_endpoint == "/embeddings":
-      body=await request.json()
-      if "model" in body and body["model"] in config.model_map:
-          body["model"] = config.model_map[body["model"]]
-      async with httpx.AsyncClient() as client:
-              try:
-                  response = await client.post(url,headers=headers, json=body)
-                  response.raise_for_status()
-                  return response.json()
-              except  httpx.HTTPError as exc :
-                  logging.error(f"httpx request embedding error :{exc}")
-                  raise HTTPException(
-                      status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                      detail=f"Request failed with status: {exc.response.status_code}, url:{exc.request.url}, detail:{exc.response.text}"
-                  )
-    elif is_stream:
-            body = await request.json()
-            async def generate_response():
-                   async with httpx.AsyncClient() as client:
-                        async with client.stream("POST",url, headers=headers, json=body) as response:
-                                response.raise_for_status()  # 检查响应状态码
-                                async for chunk in response.aiter_bytes():
-                                    if chunk:
-                                       yield chunk
-            return StreamingResponse(generate_response(), media_type="text/event-stream")
-
-    else :
-            body = await request.json()
-            async with httpx.AsyncClient() as client:
-                       try:
-                            response = await client.post(url, headers=headers, json=body)
-                            response.raise_for_status()
-                            return response.json()
-                       except httpx.HTTPError as exc :
-                                logging.error(f"httpx request error:{exc}")
-                                raise HTTPException(
-                                        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                                        detail=f"Request failed with status: {exc.response.status_code},url:{exc.request.url},detail:{exc.response.text}"
-                                )
+async def check():
+    """检查 API 密钥的余额，并更新 key_balance、key_balance_notes 和 last_updated_time."""
+    global key_balance, key_balance_notes, last_updated_time, keys
+    key_balance_notes = ""
+    new_keys = []
+    key_balance = {}
+    async with httpx.AsyncClient() as client:
+        for key in keys:
+            balance = await check_key(client, key)
+            if balance is not None and balance >= 0.1:
+                key_balance[key] = balance
+                key_balance_notes += format_key_balance_note(key, balance)
+                new_keys.append(key)
+        keys=new_keys
+    last_updated_time = datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
+    return f"更新成功：{last_updated_time}"
+
+
 @app.post("/hf/v1/chat/completions")
-async def chat_completions(request: Request) -> Any:
-        """转发 chat 完成请求"""
-        return await _forward_request(request, "/chat/completions",is_chat_stream_request(request))
+async def chat_completions(request: Request):
+    """转发聊天补全请求，处理流式和非流式响应."""
+    return await forward_request(request, "/chat/completions", is_stream=True)
 
 @app.post("/hf/v1/embeddings")
-async def embeddings(request: Request)-> Any:
-    """转发 embeddings 请求"""
-    return await _forward_request(request, "/embeddings")
-
-@app.post("/hf/v1/images/generations")
-async def image_generations(request:Request) -> Any:
-     """转发图片生成请求"""
-     return await _forward_request(request, "/images/generations")
+async def embeddings(request: Request):
+    """转发 embedding 请求."""
+    return await forward_request(request, "/embeddings")
 
 
+@app.post("/hf/v1/images/generations")
+async def images_generations(request: Request):
+     """转发图像生成请求."""
+     return await forward_request(request,"/images/generations",)