Spaces:
Running
Running
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
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(
|
430 |
"""查看账户余额"""
|
431 |
-
#
|
432 |
-
|
433 |
-
if not
|
434 |
-
raise HTTPException(status_code=
|
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 |
-
|
|
|
|
|
|
|
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 |
-
|
453 |
-
"
|
454 |
-
"
|
455 |
-
}
|
456 |
-
else:
|
457 |
-
result[token_display] = {
|
458 |
-
"status": "error",
|
459 |
-
"message": "无法获取账户余额信息"
|
460 |
-
}
|
461 |
|
462 |
-
return
|
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)
|