keycycle 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,22 @@
1
+ Metadata-Version: 2.4
2
+ Name: keycycle
3
+ Version: 0.1.0
4
+ Summary: A thread-safe key rotation and rate limiting manager for API keys.
5
+ Author-email: William Jin <jinwilliam.jin@gmail.com>
6
+ Classifier: Programming Language :: Python :: 3
7
+ Classifier: License :: OSI Approved :: MIT License
8
+ Classifier: Operating System :: OS Independent
9
+ Classifier: Development Status :: 3 - Alpha
10
+ Requires-Python: >=3.9
11
+ Description-Content-Type: text/markdown
12
+ Requires-Dist: python-dotenv
13
+ Requires-Dist: sqlalchemy
14
+ Requires-Dist: pymysql
15
+ Requires-Dist: rich
16
+ Provides-Extra: openai
17
+ Requires-Dist: openai; extra == "openai"
18
+ Provides-Extra: agno
19
+ Requires-Dist: agno; extra == "agno"
20
+ Provides-Extra: all
21
+ Requires-Dist: openai; extra == "all"
22
+ Requires-Dist: agno; extra == "all"
@@ -0,0 +1,8 @@
1
+ from .multi_provider_wrapper import RateLimits, MultiProviderWrapper, RotatingAsyncOpenAIClient, RotatingOpenAIClient
2
+ __all__ = [
3
+ "RateLimits",
4
+ "RotatingKeyManager",
5
+ "MultiProviderWrapper",
6
+ "RotatingAsyncOpenAIClient",
7
+ "RotatingOpenAIClient",
8
+ ]
@@ -0,0 +1,3 @@
1
+ from .openai_adapter import RotatingOpenAIClient, RotatingAsyncOpenAIClient
2
+
3
+ __all__ = ["RotatingOpenAIClient", "RotatingAsyncOpenAIClient"]
@@ -0,0 +1,254 @@
1
+ import logging
2
+ import time
3
+ from typing import List, Any, Optional, Union, Callable, Generator, AsyncGenerator
4
+
5
+ from ..config.dataclasses import KeyUsage, RateLimits
6
+ from ..key_rotation.rotation_manager import RotatingKeyManager
7
+
8
+ logger = logging.getLogger(__name__)
9
+
10
+ try:
11
+ import openai
12
+ from openai import OpenAI, AsyncOpenAI, RateLimitError, APIError
13
+ HAS_OPENAI = True
14
+ except ImportError:
15
+ HAS_OPENAI = False
16
+ # Generic placeholders
17
+ OpenAI = object
18
+ AsyncOpenAI = object
19
+ RateLimitError = Exception
20
+ APIError = Exception
21
+
22
+ PROVIDER_BASE_URLS = {
23
+ "openai": "https://api.openai.com/v1",
24
+ "openrouter": "https://openrouter.ai/api/v1",
25
+ "gemini": "https://generativelanguage.googleapis.com/v1beta/openai/",
26
+ "cerebras": "https://api.cerebras.ai/v1",
27
+ "groq": "https://api.groq.com/openai/v1",
28
+ "cohere": "https://api.cohere.ai/compatibility/v1",
29
+
30
+ }
31
+
32
+ class BaseRotatingClient:
33
+ def __init__(self,
34
+ manager: RotatingKeyManager,
35
+ limit_resolver: Callable[[str], RateLimits],
36
+ default_model: str,
37
+ estimated_tokens: int = 1000,
38
+ max_retries: int = 5,
39
+ base_url: Optional[str] = None,
40
+ provider: Optional[str] = None,
41
+ client_kwargs: dict = None
42
+ ):
43
+ """
44
+ Initialize the rotating client.
45
+
46
+ Args:
47
+ manager: The key rotation manager
48
+ limit_resolver: Function to resolve rate limits for a model
49
+ default_model: Default model to use
50
+ estimated_tokens: Estimated tokens per request
51
+ max_retries: Maximum number of retries on rate limit
52
+ base_url: Base URL for the API (takes precedence over provider)
53
+ provider: Provider name (openai, openrouter, gemini, cerebras, groq)
54
+ client_kwargs: Additional kwargs to pass to OpenAI client
55
+ """
56
+
57
+ if not HAS_OPENAI:
58
+ raise ImportError("The 'openai' library is required. Install with `pip install openai`.")
59
+
60
+ self.manager = manager
61
+ self.limit_resolver = limit_resolver
62
+ self.default_model = default_model
63
+ self.estimated_tokens = estimated_tokens
64
+ self.max_retries = max_retries
65
+ self.client_kwargs = client_kwargs or {}
66
+
67
+ if base_url:
68
+ self.base_url = base_url
69
+ elif provider:
70
+ provider_lower = provider.lower()
71
+ if provider_lower not in PROVIDER_BASE_URLS:
72
+ raise ValueError(
73
+ f"Unknown provider: {provider}. "
74
+ f"Valid providers: {', '.join(PROVIDER_BASE_URLS.keys())}"
75
+ )
76
+ self.base_url = PROVIDER_BASE_URLS[provider_lower]
77
+ else:
78
+ self.base_url = None # Will use OpenAI's default
79
+
80
+ if self.base_url:
81
+ self.client_kwargs['base_url'] = self.base_url
82
+
83
+ def _is_rate_limit(self, e: Exception) -> bool:
84
+ if isinstance(e, RateLimitError):
85
+ return True
86
+ err_str = str(e).lower()
87
+ return "429" in err_str or "too many requests" in err_str or "rate limit" in err_str or "resource exhausted" in err_str
88
+
89
+ def _record_usage(self, key_usage: KeyUsage, model_id: str, actual_tokens: int):
90
+ self.manager.record_usage(
91
+ key_obj=key_usage,
92
+ model_id=model_id,
93
+ actual_tokens=actual_tokens,
94
+ estimated_tokens=self.estimated_tokens
95
+ )
96
+
97
+ def _extract_usage(self, response: Any) -> int:
98
+ try:
99
+ if hasattr(response, 'usage') and response.usage:
100
+ return response.usage.total_tokens
101
+ except:
102
+ pass
103
+ return 0
104
+
105
+ # --- SYNC IMPLEMENTATION ---
106
+
107
+ class RotatingOpenAIClient(BaseRotatingClient):
108
+ def _get_fresh_client(self, api_key: str):
109
+ return OpenAI(api_key=api_key, **self.client_kwargs)
110
+
111
+ def __getattr__(self, name):
112
+ return SyncProxyHelper(self, [name])
113
+
114
+ def _execute(self, path: List[str], args, kwargs):
115
+ model_id = kwargs.get('model', self.default_model)
116
+ if 'model' not in kwargs:
117
+ kwargs['model'] = model_id
118
+ limits = self.limit_resolver(model_id)
119
+
120
+ for attempt in range(self.max_retries + 1):
121
+ key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
122
+ if not key_usage:
123
+ raise RuntimeError(f"No available keys for {model_id}")
124
+
125
+ try:
126
+ real_client = self._get_fresh_client(key_usage.api_key)
127
+
128
+ target = real_client
129
+ for p in path:
130
+ target = getattr(target, p)
131
+
132
+ result = target(*args, **kwargs)
133
+
134
+ if kwargs.get('stream', False):
135
+ return self._wrap_stream(result, key_usage, model_id)
136
+
137
+ self._record_usage(key_usage, model_id, self._extract_usage(result))
138
+ return result
139
+
140
+ except Exception as e:
141
+ if self._is_rate_limit(e) and attempt < self.max_retries:
142
+ logger.warning(f"429/RateLimit hit for {model_id} on key ...{key_usage.api_key[-8:]}. Rotating. (Attempt {attempt + 1}/{self.max_retries + 1})")
143
+ key_usage.trigger_cooldown()
144
+ self.manager.force_rotate_index()
145
+ time.sleep(0.5)
146
+ continue
147
+ self._record_usage(key_usage, model_id, 0)
148
+ raise e
149
+
150
+ def _wrap_stream(self, generator: Generator, key_usage: KeyUsage, model_id: str):
151
+ accumulated_tokens = 0
152
+ try:
153
+ for chunk in generator:
154
+ if hasattr(chunk, 'usage') and chunk.usage:
155
+ accumulated_tokens = chunk.usage.total_tokens
156
+ yield chunk
157
+ except Exception as e:
158
+ if self._is_rate_limit(e):
159
+ logger.warning(f"Rate limit hit during streaming for {model_id} on key ...{key_usage.api_key[-8:]}.")
160
+ key_usage.trigger_cooldown()
161
+ self.manager.force_rotate_index()
162
+ raise
163
+ finally:
164
+ self._record_usage(key_usage, model_id, accumulated_tokens)
165
+
166
+ class SyncProxyHelper:
167
+ def __init__(self, client: RotatingOpenAIClient, path: List[str]):
168
+ self.client = client
169
+ self.path = path
170
+
171
+ def __getattr__(self, name):
172
+ return SyncProxyHelper(self.client, self.path + [name])
173
+
174
+ def __call__(self, *args, **kwargs):
175
+ return self.client._execute(self.path, args, kwargs)
176
+
177
+
178
+ # --- ASYNC IMPLEMENTATION ---
179
+
180
+ class RotatingAsyncOpenAIClient(BaseRotatingClient):
181
+ def _get_fresh_client(self, api_key: str):
182
+ return AsyncOpenAI(api_key=api_key, **self.client_kwargs)
183
+
184
+ def __getattr__(self, name):
185
+ return AsyncProxyHelper(self, [name])
186
+
187
+ async def _execute(self, path: List[str], args, kwargs: dict):
188
+ model_id = kwargs.get('model', self.default_model)
189
+ if 'model' not in kwargs:
190
+ kwargs['model'] = model_id
191
+ limits = self.limit_resolver(model_id)
192
+
193
+ if kwargs.get('stream', False) and 'stream_options' not in kwargs:
194
+ kwargs['stream_options'] = {"include_usage": True}
195
+
196
+ if kwargs.get('stream', False) and 'stream_options' not in kwargs:
197
+ kwargs['stream_options'] = {"include_usage": True}
198
+
199
+ for attempt in range(self.max_retries + 1):
200
+ key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
201
+ if not key_usage:
202
+ raise RuntimeError(f"No available keys for {model_id}")
203
+
204
+ try:
205
+ real_client = self._get_fresh_client(key_usage.api_key)
206
+
207
+ target = real_client
208
+ for p in path:
209
+ target = getattr(target, p)
210
+
211
+ result = await target(*args, **kwargs)
212
+
213
+ if kwargs.get('stream', False):
214
+ return self._wrap_stream(result, key_usage, model_id)
215
+
216
+ self._record_usage(key_usage, model_id, self._extract_usage(result))
217
+ return result
218
+
219
+ except Exception as e:
220
+ if self._is_rate_limit(e) and attempt < self.max_retries:
221
+ logger.warning(f"429/RateLimit hit for {model_id} on key ...{key_usage.api_key[-8:]}. Rotating. (Attempt {attempt + 1}/{self.max_retries + 1})")
222
+ key_usage.trigger_cooldown()
223
+ self.manager.force_rotate_index()
224
+ time.sleep(0.5)
225
+ continue
226
+ self._record_usage(key_usage, model_id, 0)
227
+ raise e
228
+
229
+ async def _wrap_stream(self, generator: AsyncGenerator, key_usage: KeyUsage, model_id: str):
230
+ accumulated_tokens = 0
231
+ try:
232
+ async for chunk in generator:
233
+ if hasattr(chunk, 'usage') and chunk.usage:
234
+ accumulated_tokens = chunk.usage.total_tokens
235
+ yield chunk
236
+ except Exception as e:
237
+ if self._is_rate_limit(e):
238
+ logger.warning(f"Rate limit hit during streaming for {model_id} on key ...{key_usage.api_key[-8:]}.")
239
+ key_usage.trigger_cooldown()
240
+ self.manager.force_rotate_index()
241
+ raise
242
+ finally:
243
+ self._record_usage(key_usage, model_id, accumulated_tokens)
244
+
245
+ class AsyncProxyHelper:
246
+ def __init__(self, client: RotatingAsyncOpenAIClient, path: List[str]):
247
+ self.client = client
248
+ self.path = path
249
+
250
+ def __getattr__(self, name):
251
+ return AsyncProxyHelper(self.client, self.path + [name])
252
+
253
+ async def __call__(self, *args, **kwargs):
254
+ return await self.client._execute(self.path, args, kwargs)
@@ -0,0 +1,25 @@
1
+ from .dataclasses import (
2
+ RateLimits,
3
+ UsageSnapshot,
4
+ UsageBucket,
5
+ KeyUsage,
6
+ KeyDetailedStats,
7
+ KeySummary,
8
+ GlobalStats,
9
+ ModelAggregatedStats,
10
+ )
11
+ from .enums import RateLimitStrategy
12
+ from .log_config import configure_logging
13
+
14
+ __all__ = [
15
+ "RateLimits",
16
+ "UsageSnapshot",
17
+ "UsageBucket",
18
+ "KeyUsage",
19
+ "KeyDetailedStats",
20
+ "KeySummary",
21
+ "GlobalStats",
22
+ "ModelAggregatedStats",
23
+ "RateLimitStrategy",
24
+ "configure_logging",
25
+ ]
@@ -0,0 +1,56 @@
1
+ from .dataclasses import RateLimits
2
+ from .enums import RateLimitStrategy
3
+
4
+ PROVIDER_STRATEGIES = {
5
+ 'cerebras': RateLimitStrategy.PER_MODEL,
6
+ 'groq': RateLimitStrategy.PER_MODEL,
7
+ 'gemini': RateLimitStrategy.PER_MODEL,
8
+ 'openrouter': RateLimitStrategy.GLOBAL,
9
+ }
10
+
11
+ MODEL_LIMITS = {
12
+ 'cerebras': {
13
+ 'gpt-oss-120b': RateLimits(30, 900, 14400, 60000, 1000000, 1000000),
14
+ 'llama3.1-8b': RateLimits(30, 900, 14400, 60000, 1000000, 1000000),
15
+ 'llama-3.3-70b': RateLimits(30, 900, 14400, 60000, 1000000, 1000000),
16
+ 'qwen-3-32b': RateLimits(30, 900, 14400, 60000, 1000000, 1000000),
17
+ 'qwen-3-235b-a22b-instruct-2507': RateLimits(30, 900, 14400, 60000, 1000000, 1000000),
18
+ 'zai-glm-4.6': RateLimits(10, 100, 100, 150000, 1000000, 1000000),
19
+ },
20
+ 'groq': {
21
+ 'allam-2-7b': RateLimits(30, 1800, 7000, 6000, 360000, 500000),
22
+ 'groq/compound': RateLimits(30, 250, 250, 70000, None, None),
23
+ 'groq/compound-mini': RateLimits(30, 250, 250, 70000, None, None),
24
+ 'llama-3.1-8b-instant': RateLimits(30, 1800, 14400, 6000, 360000, 500000),
25
+ 'llama-3.3-70b-versatile': RateLimits(30, 1000, 1000, 12000, 720000, 100000),
26
+ 'meta-llama/llama-4-maverick-17b-128e-instruct': RateLimits(30, 1000, 1000, 6000, 360000, 500000),
27
+ 'meta-llama/llama-4-scout-17b-16e-instruct': RateLimits(30, 1000, 1000, 30000, 1800000, 500000),
28
+ 'meta-llama/llama-guard-4-12b': RateLimits(30, 1800, 14400, 15000, 900000, 500000),
29
+ 'meta-llama/llama-prompt-guard-2-22m': RateLimits(30, 1800, 14400, 15000, 900000, 500000),
30
+ 'meta-llama/llama-prompt-guard-2-86m': RateLimits(30, 1800, 14400, 15000, 900000, 500000),
31
+ 'moonshotai/kimi-k2-instruct': RateLimits(60, 1000, 1000, 10000, 600000, 300000),
32
+ 'moonshotai/kimi-k2-instruct-0905': RateLimits(60, 1000, 1000, 10000, 600000, 300000),
33
+ 'openai/gpt-oss-120b': RateLimits(30, 1000, 1000, 8000, 480000, 200000),
34
+ 'openai/gpt-oss-20b': RateLimits(30, 1000, 1000, 8000, 480000, 200000),
35
+ 'openai/gpt-oss-safeguard-20b': RateLimits(30, 1000, 1000, 8000, 480000, 200000),
36
+ 'playai-tts': RateLimits(10, 100, 100, 1200, 72000, 3600),
37
+ 'playai-tts-arabic': RateLimits(10, 100, 100, 1200, 72000, 3600),
38
+ 'qwen/qwen3-32b': RateLimits(60, 1000, 1000, 6000, 360000, 500000),
39
+ 'whisper-large-v3': RateLimits(20, 2000, 2000),
40
+ 'whisper-large-v3-turbo': RateLimits(20, 2000, 2000),
41
+ },
42
+ 'gemini': {
43
+ 'gemini-2.5-flash': RateLimits(5, 300, 20, 250000, 15000000),
44
+ 'gemini-2.5-flash-lite': RateLimits(10, 600, 20, 250000, 15000000),
45
+ 'gemini-2.5-flash-tts': RateLimits(3, 180, 10, 10000, 600000),
46
+ 'gemini-robotics-er-1.5-preview': RateLimits(10, 600, 250, 250000, 15000000),
47
+ 'gemma-3-12b': RateLimits(30, 1800, 14400, 15000, 900000),
48
+ 'gemma-3-1b': RateLimits(30, 1800, 14400, 15000, 900000),
49
+ 'gemma-3-27b': RateLimits(30, 1800, 14400, 15000, 900000),
50
+ 'gemma-3-2b': RateLimits(30, 1800, 14400, 15000, 900000),
51
+ 'gemma-3-4b': RateLimits(30, 1800, 14400, 15000, 900000),
52
+ },
53
+ 'openrouter': {
54
+ 'default': RateLimits(20, 50, 50),
55
+ },
56
+ }
@@ -0,0 +1,190 @@
1
+ from typing import Dict, List, Optional
2
+ from .enums import RateLimitStrategy
3
+ import time
4
+ from collections import deque, defaultdict
5
+ from dataclasses import dataclass, field
6
+
7
+ # --- CONFIGURATION DATA ---
8
+
9
+ @dataclass
10
+ class RateLimits:
11
+ """Rate limits for a provider"""
12
+ requests_per_minute: int
13
+ requests_per_hour: int
14
+ requests_per_day: int
15
+ tokens_per_minute: Optional[int] = None
16
+ tokens_per_hour: Optional[int] = None
17
+ tokens_per_day: Optional[int] = None
18
+
19
+ @dataclass
20
+ class UsageSnapshot:
21
+ """Standardized view of usage counters"""
22
+ rpm: int = 0
23
+ rph: int = 0
24
+ rpd: int = 0
25
+ tpm: int = 0
26
+ tph: int = 0
27
+ tpd: int = 0
28
+ total_requests: int = 0
29
+ total_tokens: int = 0
30
+
31
+ def __add__(self, other):
32
+ """Allow summing snapshots for aggregation"""
33
+ if not isinstance(other, UsageSnapshot): return NotImplemented
34
+ return UsageSnapshot(
35
+ self.rpm + other.rpm, self.rph + other.rph, self.rpd + other.rpd,
36
+ self.tpm + other.tpm, self.tph + other.tph, self.tpd + other.tpd,
37
+ self.total_requests + other.total_requests, self.total_tokens + other.total_tokens
38
+ )
39
+
40
+ # --- STATS DATA TRANSFER OBJECTS (DTOs) ---
41
+
42
+ @dataclass
43
+ class KeySummary:
44
+ index: int; suffix: str; snapshot: UsageSnapshot
45
+ @dataclass
46
+ class GlobalStats:
47
+ total: UsageSnapshot; keys: List[KeySummary]
48
+ @dataclass
49
+ class KeyDetailedStats:
50
+ index: int; suffix: str; total: UsageSnapshot; breakdown: Dict[str, UsageSnapshot]
51
+ @dataclass
52
+ class ModelAggregatedStats:
53
+ model_id: str; total: UsageSnapshot; keys: List[KeySummary]
54
+
55
+
56
+ # --- USAGE TRACKING ---
57
+ @dataclass
58
+ class UsageBucket:
59
+ """Tracks counters for a SINGLE model context"""
60
+ requests_minute: deque[float] = field(default_factory=deque)
61
+ requests_hour: deque[float] = field(default_factory=deque)
62
+ requests_day: deque[float] = field(default_factory=deque)
63
+
64
+ tokens_minute: deque[tuple[float, int]] = field(default_factory=deque)
65
+ tokens_hour: deque[tuple[float, int]] = field(default_factory=deque)
66
+ tokens_day: deque[tuple[float, int]] = field(default_factory=deque)
67
+
68
+ total_requests: int = 0
69
+ total_tokens: int = 0
70
+
71
+ pending_tokens: int = 0
72
+
73
+ def clean(self):
74
+ """Clean old entries based on current time"""
75
+ now = time.time()
76
+ cutoffs = (now - 60, now - 3600, now - 86400)
77
+
78
+ for d, cut in zip(
79
+ [self.requests_minute, self.requests_hour, self.requests_day], cutoffs
80
+ ):
81
+ while d and d[0] <= cut: d.popleft()
82
+
83
+ for d, cut in zip(
84
+ [self.tokens_minute, self.tokens_hour, self.tokens_day], cutoffs
85
+ ):
86
+ while d and d[0][0] <= cut: d.popleft()
87
+
88
+ def add(self, tokens: int, timestamp: float):
89
+ self.requests_minute.append(timestamp)
90
+ self.requests_hour.append(timestamp)
91
+ self.requests_day.append(timestamp)
92
+ self.total_requests += 1
93
+
94
+ if tokens > 0:
95
+ self.tokens_minute.append((timestamp, tokens))
96
+ self.tokens_hour.append((timestamp, tokens))
97
+ self.tokens_day.append((timestamp, tokens))
98
+ self.total_tokens += tokens
99
+
100
+ def check_limits(self, limits: RateLimits, estimated_tokens: int) -> bool:
101
+ self.clean()
102
+ if len(self.requests_minute) >= limits.requests_per_minute: return False
103
+ if len(self.requests_hour) >= limits.requests_per_hour: return False
104
+ if len(self.requests_day) >= limits.requests_per_day: return False
105
+
106
+ current_tpm = sum(t[1] for t in self.tokens_minute) + self.pending_tokens
107
+ current_tph = sum(t[1] for t in self.tokens_hour) + self.pending_tokens
108
+ current_tpd = sum(t[1] for t in self.tokens_day) + self.pending_tokens
109
+
110
+ if limits.tokens_per_minute and (current_tpm + estimated_tokens > limits.tokens_per_minute): return False
111
+ if limits.tokens_per_hour and (current_tph + estimated_tokens > limits.tokens_per_hour): return False
112
+ if limits.tokens_per_day and (current_tpd + estimated_tokens > limits.tokens_per_day): return False
113
+
114
+ return True
115
+
116
+ def reserve(self, tokens: int):
117
+ """Lock in estimated tokens"""
118
+ self.pending_tokens += tokens
119
+
120
+ def commit(self, actual_tokens: int, reserved_tokens: int, timestamp: float):
121
+ """Remove reservation and add actual usage"""
122
+ self.pending_tokens -= reserved_tokens
123
+ if self.pending_tokens < 0: self.pending_tokens = 0 # Safety floor
124
+ self.add(actual_tokens, timestamp)
125
+
126
+ def get_snapshot(self) -> UsageSnapshot:
127
+ """Return current counts as a clean snapshot"""
128
+ self.clean()
129
+ return UsageSnapshot(
130
+ rpm=len(self.requests_minute),
131
+ rph=len(self.requests_hour),
132
+ rpd=len(self.requests_day),
133
+ tpm=sum(t[1] for t in self.tokens_minute),
134
+ tph=sum(t[1] for t in self.tokens_hour),
135
+ tpd=sum(t[1] for t in self.tokens_day),
136
+ total_requests=self.total_requests,
137
+ total_tokens=self.total_tokens
138
+ )
139
+
140
+
141
+ @dataclass
142
+ class KeyUsage:
143
+ """Represents an API Key and holds multiple UsageBuckets (one per model)"""
144
+ api_key: str
145
+ strategy: RateLimitStrategy
146
+ buckets: Dict[str, UsageBucket] = field(default_factory=lambda: defaultdict(UsageBucket))
147
+ global_bucket: UsageBucket = field(default_factory=UsageBucket)
148
+ last_429: float = 0.0
149
+
150
+ def record_usage(self, model_id: str, tokens: int, timestamp: float = None):
151
+ ts = timestamp if timestamp else time.time()
152
+ self.buckets[model_id].add(tokens, ts)
153
+ if self.strategy == RateLimitStrategy.GLOBAL:
154
+ self.global_bucket.add(tokens, ts)
155
+
156
+ def can_use_model(self, model_id: str, limits: RateLimits, estimated_tokens: int = 1000) -> bool:
157
+ """Check limits based on the provider's strategy"""
158
+ if self.strategy == RateLimitStrategy.GLOBAL:
159
+ return self.global_bucket.check_limits(limits, estimated_tokens)
160
+ else: # Per-Model Limits
161
+ return self.buckets[model_id].check_limits(limits, estimated_tokens)
162
+
163
+ def get_total_snapshot(self) -> UsageSnapshot:
164
+ if self.strategy == RateLimitStrategy.GLOBAL:
165
+ return self.global_bucket.get_snapshot()
166
+ total = UsageSnapshot()
167
+ for b in self.buckets.values():
168
+ total = total + b.get_snapshot()
169
+ return total
170
+
171
+ def reserve(self, model_id: str, tokens: int):
172
+ self.buckets[model_id].reserve(tokens)
173
+ if self.strategy == RateLimitStrategy.GLOBAL:
174
+ self.global_bucket.reserve(tokens)
175
+
176
+ def commit(self, model_id: str, actual_tokens: int, reserved_tokens: int, timestamp: float = None):
177
+ ts = timestamp if timestamp else time.time()
178
+ self.buckets[model_id].commit(actual_tokens, reserved_tokens, ts)
179
+ if self.strategy == RateLimitStrategy.GLOBAL:
180
+ self.global_bucket.commit(actual_tokens, reserved_tokens, ts)
181
+
182
+
183
+ def is_cooling_down(self, cooldown_seconds: int = 30) -> bool:
184
+ """Returns True if the key is still in its 30s penalty box."""
185
+ if self.last_429 == 0: return False
186
+ return (time.time() - self.last_429) < cooldown_seconds
187
+
188
+ def trigger_cooldown(self):
189
+ """Mark this key as rate-limited."""
190
+ self.last_429 = time.time()
@@ -0,0 +1,5 @@
1
+ from enum import Enum
2
+
3
+ class RateLimitStrategy(Enum):
4
+ PER_MODEL = "per_model" # Cerebras, Groq, Gemini
5
+ GLOBAL = "global" # OpenRouter (Shared limits across all models)
@@ -0,0 +1,53 @@
1
+ import logging
2
+ import logging.config
3
+ from logging.handlers import RotatingFileHandler
4
+ from rich.logging import RichHandler
5
+
6
+ LOGGING_CONFIG = {
7
+ "version": 1,
8
+ "disable_existing_loggers": False,
9
+ "formatters": {
10
+ "standard": {
11
+ "format": "%(asctime)s | %(levelname)-8s | %(name)s:%(lineno)d | %(message)s",
12
+ "datefmt": "%Y-%m-%d %H:%M:%S",
13
+ },
14
+ },
15
+
16
+ "handlers": {
17
+ "console": {
18
+ "class": "rich.logging.RichHandler",
19
+ "level": "DEBUG",
20
+ "rich_tracebacks": True, # Beautiful traceback formatting
21
+ "markup": True, # Allow "[bold]text[/]" inside log messages!
22
+ "show_path": False # Cleaner output (hides file path column)
23
+ },
24
+ "file": {
25
+ "class": "logging.handlers.RotatingFileHandler",
26
+ "level": "INFO",
27
+ "formatter": "standard",
28
+ "filename": "app.log",
29
+ "maxBytes": 10 * 1024 * 1024, # 10 MiB per file
30
+ "backupCount": 5,
31
+ "encoding": "utf-8",
32
+ },
33
+ },
34
+
35
+ "root": {
36
+ "handlers": ["console", "file"],
37
+ "level": "WARNING",
38
+ },
39
+
40
+ "loggers": {
41
+ "key_manager": {
42
+ "handlers": ["console", "file"],
43
+ "level": "DEBUG",
44
+ "propagate": False,
45
+ },
46
+ },
47
+ }
48
+
49
+ def configure_logging():
50
+ logging.config.dictConfig(LOGGING_CONFIG)
51
+
52
+ configure_logging()
53
+ default_logger = logging.getLogger(__name__)
@@ -0,0 +1,7 @@
1
+ from .rotating_mixin import RotatingCredentialsMixin
2
+ from .rotation_manager import RotatingKeyManager
3
+
4
+ __all__ = [
5
+ "RotatingKeyManager",
6
+ "RotatingCredentialsMixin"
7
+ ]