Skip to content

Commit 45cee6e

Browse files
committed
feat(quota): ✨ add Firmware.ai quota tracking with 5-hour rolling window
Implement quota tracking for Firmware.ai provider using their /api/v1/quota endpoint. The provider tracks a 5-hour rolling window quota where `used` is already a 0-1 ratio from the API. - Add FirmwareQuotaTracker mixin with configurable api_base - Add FirmwareProvider with background job for periodic quota refresh - Parse ISO 8601 reset timestamps with proper Z suffix handling - Validate API response types and clamp remaining fraction to 0.0-1.0 - Support FIRMWARE_QUOTA_REFRESH_INTERVAL env var (default: 300s)
1 parent b7c5dad commit 45cee6e

File tree

4 files changed

+503
-5
lines changed

4 files changed

+503
-5
lines changed

src/rotator_library/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3122,7 +3122,9 @@ async def get_quota_stats(
31223122
)
31233123
else:
31243124
group_stats["total_requests_remaining"] = 0
3125-
group_stats["total_remaining_pct"] = None
3125+
# Fallback to avg_remaining_pct when max_requests unavailable
3126+
# This handles providers like Firmware that only provide percentage
3127+
group_stats["total_remaining_pct"] = group_stats.get("avg_remaining_pct")
31263128

31273129
prov_stats["quota_groups"][group_name] = group_stats
31283130

@@ -3188,14 +3190,22 @@ async def get_quota_stats(
31883190
requests_remaining = (
31893191
max(0, max_req - req_count) if max_req else 0
31903192
)
3193+
3194+
# Determine display format
3195+
# Priority: requests (if max known) > percentage (if baseline available) > unknown
3196+
if max_req:
3197+
display = f"{requests_remaining}/{max_req}"
3198+
elif remaining_pct is not None:
3199+
display = f"{remaining_pct}%"
3200+
else:
3201+
display = "?/?"
3202+
31913203
cred["model_groups"][group_name] = {
31923204
"remaining_pct": remaining_pct,
31933205
"requests_used": req_count,
31943206
"requests_remaining": requests_remaining,
31953207
"requests_max": max_req,
3196-
"display": f"{requests_remaining}/{max_req}"
3197-
if max_req
3198-
else f"?/?",
3208+
"display": display,
31993209
"is_exhausted": is_exhausted,
32003210
"reset_time_iso": reset_iso,
32013211
"models": group_models,
Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
"""
2+
Firmware.ai Provider with Quota Tracking
3+
4+
Provider implementation for the Firmware.ai API with 5-hour rolling window quota tracking.
5+
Uses the FirmwareQuotaTracker mixin to fetch quota usage from their API.
6+
7+
Environment variables:
8+
FIRMWARE_API_BASE: API base URL (default: https://app.firmware.ai/api/v1)
9+
FIRMWARE_API_KEY: API key for authentication
10+
FIRMWARE_QUOTA_REFRESH_INTERVAL: Quota refresh interval in seconds (default: 300)
11+
"""
12+
13+
import asyncio
14+
import httpx
15+
import os
16+
from typing import Any, Dict, List, Optional, TYPE_CHECKING
17+
18+
from .provider_interface import ProviderInterface
19+
from .utilities.firmware_quota_tracker import FirmwareQuotaTracker
20+
21+
if TYPE_CHECKING:
22+
from ..usage_manager import UsageManager
23+
24+
import logging
25+
26+
lib_logger = logging.getLogger("rotator_library")
27+
28+
# Concurrency limit for parallel quota fetches
29+
QUOTA_FETCH_CONCURRENCY = 5
30+
31+
32+
class FirmwareProvider(FirmwareQuotaTracker, ProviderInterface):
33+
"""
34+
Provider implementation for the Firmware.ai API with quota tracking.
35+
"""
36+
37+
# Quota groups for tracking 5-hour rolling window limits
38+
# Uses a virtual model "_quota" for credential-level quota tracking
39+
model_quota_groups = {
40+
"firmware_global": ["_quota"],
41+
}
42+
43+
def __init__(self, *args, **kwargs):
44+
"""Initialize FirmwareProvider with quota tracking."""
45+
super().__init__(*args, **kwargs)
46+
47+
# Quota tracking cache and refresh interval
48+
self._quota_cache: Dict[str, Dict[str, Any]] = {}
49+
try:
50+
self._quota_refresh_interval = int(
51+
os.environ.get("FIRMWARE_QUOTA_REFRESH_INTERVAL", "300")
52+
)
53+
except ValueError:
54+
lib_logger.warning(
55+
"Invalid FIRMWARE_QUOTA_REFRESH_INTERVAL value, using default 300"
56+
)
57+
self._quota_refresh_interval = 300
58+
59+
# API base URL (default to Firmware.ai)
60+
self.api_base = os.environ.get(
61+
"FIRMWARE_API_BASE", "https://app.firmware.ai/api/v1"
62+
)
63+
64+
def get_model_quota_group(self, model: str) -> Optional[str]:
65+
"""
66+
Get the quota group for a model.
67+
68+
All Firmware.ai models share the same credential-level quota pool,
69+
so they all belong to the same quota group.
70+
71+
Args:
72+
model: Model name (ignored - all models share quota)
73+
74+
Returns:
75+
Quota group identifier for shared credential-level tracking
76+
"""
77+
return "firmware_global"
78+
79+
def get_models_in_quota_group(self, group: str) -> List[str]:
80+
"""
81+
Get all models in a quota group.
82+
83+
For Firmware.ai, we use a virtual model "_quota" to track the
84+
credential-level 5-hour rolling window quota.
85+
86+
Args:
87+
group: Quota group name
88+
89+
Returns:
90+
List of model names in the group
91+
"""
92+
if group == "firmware_global":
93+
return ["_quota"]
94+
return []
95+
96+
def get_usage_reset_config(self, credential: str) -> Optional[Dict[str, Any]]:
97+
"""
98+
Return usage reset configuration for Firmware.ai credentials.
99+
100+
Firmware.ai uses per_model mode to track usage at the model level,
101+
with 5-hour rolling window quotas managed via the background job.
102+
103+
Args:
104+
credential: The API key (unused, same config for all)
105+
106+
Returns:
107+
Configuration with per_model mode and 5-hour window
108+
"""
109+
return {
110+
"mode": "per_model",
111+
"window_seconds": 18000, # 5 hours (5-hour rolling window)
112+
"field_name": "models",
113+
}
114+
115+
async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]:
116+
"""
117+
Fetch available models from the Firmware.ai API.
118+
119+
Args:
120+
api_key: Firmware.ai API key
121+
client: HTTP client
122+
123+
Returns:
124+
List of model names prefixed with 'firmware/'
125+
"""
126+
try:
127+
response = await client.get(
128+
f"{self.api_base.rstrip('/')}/models",
129+
headers={"Authorization": f"Bearer {api_key}"},
130+
)
131+
response.raise_for_status()
132+
return [
133+
f"firmware/{model['id']}" for model in response.json().get("data", [])
134+
]
135+
except (httpx.RequestError, httpx.HTTPStatusError) as e:
136+
lib_logger.error(f"Failed to fetch Firmware.ai models: {e}")
137+
return []
138+
139+
# =========================================================================
140+
# BACKGROUND JOB CONFIGURATION
141+
# =========================================================================
142+
143+
def get_background_job_config(self) -> Optional[Dict[str, Any]]:
144+
"""
145+
Configure periodic quota usage refresh.
146+
147+
Returns:
148+
Background job configuration for quota refresh
149+
"""
150+
return {
151+
"interval": self._quota_refresh_interval,
152+
"name": "firmware_quota_refresh",
153+
"run_on_start": True,
154+
}
155+
156+
async def run_background_job(
157+
self,
158+
usage_manager: "UsageManager",
159+
credentials: List[str],
160+
) -> None:
161+
"""
162+
Refresh quota usage for all credentials in parallel.
163+
164+
Args:
165+
usage_manager: UsageManager instance
166+
credentials: List of API keys
167+
"""
168+
semaphore = asyncio.Semaphore(QUOTA_FETCH_CONCURRENCY)
169+
170+
async def refresh_single_credential(
171+
api_key: str, client: httpx.AsyncClient
172+
) -> None:
173+
async with semaphore:
174+
try:
175+
usage_data = await self.fetch_quota_usage(api_key, client)
176+
177+
if usage_data.get("status") == "success":
178+
# Update quota cache
179+
self._quota_cache[api_key] = usage_data
180+
181+
# Calculate values for usage manager
182+
remaining_fraction = usage_data.get("remaining_fraction", 0.0)
183+
reset_ts = usage_data.get("reset_at")
184+
185+
# Store baseline in usage manager
186+
# Since Firmware.ai uses credential-level quota, we use a virtual model name
187+
await usage_manager.update_quota_baseline(
188+
api_key,
189+
"firmware/_quota", # Virtual model for credential-level tracking
190+
remaining_fraction,
191+
# No max_requests - Firmware.ai doesn't expose this
192+
reset_timestamp=reset_ts,
193+
)
194+
195+
lib_logger.debug(
196+
f"Updated Firmware.ai quota baseline: "
197+
f"{remaining_fraction * 100:.1f}% remaining, "
198+
f"active_window={usage_data.get('has_active_window', False)}"
199+
)
200+
201+
except Exception as e:
202+
lib_logger.warning(f"Failed to refresh Firmware.ai quota usage: {e}")
203+
204+
# Fetch all credentials in parallel with shared HTTP client
205+
async with httpx.AsyncClient(timeout=30.0) as client:
206+
tasks = [
207+
refresh_single_credential(api_key, client) for api_key in credentials
208+
]
209+
await asyncio.gather(*tasks, return_exceptions=True)

0 commit comments

Comments
 (0)