Refactor get_account_balance function to remove admin key check and improve error handling for API key retrieval; return total quota information for all tokens.

app.py

@@ -426,40 +426,29 @@ async def create_chat_completion(
         raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")
 
 @app.get("/admin/balance")
-async def get_account_balance(request: Request, admin_key: str = Header(None, alias="X-Admin-Key")):
+async def get_account_balance():
     """查看账户余额"""
-    # 简单的管理密钥检查
-    expected_admin_key = os.getenv("ADMIN_KEY", "admin")
-    if not admin_key or admin_key != expected_admin_key:
-        raise HTTPException(status_code=403, detail="Unauthorized")
+    # 从环境变量获取API密钥
+    api_key = os.getenv("OPENAI_API_KEY", "")
+    if not api_key:
+        raise HTTPException(status_code=500, detail="未配置 OPENAI_API_KEY 环境变量")
     
-    # 从请求头中获取API密钥
-    auth_header = request.headers.get("Authorization", "")
-    if not auth_header or not auth_header.startswith("Bearer "):
-        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
-    
-    api_key = auth_header.replace("Bearer ", "")
     tokens = api_key.split(',')
     
-    result = {}
+    total_quota = {
+        "total": 0,
+        "available": 0
+    }
     
-    # 获取所有token的余额信息
+    # 获取所有token的余额信息并计算总和
     for i, token in enumerate(tokens):
-        token_display = f"token_{i+1}"
         success, quota_info = await check_account_balance(api_key, i)
-        
         if success:
-            result[token_display] = {
-                "status": "success",
-                "quota": quota_info
-            }
-        else:
-            result[token_display] = {
-                "status": "error",
-                "message": "无法获取账户余额信息"
-            }
+            for quota_type, info in quota_info.items():
+                total_quota["total"] += info.get("total", 0)
+                total_quota["available"] += info.get("available", 0)
     
-    return result
+    return total_quota
 
 # 错误处理器
 @app.exception_handler(404)