""" API Gateway Models - Groq/OpenRouter-style account and chat management """ from dataclasses import dataclass, field from datetime import datetime, timedelta from typing import List, Optional, Dict from enum import Enum import secrets import hashlib class AccountStatus(str, Enum): ACTIVE = "active" SUSPENDED = "suspended" TRIAL = "trial" class Role(str, Enum): USER = "user" ADMIN = "admin" @dataclass class APIKey: key_id: str key_hash: str account_id: str name: str created_at: datetime = field(default_factory=datetime.utcnow) last_used: Optional[datetime] = None expires_at: Optional[datetime] = None is_active: bool = True rate_limit_per_minute: int = 60 rate_limit_per_day: int = 1000 @staticmethod def generate_key() -> str: """Generate a secure API key.""" return f"m5_{secrets.token_urlsafe(32)}" @staticmethod def hash_key(key: str) -> str: """Hash an API key for storage.""" return hashlib.sha256(key.encode()).hexdigest() @dataclass class ChatMessage: role: str # "user", "assistant", "system" content: str timestamp: datetime = field(default_factory=datetime.utcnow) tokens_used: Optional[int] = None model: Optional[str] = None latency_ms: Optional[int] = None @dataclass class ChatSession: session_id: str account_id: str title: str messages: List[ChatMessage] = field(default_factory=list) created_at: datetime = field(default_factory=datetime.utcnow) updated_at: datetime = field(default_factory=datetime.utcnow) model: Optional[str] = None total_tokens: int = 0 total_cost_usd: float = 0.0 def add_message(self, message: ChatMessage): self.messages.append(message) self.updated_at = datetime.utcnow() if message.tokens_used: self.total_tokens += message.tokens_used @dataclass class Account: account_id: str email: str password_hash: str name: str status: AccountStatus = AccountStatus.TRIAL role: Role = Role.USER created_at: datetime = field(default_factory=datetime.utcnow) api_keys: List[APIKey] = field(default_factory=list) chat_sessions: List[ChatSession] = field(default_factory=list) balance_usd: float = 10.0 # Free trial credits total_requests: int = 0 total_tokens: int = 0 @staticmethod def hash_password(password: str) -> str: """Hash a password for storage.""" return hashlib.sha256(password.encode()).hexdigest() def verify_password(self, password: str) -> bool: """Verify a password against the hash.""" return self.password_hash == self.hash_password(password) def add_api_key(self, name: str, expires_days: Optional[int] = None) -> APIKey: """Add a new API key to the account.""" key = APIKey.generate_key() key_hash = APIKey.hash_key(key) expires_at = datetime.utcnow() + timedelta(days=expires_days) if expires_days else None api_key = APIKey( key_id=secrets.token_hex(8), key_hash=key_hash, account_id=self.account_id, name=name, expires_at=expires_at ) self.api_keys.append(api_key) return api_key def revoke_api_key(self, key_id: str) -> bool: """Revoke an API key.""" for key in self.api_keys: if key.key_id == key_id: key.is_active = False return True return False def create_chat_session(self, title: str) -> ChatSession: """Create a new chat session.""" session = ChatSession( session_id=secrets.token_hex(16), account_id=self.account_id, title=title ) self.chat_sessions.append(session) return session def deduct_balance(self, amount_usd: float) -> bool: """Deduct balance from account.""" if self.balance_usd >= amount_usd: self.balance_usd -= amount_usd return True return False @dataclass class UsageMetrics: account_id: str date: datetime requests_count: int = 0 tokens_count: int = 0 cost_usd: float = 0.0 model_usage: Dict[str, int] = field(default_factory=dict) def record_request(self, tokens: int, model: str, cost_usd: float): self.requests_count += 1 self.tokens_count += tokens self.cost_usd += cost_usd self.model_usage[model] = self.model_usage.get(model, 0) + tokens # In-memory storage (replace with database in production) _accounts: Dict[str, Account] = {} _api_keys: Dict[str, APIKey] = {} # key_hash -> APIKey _rate_limits: Dict[str, List[datetime]] = {} # key_id -> request timestamps class AccountManager: """Manage accounts and API keys.""" @staticmethod def create_account(email: str, password: str, name: str) -> Account: """Create a new account.""" account_id = secrets.token_hex(16) account = Account( account_id=account_id, email=email, password_hash=Account.hash_password(password), name=name ) _accounts[account_id] = account return account @staticmethod def get_account_by_email(email: str) -> Optional[Account]: """Get account by email.""" for account in _accounts.values(): if account.email == email: return account return None @staticmethod def get_account_by_id(account_id: str) -> Optional[Account]: """Get account by ID.""" return _accounts.get(account_id) @staticmethod def verify_api_key(api_key: str) -> Optional[Account]: """Verify an API key and return the associated account.""" key_hash = APIKey.hash_key(api_key) for account in _accounts.values(): for key in account.api_keys: if key.key_hash == key_hash and key.is_active: if key.expires_at and datetime.utcnow() > key.expires_at: key.is_active = False return None key.last_used = datetime.utcnow() return account return None @staticmethod def check_rate_limit(api_key: str) -> bool: """Check if API key is within rate limits.""" key_hash = APIKey.hash_key(api_key) for account in _accounts.values(): for key in account.api_keys: if key.key_hash == key_hash: now = datetime.utcnow() minute_ago = now - timedelta(minutes=1) day_ago = now - timedelta(days=1) if key.key_id not in _rate_limits: _rate_limits[key.key_id] = [] # Clean old timestamps _rate_limits[key.key_id] = [ ts for ts in _rate_limits[key.key_id] if ts > day_ago ] # Check limits recent_minute = sum(1 for ts in _rate_limits[key.key_id] if ts > minute_ago) recent_day = len(_rate_limits[key.key_id]) if recent_minute >= key.rate_limit_per_minute: return False if recent_day >= key.rate_limit_per_day: return False # Record this request _rate_limits[key.key_id].append(now) return True return False @staticmethod def get_all_accounts() -> List[Account]: """Get all accounts (admin only).""" return list(_accounts.values())