Update app.py

app.py

@@ -1,4 +1,6 @@
-from fastapi import FastAPI, Request, HTTPException
+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.middleware.cors import CORSMiddleware
 import httpx
@@ -12,9 +14,11 @@ 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")
-MODEL_MAP = {
+@dataclass
+class ApiConfig:
+    """API 配置数据类"""
+    base_url: str = os.environ.get("BASE_URL", "https://api.siliconflow.cn/v1")
+    model_map: Dict[str, str] = {
     "qwen-72b": "Qwen/Qwen2.5-72B-Instruct",
     "qwen-32b": "Qwen/Qwen2.5-32B-Instruct",
     "qwen-14b": "Qwen/Qwen2.5-14B-Instruct",
@@ -34,24 +38,20 @@ MODEL_MAP = {
     "flux-s":"black-forest-labs/FLUX.1-schnell",
     "flux-d":"black-forest-labs/FLUX.1-dev"
 }
-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")
+    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"))
 
-KEYS = KEY_STR.split(",")
-KEY_BALANCE = {}
-KEY_BALANCE_NOTES = ""
-# 创建一个东八区的时区对象
-TIMEZONE = pytz.timezone("Asia/Shanghai")
-LAST_UPDATE_TIME = ""
+config = ApiConfig()
 
 app = FastAPI()
-
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -62,16 +62,16 @@ app.add_middleware(
 
 # API Key 管理类
 class ApiKeyManager:
-    def __init__(self, keys):
+    def __init__(self, keys:list[str]):
         self.keys = keys
-        self.key_balance = {}
+        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=500, detail="No available API keys")
+            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="No available API keys")
         random.shuffle(self.keys)
         return self.keys[0]
     
@@ -104,7 +104,7 @@ class ApiKeyManager:
                 key_to_remove.append(key)
         for remove_key in key_to_remove:
             self.remove_key(remove_key)
-        self.last_update_time = datetime.now(TIMEZONE)
+        self.last_update_time = datetime.now(pytz.timezone(config.timezone))
 
     async def _fetch_balance(self, key: str) -> float:
         """发送 API 请求，获取 API Key 的余额"""
@@ -118,19 +118,21 @@ class ApiKeyManager:
                 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():
+                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 MODEL_MAP.items():
+    for key, value in config.model_map.items():
         models += f"<h2>{key}————{value}</h2>"
     return models
 
 @app.get("/", response_class=HTMLResponse)
-async def root():
+async def root() -> HTMLResponse:
     """根路由，返回 HTML 页面，展示模型信息和更新时间"""
     models_info = await get_model_info()
-    return f"""
+    return HTMLResponse(f"""
     <html>
         <head>
             <title>富文本示例</title>
@@ -142,98 +144,87 @@ async def root():
             {key_manager.key_balance_notes}
         </body>
     </html>
-    """
+    """)
 
 
 @app.get("/check")
-async def check():
+async def check() -> str:
     """手动触发检查 API Key 余额的路由"""
     await key_manager.check_keys_balance()
     return f"更新成功：{key_manager.last_update_time}"
 
-async def _forward_request(request: Request, api_type: str, image_generations:bool = False):
+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_type}")
+    logging.info(f"using key {key[0:4]}***{key[-4:]} to {api_endpoint}")
     headers = {"Authorization": f"Bearer {key}"}
-    if  api_type == "chat" and not image_generations :
+    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()
-            # 处理模型映射
-            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 def generate_response():
                    async with httpx.AsyncClient() as client:
-                        async with client.stream("POST", f"{BASE_URL}/chat/completions" , headers=headers, json=body) as response:
+                        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 :
-                async with httpx.AsyncClient() as client:
+            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(f"{BASE_URL}/chat/completions", headers=headers, json=body)
+                            response = await client.post(url, headers=headers, json=body)
                             response.raise_for_status()
                             return response.json()
                        except httpx.HTTPError as exc :
-                                logging.error("httpx request error:" + str(exc))
+                                logging.error(f"httpx request error:{exc}")
                                 raise HTTPException(
-                                        status_code=500,
+                                        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 image_generations:
-            url =f"{BASE_URL}/images/generations"
-            body = await request.json()
-            headers["Content-Type"] = "application/json"
-            async with httpx.AsyncClient() as client:
-                
-                 try:
-                         response = await client.post(url, headers=headers,json=body)
-                         response.raise_for_status()
-                         return JSONResponse(content=response.json(), status_code=response.status_code)
-                 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}"
-                           )
-    elif api_type =="embedding":
-          url=f"{BASE_URL}/embeddings"
-          body=await request.json()
-          if "model" in body and body["model"] in MODEL_MAP:
-              body["model"] = MODEL_MAP[body["model"]]
-          async with httpx.AsyncClient() as client:
-              try:
-                        response = await client.post(url, json=body,headers=headers)
-                        response.raise_for_status()
-                        return response.json()
-              except  httpx.HTTPError as exc :
-                           logging.error(f"httpx request embedding error :{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):
+async def chat_completions(request: Request) -> Any:
         """转发 chat 完成请求"""
-        return await _forward_request(request, "chat")
+        return await _forward_request(request, "/chat/completions",is_chat_stream_request(request))
 
 @app.post("/hf/v1/embeddings")
-async def embeddings(request: Request):
+async def embeddings(request: Request)-> Any:
     """转发 embeddings 请求"""
-    return await _forward_request(request, "embedding")
+    return await _forward_request(request, "/embeddings")
 
 @app.post("/hf/v1/images/generations")
-async def image_generations(request:Request):
+async def image_generations(request:Request) -> Any:
      """转发图片生成请求"""
-     return await _forward_request(request, api_type="image",image_generations=True)
+     return await _forward_request(request, "/images/generations")
+