OraCatQAQ commited on
Commit
58d32e7
·
1 Parent(s): 92a01b7

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.

Browse files
Files changed (1) hide show
  1. app.py +14 -25
app.py CHANGED
@@ -426,40 +426,29 @@ async def create_chat_completion(
426
  raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")
427
 
428
  @app.get("/admin/balance")
429
- async def get_account_balance(request: Request, admin_key: str = Header(None, alias="X-Admin-Key")):
430
  """查看账户余额"""
431
- # 简单的管理密钥检查
432
- expected_admin_key = os.getenv("ADMIN_KEY", "admin")
433
- if not admin_key or admin_key != expected_admin_key:
434
- raise HTTPException(status_code=403, detail="Unauthorized")
435
 
436
- # 从请求头中获取API密钥
437
- auth_header = request.headers.get("Authorization", "")
438
- if not auth_header or not auth_header.startswith("Bearer "):
439
- raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
440
-
441
- api_key = auth_header.replace("Bearer ", "")
442
  tokens = api_key.split(',')
443
 
444
- result = {}
 
 
 
445
 
446
- # 获取所有token的余额信息
447
  for i, token in enumerate(tokens):
448
- token_display = f"token_{i+1}"
449
  success, quota_info = await check_account_balance(api_key, i)
450
-
451
  if success:
452
- result[token_display] = {
453
- "status": "success",
454
- "quota": quota_info
455
- }
456
- else:
457
- result[token_display] = {
458
- "status": "error",
459
- "message": "无法获取账户余额信息"
460
- }
461
 
462
- return result
463
 
464
  # 错误处理器
465
  @app.exception_handler(404)
 
426
  raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")
427
 
428
  @app.get("/admin/balance")
429
+ async def get_account_balance():
430
  """查看账户余额"""
431
+ # 从环境变量获取API密钥
432
+ api_key = os.getenv("OPENAI_API_KEY", "")
433
+ if not api_key:
434
+ raise HTTPException(status_code=500, detail="未配置 OPENAI_API_KEY 环境变量")
435
 
 
 
 
 
 
 
436
  tokens = api_key.split(',')
437
 
438
+ total_quota = {
439
+ "total": 0,
440
+ "available": 0
441
+ }
442
 
443
+ # 获取所有token的余额信息并计算总和
444
  for i, token in enumerate(tokens):
 
445
  success, quota_info = await check_account_balance(api_key, i)
 
446
  if success:
447
+ for quota_type, info in quota_info.items():
448
+ total_quota["total"] += info.get("total", 0)
449
+ total_quota["available"] += info.get("available", 0)
 
 
 
 
 
 
450
 
451
+ return total_quota
452
 
453
  # 错误处理器
454
  @app.exception_handler(404)