| """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, |
| 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) |
|
|
