File size: 2,077 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 | """Shared API runtime state."""
from __future__ import annotations
import asyncio
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from aistudio_api.application.account_rotator import AccountRotator
from aistudio_api.application.account_service import AccountService
from aistudio_api.infrastructure.cache.snapshot_cache import SnapshotCache
from aistudio_api.infrastructure.gateway.client import AIStudioClient
@dataclass
class ModelStatsItem:
requests: int = 0
success: int = 0
rate_limited: int = 0
errors: int = 0
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
last_used: str | None = None
@dataclass
class RuntimeState:
client: AIStudioClient | None = None
busy_lock: asyncio.Semaphore | None = None
browser_port: int = 9222
snapshot_cache: SnapshotCache | None = None
account_service: AccountService | None = None
rotator: AccountRotator | None = None
model_stats: dict[str, ModelStatsItem] = field(
default_factory=lambda: defaultdict(ModelStatsItem)
)
def record(self, model: str, event: str, usage: dict[str, object] | None = None) -> None:
stats = self.model_stats[model]
stats.requests += 1
if event == "success":
stats.success += 1
elif event == "rate_limited":
stats.rate_limited += 1
elif event == "errors":
stats.errors += 1
stats.last_used = datetime.now(timezone(timedelta(hours=8))).isoformat()
if usage and event == "success":
pt = usage.get("prompt_tokens", 0)
ct = usage.get("completion_tokens", 0)
tt = usage.get("total_tokens", 0)
stats.prompt_tokens += pt if isinstance(pt, int) else 0
stats.completion_tokens += ct if isinstance(ct, int) else 0
stats.total_tokens += tt if isinstance(tt, int) else 0
runtime_state = RuntimeState()
|