tiphys.auth.profiles

 1"""
 2Auth profile management for LLM providers.
 3
 4Supports multiple API keys with rotation and cooldown tracking.
 5"""
 6
 7from dataclasses import dataclass, field
 8from datetime import datetime, timedelta
 9
10
11@dataclass
12class AuthProfile:
13    """A single authentication profile (API key)."""
14
15    id: str
16    provider: str
17    api_key: str
18    models: list[str] = field(default_factory=list)  # Models this key can access
19    priority: int = 0  # Lower = preferred
20
21    # State
22    cooldown_until: datetime | None = None
23    failure_count: int = 0
24    last_used: datetime | None = None
25
26    def is_available(self, now: datetime | None = None) -> bool:
27        """Check if the profile is currently available."""
28        if now is None:
29            now = datetime.utcnow()
30
31        return not (self.cooldown_until and self.cooldown_until > now)
32
33
34class AuthProfileManager:
35    """
36    Manages multiple auth profiles for rotation and failover.
37    """
38
39    def __init__(self) -> None:
40        """Initialize the auth profile manager."""
41        self._profiles: list[AuthProfile] = []
42
43    def add_profile(self, profile: AuthProfile) -> None:
44        """Add a new auth profile."""
45        self._profiles.append(profile)
46        # Sort by priority
47        self._profiles.sort(key=lambda p: p.priority)
48
49    def get_active(self, provider: str, model: str | None = None) -> AuthProfile | None:
50        """
51        Get an active auth profile for the provider and model.
52
53        Args:
54            provider: LLM provider (e.g., 'openai', 'anthropic').
55            model: Optional model identifier.
56
57        Returns:
58            The best available profile or None.
59        """
60        now = datetime.utcnow()
61
62        # Filter by provider and availability
63        candidates = [p for p in self._profiles if p.provider == provider and p.is_available(now)]
64
65        # Filter by model if provided
66        if model:
67            model_candidates = [p for p in candidates if not p.models or model in p.models]
68            if model_candidates:
69                candidates = model_candidates
70
71        if not candidates:
72            return None
73
74        # If multiple candidates with same priority, pick one randomly or use last_used
75        # For now, just pick the one that was used longest ago
76        candidates.sort(key=lambda p: (p.priority, p.last_used or datetime.min))
77
78        profile = candidates[0]
79        profile.last_used = now
80        return profile
81
82    def mark_rate_limited(self, profile_id: str, cooldown_seconds: int = 60) -> None:
83        """Mark a profile as rate limited."""
84        for p in self._profiles:
85            if p.id == profile_id:
86                p.cooldown_until = datetime.utcnow() + timedelta(seconds=cooldown_seconds)
87                p.failure_count += 1
88                break
89
90    def mark_failure(self, profile_id: str) -> None:
91        """Mark a profile as failed."""
92        for p in self._profiles:
93            if p.id == profile_id:
94                p.failure_count += 1
95                # Small cooldown on failure
96                p.cooldown_until = datetime.utcnow() + timedelta(seconds=10)
97                break