File size: 3,594 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """Shared helpers for API service handlers."""
from __future__ import annotations
import logging
from fastapi import HTTPException
from aistudio_api.api.response_models import (
HealthResponse,
ModelStatsResponse,
StatsResponse,
StatsTotalsResponse,
)
from aistudio_api.api.state import runtime_state
logger = logging.getLogger("aistudio.server")
MAX_RETRIES = 3
async def try_switch_account() -> bool:
"""尝试切换到下一个可用账号。返回是否成功切换。"""
rotator = runtime_state.rotator
if rotator is None:
return False
next_account = await rotator.get_next_account()
if next_account is None:
return False
account_service = runtime_state.account_service
client = runtime_state.client
if account_service is None or client is None or client._session is None:
return False
result = await account_service.activate_account(
next_account.id,
client._session,
runtime_state.snapshot_cache,
None, # skip lock — caller already holds it
keep_snapshot_cache=False,
)
return result is not None
def require_busy_lock():
busy_lock = runtime_state.busy_lock
if busy_lock is None:
raise HTTPException(
503, detail={"message": "Server not ready", "type": "service_unavailable"}
)
if busy_lock.locked():
raise HTTPException(
503,
detail={
"message": "Server is busy with maximum concurrent requests",
"type": "service_unavailable",
},
headers={"Retry-After": "2"},
)
return busy_lock
async def ensure_active_account(attempt: int) -> None:
if attempt != 0:
return
account_svc = runtime_state.account_service
if account_svc and not account_svc.get_active_account():
await try_switch_account()
def record_rotator_event(event: str) -> None:
rotator = runtime_state.rotator
account_service = runtime_state.account_service
account = account_service.get_active_account() if account_service else None
if not rotator or account is None:
return
if event == "success":
rotator.record_success(account.id)
elif event == "rate_limited":
rotator.record_rate_limited(account.id)
elif event == "error":
rotator.record_error(account.id)
def health_response() -> HealthResponse:
busy_lock = runtime_state.busy_lock
return HealthResponse(status="ok", busy=busy_lock.locked() if busy_lock else False)
def stats_response() -> StatsResponse:
stats = dict(runtime_state.model_stats)
totals = StatsTotalsResponse(
requests=sum(s.requests for s in stats.values()),
success=sum(s.success for s in stats.values()),
rate_limited=sum(s.rate_limited for s in stats.values()),
errors=sum(s.errors for s in stats.values()),
prompt_tokens=sum(s.prompt_tokens for s in stats.values()),
completion_tokens=sum(s.completion_tokens for s in stats.values()),
total_tokens=sum(s.total_tokens for s in stats.values()),
)
models = {
name: ModelStatsResponse(
requests=s.requests,
success=s.success,
rate_limited=s.rate_limited,
errors=s.errors,
prompt_tokens=s.prompt_tokens,
completion_tokens=s.completion_tokens,
total_tokens=s.total_tokens,
last_used=s.last_used,
)
for name, s in stats.items()
}
return StatsResponse(models=models, totals=totals)
|