keycycle 0.3.2__py3-none-any.whl
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.
- keycycle/__init__.py +48 -0
- keycycle/adapters/__init__.py +22 -0
- keycycle/adapters/generic_adapter.py +663 -0
- keycycle/adapters/openai_adapter.py +381 -0
- keycycle/config/__init__.py +25 -0
- keycycle/config/constants.py +35 -0
- keycycle/config/dataclasses.py +225 -0
- keycycle/config/enums.py +5 -0
- keycycle/config/loader.py +54 -0
- keycycle/config/log_config.py +53 -0
- keycycle/config/models/cerebras.yaml +36 -0
- keycycle/config/models/cohere.yaml +14 -0
- keycycle/config/models/gemini.yaml +69 -0
- keycycle/config/models/groq.yaml +161 -0
- keycycle/config/models/mistral.yaml +13 -0
- keycycle/config/models/moonshot.yaml +59 -0
- keycycle/config/models/openrouter.yaml +4 -0
- keycycle/config/models/openrouter_models.yaml +97 -0
- keycycle/config/models.py +53 -0
- keycycle/core/__init__.py +40 -0
- keycycle/core/backoff.py +79 -0
- keycycle/core/exceptions.py +94 -0
- keycycle/core/utils.py +480 -0
- keycycle/key_rotation/__init__.py +7 -0
- keycycle/key_rotation/rotating_mixin.py +392 -0
- keycycle/key_rotation/rotation_manager.py +231 -0
- keycycle/legacy_multi_provider_wrapper.py +734 -0
- keycycle/multi_client_wrapper.py +372 -0
- keycycle/py.typed +0 -0
- keycycle/usage/__init__.py +6 -0
- keycycle/usage/db_logic.py +98 -0
- keycycle/usage/usage_logger.py +89 -0
- keycycle/utils.py +27 -0
- keycycle-0.3.2.dist-info/METADATA +44 -0
- keycycle-0.3.2.dist-info/RECORD +37 -0
- keycycle-0.3.2.dist-info/WHEEL +5 -0
- keycycle-0.3.2.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,372 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Multi-client wrapper for managing multiple providers with rotating API keys.
|
|
3
|
+
|
|
4
|
+
This module provides the primary interface for keycycle, supporting:
|
|
5
|
+
- Multi-parameter key rotation (api_key + index_id, etc.)
|
|
6
|
+
- Multiple providers in a single wrapper
|
|
7
|
+
- Environment-based configuration with indexed params
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
import os
|
|
11
|
+
from dataclasses import dataclass, field
|
|
12
|
+
from pathlib import Path
|
|
13
|
+
from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
|
|
14
|
+
|
|
15
|
+
from dotenv import load_dotenv
|
|
16
|
+
|
|
17
|
+
from .key_rotation.rotation_manager import RotatingKeyManager
|
|
18
|
+
from .config.dataclasses import RateLimits, KeyLimitOverride
|
|
19
|
+
from .config.enums import RateLimitStrategy
|
|
20
|
+
from .config.models import MODEL_LIMITS, PROVIDER_STRATEGIES
|
|
21
|
+
from .core.utils import (
|
|
22
|
+
KeyEntry,
|
|
23
|
+
get_key_suffix,
|
|
24
|
+
load_api_keys as _load_api_keys,
|
|
25
|
+
normalize_key_limits as _normalize_key_limits,
|
|
26
|
+
)
|
|
27
|
+
from .usage.db_logic import UsageDatabase
|
|
28
|
+
from .adapters.generic_adapter import (
|
|
29
|
+
create_rotating_client,
|
|
30
|
+
SyncGenericRotatingClient,
|
|
31
|
+
AsyncGenericRotatingClient,
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
T = TypeVar("T")
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
@dataclass
|
|
38
|
+
class ProviderConfig:
|
|
39
|
+
"""Internal config for a registered provider."""
|
|
40
|
+
default_model: Optional[str] = None
|
|
41
|
+
limits: Optional[Dict[str, RateLimits]] = None
|
|
42
|
+
excluded_kwargs: Optional[List[str]] = None
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
@dataclass
|
|
46
|
+
class ProviderEnvConfig:
|
|
47
|
+
"""
|
|
48
|
+
Configuration for loading a provider from environment variables.
|
|
49
|
+
|
|
50
|
+
Attributes:
|
|
51
|
+
default_model: Default model ID for this provider
|
|
52
|
+
extra_params: List of extra parameter names to load alongside API keys.
|
|
53
|
+
For each param, loads {PROVIDER}_{PARAM}_N environment variables.
|
|
54
|
+
strategy: Rate limit strategy (per-model or global)
|
|
55
|
+
limits: Custom rate limits per model
|
|
56
|
+
api_key_param: Name of the API key parameter (default: "api_key")
|
|
57
|
+
excluded_kwargs: List of kwarg names to exclude from client constructor.
|
|
58
|
+
Useful for clients that don't accept certain params (e.g., TwelveLabs doesn't accept 'model').
|
|
59
|
+
"""
|
|
60
|
+
default_model: Optional[str] = None
|
|
61
|
+
extra_params: Optional[List[str]] = None
|
|
62
|
+
strategy: RateLimitStrategy = RateLimitStrategy.PER_MODEL
|
|
63
|
+
limits: Optional[Dict[str, RateLimits]] = None
|
|
64
|
+
api_key_param: str = "api_key"
|
|
65
|
+
excluded_kwargs: Optional[List[str]] = None
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
class MultiClientWrapper:
|
|
69
|
+
"""
|
|
70
|
+
Manages multiple providers with rotating API keys.
|
|
71
|
+
|
|
72
|
+
This is the primary wrapper for keycycle, supporting:
|
|
73
|
+
- Multiple providers in a single instance
|
|
74
|
+
- Multi-parameter rotation (api_key + other params)
|
|
75
|
+
- Environment-based configuration
|
|
76
|
+
|
|
77
|
+
Example:
|
|
78
|
+
>>> wrapper = MultiClientWrapper.from_env({
|
|
79
|
+
... "twelvelabs": ProviderEnvConfig(
|
|
80
|
+
... default_model="pegasus-1",
|
|
81
|
+
... extra_params=["index_id"],
|
|
82
|
+
... ),
|
|
83
|
+
... "anthropic": ProviderEnvConfig(
|
|
84
|
+
... default_model="claude-3-sonnet",
|
|
85
|
+
... ),
|
|
86
|
+
... })
|
|
87
|
+
>>> tl_client = wrapper.get_rotating_client("twelvelabs", TwelveLabs)
|
|
88
|
+
>>> claude_client = wrapper.get_rotating_client("anthropic", Anthropic)
|
|
89
|
+
"""
|
|
90
|
+
|
|
91
|
+
MODEL_LIMITS = MODEL_LIMITS
|
|
92
|
+
PROVIDER_STRATEGIES = PROVIDER_STRATEGIES
|
|
93
|
+
|
|
94
|
+
def __init__(self, db_url: Optional[str] = None, db_env_var: str = "TIDB_DB_URL"):
|
|
95
|
+
"""
|
|
96
|
+
Initialize a MultiClientWrapper.
|
|
97
|
+
|
|
98
|
+
Args:
|
|
99
|
+
db_url: Database URL for usage persistence (optional)
|
|
100
|
+
db_env_var: Environment variable name for database URL
|
|
101
|
+
"""
|
|
102
|
+
self.db = UsageDatabase(db_url, db_env_var)
|
|
103
|
+
self._managers: Dict[str, RotatingKeyManager] = {}
|
|
104
|
+
self._configs: Dict[str, ProviderConfig] = {}
|
|
105
|
+
self._key_limits: Dict[str, Dict[str, KeyLimitOverride]] = {}
|
|
106
|
+
|
|
107
|
+
def register_provider(
|
|
108
|
+
self,
|
|
109
|
+
provider: str,
|
|
110
|
+
keys: List[KeyEntry],
|
|
111
|
+
strategy: Optional[RateLimitStrategy] = None,
|
|
112
|
+
default_model: Optional[str] = None,
|
|
113
|
+
limits: Optional[Dict[str, RateLimits]] = None,
|
|
114
|
+
api_key_param: str = "api_key",
|
|
115
|
+
key_limits: Optional[Dict[Union[int, str], KeyLimitOverride]] = None,
|
|
116
|
+
**kwargs
|
|
117
|
+
) -> "MultiClientWrapper":
|
|
118
|
+
"""
|
|
119
|
+
Register a provider with its keys.
|
|
120
|
+
|
|
121
|
+
Args:
|
|
122
|
+
provider: Provider name (e.g., "twelvelabs", "anthropic")
|
|
123
|
+
keys: List of key entries (strings or dicts with params)
|
|
124
|
+
strategy: Rate limit strategy (defaults to provider's default)
|
|
125
|
+
default_model: Default model ID
|
|
126
|
+
limits: Custom rate limits per model
|
|
127
|
+
api_key_param: Name of the API key parameter
|
|
128
|
+
key_limits: Per-key rate limit overrides
|
|
129
|
+
**kwargs: Additional arguments for RotatingKeyManager
|
|
130
|
+
|
|
131
|
+
Returns:
|
|
132
|
+
Self for method chaining
|
|
133
|
+
"""
|
|
134
|
+
provider = provider.lower()
|
|
135
|
+
|
|
136
|
+
# Use provider-specific strategy if not specified
|
|
137
|
+
if strategy is None:
|
|
138
|
+
strategy = self.PROVIDER_STRATEGIES.get(provider, RateLimitStrategy.PER_MODEL)
|
|
139
|
+
|
|
140
|
+
# Normalize key_limits to suffix-based
|
|
141
|
+
normalized_key_limits = self._normalize_key_limits_from_entries(keys, key_limits, api_key_param)
|
|
142
|
+
self._key_limits[provider] = normalized_key_limits
|
|
143
|
+
|
|
144
|
+
manager = RotatingKeyManager(
|
|
145
|
+
api_keys=keys,
|
|
146
|
+
provider_name=provider,
|
|
147
|
+
strategy=strategy,
|
|
148
|
+
db=self.db,
|
|
149
|
+
api_key_param=api_key_param,
|
|
150
|
+
limit_resolver=lambda m, k, p=provider: self._resolve_limits(p, m, k),
|
|
151
|
+
**kwargs
|
|
152
|
+
)
|
|
153
|
+
self._managers[provider] = manager
|
|
154
|
+
self._configs[provider] = ProviderConfig(default_model, limits)
|
|
155
|
+
return self
|
|
156
|
+
|
|
157
|
+
def _normalize_key_limits_from_entries(
|
|
158
|
+
self,
|
|
159
|
+
keys: List[KeyEntry],
|
|
160
|
+
key_limits: Optional[Dict[Union[int, str], KeyLimitOverride]],
|
|
161
|
+
api_key_param: str = "api_key",
|
|
162
|
+
) -> Dict[str, KeyLimitOverride]:
|
|
163
|
+
"""Normalize key_limits to suffix-based lookup."""
|
|
164
|
+
if not key_limits:
|
|
165
|
+
return {}
|
|
166
|
+
|
|
167
|
+
# Extract primary keys
|
|
168
|
+
primary_keys = []
|
|
169
|
+
for entry in keys:
|
|
170
|
+
if isinstance(entry, str):
|
|
171
|
+
primary_keys.append(entry)
|
|
172
|
+
else:
|
|
173
|
+
primary_keys.append(entry.get(api_key_param, ""))
|
|
174
|
+
|
|
175
|
+
return _normalize_key_limits(primary_keys, key_limits)
|
|
176
|
+
|
|
177
|
+
def _resolve_limits(
|
|
178
|
+
self,
|
|
179
|
+
provider: str,
|
|
180
|
+
model_id: str,
|
|
181
|
+
key_suffix: Optional[str] = None
|
|
182
|
+
) -> RateLimits:
|
|
183
|
+
"""Resolve rate limits for a provider/model/key combination."""
|
|
184
|
+
config = self._configs.get(provider)
|
|
185
|
+
provider_key_limits = self._key_limits.get(provider, {})
|
|
186
|
+
|
|
187
|
+
# Check key-specific overrides first
|
|
188
|
+
if key_suffix and provider_key_limits:
|
|
189
|
+
override = provider_key_limits.get(key_suffix)
|
|
190
|
+
if override:
|
|
191
|
+
if isinstance(override, RateLimits):
|
|
192
|
+
return override
|
|
193
|
+
elif isinstance(override, dict):
|
|
194
|
+
if model_id in override:
|
|
195
|
+
return override[model_id]
|
|
196
|
+
if '__default__' in override:
|
|
197
|
+
return override['__default__']
|
|
198
|
+
|
|
199
|
+
# Check provider-specific limits from config
|
|
200
|
+
if config and config.limits and model_id in config.limits:
|
|
201
|
+
return config.limits[model_id]
|
|
202
|
+
|
|
203
|
+
# Fall back to global model limits
|
|
204
|
+
provider_limits = self.MODEL_LIMITS.get(provider, {})
|
|
205
|
+
return provider_limits.get(model_id, provider_limits.get('default', RateLimits(10, 100, 1000)))
|
|
206
|
+
|
|
207
|
+
def get_rotating_client(
|
|
208
|
+
self,
|
|
209
|
+
provider: str,
|
|
210
|
+
client_class: Type[T],
|
|
211
|
+
model: Optional[str] = None,
|
|
212
|
+
api_key_param: Optional[str] = None,
|
|
213
|
+
is_async: Optional[bool] = None,
|
|
214
|
+
usage_extractor: Optional[Callable[[Any], int]] = None,
|
|
215
|
+
estimated_tokens: int = 1000,
|
|
216
|
+
max_retries: int = 5,
|
|
217
|
+
model_param: str = "model",
|
|
218
|
+
excluded_kwargs: Optional[List[str]] = None,
|
|
219
|
+
**client_kwargs
|
|
220
|
+
) -> Union[SyncGenericRotatingClient[T], AsyncGenericRotatingClient[T]]:
|
|
221
|
+
"""
|
|
222
|
+
Get a rotating client for a registered provider.
|
|
223
|
+
|
|
224
|
+
Args:
|
|
225
|
+
provider: Registered provider name
|
|
226
|
+
client_class: The client class to wrap
|
|
227
|
+
model: Model ID (defaults to provider's default_model)
|
|
228
|
+
api_key_param: Override the API key parameter name
|
|
229
|
+
is_async: Whether the client is async (auto-detected if None)
|
|
230
|
+
usage_extractor: Custom function to extract token usage
|
|
231
|
+
estimated_tokens: Estimated tokens per request
|
|
232
|
+
max_retries: Max key rotations on rate limit errors
|
|
233
|
+
model_param: Name of the model parameter in API calls
|
|
234
|
+
excluded_kwargs: List of kwarg names to exclude from client constructor.
|
|
235
|
+
Useful for clients that don't accept certain params (e.g., 'model' for TwelveLabs).
|
|
236
|
+
**client_kwargs: Additional kwargs for client constructor
|
|
237
|
+
|
|
238
|
+
Returns:
|
|
239
|
+
A rotating client wrapper
|
|
240
|
+
|
|
241
|
+
Raises:
|
|
242
|
+
ValueError: If provider is not registered
|
|
243
|
+
"""
|
|
244
|
+
provider = provider.lower()
|
|
245
|
+
manager = self._managers.get(provider)
|
|
246
|
+
if not manager:
|
|
247
|
+
raise ValueError(f"Provider '{provider}' not registered. Call register_provider() first.")
|
|
248
|
+
|
|
249
|
+
config = self._configs[provider]
|
|
250
|
+
default_model = model or config.default_model or "default"
|
|
251
|
+
|
|
252
|
+
# Use the manager's api_key_param if not overridden
|
|
253
|
+
if api_key_param is None:
|
|
254
|
+
api_key_param = manager.api_key_param
|
|
255
|
+
|
|
256
|
+
# Use config's excluded_kwargs as default if not specified
|
|
257
|
+
if excluded_kwargs is None:
|
|
258
|
+
excluded_kwargs = config.excluded_kwargs
|
|
259
|
+
|
|
260
|
+
return create_rotating_client(
|
|
261
|
+
client_class=client_class,
|
|
262
|
+
manager=manager,
|
|
263
|
+
limit_resolver=lambda m, k: self._resolve_limits(provider, m, k),
|
|
264
|
+
default_model=default_model,
|
|
265
|
+
api_key_param=api_key_param,
|
|
266
|
+
is_async=is_async,
|
|
267
|
+
usage_extractor=usage_extractor,
|
|
268
|
+
estimated_tokens=estimated_tokens,
|
|
269
|
+
max_retries=max_retries,
|
|
270
|
+
model_param=model_param,
|
|
271
|
+
excluded_kwargs=excluded_kwargs,
|
|
272
|
+
**client_kwargs,
|
|
273
|
+
)
|
|
274
|
+
|
|
275
|
+
def get_manager(self, provider: str) -> RotatingKeyManager:
|
|
276
|
+
"""Get the RotatingKeyManager for a provider."""
|
|
277
|
+
provider = provider.lower()
|
|
278
|
+
manager = self._managers.get(provider)
|
|
279
|
+
if not manager:
|
|
280
|
+
raise ValueError(f"Provider '{provider}' not registered.")
|
|
281
|
+
return manager
|
|
282
|
+
|
|
283
|
+
@classmethod
|
|
284
|
+
def from_env(
|
|
285
|
+
cls,
|
|
286
|
+
providers: Dict[str, ProviderEnvConfig],
|
|
287
|
+
env_file: Optional[str] = None,
|
|
288
|
+
db_url: Optional[str] = None,
|
|
289
|
+
db_env_var: str = "TIDB_DB_URL",
|
|
290
|
+
) -> "MultiClientWrapper":
|
|
291
|
+
"""
|
|
292
|
+
Create a wrapper from environment variables.
|
|
293
|
+
|
|
294
|
+
Args:
|
|
295
|
+
providers: Dict mapping provider names to their env configs
|
|
296
|
+
env_file: Path to .env file (optional)
|
|
297
|
+
db_url: Database URL for usage persistence
|
|
298
|
+
db_env_var: Environment variable name for database URL
|
|
299
|
+
|
|
300
|
+
Returns:
|
|
301
|
+
Configured MultiClientWrapper instance
|
|
302
|
+
|
|
303
|
+
Example:
|
|
304
|
+
>>> wrapper = MultiClientWrapper.from_env({
|
|
305
|
+
... "twelvelabs": ProviderEnvConfig(
|
|
306
|
+
... default_model="pegasus-1",
|
|
307
|
+
... extra_params=["index_id"],
|
|
308
|
+
... ),
|
|
309
|
+
... "anthropic": ProviderEnvConfig(
|
|
310
|
+
... default_model="claude-3-sonnet",
|
|
311
|
+
... ),
|
|
312
|
+
... })
|
|
313
|
+
|
|
314
|
+
Environment variables format:
|
|
315
|
+
NUM_TWELVELABS=2
|
|
316
|
+
TWELVELABS_API_KEY_1=key1
|
|
317
|
+
TWELVELABS_INDEX_ID_1=idx_abc
|
|
318
|
+
TWELVELABS_API_KEY_2=key2
|
|
319
|
+
TWELVELABS_INDEX_ID_2=idx_xyz
|
|
320
|
+
"""
|
|
321
|
+
instance = cls(db_url=db_url, db_env_var=db_env_var)
|
|
322
|
+
|
|
323
|
+
for provider, config in providers.items():
|
|
324
|
+
keys = cls.load_api_keys(
|
|
325
|
+
provider,
|
|
326
|
+
env_file=env_file,
|
|
327
|
+
extra_params=config.extra_params,
|
|
328
|
+
api_key_param=config.api_key_param,
|
|
329
|
+
)
|
|
330
|
+
instance.register_provider(
|
|
331
|
+
provider=provider,
|
|
332
|
+
keys=keys,
|
|
333
|
+
default_model=config.default_model,
|
|
334
|
+
strategy=config.strategy,
|
|
335
|
+
limits=config.limits,
|
|
336
|
+
api_key_param=config.api_key_param,
|
|
337
|
+
)
|
|
338
|
+
# Store excluded_kwargs from env config
|
|
339
|
+
if config.excluded_kwargs:
|
|
340
|
+
instance._configs[provider.lower()].excluded_kwargs = config.excluded_kwargs
|
|
341
|
+
return instance
|
|
342
|
+
|
|
343
|
+
@staticmethod
|
|
344
|
+
def load_api_keys(
|
|
345
|
+
provider: str,
|
|
346
|
+
env_file: Optional[str] = None,
|
|
347
|
+
extra_params: Optional[List[str]] = None,
|
|
348
|
+
api_key_param: str = "api_key",
|
|
349
|
+
) -> List[KeyEntry]:
|
|
350
|
+
"""
|
|
351
|
+
Load keys from environment variables with indexed pattern.
|
|
352
|
+
|
|
353
|
+
Args:
|
|
354
|
+
provider: Provider name (e.g., "twelvelabs")
|
|
355
|
+
env_file: Path to .env file (optional)
|
|
356
|
+
extra_params: List of extra param names to load
|
|
357
|
+
api_key_param: Name for the API key in returned dicts
|
|
358
|
+
|
|
359
|
+
Returns:
|
|
360
|
+
List of key entries
|
|
361
|
+
|
|
362
|
+
Environment variables format:
|
|
363
|
+
NUM_{PROVIDER}=N
|
|
364
|
+
{PROVIDER}_API_KEY_1, {PROVIDER}_API_KEY_2, ...
|
|
365
|
+
{PROVIDER}_{PARAM}_1, {PROVIDER}_{PARAM}_2, ... (for each extra_param)
|
|
366
|
+
"""
|
|
367
|
+
return _load_api_keys(provider, env_file, extra_params, api_key_param)
|
|
368
|
+
|
|
369
|
+
def stop(self) -> None:
|
|
370
|
+
"""Stop all managers and flush logs."""
|
|
371
|
+
for manager in self._managers.values():
|
|
372
|
+
manager.stop()
|
keycycle/py.typed
ADDED
|
File without changes
|
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import time
|
|
3
|
+
from typing import Optional, List
|
|
4
|
+
from sqlalchemy import (
|
|
5
|
+
create_engine, select, and_, Table, Column,
|
|
6
|
+
Integer, String, Float, MetaData, Index, delete,
|
|
7
|
+
URL
|
|
8
|
+
)
|
|
9
|
+
|
|
10
|
+
from ..core.utils import get_key_suffix
|
|
11
|
+
from ..config.constants import SECONDS_PER_DAY
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
# --- DATABASE LAYER ---
|
|
15
|
+
|
|
16
|
+
class UsageDatabase:
|
|
17
|
+
"""Handles Online persistence for API usage"""
|
|
18
|
+
def __init__(self, db_url: Optional[str] = None, db_env_var: str = "TIDB_DB_URL"):
|
|
19
|
+
self.db_url = db_url or os.getenv(db_env_var)
|
|
20
|
+
if not self.db_url:
|
|
21
|
+
raise ValueError(f"Database URL not provided and {db_env_var} not set.")
|
|
22
|
+
self.engine = create_engine(
|
|
23
|
+
self.db_url,
|
|
24
|
+
pool_recycle = 300
|
|
25
|
+
)
|
|
26
|
+
self._init_db()
|
|
27
|
+
|
|
28
|
+
def _init_db(self):
|
|
29
|
+
metadata = MetaData()
|
|
30
|
+
self.usage_logs = Table(
|
|
31
|
+
'usage_logs',
|
|
32
|
+
metadata,
|
|
33
|
+
Column('id', Integer, primary_key=True),
|
|
34
|
+
Column('provider', String(100)),
|
|
35
|
+
Column('model', String(100)),
|
|
36
|
+
Column('api_key_suffix', String(50)),
|
|
37
|
+
Column('timestamp', Float),
|
|
38
|
+
Column('tokens', Integer),
|
|
39
|
+
|
|
40
|
+
Index('idx_key_usage', 'provider', 'api_key_suffix', 'timestamp'),
|
|
41
|
+
Index('idx_cleanup', 'timestamp'),
|
|
42
|
+
Index('idx_model_reporting', 'provider', 'model', 'timestamp')
|
|
43
|
+
)
|
|
44
|
+
metadata.create_all(self.engine)
|
|
45
|
+
|
|
46
|
+
def load_history(self, provider: str, api_key: str, seconds_lookback: int) -> List[tuple[str, float, int]]:
|
|
47
|
+
"""Load history SPECIFIC to this Provider + Model combination"""
|
|
48
|
+
suffix = get_key_suffix(api_key)
|
|
49
|
+
cutoff = time.time() - seconds_lookback
|
|
50
|
+
|
|
51
|
+
stmt = (
|
|
52
|
+
select(
|
|
53
|
+
self.usage_logs.c.model,
|
|
54
|
+
self.usage_logs.c.timestamp,
|
|
55
|
+
self.usage_logs.c.tokens,
|
|
56
|
+
)
|
|
57
|
+
.where(
|
|
58
|
+
and_(
|
|
59
|
+
self.usage_logs.c.provider == provider,
|
|
60
|
+
self.usage_logs.c.api_key_suffix == suffix,
|
|
61
|
+
self.usage_logs.c.timestamp > cutoff,
|
|
62
|
+
)
|
|
63
|
+
)
|
|
64
|
+
.order_by(self.usage_logs.c.timestamp.asc())
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
with self.engine.connect() as conn:
|
|
68
|
+
return conn.execute(stmt).all()
|
|
69
|
+
|
|
70
|
+
def load_provider_history(self, provider: str, seconds_lookback: int):
|
|
71
|
+
"""Optimization: Load everything for the provider in ONE call"""
|
|
72
|
+
cutoff = time.time() - seconds_lookback
|
|
73
|
+
|
|
74
|
+
stmt = (
|
|
75
|
+
select(
|
|
76
|
+
self.usage_logs.c.api_key_suffix,
|
|
77
|
+
self.usage_logs.c.model,
|
|
78
|
+
self.usage_logs.c.timestamp,
|
|
79
|
+
self.usage_logs.c.tokens,
|
|
80
|
+
)
|
|
81
|
+
.where(
|
|
82
|
+
self.usage_logs.c.provider == provider,
|
|
83
|
+
self.usage_logs.c.timestamp > cutoff,
|
|
84
|
+
)
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
with self.engine.connect() as conn:
|
|
88
|
+
return conn.execute(stmt).all()
|
|
89
|
+
|
|
90
|
+
def prune_old_records(self, days_retention: int = 3) -> None:
|
|
91
|
+
"""Delete records older than retention period to keep DB small (3 days)"""
|
|
92
|
+
cutoff = time.time() - (days_retention * SECONDS_PER_DAY)
|
|
93
|
+
with self.engine.connect() as conn:
|
|
94
|
+
conn.execute(
|
|
95
|
+
delete(self.usage_logs).where(
|
|
96
|
+
self.usage_logs.c.timestamp < cutoff
|
|
97
|
+
))
|
|
98
|
+
conn.commit()
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
import atexit
|
|
2
|
+
import logging
|
|
3
|
+
import queue
|
|
4
|
+
import threading
|
|
5
|
+
import time
|
|
6
|
+
from typing import Optional
|
|
7
|
+
|
|
8
|
+
from .db_logic import UsageDatabase
|
|
9
|
+
from ..config.log_config import default_logger
|
|
10
|
+
from ..config.constants import USAGE_LOG_BATCH_SIZE
|
|
11
|
+
from ..core.utils import get_key_suffix
|
|
12
|
+
|
|
13
|
+
# --- ASYNC LOGGER ---
|
|
14
|
+
class AsyncUsageLogger:
|
|
15
|
+
"""Decouples Turso DB writes from the main thread using batching."""
|
|
16
|
+
def __init__(self, db: UsageDatabase, logger: Optional[logging.Logger] = None):
|
|
17
|
+
self.db = db
|
|
18
|
+
self.queue = queue.Queue()
|
|
19
|
+
self._stop_event = threading.Event()
|
|
20
|
+
self._thread = threading.Thread(target=self._writer_loop, daemon=True)
|
|
21
|
+
self._thread.start()
|
|
22
|
+
atexit.register(self.stop)
|
|
23
|
+
|
|
24
|
+
self.logger = logger or default_logger
|
|
25
|
+
|
|
26
|
+
def log(self, provider: str, model: str, api_key: str, tokens: int):
|
|
27
|
+
self.queue.put((provider, model, api_key, time.time(), tokens))
|
|
28
|
+
|
|
29
|
+
def _writer_loop(self):
|
|
30
|
+
batch = []
|
|
31
|
+
insert_stmt = self.db.usage_logs.insert()
|
|
32
|
+
|
|
33
|
+
while not self._stop_event.is_set() or not self.queue.empty(): #always empty queue
|
|
34
|
+
try:
|
|
35
|
+
record = self.queue.get(timeout=1.0)
|
|
36
|
+
# Parse record for batch formatting
|
|
37
|
+
provider, model, full_key, ts, tokens = record
|
|
38
|
+
suffix = get_key_suffix(full_key)
|
|
39
|
+
batch.append({
|
|
40
|
+
"provider": provider,
|
|
41
|
+
"model": model,
|
|
42
|
+
"api_key_suffix": suffix,
|
|
43
|
+
"timestamp": ts,
|
|
44
|
+
"tokens": tokens
|
|
45
|
+
})
|
|
46
|
+
|
|
47
|
+
# Drain queue up to USAGE_LOG_BATCH_SIZE items to batch write
|
|
48
|
+
while len(batch) < USAGE_LOG_BATCH_SIZE:
|
|
49
|
+
try:
|
|
50
|
+
r = self.queue.get_nowait()
|
|
51
|
+
p, m, k, t, tok = r
|
|
52
|
+
s = get_key_suffix(k)
|
|
53
|
+
batch.append({
|
|
54
|
+
"provider": p,
|
|
55
|
+
"model": m,
|
|
56
|
+
"api_key_suffix": s,
|
|
57
|
+
"timestamp": t,
|
|
58
|
+
"tokens": tok
|
|
59
|
+
})
|
|
60
|
+
except queue.Empty:
|
|
61
|
+
break
|
|
62
|
+
|
|
63
|
+
if batch:
|
|
64
|
+
with self.db.engine.connect() as conn:
|
|
65
|
+
conn.execute(insert_stmt, batch)
|
|
66
|
+
conn.commit()
|
|
67
|
+
batch.clear()
|
|
68
|
+
|
|
69
|
+
except queue.Empty:
|
|
70
|
+
continue
|
|
71
|
+
except Exception as e:
|
|
72
|
+
# print(f"Logging thread error: {e}")
|
|
73
|
+
self.logger.exception("Logging thread error", exc_info=e)
|
|
74
|
+
time.sleep(2)
|
|
75
|
+
|
|
76
|
+
if batch:
|
|
77
|
+
try:
|
|
78
|
+
with self.db.engine.connect() as conn:
|
|
79
|
+
conn.execute(insert_stmt, batch)
|
|
80
|
+
conn.commit()
|
|
81
|
+
except Exception as e:
|
|
82
|
+
self.logger.exception("Logging thread error on exit", exc_info=e)
|
|
83
|
+
|
|
84
|
+
def stop(self):
|
|
85
|
+
self._stop_event.set()
|
|
86
|
+
if self._thread.is_alive():
|
|
87
|
+
self._thread.join(timeout=10)
|
|
88
|
+
if self._thread.is_alive():
|
|
89
|
+
self.logger.warning("AsyncUsageLogger thread did not exit cleanly within timeout.")
|
keycycle/utils.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
import importlib
|
|
2
|
+
|
|
3
|
+
def get_agno_model_class(provider: str):
|
|
4
|
+
"""
|
|
5
|
+
Dynamically maps a provider string to the actual Agno model class.
|
|
6
|
+
"""
|
|
7
|
+
p_low = provider.lower()
|
|
8
|
+
|
|
9
|
+
overrides = {
|
|
10
|
+
"openai": "OpenAI",
|
|
11
|
+
"google": "Gemini",
|
|
12
|
+
"gemini": "Gemini",
|
|
13
|
+
"azure": "AzureOpenAI",
|
|
14
|
+
"aws": "Bedrock",
|
|
15
|
+
"openrouter": "OpenRouter",
|
|
16
|
+
"moonshot": "MoonShot"
|
|
17
|
+
}
|
|
18
|
+
|
|
19
|
+
class_name = overrides.get(p_low, p_low.capitalize())
|
|
20
|
+
|
|
21
|
+
module_path = "google" if p_low in ["google", "gemini"] else p_low
|
|
22
|
+
|
|
23
|
+
try:
|
|
24
|
+
module = importlib.import_module(f"agno.models.{module_path}")
|
|
25
|
+
return getattr(module, class_name)
|
|
26
|
+
except (ImportError, AttributeError) as e:
|
|
27
|
+
raise ValueError(f"Agno class '{class_name}' not found for provider '{provider}': {e}")
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: keycycle
|
|
3
|
+
Version: 0.3.2
|
|
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
|
+
Requires-Dist: PyYAML
|
|
17
|
+
Provides-Extra: openai
|
|
18
|
+
Requires-Dist: openai; extra == "openai"
|
|
19
|
+
Provides-Extra: agno
|
|
20
|
+
Requires-Dist: agno; extra == "agno"
|
|
21
|
+
Provides-Extra: cohere
|
|
22
|
+
Requires-Dist: cohere; extra == "cohere"
|
|
23
|
+
Provides-Extra: cerebras
|
|
24
|
+
Requires-Dist: cerebras-cloud-sdk; extra == "cerebras"
|
|
25
|
+
Provides-Extra: groq
|
|
26
|
+
Requires-Dist: groq; extra == "groq"
|
|
27
|
+
Provides-Extra: openrouter
|
|
28
|
+
Requires-Dist: openrouter; extra == "openrouter"
|
|
29
|
+
Provides-Extra: all
|
|
30
|
+
Requires-Dist: openai; extra == "all"
|
|
31
|
+
Requires-Dist: agno; extra == "all"
|
|
32
|
+
Requires-Dist: cohere; extra == "all"
|
|
33
|
+
Requires-Dist: cerebras-cloud-sdk; extra == "all"
|
|
34
|
+
Requires-Dist: groq; extra == "all"
|
|
35
|
+
Requires-Dist: openrouter; extra == "all"
|
|
36
|
+
Provides-Extra: test
|
|
37
|
+
Requires-Dist: pytest>=7.4.0; extra == "test"
|
|
38
|
+
Requires-Dist: pytest-cov>=4.1.0; extra == "test"
|
|
39
|
+
Requires-Dist: pytest-asyncio>=0.21.0; extra == "test"
|
|
40
|
+
Requires-Dist: pytest-timeout>=2.2.0; extra == "test"
|
|
41
|
+
Provides-Extra: dev
|
|
42
|
+
Requires-Dist: keycycle[test]; extra == "dev"
|
|
43
|
+
Requires-Dist: mypy>=1.5.0; extra == "dev"
|
|
44
|
+
Requires-Dist: types-PyYAML; extra == "dev"
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
keycycle/__init__.py,sha256=IRm7jD151rZk8KsdzF5527n7ve4NRPlYDGEp0z371Kk,1357
|
|
2
|
+
keycycle/legacy_multi_provider_wrapper.py,sha256=u_8EW1eE8XfF84MBHKKK3KYQuPpPxdf4XFO4KPtLz8U,28820
|
|
3
|
+
keycycle/multi_client_wrapper.py,sha256=vnCU2Tz3Tv75vMm-mALIZgj_SuLd-e2eozRLSW30NUo,13767
|
|
4
|
+
keycycle/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
5
|
+
keycycle/utils.py,sha256=4nNzR35AG73HfaEdjIoct1yYhBYtHO58tLUJTupeaLQ,811
|
|
6
|
+
keycycle/adapters/__init__.py,sha256=YclnV2JDpRBDaE5htVChjkSgbx_6o6xFK-pBDfNIxWk,540
|
|
7
|
+
keycycle/adapters/generic_adapter.py,sha256=3iWAlMlvmZg83YP8I23ORzX0hzbD996oQK6GXffMJqg,25984
|
|
8
|
+
keycycle/adapters/openai_adapter.py,sha256=CoarZIHfzVLdVfj1PBXoj46lFw8SX_itg0_RXeV9PEk,16133
|
|
9
|
+
keycycle/config/__init__.py,sha256=ZTqtkbw3PBBc-tc4uvjhQkgpU6azATP7cmdBCQ5HHVw,483
|
|
10
|
+
keycycle/config/constants.py,sha256=YBTedy-pDo1KTAKsTwot93dVdYh0b8Rbfoyd4LlpoAc,856
|
|
11
|
+
keycycle/config/dataclasses.py,sha256=l-8wEixdd7eVIyXvblt2BFW93q7d36R_OxM7vYoY54Y,8769
|
|
12
|
+
keycycle/config/enums.py,sha256=J8taDqXN0hT9-mBv8alihCmDoNRGtZuKtFMa87AOBvI,184
|
|
13
|
+
keycycle/config/loader.py,sha256=cei34gYic0rz9PAdIE1LSbx6i8Xhg_I4Mb-fp2DlpcA,2008
|
|
14
|
+
keycycle/config/log_config.py,sha256=tVVOwH4uekrdACOp1I-6NfOs0dxQ9ToBGg0X7vmQ344,1488
|
|
15
|
+
keycycle/config/models.py,sha256=SLfJ8Myedpar7AcP118WLjQDN5MeU_8y3TLs5t3r1sI,1924
|
|
16
|
+
keycycle/config/models/cerebras.yaml,sha256=frD4wiIUw4LaAmpDglS-zuJElTd7_wUxUX682ykcr_0,1441
|
|
17
|
+
keycycle/config/models/cohere.yaml,sha256=itL2P0rwGHMBZh0vZJ9OOgBres-FTE1zKp8cpm3I2VU,268
|
|
18
|
+
keycycle/config/models/gemini.yaml,sha256=vjc1bm_hPhGNj9dieFJTIcDpSAMNcsOvtpjfBi7SjV8,1505
|
|
19
|
+
keycycle/config/models/groq.yaml,sha256=r-2-BqCloxK7O74t3Zq71ieJyP_zT1j8YCHS09CxqtI,3624
|
|
20
|
+
keycycle/config/models/mistral.yaml,sha256=SHDDVHnU1LQyqBD2XaNP4ixWBfzWzzQ0-MPO_xfeTcA,563
|
|
21
|
+
keycycle/config/models/moonshot.yaml,sha256=wUcsAwIU9hwYOrPqXdutBVR3q2Y93rZ_Yz3Fc6x9-pQ,1493
|
|
22
|
+
keycycle/config/models/openrouter.yaml,sha256=hI0EZc2KR4YtCzJ7NubYNL-f5rux1y-NvxuYzgy_Fds,82
|
|
23
|
+
keycycle/config/models/openrouter_models.yaml,sha256=qd43vjFAKh-9Uz2aeSvOqQtTD4bDbKBmc0AOrzJ02VM,2936
|
|
24
|
+
keycycle/core/__init__.py,sha256=l7c_73JyYb9ao6MfcF360VnpEAWRWUq0dgASMEqF6_s,905
|
|
25
|
+
keycycle/core/backoff.py,sha256=SkvJzXF0ZbUrdjbwOZk6csCbcz3QwKlwgRQYQ4rpcd8,2209
|
|
26
|
+
keycycle/core/exceptions.py,sha256=TTzvsFvPv2aW38hGrxnEGwZQxLQ8o6eqMSwJy8vI3z0,2814
|
|
27
|
+
keycycle/core/utils.py,sha256=KfHIyl0RbO2oLbV_0-Fs8ZmGPGH8QZYV8cZt04CXlWY,16085
|
|
28
|
+
keycycle/key_rotation/__init__.py,sha256=Gk9Yk-Dkbw_OlUqQZbzGhJlkkIKRHb42hd0f_RE09vo,173
|
|
29
|
+
keycycle/key_rotation/rotating_mixin.py,sha256=nh53FWb9uAAWQ6_yJMGaChlhOG3BOvIixxRksRv1zyg,17465
|
|
30
|
+
keycycle/key_rotation/rotation_manager.py,sha256=sT6LKEZ_bz_kQdVP6Ie9J3AcpRYTEk7JkHvK1qRyVUQ,9388
|
|
31
|
+
keycycle/usage/__init__.py,sha256=y7lJv-9VSpUPWU5cC-LznS1c4vR8Kec_SHKR7ZTpfAQ,136
|
|
32
|
+
keycycle/usage/db_logic.py,sha256=Kd7tWKzr6kqJIPq3AATF2y4T5JAH31Ko7WhSgSOcQs4,3386
|
|
33
|
+
keycycle/usage/usage_logger.py,sha256=TfTmD37xwKZpqeqzxhNNThX_UhA8NG3QtTyfBrBXuw0,3351
|
|
34
|
+
keycycle-0.3.2.dist-info/METADATA,sha256=wm5OFHZfH2-NFro_32tRnG03T_ZQTVTlH6ANv-MVl4A,1593
|
|
35
|
+
keycycle-0.3.2.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
|
|
36
|
+
keycycle-0.3.2.dist-info/top_level.txt,sha256=Vfq7VDf-qh5oxQpTVd4qKB0m5CTpO9LscJ-oGu30DvA,9
|
|
37
|
+
keycycle-0.3.2.dist-info/RECORD,,
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
keycycle
|