Update app.py

app.py

@@ -1,137 +1,192 @@
-from fastapi import FastAPI,Request,HTTPException
-from fastapi.responses import StreamingResponse
-from fastapi.responses import HTMLResponse
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 import httpx
-import json,os,random
+import json
+import os
+import random
 from datetime import datetime
 import pytz
-def get_si_key():
-    random.shuffle(keys)
-    return keys[0]
+import logging
+
+# 配置日志
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# 从环境变量或者文件中加载配置
+BASE_URL = os.environ.get("BASE_URL", "https://api.siliconflow.cn/v1/chat/completions")
+MODEL_MAP = {
+    "qwen-72b": "Qwen/Qwen2.5-72B-Instruct",
+    "qwen-32b": "Qwen/Qwen2.5-32B-Instruct",
+    "qwen-14b": "Qwen/Qwen2.5-14B-Instruct",
+    "qwen-7b": "Qwen/Qwen2.5-7B-Instruct",
+    "qwen-vl": "Qwen/Qwen2-VL-72B-Instruct",
+    "qwen-coder": "Qwen/Qwen2.5-Coder-32B-Instruct",
+    "qwq": "Qwen/QwQ-32B-Preview",
+    "o1": "AIDC-AI/Marco-o1",
+    "deepseek": "deepseek-ai/DeepSeek-V2.5",
+    "deepseek-vl": "deepseek-ai/deepseek-vl2",
+    "glm-9b": "THUDM/glm-4-9b-chat",
+    "bce": "netease-youdao/bce-embedding-base_v1",
+    "bge-m3": "BAAI/bge-m3",
+    "bge-zh": "BAAI/bge-large-zh-v1.5",
+}
+if os.environ.get("MODEL_MAP"):
+    MODEL_MAP = json.loads(os.environ.get("MODEL_MAP"))
+
+KEY_STR = os.environ.get("SI_KEY")
+if KEY_STR is None:
+    logging.error("SI_KEY not found in env")
+    raise EnvironmentError("SI_KEY not found in env")
+
+KEYS = KEY_STR.split(",")
+KEY_BALANCE = {}
+KEY_BALANCE_NOTES = ""
+# 创建一个东八区的时区对象
+TIMEZONE = pytz.timezone("Asia/Shanghai")
+LAST_UPDATE_TIME = ""
+
 app = FastAPI()
+
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # 允许的源列表
-    allow_credentials=True,  # 允许发送 cookies
-    allow_methods=["*"],  # 允许所有方法
-    allow_headers=["*"],  # 允许所有头部
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
-base_url = "https://api.siliconflow.cn/v1/chat/completions"
-if os.environ.get("MODEL_MAP"): 
-    model_map=json.loads(os.environ.get("MODEL_MAP"))
-else:
-    model_map={
-        "qwen-72b":"Qwen/Qwen2.5-72B-Instruct",
-        "qwen-32b":"Qwen/Qwen2.5-32B-Instruct",
-        "qwen-14b":"Qwen/Qwen2.5-14B-Instruct",
-        "qwen-7b":"Qwen/Qwen2.5-7B-Instruct",
-        "qwen-vl":"Qwen/Qwen2-VL-72B-Instruct",
-        "qwen-coder":"Qwen/Qwen2.5-Coder-32B-Instruct",
-        "qwq":"Qwen/QwQ-32B-Preview",
-        "o1":"AIDC-AI/Marco-o1",
-        "deepseek":"deepseek-ai/DeepSeek-V2.5",
-        "deepseek-vl":"deepseek-ai/deepseek-vl2",
-        "glm-9b":"THUDM/glm-4-9b-chat",
-        "bce":"netease-youdao/bce-embedding-base_v1",
-        "bge-m3":"BAAI/bge-m3",
-        "bge-zh":"BAAI/bge-large-zh-v1.5"
-    }
-if os.environ.get("SI_KEY"):
-    keys=os.environ.get("SI_KEY").split(",")
-else:
-    keys=["sk-jopguysfqvrlciqjvlydkbxtynqagxqdhmrprfjivupuutfk","sk-bvkqkygxqrusyhdfrqoyoctqfhnxylpuoajgfzbtkhiecffo","sk-eeoefacvytttokuwhslnvjazspigjhrdkhxuzdrxrizrpeep"]
-key_balacnce={}
-key_balacnce_notes=""
-# 创建一个东八区的时区对象
-tz = pytz.timezone('Asia/Shanghai')
 
-now=""
-@app.get("/",response_class=HTMLResponse)
+# API Key 管理类
+class ApiKeyManager:
+    def __init__(self, keys):
+        self.keys = keys
+        self.key_balance = {}
+        self.key_balance_notes = ""
+        self.last_update_time = ""
+
+    def get_key(self) -> str:
+        """随机获取一个可用 API Key"""
+        if not self.keys:
+            raise HTTPException(status_code=500, detail="No available API keys")
+        random.shuffle(self.keys)
+        return self.keys[0]
+    
+    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(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()}"}
+        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=500, detail=f"Check balance failed with status:{exc.response.status_code},url:{exc.request.url}")
+key_manager = ApiKeyManager(KEYS)
+async def get_model_info():
+    models = ""
+    for key, value in MODEL_MAP.items():
+        models += f"<h2>{key}————{value}</h2>"
+    return models
+
+
+@app.get("/", response_class=HTMLResponse)
 async def root():
-    # 获取当前时间，并设置时区为东八区
-    models=""
-    for key in model_map:
-        models+=f"<h2>{key}————{model_map[key]}</h2>"
-    global now
-    return  f"""
+    """根路由，返回 HTML 页面，展示模型信息和更新时间"""
+    models_info = await get_model_info()
+    return f"""
     <html>
         <head>
             <title>富文本示例</title>
         </head>
         <body>
-            <h1>有效key数量:{len(keys)}</h1>
-            {models}
-            <h1>最后更新时间：{now}</h1>
-            {key_balacnce_notes}
+            <h1>有效key数量:{len(key_manager.keys)}</h1>
+            {models_info}
+            <h1>最后更新时间：{key_manager.last_update_time}</h1>
+            {key_manager.key_balance_notes}
         </body>
     </html>
     """
+
 @app.get("/check")
 async def check():
-    global key_balacnce,key_balacnce_notes,now,keys
-    key_balacnce_notes=""
-    for i,key in enumerate(keys):
-        url = "https://api.siliconflow.cn/v1/user/info"
-        headers={
-            "Authorization":f"Bearer {key.strip()}"
-        }
-        async with httpx.AsyncClient() as client:
-            res=await client.get(url,headers=headers)
-            if res.status_code==200:
-                balance=res.json()['data']['balance']
-                if float(balance)<0.1:
-                    keys.pop(i)
-                    continue
-                va=f'''<h2>{key.strip()[0:4]}****{key.strip()[-4:]}————{balance}</h2>'''
-                key_balacnce[key.strip()]=balance
-                key_balacnce_notes+=va
-    
-    now = datetime.now(tz)
-    
-    return f"更新成功：{now}"            
-@app.post("/hf/v1/chat/completions")
-async def reforword(request:Request):
+    """手动触发检查 API Key 余额的路由"""
+    await key_manager.check_keys_balance()
+    return f"更新成功：{key_manager.last_update_time}"
+
+async def _forward_request(request: Request, api_type: str):
+    """转发请求到硅流 API"""
     body = await request.json()
-    # 获取 API 密钥
-    key = get_si_key()
-    print(key)
-    headers = {
-            "Authorization": f"Bearer {key}"
-        }
-    body_map=model_map
-    if "model" in body.keys() and body['model'] in body_map.keys():
-        body['model']=body_map[body['model']]
-    if "stream" in body.keys() and body['stream']:    
-        # 使用 httpx 发送流式请求
+    try:
+          key = key_manager.get_key()
+    except HTTPException as e:
+        return e
+    logging.info(f"using key {key[0:4]}***{key[-4:]} to {api_type}")
+    headers = {"Authorization": f"Bearer {key}"}
+    # 处理模型映射
+    if "model" in body and body["model"] in MODEL_MAP:
+        body["model"] = MODEL_MAP[body["model"]]
+
+    if api_type == "chat" and "stream" in body and body["stream"]:
         async def generate_response():
             async with httpx.AsyncClient() as client:
-                async with client.stream("POST", base_url, headers=headers, json=body) as response:
-                    response.raise_for_status()  # 检查响应状态码
-                    async for chunk in response.aiter_bytes():
-                        if chunk:
-                            yield chunk
-
+                async with client.stream("POST", BASE_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:
-        # 发送 POST 请求
         async with httpx.AsyncClient() as client:
-            response = await client.post(base_url, headers=headers, json=body)
-            response.raise_for_status()
-            return response.json()
+            try:
+                response = await client.post(BASE_URL, headers=headers, json=body)
+                response.raise_for_status()
+                return response.json()
+            except httpx.HTTPError as exc :
+                 logging.error("httpx request error:" + str(exc))
+                 raise HTTPException(
+                     status_code=500,
+                     detail=f"Request failed with status: {exc.response.status_code},url:{exc.request.url},detail:{exc.response.text}"
+                 )
+
+@app.post("/hf/v1/chat/completions")
+async def chat_completions(request: Request):
+        """转发 chat 完成请求"""
+        return await _forward_request(request, "chat")
 @app.post("/hf/v1/embeddings")
-async def embedding(request:Request):
-    body=await request.json()
-    body_map=model_map
-    if "model" in body.keys() and body['model'] in body_map.keys():
-        body['model']=body_map[body['model']]
-    # 获取 API 密钥
-    key = get_si_key()
-    print(key)
-    headers = {
-        "Authorization": f"Bearer {key}"
-    }
-    async with httpx.AsyncClient() as client:
-        response = await client.post(base_url, json=body, headers=headers)
-        response.raise_for_status()  # 检查请求是否成功
-        return response.json()
+async def embeddings(request: Request):
+    """转发 embeddings 请求"""
+    return await _forward_request(request, "embedding")
+