Spaces:
Paused
feat(quota-viewer): ✨ add quota and usage statistics viewer system
Browse filesThis introduces a comprehensive quota and usage statistics monitoring system with API endpoints and a rich terminal UI for real-time credential consumption visibility.
Key components:
- Add GET/POST /v1/quota-stats API endpoints to main proxy for retrieving and refreshing statistics
- Implement full-featured TUI viewer (quota_viewer.py) with provider summaries, credential details, progress bars, and remote management
- Add configuration system (quota_viewer_config.py) supporting multiple remote proxies with API key management
- Extend RotatingClient with get_quota_stats(), reload_usage_from_disk(), and force_refresh_quota() methods
- Enhance usage_manager with get_stats_for_endpoint() for aggregated statistics
- Integrate viewer into launcher TUI as new menu option 5
Features include toggle between current and global/lifetime stats, token formatting, cooldown timers, quota group visualization, and estimated cost tracking. Supports Antigravity quota group enrichment and live API refresh.
Also in this commit:
- fix(antigravity): increase empty response retry attempts from 4 to 6
- src/proxy_app/launcher_tui.py +23 -6
- src/proxy_app/main.py +139 -0
- src/proxy_app/quota_viewer.py +1086 -0
- src/proxy_app/quota_viewer_config.py +288 -0
- src/rotator_library/client.py +277 -0
- src/rotator_library/providers/antigravity_provider.py +1 -1
- src/rotator_library/usage_manager.py +476 -0
|
@@ -429,9 +429,10 @@ class LauncherTUI:
|
|
| 429 |
self.console.print(" 3. 🔑 Manage Credentials")
|
| 430 |
|
| 431 |
self.console.print(" 4. 📊 View Provider & Advanced Settings")
|
| 432 |
-
self.console.print(" 5.
|
| 433 |
-
self.console.print(" 6.
|
| 434 |
-
self.console.print(" 7.
|
|
|
|
| 435 |
|
| 436 |
self.console.print()
|
| 437 |
self.console.print("━" * 70)
|
|
@@ -439,7 +440,7 @@ class LauncherTUI:
|
|
| 439 |
|
| 440 |
choice = Prompt.ask(
|
| 441 |
"Select option",
|
| 442 |
-
choices=["1", "2", "3", "4", "5", "6", "7"],
|
| 443 |
show_choices=False,
|
| 444 |
)
|
| 445 |
|
|
@@ -452,12 +453,14 @@ class LauncherTUI:
|
|
| 452 |
elif choice == "4":
|
| 453 |
self.show_provider_settings_menu()
|
| 454 |
elif choice == "5":
|
|
|
|
|
|
|
| 455 |
load_dotenv(dotenv_path=_get_env_file(), override=True)
|
| 456 |
self.config = LauncherConfig() # Reload config
|
| 457 |
self.console.print("\n[green]✅ Configuration reloaded![/green]")
|
| 458 |
-
elif choice == "6":
|
| 459 |
-
self.show_about()
|
| 460 |
elif choice == "7":
|
|
|
|
|
|
|
| 461 |
self.running = False
|
| 462 |
sys.exit(0)
|
| 463 |
|
|
@@ -874,6 +877,20 @@ class LauncherTUI:
|
|
| 874 |
# Reload environment after settings tool
|
| 875 |
load_dotenv(dotenv_path=_get_env_file(), override=True)
|
| 876 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 877 |
def show_about(self):
|
| 878 |
"""Display About page with project information"""
|
| 879 |
clear_screen()
|
|
|
|
| 429 |
self.console.print(" 3. 🔑 Manage Credentials")
|
| 430 |
|
| 431 |
self.console.print(" 4. 📊 View Provider & Advanced Settings")
|
| 432 |
+
self.console.print(" 5. 📈 View Quota & Usage Stats")
|
| 433 |
+
self.console.print(" 6. 🔄 Reload Configuration")
|
| 434 |
+
self.console.print(" 7. ℹ️ About")
|
| 435 |
+
self.console.print(" 8. 🚪 Exit")
|
| 436 |
|
| 437 |
self.console.print()
|
| 438 |
self.console.print("━" * 70)
|
|
|
|
| 440 |
|
| 441 |
choice = Prompt.ask(
|
| 442 |
"Select option",
|
| 443 |
+
choices=["1", "2", "3", "4", "5", "6", "7", "8"],
|
| 444 |
show_choices=False,
|
| 445 |
)
|
| 446 |
|
|
|
|
| 453 |
elif choice == "4":
|
| 454 |
self.show_provider_settings_menu()
|
| 455 |
elif choice == "5":
|
| 456 |
+
self.launch_quota_viewer()
|
| 457 |
+
elif choice == "6":
|
| 458 |
load_dotenv(dotenv_path=_get_env_file(), override=True)
|
| 459 |
self.config = LauncherConfig() # Reload config
|
| 460 |
self.console.print("\n[green]✅ Configuration reloaded![/green]")
|
|
|
|
|
|
|
| 461 |
elif choice == "7":
|
| 462 |
+
self.show_about()
|
| 463 |
+
elif choice == "8":
|
| 464 |
self.running = False
|
| 465 |
sys.exit(0)
|
| 466 |
|
|
|
|
| 877 |
# Reload environment after settings tool
|
| 878 |
load_dotenv(dotenv_path=_get_env_file(), override=True)
|
| 879 |
|
| 880 |
+
def launch_quota_viewer(self):
|
| 881 |
+
"""Launch the quota stats viewer"""
|
| 882 |
+
clear_screen()
|
| 883 |
+
|
| 884 |
+
self.console.print("━" * 70)
|
| 885 |
+
self.console.print("Quota & Usage Statistics Viewer")
|
| 886 |
+
self.console.print("━" * 70)
|
| 887 |
+
self.console.print()
|
| 888 |
+
|
| 889 |
+
# Import the lightweight viewer (no heavy imports)
|
| 890 |
+
from proxy_app.quota_viewer import run_quota_viewer
|
| 891 |
+
|
| 892 |
+
run_quota_viewer()
|
| 893 |
+
|
| 894 |
def show_about(self):
|
| 895 |
"""Display About page with project information"""
|
| 896 |
clear_screen()
|
|
@@ -1148,6 +1148,145 @@ async def list_providers(_=Depends(verify_api_key)):
|
|
| 1148 |
return list(PROVIDER_PLUGINS.keys())
|
| 1149 |
|
| 1150 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1151 |
@app.post("/v1/token-count")
|
| 1152 |
async def token_count(
|
| 1153 |
request: Request,
|
|
|
|
| 1148 |
return list(PROVIDER_PLUGINS.keys())
|
| 1149 |
|
| 1150 |
|
| 1151 |
+
@app.get("/v1/quota-stats")
|
| 1152 |
+
async def get_quota_stats(
|
| 1153 |
+
request: Request,
|
| 1154 |
+
client: RotatingClient = Depends(get_rotating_client),
|
| 1155 |
+
_=Depends(verify_api_key),
|
| 1156 |
+
provider: str = None,
|
| 1157 |
+
):
|
| 1158 |
+
"""
|
| 1159 |
+
Returns quota and usage statistics for all credentials.
|
| 1160 |
+
|
| 1161 |
+
This returns cached data from the proxy without making external API calls.
|
| 1162 |
+
Use POST to reload from disk or force refresh from external APIs.
|
| 1163 |
+
|
| 1164 |
+
Query Parameters:
|
| 1165 |
+
provider: Optional filter to return stats for a specific provider only
|
| 1166 |
+
|
| 1167 |
+
Returns:
|
| 1168 |
+
{
|
| 1169 |
+
"providers": {
|
| 1170 |
+
"provider_name": {
|
| 1171 |
+
"credential_count": int,
|
| 1172 |
+
"active_count": int,
|
| 1173 |
+
"on_cooldown_count": int,
|
| 1174 |
+
"exhausted_count": int,
|
| 1175 |
+
"total_requests": int,
|
| 1176 |
+
"tokens": {...},
|
| 1177 |
+
"approx_cost": float | null,
|
| 1178 |
+
"quota_groups": {...}, // For Antigravity
|
| 1179 |
+
"credentials": [...]
|
| 1180 |
+
}
|
| 1181 |
+
},
|
| 1182 |
+
"summary": {...},
|
| 1183 |
+
"data_source": "cache",
|
| 1184 |
+
"timestamp": float
|
| 1185 |
+
}
|
| 1186 |
+
"""
|
| 1187 |
+
try:
|
| 1188 |
+
stats = await client.get_quota_stats(provider_filter=provider)
|
| 1189 |
+
return stats
|
| 1190 |
+
except Exception as e:
|
| 1191 |
+
logging.error(f"Failed to get quota stats: {e}")
|
| 1192 |
+
raise HTTPException(status_code=500, detail=str(e))
|
| 1193 |
+
|
| 1194 |
+
|
| 1195 |
+
@app.post("/v1/quota-stats")
|
| 1196 |
+
async def refresh_quota_stats(
|
| 1197 |
+
request: Request,
|
| 1198 |
+
client: RotatingClient = Depends(get_rotating_client),
|
| 1199 |
+
_=Depends(verify_api_key),
|
| 1200 |
+
):
|
| 1201 |
+
"""
|
| 1202 |
+
Refresh quota and usage statistics.
|
| 1203 |
+
|
| 1204 |
+
Request body:
|
| 1205 |
+
{
|
| 1206 |
+
"action": "reload" | "force_refresh",
|
| 1207 |
+
"scope": "all" | "provider" | "credential",
|
| 1208 |
+
"provider": "antigravity", // required if scope != "all"
|
| 1209 |
+
"credential": "antigravity_oauth_1.json" // required if scope == "credential"
|
| 1210 |
+
}
|
| 1211 |
+
|
| 1212 |
+
Actions:
|
| 1213 |
+
- reload: Re-read data from disk (no external API calls)
|
| 1214 |
+
- force_refresh: For Antigravity, fetch live quota from API.
|
| 1215 |
+
For other providers, same as reload.
|
| 1216 |
+
|
| 1217 |
+
Returns:
|
| 1218 |
+
Same as GET, plus a "refresh_result" field with operation details.
|
| 1219 |
+
"""
|
| 1220 |
+
try:
|
| 1221 |
+
data = await request.json()
|
| 1222 |
+
action = data.get("action", "reload")
|
| 1223 |
+
scope = data.get("scope", "all")
|
| 1224 |
+
provider = data.get("provider")
|
| 1225 |
+
credential = data.get("credential")
|
| 1226 |
+
|
| 1227 |
+
# Validate parameters
|
| 1228 |
+
if action not in ("reload", "force_refresh"):
|
| 1229 |
+
raise HTTPException(
|
| 1230 |
+
status_code=400,
|
| 1231 |
+
detail="action must be 'reload' or 'force_refresh'",
|
| 1232 |
+
)
|
| 1233 |
+
|
| 1234 |
+
if scope not in ("all", "provider", "credential"):
|
| 1235 |
+
raise HTTPException(
|
| 1236 |
+
status_code=400,
|
| 1237 |
+
detail="scope must be 'all', 'provider', or 'credential'",
|
| 1238 |
+
)
|
| 1239 |
+
|
| 1240 |
+
if scope in ("provider", "credential") and not provider:
|
| 1241 |
+
raise HTTPException(
|
| 1242 |
+
status_code=400,
|
| 1243 |
+
detail="'provider' is required when scope is 'provider' or 'credential'",
|
| 1244 |
+
)
|
| 1245 |
+
|
| 1246 |
+
if scope == "credential" and not credential:
|
| 1247 |
+
raise HTTPException(
|
| 1248 |
+
status_code=400,
|
| 1249 |
+
detail="'credential' is required when scope is 'credential'",
|
| 1250 |
+
)
|
| 1251 |
+
|
| 1252 |
+
refresh_result = {
|
| 1253 |
+
"action": action,
|
| 1254 |
+
"scope": scope,
|
| 1255 |
+
"provider": provider,
|
| 1256 |
+
"credential": credential,
|
| 1257 |
+
}
|
| 1258 |
+
|
| 1259 |
+
if action == "reload":
|
| 1260 |
+
# Just reload from disk
|
| 1261 |
+
start_time = time.time()
|
| 1262 |
+
await client.reload_usage_from_disk()
|
| 1263 |
+
refresh_result["duration_ms"] = int((time.time() - start_time) * 1000)
|
| 1264 |
+
refresh_result["success"] = True
|
| 1265 |
+
refresh_result["message"] = "Reloaded usage data from disk"
|
| 1266 |
+
|
| 1267 |
+
elif action == "force_refresh":
|
| 1268 |
+
# Force refresh from external API (for supported providers like Antigravity)
|
| 1269 |
+
result = await client.force_refresh_quota(
|
| 1270 |
+
provider=provider if scope in ("provider", "credential") else None,
|
| 1271 |
+
credential=credential if scope == "credential" else None,
|
| 1272 |
+
)
|
| 1273 |
+
refresh_result.update(result)
|
| 1274 |
+
refresh_result["success"] = result["failed_count"] == 0
|
| 1275 |
+
|
| 1276 |
+
# Get updated stats
|
| 1277 |
+
stats = await client.get_quota_stats(provider_filter=provider)
|
| 1278 |
+
stats["refresh_result"] = refresh_result
|
| 1279 |
+
stats["data_source"] = "refreshed"
|
| 1280 |
+
|
| 1281 |
+
return stats
|
| 1282 |
+
|
| 1283 |
+
except HTTPException:
|
| 1284 |
+
raise
|
| 1285 |
+
except Exception as e:
|
| 1286 |
+
logging.error(f"Failed to refresh quota stats: {e}")
|
| 1287 |
+
raise HTTPException(status_code=500, detail=str(e))
|
| 1288 |
+
|
| 1289 |
+
|
| 1290 |
@app.post("/v1/token-count")
|
| 1291 |
async def token_count(
|
| 1292 |
request: Request,
|
|
@@ -0,0 +1,1086 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Lightweight Quota Stats Viewer TUI.
|
| 3 |
+
|
| 4 |
+
Connects to a running proxy to display quota and usage statistics.
|
| 5 |
+
Uses only httpx + rich (no heavy rotator_library imports).
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import os
|
| 9 |
+
import sys
|
| 10 |
+
import time
|
| 11 |
+
from datetime import datetime, timezone
|
| 12 |
+
from typing import Any, Dict, List, Optional, Tuple
|
| 13 |
+
|
| 14 |
+
import httpx
|
| 15 |
+
from rich.console import Console
|
| 16 |
+
from rich.panel import Panel
|
| 17 |
+
from rich.progress import BarColumn, Progress, TextColumn
|
| 18 |
+
from rich.prompt import Prompt
|
| 19 |
+
from rich.table import Table
|
| 20 |
+
from rich.text import Text
|
| 21 |
+
|
| 22 |
+
from .quota_viewer_config import QuotaViewerConfig
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def clear_screen():
|
| 26 |
+
"""Clear the terminal screen."""
|
| 27 |
+
os.system("cls" if os.name == "nt" else "clear")
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def format_tokens(count: int) -> str:
|
| 31 |
+
"""Format token count for display (e.g., 125000 -> 125k)."""
|
| 32 |
+
if count >= 1_000_000:
|
| 33 |
+
return f"{count / 1_000_000:.1f}M"
|
| 34 |
+
elif count >= 1_000:
|
| 35 |
+
return f"{count / 1_000:.0f}k"
|
| 36 |
+
return str(count)
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def format_cost(cost: Optional[float]) -> str:
|
| 40 |
+
"""Format cost for display."""
|
| 41 |
+
if cost is None or cost == 0:
|
| 42 |
+
return "-"
|
| 43 |
+
if cost < 0.01:
|
| 44 |
+
return f"${cost:.4f}"
|
| 45 |
+
return f"${cost:.2f}"
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def format_time_ago(timestamp: Optional[float]) -> str:
|
| 49 |
+
"""Format timestamp as relative time (e.g., '5 min ago')."""
|
| 50 |
+
if not timestamp:
|
| 51 |
+
return "Never"
|
| 52 |
+
try:
|
| 53 |
+
delta = time.time() - timestamp
|
| 54 |
+
if delta < 60:
|
| 55 |
+
return f"{int(delta)}s ago"
|
| 56 |
+
elif delta < 3600:
|
| 57 |
+
return f"{int(delta / 60)} min ago"
|
| 58 |
+
elif delta < 86400:
|
| 59 |
+
return f"{int(delta / 3600)}h ago"
|
| 60 |
+
else:
|
| 61 |
+
return f"{int(delta / 86400)}d ago"
|
| 62 |
+
except (ValueError, OSError):
|
| 63 |
+
return "Unknown"
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def format_reset_time(iso_time: Optional[str]) -> str:
|
| 67 |
+
"""Format ISO time string for display."""
|
| 68 |
+
if not iso_time:
|
| 69 |
+
return "-"
|
| 70 |
+
try:
|
| 71 |
+
dt = datetime.fromisoformat(iso_time.replace("Z", "+00:00"))
|
| 72 |
+
# Convert to local time
|
| 73 |
+
local_dt = dt.astimezone()
|
| 74 |
+
return local_dt.strftime("%b %d %H:%M")
|
| 75 |
+
except (ValueError, AttributeError):
|
| 76 |
+
return iso_time[:16] if iso_time else "-"
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def create_progress_bar(percent: Optional[int], width: int = 10) -> str:
|
| 80 |
+
"""Create a text-based progress bar."""
|
| 81 |
+
if percent is None:
|
| 82 |
+
return "░" * width
|
| 83 |
+
filled = int(percent / 100 * width)
|
| 84 |
+
return "▓" * filled + "░" * (width - filled)
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def format_cooldown(seconds: int) -> str:
|
| 88 |
+
"""Format cooldown seconds as human-readable string."""
|
| 89 |
+
if seconds < 60:
|
| 90 |
+
return f"{seconds}s"
|
| 91 |
+
elif seconds < 3600:
|
| 92 |
+
mins = seconds // 60
|
| 93 |
+
secs = seconds % 60
|
| 94 |
+
return f"{mins}m {secs}s" if secs > 0 else f"{mins}m"
|
| 95 |
+
else:
|
| 96 |
+
hours = seconds // 3600
|
| 97 |
+
mins = (seconds % 3600) // 60
|
| 98 |
+
return f"{hours}h {mins}m" if mins > 0 else f"{hours}h"
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
class QuotaViewer:
|
| 102 |
+
"""Main Quota Viewer TUI class."""
|
| 103 |
+
|
| 104 |
+
def __init__(self, config: Optional[QuotaViewerConfig] = None):
|
| 105 |
+
"""
|
| 106 |
+
Initialize the viewer.
|
| 107 |
+
|
| 108 |
+
Args:
|
| 109 |
+
config: Optional config object. If not provided, one will be created.
|
| 110 |
+
"""
|
| 111 |
+
self.console = Console()
|
| 112 |
+
self.config = config or QuotaViewerConfig()
|
| 113 |
+
self.config.sync_with_launcher_config()
|
| 114 |
+
|
| 115 |
+
self.current_remote: Optional[Dict[str, Any]] = None
|
| 116 |
+
self.cached_stats: Optional[Dict[str, Any]] = None
|
| 117 |
+
self.last_error: Optional[str] = None
|
| 118 |
+
self.running = True
|
| 119 |
+
self.view_mode = "current" # "current" or "global"
|
| 120 |
+
|
| 121 |
+
def _get_headers(self) -> Dict[str, str]:
|
| 122 |
+
"""Get HTTP headers including auth if configured."""
|
| 123 |
+
headers = {}
|
| 124 |
+
if self.current_remote and self.current_remote.get("api_key"):
|
| 125 |
+
headers["Authorization"] = f"Bearer {self.current_remote['api_key']}"
|
| 126 |
+
return headers
|
| 127 |
+
|
| 128 |
+
def _get_base_url(self) -> str:
|
| 129 |
+
"""Get base URL for the current remote."""
|
| 130 |
+
if not self.current_remote:
|
| 131 |
+
return "http://127.0.0.1:8000"
|
| 132 |
+
host = self.current_remote.get("host", "127.0.0.1")
|
| 133 |
+
port = self.current_remote.get("port", 8000)
|
| 134 |
+
# Use https if port is 443 or host looks like a domain
|
| 135 |
+
scheme = "https" if port == 443 or "." in host else "http"
|
| 136 |
+
return f"{scheme}://{host}:{port}"
|
| 137 |
+
|
| 138 |
+
def check_connection(
|
| 139 |
+
self, remote: Dict[str, Any], timeout: float = 3.0
|
| 140 |
+
) -> Tuple[bool, str]:
|
| 141 |
+
"""
|
| 142 |
+
Check if a remote proxy is reachable.
|
| 143 |
+
|
| 144 |
+
Args:
|
| 145 |
+
remote: Remote configuration dict
|
| 146 |
+
timeout: Connection timeout in seconds
|
| 147 |
+
|
| 148 |
+
Returns:
|
| 149 |
+
Tuple of (is_online, status_message)
|
| 150 |
+
"""
|
| 151 |
+
host = remote.get("host", "127.0.0.1")
|
| 152 |
+
port = remote.get("port", 8000)
|
| 153 |
+
scheme = "https" if port == 443 or "." in host else "http"
|
| 154 |
+
url = f"{scheme}://{host}:{port}/"
|
| 155 |
+
|
| 156 |
+
headers = {}
|
| 157 |
+
if remote.get("api_key"):
|
| 158 |
+
headers["Authorization"] = f"Bearer {remote['api_key']}"
|
| 159 |
+
|
| 160 |
+
try:
|
| 161 |
+
with httpx.Client(timeout=timeout) as client:
|
| 162 |
+
response = client.get(url, headers=headers)
|
| 163 |
+
if response.status_code == 200:
|
| 164 |
+
return True, "Online"
|
| 165 |
+
elif response.status_code == 401:
|
| 166 |
+
return False, "Auth failed"
|
| 167 |
+
else:
|
| 168 |
+
return False, f"HTTP {response.status_code}"
|
| 169 |
+
except httpx.ConnectError:
|
| 170 |
+
return False, "Offline"
|
| 171 |
+
except httpx.TimeoutException:
|
| 172 |
+
return False, "Timeout"
|
| 173 |
+
except Exception as e:
|
| 174 |
+
return False, str(e)[:20]
|
| 175 |
+
|
| 176 |
+
def fetch_stats(self, provider: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
| 177 |
+
"""
|
| 178 |
+
Fetch quota stats from the current remote.
|
| 179 |
+
|
| 180 |
+
Args:
|
| 181 |
+
provider: Optional provider filter
|
| 182 |
+
|
| 183 |
+
Returns:
|
| 184 |
+
Stats dict or None on failure
|
| 185 |
+
"""
|
| 186 |
+
url = f"{self._get_base_url()}/v1/quota-stats"
|
| 187 |
+
if provider:
|
| 188 |
+
url += f"?provider={provider}"
|
| 189 |
+
|
| 190 |
+
try:
|
| 191 |
+
with httpx.Client(timeout=30.0) as client:
|
| 192 |
+
response = client.get(url, headers=self._get_headers())
|
| 193 |
+
|
| 194 |
+
if response.status_code == 401:
|
| 195 |
+
self.last_error = "Authentication failed. Check API key."
|
| 196 |
+
return None
|
| 197 |
+
elif response.status_code != 200:
|
| 198 |
+
self.last_error = (
|
| 199 |
+
f"HTTP {response.status_code}: {response.text[:100]}"
|
| 200 |
+
)
|
| 201 |
+
return None
|
| 202 |
+
|
| 203 |
+
self.cached_stats = response.json()
|
| 204 |
+
self.last_error = None
|
| 205 |
+
return self.cached_stats
|
| 206 |
+
|
| 207 |
+
except httpx.ConnectError:
|
| 208 |
+
self.last_error = "Connection failed. Is the proxy running?"
|
| 209 |
+
return None
|
| 210 |
+
except httpx.TimeoutException:
|
| 211 |
+
self.last_error = "Request timed out."
|
| 212 |
+
return None
|
| 213 |
+
except Exception as e:
|
| 214 |
+
self.last_error = str(e)
|
| 215 |
+
return None
|
| 216 |
+
|
| 217 |
+
def post_action(
|
| 218 |
+
self,
|
| 219 |
+
action: str,
|
| 220 |
+
scope: str = "all",
|
| 221 |
+
provider: Optional[str] = None,
|
| 222 |
+
credential: Optional[str] = None,
|
| 223 |
+
) -> Optional[Dict[str, Any]]:
|
| 224 |
+
"""
|
| 225 |
+
Post a refresh action to the proxy.
|
| 226 |
+
|
| 227 |
+
Args:
|
| 228 |
+
action: "reload" or "force_refresh"
|
| 229 |
+
scope: "all", "provider", or "credential"
|
| 230 |
+
provider: Provider name (required for scope != "all")
|
| 231 |
+
credential: Credential identifier (required for scope == "credential")
|
| 232 |
+
|
| 233 |
+
Returns:
|
| 234 |
+
Response dict or None on failure
|
| 235 |
+
"""
|
| 236 |
+
url = f"{self._get_base_url()}/v1/quota-stats"
|
| 237 |
+
payload = {
|
| 238 |
+
"action": action,
|
| 239 |
+
"scope": scope,
|
| 240 |
+
}
|
| 241 |
+
if provider:
|
| 242 |
+
payload["provider"] = provider
|
| 243 |
+
if credential:
|
| 244 |
+
payload["credential"] = credential
|
| 245 |
+
|
| 246 |
+
try:
|
| 247 |
+
with httpx.Client(timeout=60.0) as client:
|
| 248 |
+
response = client.post(url, headers=self._get_headers(), json=payload)
|
| 249 |
+
|
| 250 |
+
if response.status_code == 401:
|
| 251 |
+
self.last_error = "Authentication failed. Check API key."
|
| 252 |
+
return None
|
| 253 |
+
elif response.status_code != 200:
|
| 254 |
+
self.last_error = (
|
| 255 |
+
f"HTTP {response.status_code}: {response.text[:100]}"
|
| 256 |
+
)
|
| 257 |
+
return None
|
| 258 |
+
|
| 259 |
+
result = response.json()
|
| 260 |
+
self.cached_stats = result
|
| 261 |
+
self.last_error = None
|
| 262 |
+
return result
|
| 263 |
+
|
| 264 |
+
except httpx.ConnectError:
|
| 265 |
+
self.last_error = "Connection failed. Is the proxy running?"
|
| 266 |
+
return None
|
| 267 |
+
except httpx.TimeoutException:
|
| 268 |
+
self.last_error = "Request timed out."
|
| 269 |
+
return None
|
| 270 |
+
except Exception as e:
|
| 271 |
+
self.last_error = str(e)
|
| 272 |
+
return None
|
| 273 |
+
|
| 274 |
+
# =========================================================================
|
| 275 |
+
# DISPLAY SCREENS
|
| 276 |
+
# =========================================================================
|
| 277 |
+
|
| 278 |
+
def show_connection_error(self):
|
| 279 |
+
"""Display connection error screen."""
|
| 280 |
+
clear_screen()
|
| 281 |
+
self.console.print(
|
| 282 |
+
Panel(
|
| 283 |
+
Text.from_markup(
|
| 284 |
+
"[bold red]Connection Error[/bold red]\n\n"
|
| 285 |
+
f"{self.last_error or 'Unknown error'}\n\n"
|
| 286 |
+
"[bold]This tool requires the proxy to be running.[/bold]\n"
|
| 287 |
+
"Start the proxy first, then try again.\n\n"
|
| 288 |
+
"[dim]Tip: Select option 1 from the main menu to run the proxy.[/dim]"
|
| 289 |
+
),
|
| 290 |
+
border_style="red",
|
| 291 |
+
expand=False,
|
| 292 |
+
)
|
| 293 |
+
)
|
| 294 |
+
Prompt.ask("\nPress Enter to return to main menu", default="")
|
| 295 |
+
|
| 296 |
+
def show_summary_screen(self):
|
| 297 |
+
"""Display the main summary screen with all providers."""
|
| 298 |
+
clear_screen()
|
| 299 |
+
|
| 300 |
+
# Header
|
| 301 |
+
remote_name = (
|
| 302 |
+
self.current_remote.get("name", "Unknown")
|
| 303 |
+
if self.current_remote
|
| 304 |
+
else "None"
|
| 305 |
+
)
|
| 306 |
+
remote_host = self.current_remote.get("host", "") if self.current_remote else ""
|
| 307 |
+
remote_port = self.current_remote.get("port", "") if self.current_remote else ""
|
| 308 |
+
|
| 309 |
+
# Calculate data age
|
| 310 |
+
data_age = ""
|
| 311 |
+
if self.cached_stats and self.cached_stats.get("timestamp"):
|
| 312 |
+
age_seconds = int(time.time() - self.cached_stats["timestamp"])
|
| 313 |
+
data_age = f"Data age: {age_seconds}s"
|
| 314 |
+
|
| 315 |
+
# View mode indicator
|
| 316 |
+
if self.view_mode == "global":
|
| 317 |
+
view_label = "[magenta]📊 Global/Lifetime[/magenta]"
|
| 318 |
+
else:
|
| 319 |
+
view_label = "[cyan]📈 Current Period[/cyan]"
|
| 320 |
+
|
| 321 |
+
self.console.print("━" * 78)
|
| 322 |
+
self.console.print(
|
| 323 |
+
f"[bold cyan]📈 Quota & Usage Statistics[/bold cyan] | {view_label}"
|
| 324 |
+
)
|
| 325 |
+
self.console.print("━" * 78)
|
| 326 |
+
self.console.print(
|
| 327 |
+
f"Connected to: [bold]{remote_name}[/bold] ({remote_host}:{remote_port}) "
|
| 328 |
+
f"[green]✅[/green] | {data_age}"
|
| 329 |
+
)
|
| 330 |
+
self.console.print()
|
| 331 |
+
|
| 332 |
+
if not self.cached_stats:
|
| 333 |
+
self.console.print("[yellow]No data available. Press R to reload.[/yellow]")
|
| 334 |
+
else:
|
| 335 |
+
# Build provider table
|
| 336 |
+
table = Table(box=None, show_header=True, header_style="bold")
|
| 337 |
+
table.add_column("Provider", style="cyan", min_width=12)
|
| 338 |
+
table.add_column("Creds", justify="center", min_width=6)
|
| 339 |
+
table.add_column("Quota Status", min_width=28)
|
| 340 |
+
table.add_column("Requests", justify="right", min_width=9)
|
| 341 |
+
table.add_column("Tokens (in/out)", min_width=22)
|
| 342 |
+
table.add_column("Cost", justify="right", min_width=8)
|
| 343 |
+
|
| 344 |
+
providers = self.cached_stats.get("providers", {})
|
| 345 |
+
provider_list = list(providers.keys())
|
| 346 |
+
|
| 347 |
+
for idx, (provider, prov_stats) in enumerate(providers.items(), 1):
|
| 348 |
+
cred_count = prov_stats.get("credential_count", 0)
|
| 349 |
+
|
| 350 |
+
# Use global stats if in global mode
|
| 351 |
+
if self.view_mode == "global":
|
| 352 |
+
stats_source = prov_stats.get("global", prov_stats)
|
| 353 |
+
total_requests = stats_source.get("total_requests", 0)
|
| 354 |
+
tokens = stats_source.get("tokens", {})
|
| 355 |
+
cost_value = stats_source.get("approx_cost")
|
| 356 |
+
else:
|
| 357 |
+
total_requests = prov_stats.get("total_requests", 0)
|
| 358 |
+
tokens = prov_stats.get("tokens", {})
|
| 359 |
+
cost_value = prov_stats.get("approx_cost")
|
| 360 |
+
|
| 361 |
+
# Format tokens
|
| 362 |
+
input_total = tokens.get("input_cached", 0) + tokens.get(
|
| 363 |
+
"input_uncached", 0
|
| 364 |
+
)
|
| 365 |
+
output = tokens.get("output", 0)
|
| 366 |
+
cache_pct = tokens.get("input_cache_pct", 0)
|
| 367 |
+
token_str = f"{format_tokens(input_total)}/{format_tokens(output)} ({cache_pct}% cached)"
|
| 368 |
+
|
| 369 |
+
# Format cost
|
| 370 |
+
cost_str = format_cost(cost_value)
|
| 371 |
+
|
| 372 |
+
# Build quota status string (for providers with quota groups)
|
| 373 |
+
quota_groups = prov_stats.get("quota_groups", {})
|
| 374 |
+
if quota_groups:
|
| 375 |
+
quota_lines = []
|
| 376 |
+
for group_name, group_stats in quota_groups.items():
|
| 377 |
+
avg_pct = group_stats.get("avg_remaining_pct", 0)
|
| 378 |
+
exhausted = group_stats.get("credentials_exhausted", 0)
|
| 379 |
+
total = group_stats.get("credentials_total", 0)
|
| 380 |
+
|
| 381 |
+
# Determine color based on remaining
|
| 382 |
+
if exhausted > 0:
|
| 383 |
+
color = "red"
|
| 384 |
+
status = f"({exhausted}/{total} exhausted)"
|
| 385 |
+
elif avg_pct < 20:
|
| 386 |
+
color = "yellow"
|
| 387 |
+
status = ""
|
| 388 |
+
else:
|
| 389 |
+
color = "green"
|
| 390 |
+
status = ""
|
| 391 |
+
|
| 392 |
+
bar = create_progress_bar(avg_pct)
|
| 393 |
+
display_name = group_name[:10]
|
| 394 |
+
quota_lines.append(
|
| 395 |
+
f"[{color}]{display_name}: {avg_pct}% {bar}[/{color}] {status}"
|
| 396 |
+
)
|
| 397 |
+
|
| 398 |
+
# First line goes in the main row
|
| 399 |
+
first_quota = quota_lines[0] if quota_lines else "-"
|
| 400 |
+
table.add_row(
|
| 401 |
+
provider,
|
| 402 |
+
str(cred_count),
|
| 403 |
+
first_quota,
|
| 404 |
+
str(total_requests),
|
| 405 |
+
token_str,
|
| 406 |
+
cost_str,
|
| 407 |
+
)
|
| 408 |
+
# Additional quota lines as sub-rows
|
| 409 |
+
for quota_line in quota_lines[1:]:
|
| 410 |
+
table.add_row("", "", quota_line, "", "", "")
|
| 411 |
+
else:
|
| 412 |
+
# No quota groups
|
| 413 |
+
table.add_row(
|
| 414 |
+
provider,
|
| 415 |
+
str(cred_count),
|
| 416 |
+
"-",
|
| 417 |
+
str(total_requests),
|
| 418 |
+
token_str,
|
| 419 |
+
cost_str,
|
| 420 |
+
)
|
| 421 |
+
|
| 422 |
+
# Add separator between providers (except last)
|
| 423 |
+
if idx < len(providers):
|
| 424 |
+
table.add_row(
|
| 425 |
+
"─" * 10, "─" * 4, "─" * 26, "─" * 7, "─" * 20, "─" * 6
|
| 426 |
+
)
|
| 427 |
+
|
| 428 |
+
self.console.print(table)
|
| 429 |
+
|
| 430 |
+
# Summary line - use global_summary if in global mode
|
| 431 |
+
if self.view_mode == "global":
|
| 432 |
+
summary = self.cached_stats.get(
|
| 433 |
+
"global_summary", self.cached_stats.get("summary", {})
|
| 434 |
+
)
|
| 435 |
+
else:
|
| 436 |
+
summary = self.cached_stats.get("summary", {})
|
| 437 |
+
|
| 438 |
+
total_creds = summary.get("total_credentials", 0)
|
| 439 |
+
total_requests = summary.get("total_requests", 0)
|
| 440 |
+
total_tokens = summary.get("tokens", {})
|
| 441 |
+
total_input = total_tokens.get("input_cached", 0) + total_tokens.get(
|
| 442 |
+
"input_uncached", 0
|
| 443 |
+
)
|
| 444 |
+
total_output = total_tokens.get("output", 0)
|
| 445 |
+
total_cost = format_cost(summary.get("approx_total_cost"))
|
| 446 |
+
|
| 447 |
+
self.console.print()
|
| 448 |
+
self.console.print(
|
| 449 |
+
f"[bold]Total:[/bold] {total_creds} credentials | "
|
| 450 |
+
f"{total_requests} requests | "
|
| 451 |
+
f"{format_tokens(total_input)}/{format_tokens(total_output)} tokens | "
|
| 452 |
+
f"{total_cost} cost"
|
| 453 |
+
)
|
| 454 |
+
|
| 455 |
+
# Menu
|
| 456 |
+
self.console.print()
|
| 457 |
+
self.console.print("━" * 78)
|
| 458 |
+
self.console.print()
|
| 459 |
+
|
| 460 |
+
# Build provider menu options
|
| 461 |
+
providers = self.cached_stats.get("providers", {}) if self.cached_stats else {}
|
| 462 |
+
provider_list = list(providers.keys())
|
| 463 |
+
|
| 464 |
+
for idx, provider in enumerate(provider_list, 1):
|
| 465 |
+
self.console.print(f" {idx}. View [cyan]{provider}[/cyan] details")
|
| 466 |
+
|
| 467 |
+
self.console.print()
|
| 468 |
+
self.console.print(" G. Toggle view mode (current/global)")
|
| 469 |
+
self.console.print(" R. Reload all stats (re-read from proxy)")
|
| 470 |
+
self.console.print(" S. Switch remote")
|
| 471 |
+
self.console.print(" M. Manage remotes")
|
| 472 |
+
self.console.print(" B. Back to main menu")
|
| 473 |
+
self.console.print()
|
| 474 |
+
self.console.print("━" * 78)
|
| 475 |
+
|
| 476 |
+
# Get input
|
| 477 |
+
valid_choices = [str(i) for i in range(1, len(provider_list) + 1)]
|
| 478 |
+
valid_choices.extend(["r", "R", "s", "S", "m", "M", "b", "B", "g", "G"])
|
| 479 |
+
|
| 480 |
+
choice = Prompt.ask("Select option", default="B").strip()
|
| 481 |
+
|
| 482 |
+
if choice.lower() == "b":
|
| 483 |
+
self.running = False
|
| 484 |
+
elif choice.lower() == "g":
|
| 485 |
+
# Toggle view mode
|
| 486 |
+
self.view_mode = "global" if self.view_mode == "current" else "current"
|
| 487 |
+
elif choice.lower() == "r":
|
| 488 |
+
with self.console.status("[bold]Reloading stats...", spinner="dots"):
|
| 489 |
+
self.post_action("reload", scope="all")
|
| 490 |
+
elif choice.lower() == "s":
|
| 491 |
+
self.show_switch_remote_screen()
|
| 492 |
+
elif choice.lower() == "m":
|
| 493 |
+
self.show_manage_remotes_screen()
|
| 494 |
+
elif choice.isdigit() and 1 <= int(choice) <= len(provider_list):
|
| 495 |
+
provider = provider_list[int(choice) - 1]
|
| 496 |
+
self.show_provider_detail_screen(provider)
|
| 497 |
+
|
| 498 |
+
def show_provider_detail_screen(self, provider: str):
|
| 499 |
+
"""Display detailed stats for a specific provider."""
|
| 500 |
+
while True:
|
| 501 |
+
clear_screen()
|
| 502 |
+
|
| 503 |
+
# View mode indicator
|
| 504 |
+
if self.view_mode == "global":
|
| 505 |
+
view_label = "[magenta]Global/Lifetime[/magenta]"
|
| 506 |
+
else:
|
| 507 |
+
view_label = "[cyan]Current Period[/cyan]"
|
| 508 |
+
|
| 509 |
+
self.console.print("━" * 78)
|
| 510 |
+
self.console.print(
|
| 511 |
+
f"[bold cyan]📊 {provider.title()} - Detailed Stats[/bold cyan] | {view_label}"
|
| 512 |
+
)
|
| 513 |
+
self.console.print("━" * 78)
|
| 514 |
+
self.console.print()
|
| 515 |
+
|
| 516 |
+
if not self.cached_stats:
|
| 517 |
+
self.console.print("[yellow]No data available.[/yellow]")
|
| 518 |
+
else:
|
| 519 |
+
prov_stats = self.cached_stats.get("providers", {}).get(provider, {})
|
| 520 |
+
credentials = prov_stats.get("credentials", [])
|
| 521 |
+
|
| 522 |
+
if not credentials:
|
| 523 |
+
self.console.print(
|
| 524 |
+
"[dim]No credentials configured for this provider.[/dim]"
|
| 525 |
+
)
|
| 526 |
+
else:
|
| 527 |
+
for idx, cred in enumerate(credentials, 1):
|
| 528 |
+
self._render_credential_panel(idx, cred, provider)
|
| 529 |
+
self.console.print()
|
| 530 |
+
|
| 531 |
+
# Menu
|
| 532 |
+
self.console.print("━" * 78)
|
| 533 |
+
self.console.print()
|
| 534 |
+
self.console.print(" G. Toggle view mode (current/global)")
|
| 535 |
+
self.console.print(" R. Reload stats (from proxy cache)")
|
| 536 |
+
self.console.print(" RA. Reload all stats")
|
| 537 |
+
|
| 538 |
+
# Force refresh options (only for providers that support it)
|
| 539 |
+
has_quota_groups = bool(
|
| 540 |
+
self.cached_stats
|
| 541 |
+
and self.cached_stats.get("providers", {})
|
| 542 |
+
.get(provider, {})
|
| 543 |
+
.get("quota_groups")
|
| 544 |
+
)
|
| 545 |
+
|
| 546 |
+
if has_quota_groups:
|
| 547 |
+
self.console.print()
|
| 548 |
+
self.console.print(
|
| 549 |
+
f" F. [yellow]Force refresh ALL {provider} quotas from API[/yellow]"
|
| 550 |
+
)
|
| 551 |
+
credentials = (
|
| 552 |
+
self.cached_stats.get("providers", {})
|
| 553 |
+
.get(provider, {})
|
| 554 |
+
.get("credentials", [])
|
| 555 |
+
if self.cached_stats
|
| 556 |
+
else []
|
| 557 |
+
)
|
| 558 |
+
for idx, cred in enumerate(credentials, 1):
|
| 559 |
+
identifier = cred.get("identifier", f"credential {idx}")
|
| 560 |
+
email = cred.get("email", identifier)
|
| 561 |
+
self.console.print(
|
| 562 |
+
f" F{idx}. Force refresh [{idx}] only ({email})"
|
| 563 |
+
)
|
| 564 |
+
|
| 565 |
+
self.console.print()
|
| 566 |
+
self.console.print(" B. Back to summary")
|
| 567 |
+
self.console.print()
|
| 568 |
+
self.console.print("━" * 78)
|
| 569 |
+
|
| 570 |
+
choice = Prompt.ask("Select option", default="B").strip().upper()
|
| 571 |
+
|
| 572 |
+
if choice == "B":
|
| 573 |
+
break
|
| 574 |
+
elif choice == "G":
|
| 575 |
+
# Toggle view mode
|
| 576 |
+
self.view_mode = "global" if self.view_mode == "current" else "current"
|
| 577 |
+
elif choice == "R":
|
| 578 |
+
with self.console.status(
|
| 579 |
+
f"[bold]Reloading {provider} stats...", spinner="dots"
|
| 580 |
+
):
|
| 581 |
+
self.post_action("reload", scope="provider", provider=provider)
|
| 582 |
+
elif choice == "RA":
|
| 583 |
+
with self.console.status(
|
| 584 |
+
"[bold]Reloading all stats...", spinner="dots"
|
| 585 |
+
):
|
| 586 |
+
self.post_action("reload", scope="all")
|
| 587 |
+
elif choice == "F" and has_quota_groups:
|
| 588 |
+
with self.console.status(
|
| 589 |
+
f"[bold]Fetching live quota for ALL {provider} credentials...",
|
| 590 |
+
spinner="dots",
|
| 591 |
+
):
|
| 592 |
+
result = self.post_action(
|
| 593 |
+
"force_refresh", scope="provider", provider=provider
|
| 594 |
+
)
|
| 595 |
+
if result and result.get("refresh_result"):
|
| 596 |
+
rr = result["refresh_result"]
|
| 597 |
+
self.console.print(
|
| 598 |
+
f"\n[green]Refreshed {rr.get('credentials_refreshed', 0)} credentials "
|
| 599 |
+
f"in {rr.get('duration_ms', 0)}ms[/green]"
|
| 600 |
+
)
|
| 601 |
+
if rr.get("errors"):
|
| 602 |
+
for err in rr["errors"]:
|
| 603 |
+
self.console.print(f"[red] Error: {err}[/red]")
|
| 604 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 605 |
+
elif choice.startswith("F") and choice[1:].isdigit() and has_quota_groups:
|
| 606 |
+
idx = int(choice[1:])
|
| 607 |
+
credentials = (
|
| 608 |
+
self.cached_stats.get("providers", {})
|
| 609 |
+
.get(provider, {})
|
| 610 |
+
.get("credentials", [])
|
| 611 |
+
if self.cached_stats
|
| 612 |
+
else []
|
| 613 |
+
)
|
| 614 |
+
if 1 <= idx <= len(credentials):
|
| 615 |
+
cred = credentials[idx - 1]
|
| 616 |
+
cred_id = cred.get("identifier", "")
|
| 617 |
+
email = cred.get("email", cred_id)
|
| 618 |
+
with self.console.status(
|
| 619 |
+
f"[bold]Fetching live quota for {email}...", spinner="dots"
|
| 620 |
+
):
|
| 621 |
+
result = self.post_action(
|
| 622 |
+
"force_refresh",
|
| 623 |
+
scope="credential",
|
| 624 |
+
provider=provider,
|
| 625 |
+
credential=cred_id,
|
| 626 |
+
)
|
| 627 |
+
if result and result.get("refresh_result"):
|
| 628 |
+
rr = result["refresh_result"]
|
| 629 |
+
self.console.print(
|
| 630 |
+
f"\n[green]Refreshed in {rr.get('duration_ms', 0)}ms[/green]"
|
| 631 |
+
)
|
| 632 |
+
if rr.get("errors"):
|
| 633 |
+
for err in rr["errors"]:
|
| 634 |
+
self.console.print(f"[red] Error: {err}[/red]")
|
| 635 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 636 |
+
|
| 637 |
+
def _render_credential_panel(self, idx: int, cred: Dict[str, Any], provider: str):
|
| 638 |
+
"""Render a single credential as a panel."""
|
| 639 |
+
identifier = cred.get("identifier", f"credential {idx}")
|
| 640 |
+
email = cred.get("email")
|
| 641 |
+
tier = cred.get("tier", "")
|
| 642 |
+
status = cred.get("status", "unknown")
|
| 643 |
+
|
| 644 |
+
# Check for active cooldowns
|
| 645 |
+
key_cooldown = cred.get("key_cooldown_remaining")
|
| 646 |
+
model_cooldowns = cred.get("model_cooldowns", {})
|
| 647 |
+
has_cooldown = key_cooldown or model_cooldowns
|
| 648 |
+
|
| 649 |
+
# Status indicator
|
| 650 |
+
if status == "exhausted":
|
| 651 |
+
status_icon = "[red]⛔ Exhausted[/red]"
|
| 652 |
+
elif status == "cooldown" or has_cooldown:
|
| 653 |
+
if key_cooldown:
|
| 654 |
+
status_icon = f"[yellow]⚠️ Cooldown ({format_cooldown(int(key_cooldown))})[/yellow]"
|
| 655 |
+
else:
|
| 656 |
+
status_icon = "[yellow]⚠️ Cooldown[/yellow]"
|
| 657 |
+
else:
|
| 658 |
+
status_icon = "[green]✅ Active[/green]"
|
| 659 |
+
|
| 660 |
+
# Header line
|
| 661 |
+
display_name = email if email else identifier
|
| 662 |
+
tier_str = f" ({tier})" if tier else ""
|
| 663 |
+
header = f"[{idx}] {display_name}{tier_str} {status_icon}"
|
| 664 |
+
|
| 665 |
+
# Use global stats if in global mode
|
| 666 |
+
if self.view_mode == "global":
|
| 667 |
+
stats_source = cred.get("global", cred)
|
| 668 |
+
else:
|
| 669 |
+
stats_source = cred
|
| 670 |
+
|
| 671 |
+
# Stats line
|
| 672 |
+
last_used = format_time_ago(cred.get("last_used_ts")) # Always from current
|
| 673 |
+
requests = stats_source.get("requests", 0)
|
| 674 |
+
tokens = stats_source.get("tokens", {})
|
| 675 |
+
input_total = tokens.get("input_cached", 0) + tokens.get("input_uncached", 0)
|
| 676 |
+
output = tokens.get("output", 0)
|
| 677 |
+
cost = format_cost(stats_source.get("approx_cost"))
|
| 678 |
+
|
| 679 |
+
stats_line = (
|
| 680 |
+
f"Last used: {last_used} | Requests: {requests} | "
|
| 681 |
+
f"Tokens: {format_tokens(input_total)}/{format_tokens(output)}"
|
| 682 |
+
)
|
| 683 |
+
if cost != "-":
|
| 684 |
+
stats_line += f" | Cost: {cost}"
|
| 685 |
+
|
| 686 |
+
# Build panel content
|
| 687 |
+
content_lines = [
|
| 688 |
+
f"[dim]{stats_line}[/dim]",
|
| 689 |
+
]
|
| 690 |
+
|
| 691 |
+
# Show model cooldowns if any
|
| 692 |
+
if model_cooldowns:
|
| 693 |
+
content_lines.append("")
|
| 694 |
+
content_lines.append("[yellow]Active Cooldowns:[/yellow]")
|
| 695 |
+
for model_name, cooldown_info in model_cooldowns.items():
|
| 696 |
+
remaining = cooldown_info.get("remaining_seconds", 0)
|
| 697 |
+
if remaining > 0:
|
| 698 |
+
# Shorten model name for display
|
| 699 |
+
short_model = model_name.split("/")[-1][:35]
|
| 700 |
+
content_lines.append(
|
| 701 |
+
f" [yellow]⏱️ {short_model}: {format_cooldown(int(remaining))}[/yellow]"
|
| 702 |
+
)
|
| 703 |
+
|
| 704 |
+
# Model groups (for providers with quota tracking)
|
| 705 |
+
model_groups = cred.get("model_groups", {})
|
| 706 |
+
if model_groups:
|
| 707 |
+
content_lines.append("")
|
| 708 |
+
for group_name, group_stats in model_groups.items():
|
| 709 |
+
remaining_pct = group_stats.get("remaining_pct")
|
| 710 |
+
requests_used = group_stats.get("requests_used", 0)
|
| 711 |
+
requests_max = group_stats.get("requests_max")
|
| 712 |
+
is_exhausted = group_stats.get("is_exhausted", False)
|
| 713 |
+
reset_time = format_reset_time(group_stats.get("reset_time_iso"))
|
| 714 |
+
confidence = group_stats.get("confidence", "low")
|
| 715 |
+
|
| 716 |
+
# Format display
|
| 717 |
+
display = group_stats.get("display", f"{requests_used}/?")
|
| 718 |
+
bar = create_progress_bar(remaining_pct)
|
| 719 |
+
|
| 720 |
+
# Color based on status
|
| 721 |
+
if is_exhausted:
|
| 722 |
+
color = "red"
|
| 723 |
+
status_text = "⛔ EXHAUSTED"
|
| 724 |
+
elif remaining_pct is not None and remaining_pct < 20:
|
| 725 |
+
color = "yellow"
|
| 726 |
+
status_text = "⚠️ LOW"
|
| 727 |
+
else:
|
| 728 |
+
color = "green"
|
| 729 |
+
status_text = f"Resets: {reset_time}"
|
| 730 |
+
|
| 731 |
+
# Confidence indicator
|
| 732 |
+
conf_indicator = ""
|
| 733 |
+
if confidence == "low":
|
| 734 |
+
conf_indicator = " [dim](~)[/dim]"
|
| 735 |
+
elif confidence == "medium":
|
| 736 |
+
conf_indicator = " [dim](?)[/dim]"
|
| 737 |
+
|
| 738 |
+
pct_str = f"{remaining_pct}%" if remaining_pct is not None else "?%"
|
| 739 |
+
content_lines.append(
|
| 740 |
+
f" [{color}]{group_name:<18} {display:<10} {pct_str:>4} {bar}[/{color}] {status_text}{conf_indicator}"
|
| 741 |
+
)
|
| 742 |
+
else:
|
| 743 |
+
# For providers without quota groups, show model breakdown if available
|
| 744 |
+
models = cred.get("models", {})
|
| 745 |
+
if models:
|
| 746 |
+
content_lines.append("")
|
| 747 |
+
content_lines.append(" [dim]Models used:[/dim]")
|
| 748 |
+
for model_name, model_stats in models.items():
|
| 749 |
+
req_count = model_stats.get("success_count", 0)
|
| 750 |
+
model_cost = format_cost(model_stats.get("approx_cost"))
|
| 751 |
+
# Shorten model name for display
|
| 752 |
+
short_name = model_name.split("/")[-1][:30]
|
| 753 |
+
content_lines.append(
|
| 754 |
+
f" {short_name}: {req_count} requests, {model_cost}"
|
| 755 |
+
)
|
| 756 |
+
|
| 757 |
+
self.console.print(
|
| 758 |
+
Panel(
|
| 759 |
+
"\n".join(content_lines),
|
| 760 |
+
title=header,
|
| 761 |
+
title_align="left",
|
| 762 |
+
border_style="dim",
|
| 763 |
+
expand=True,
|
| 764 |
+
)
|
| 765 |
+
)
|
| 766 |
+
|
| 767 |
+
def show_switch_remote_screen(self):
|
| 768 |
+
"""Display remote selection screen."""
|
| 769 |
+
clear_screen()
|
| 770 |
+
|
| 771 |
+
self.console.print("━" * 78)
|
| 772 |
+
self.console.print("[bold cyan]🔄 Switch Remote[/bold cyan]")
|
| 773 |
+
self.console.print("━" * 78)
|
| 774 |
+
self.console.print()
|
| 775 |
+
|
| 776 |
+
current_name = self.current_remote.get("name") if self.current_remote else None
|
| 777 |
+
self.console.print(f"Current: [bold]{current_name}[/bold]")
|
| 778 |
+
self.console.print()
|
| 779 |
+
self.console.print("Available remotes:")
|
| 780 |
+
|
| 781 |
+
remotes = self.config.get_remotes()
|
| 782 |
+
remote_status: List[Tuple[Dict, bool, str]] = []
|
| 783 |
+
|
| 784 |
+
# Check status of all remotes
|
| 785 |
+
with self.console.status("[dim]Checking remote status...", spinner="dots"):
|
| 786 |
+
for remote in remotes:
|
| 787 |
+
is_online, status_msg = self.check_connection(remote)
|
| 788 |
+
remote_status.append((remote, is_online, status_msg))
|
| 789 |
+
|
| 790 |
+
for idx, (remote, is_online, status_msg) in enumerate(remote_status, 1):
|
| 791 |
+
name = remote.get("name", "Unknown")
|
| 792 |
+
host = remote.get("host", "")
|
| 793 |
+
port = remote.get("port", 8000)
|
| 794 |
+
|
| 795 |
+
is_current = name == current_name
|
| 796 |
+
current_marker = " (current)" if is_current else ""
|
| 797 |
+
|
| 798 |
+
if is_online:
|
| 799 |
+
status_icon = "[green]✅ Online[/green]"
|
| 800 |
+
else:
|
| 801 |
+
status_icon = f"[red]⚠️ {status_msg}[/red]"
|
| 802 |
+
|
| 803 |
+
self.console.print(
|
| 804 |
+
f" {idx}. {name:<20} {host}:{port:<6} {status_icon}{current_marker}"
|
| 805 |
+
)
|
| 806 |
+
|
| 807 |
+
self.console.print()
|
| 808 |
+
self.console.print("━" * 78)
|
| 809 |
+
self.console.print()
|
| 810 |
+
|
| 811 |
+
choice = Prompt.ask(
|
| 812 |
+
f"Select remote (1-{len(remotes)}) or B to go back", default="B"
|
| 813 |
+
).strip()
|
| 814 |
+
|
| 815 |
+
if choice.lower() == "b":
|
| 816 |
+
return
|
| 817 |
+
|
| 818 |
+
if choice.isdigit() and 1 <= int(choice) <= len(remotes):
|
| 819 |
+
selected = remotes[int(choice) - 1]
|
| 820 |
+
self.current_remote = selected
|
| 821 |
+
self.config.set_last_used(selected["name"])
|
| 822 |
+
self.cached_stats = None # Clear cache
|
| 823 |
+
|
| 824 |
+
# Try to fetch stats from new remote
|
| 825 |
+
with self.console.status("[bold]Connecting...", spinner="dots"):
|
| 826 |
+
stats = self.fetch_stats()
|
| 827 |
+
if stats is None:
|
| 828 |
+
# Try with API key from .env for Local
|
| 829 |
+
if selected["name"] == "Local" and not selected.get("api_key"):
|
| 830 |
+
env_key = self.config.get_api_key_from_env()
|
| 831 |
+
if env_key:
|
| 832 |
+
self.current_remote["api_key"] = env_key
|
| 833 |
+
stats = self.fetch_stats()
|
| 834 |
+
|
| 835 |
+
if stats is None:
|
| 836 |
+
self.show_api_key_prompt()
|
| 837 |
+
|
| 838 |
+
def show_api_key_prompt(self):
|
| 839 |
+
"""Prompt for API key when authentication fails."""
|
| 840 |
+
self.console.print()
|
| 841 |
+
self.console.print(
|
| 842 |
+
"[yellow]Authentication required or connection failed.[/yellow]"
|
| 843 |
+
)
|
| 844 |
+
self.console.print(f"Error: {self.last_error}")
|
| 845 |
+
self.console.print()
|
| 846 |
+
|
| 847 |
+
api_key = Prompt.ask(
|
| 848 |
+
"Enter API key (or press Enter to cancel)", default=""
|
| 849 |
+
).strip()
|
| 850 |
+
|
| 851 |
+
if api_key:
|
| 852 |
+
self.current_remote["api_key"] = api_key
|
| 853 |
+
# Update config with new API key
|
| 854 |
+
self.config.update_remote(self.current_remote["name"], api_key=api_key)
|
| 855 |
+
|
| 856 |
+
# Try again
|
| 857 |
+
with self.console.status("[bold]Reconnecting...", spinner="dots"):
|
| 858 |
+
if self.fetch_stats() is None:
|
| 859 |
+
self.console.print(f"[red]Still failed: {self.last_error}[/red]")
|
| 860 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 861 |
+
else:
|
| 862 |
+
self.console.print("[dim]Cancelled.[/dim]")
|
| 863 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 864 |
+
|
| 865 |
+
def show_manage_remotes_screen(self):
|
| 866 |
+
"""Display remote management screen."""
|
| 867 |
+
while True:
|
| 868 |
+
clear_screen()
|
| 869 |
+
|
| 870 |
+
self.console.print("━" * 78)
|
| 871 |
+
self.console.print("[bold cyan]⚙️ Manage Remotes[/bold cyan]")
|
| 872 |
+
self.console.print("━" * 78)
|
| 873 |
+
self.console.print()
|
| 874 |
+
|
| 875 |
+
remotes = self.config.get_remotes()
|
| 876 |
+
|
| 877 |
+
table = Table(box=None, show_header=True, header_style="bold")
|
| 878 |
+
table.add_column("#", style="dim", width=3)
|
| 879 |
+
table.add_column("Name", min_width=16)
|
| 880 |
+
table.add_column("Host", min_width=24)
|
| 881 |
+
table.add_column("Port", justify="right", width=6)
|
| 882 |
+
table.add_column("Default", width=8)
|
| 883 |
+
|
| 884 |
+
for idx, remote in enumerate(remotes, 1):
|
| 885 |
+
is_default = "★" if remote.get("is_default") else ""
|
| 886 |
+
table.add_row(
|
| 887 |
+
str(idx),
|
| 888 |
+
remote.get("name", ""),
|
| 889 |
+
remote.get("host", ""),
|
| 890 |
+
str(remote.get("port", 8000)),
|
| 891 |
+
is_default,
|
| 892 |
+
)
|
| 893 |
+
|
| 894 |
+
self.console.print(table)
|
| 895 |
+
|
| 896 |
+
self.console.print()
|
| 897 |
+
self.console.print("━" * 78)
|
| 898 |
+
self.console.print()
|
| 899 |
+
self.console.print(" A. Add new remote")
|
| 900 |
+
self.console.print(" E. Edit remote (enter number, e.g., E1)")
|
| 901 |
+
self.console.print(" D. Delete remote (enter number, e.g., D1)")
|
| 902 |
+
self.console.print(" S. Set default remote")
|
| 903 |
+
self.console.print(" B. Back")
|
| 904 |
+
self.console.print()
|
| 905 |
+
self.console.print("━" * 78)
|
| 906 |
+
|
| 907 |
+
choice = Prompt.ask("Select option", default="B").strip().upper()
|
| 908 |
+
|
| 909 |
+
if choice == "B":
|
| 910 |
+
break
|
| 911 |
+
elif choice == "A":
|
| 912 |
+
self._add_remote_dialog()
|
| 913 |
+
elif choice == "S":
|
| 914 |
+
self._set_default_dialog(remotes)
|
| 915 |
+
elif choice.startswith("E") and choice[1:].isdigit():
|
| 916 |
+
idx = int(choice[1:])
|
| 917 |
+
if 1 <= idx <= len(remotes):
|
| 918 |
+
self._edit_remote_dialog(remotes[idx - 1])
|
| 919 |
+
elif choice.startswith("D") and choice[1:].isdigit():
|
| 920 |
+
idx = int(choice[1:])
|
| 921 |
+
if 1 <= idx <= len(remotes):
|
| 922 |
+
self._delete_remote_dialog(remotes[idx - 1])
|
| 923 |
+
|
| 924 |
+
def _add_remote_dialog(self):
|
| 925 |
+
"""Dialog to add a new remote."""
|
| 926 |
+
self.console.print()
|
| 927 |
+
self.console.print("[bold]Add New Remote[/bold]")
|
| 928 |
+
self.console.print()
|
| 929 |
+
|
| 930 |
+
name = Prompt.ask("Name", default="").strip()
|
| 931 |
+
if not name:
|
| 932 |
+
self.console.print("[dim]Cancelled.[/dim]")
|
| 933 |
+
return
|
| 934 |
+
|
| 935 |
+
host = Prompt.ask("Host", default="").strip()
|
| 936 |
+
if not host:
|
| 937 |
+
self.console.print("[dim]Cancelled.[/dim]")
|
| 938 |
+
return
|
| 939 |
+
|
| 940 |
+
port_str = Prompt.ask("Port", default="8000").strip()
|
| 941 |
+
try:
|
| 942 |
+
port = int(port_str)
|
| 943 |
+
except ValueError:
|
| 944 |
+
port = 8000
|
| 945 |
+
|
| 946 |
+
api_key = Prompt.ask("API Key (optional)", default="").strip() or None
|
| 947 |
+
|
| 948 |
+
if self.config.add_remote(name, host, port, api_key):
|
| 949 |
+
self.console.print(f"[green]Added remote '{name}'.[/green]")
|
| 950 |
+
else:
|
| 951 |
+
self.console.print(f"[red]Remote '{name}' already exists.[/red]")
|
| 952 |
+
|
| 953 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 954 |
+
|
| 955 |
+
def _edit_remote_dialog(self, remote: Dict[str, Any]):
|
| 956 |
+
"""Dialog to edit an existing remote."""
|
| 957 |
+
self.console.print()
|
| 958 |
+
self.console.print(f"[bold]Edit Remote: {remote['name']}[/bold]")
|
| 959 |
+
self.console.print("[dim]Press Enter to keep current value[/dim]")
|
| 960 |
+
self.console.print()
|
| 961 |
+
|
| 962 |
+
new_name = Prompt.ask("Name", default=remote["name"]).strip()
|
| 963 |
+
new_host = Prompt.ask("Host", default=remote.get("host", "")).strip()
|
| 964 |
+
new_port_str = Prompt.ask("Port", default=str(remote.get("port", 8000))).strip()
|
| 965 |
+
try:
|
| 966 |
+
new_port = int(new_port_str)
|
| 967 |
+
except ValueError:
|
| 968 |
+
new_port = remote.get("port", 8000)
|
| 969 |
+
|
| 970 |
+
current_key = remote.get("api_key", "") or ""
|
| 971 |
+
display_key = f"{current_key[:8]}..." if len(current_key) > 8 else current_key
|
| 972 |
+
new_key = Prompt.ask(
|
| 973 |
+
f"API Key (current: {display_key or 'none'})", default=""
|
| 974 |
+
).strip()
|
| 975 |
+
|
| 976 |
+
updates = {}
|
| 977 |
+
if new_name != remote["name"]:
|
| 978 |
+
updates["new_name"] = new_name
|
| 979 |
+
if new_host != remote.get("host"):
|
| 980 |
+
updates["host"] = new_host
|
| 981 |
+
if new_port != remote.get("port"):
|
| 982 |
+
updates["port"] = new_port
|
| 983 |
+
if new_key:
|
| 984 |
+
updates["api_key"] = new_key
|
| 985 |
+
|
| 986 |
+
if updates:
|
| 987 |
+
if self.config.update_remote(remote["name"], **updates):
|
| 988 |
+
self.console.print("[green]Remote updated.[/green]")
|
| 989 |
+
# Update current_remote if it was the one being edited
|
| 990 |
+
if (
|
| 991 |
+
self.current_remote
|
| 992 |
+
and self.current_remote["name"] == remote["name"]
|
| 993 |
+
):
|
| 994 |
+
self.current_remote.update(updates)
|
| 995 |
+
if "new_name" in updates:
|
| 996 |
+
self.current_remote["name"] = updates["new_name"]
|
| 997 |
+
else:
|
| 998 |
+
self.console.print("[red]Failed to update remote.[/red]")
|
| 999 |
+
else:
|
| 1000 |
+
self.console.print("[dim]No changes made.[/dim]")
|
| 1001 |
+
|
| 1002 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 1003 |
+
|
| 1004 |
+
def _delete_remote_dialog(self, remote: Dict[str, Any]):
|
| 1005 |
+
"""Dialog to delete a remote."""
|
| 1006 |
+
self.console.print()
|
| 1007 |
+
self.console.print(f"[yellow]Delete remote '{remote['name']}'?[/yellow]")
|
| 1008 |
+
|
| 1009 |
+
confirm = Prompt.ask("Type 'yes' to confirm", default="no").strip().lower()
|
| 1010 |
+
|
| 1011 |
+
if confirm == "yes":
|
| 1012 |
+
if self.config.delete_remote(remote["name"]):
|
| 1013 |
+
self.console.print(f"[green]Deleted remote '{remote['name']}'.[/green]")
|
| 1014 |
+
# If deleted current remote, switch to another
|
| 1015 |
+
if (
|
| 1016 |
+
self.current_remote
|
| 1017 |
+
and self.current_remote["name"] == remote["name"]
|
| 1018 |
+
):
|
| 1019 |
+
self.current_remote = self.config.get_default_remote()
|
| 1020 |
+
self.cached_stats = None
|
| 1021 |
+
else:
|
| 1022 |
+
self.console.print(
|
| 1023 |
+
"[red]Cannot delete. At least one remote must exist.[/red]"
|
| 1024 |
+
)
|
| 1025 |
+
else:
|
| 1026 |
+
self.console.print("[dim]Cancelled.[/dim]")
|
| 1027 |
+
|
| 1028 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 1029 |
+
|
| 1030 |
+
def _set_default_dialog(self, remotes: List[Dict[str, Any]]):
|
| 1031 |
+
"""Dialog to set the default remote."""
|
| 1032 |
+
self.console.print()
|
| 1033 |
+
choice = Prompt.ask(f"Set default (1-{len(remotes)})", default="").strip()
|
| 1034 |
+
|
| 1035 |
+
if choice.isdigit() and 1 <= int(choice) <= len(remotes):
|
| 1036 |
+
remote = remotes[int(choice) - 1]
|
| 1037 |
+
if self.config.set_default_remote(remote["name"]):
|
| 1038 |
+
self.console.print(
|
| 1039 |
+
f"[green]'{remote['name']}' is now the default.[/green]"
|
| 1040 |
+
)
|
| 1041 |
+
else:
|
| 1042 |
+
self.console.print("[red]Failed to set default.[/red]")
|
| 1043 |
+
Prompt.ask("Press Enter to continue", default="")
|
| 1044 |
+
|
| 1045 |
+
# =========================================================================
|
| 1046 |
+
# MAIN LOOP
|
| 1047 |
+
# =========================================================================
|
| 1048 |
+
|
| 1049 |
+
def run(self):
|
| 1050 |
+
"""Main viewer loop."""
|
| 1051 |
+
# Get initial remote
|
| 1052 |
+
self.current_remote = self.config.get_last_used_remote()
|
| 1053 |
+
|
| 1054 |
+
if not self.current_remote:
|
| 1055 |
+
self.console.print("[red]No remotes configured.[/red]")
|
| 1056 |
+
return
|
| 1057 |
+
|
| 1058 |
+
# For Local remote, try to get API key from .env if not set
|
| 1059 |
+
if self.current_remote["name"] == "Local" and not self.current_remote.get(
|
| 1060 |
+
"api_key"
|
| 1061 |
+
):
|
| 1062 |
+
env_key = self.config.get_api_key_from_env()
|
| 1063 |
+
if env_key:
|
| 1064 |
+
self.current_remote["api_key"] = env_key
|
| 1065 |
+
|
| 1066 |
+
# Initial fetch
|
| 1067 |
+
with self.console.status("[bold]Connecting to proxy...", spinner="dots"):
|
| 1068 |
+
stats = self.fetch_stats()
|
| 1069 |
+
|
| 1070 |
+
if stats is None:
|
| 1071 |
+
self.show_connection_error()
|
| 1072 |
+
return
|
| 1073 |
+
|
| 1074 |
+
# Main loop
|
| 1075 |
+
while self.running:
|
| 1076 |
+
self.show_summary_screen()
|
| 1077 |
+
|
| 1078 |
+
|
| 1079 |
+
def run_quota_viewer():
|
| 1080 |
+
"""Entry point for the quota viewer."""
|
| 1081 |
+
viewer = QuotaViewer()
|
| 1082 |
+
viewer.run()
|
| 1083 |
+
|
| 1084 |
+
|
| 1085 |
+
if __name__ == "__main__":
|
| 1086 |
+
run_quota_viewer()
|
|
@@ -0,0 +1,288 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Configuration management for the Quota Viewer.
|
| 3 |
+
|
| 4 |
+
Handles remote proxy configurations including:
|
| 5 |
+
- Multiple remote proxies (local, VPS, etc.)
|
| 6 |
+
- API key storage per remote
|
| 7 |
+
- Default and last-used remote tracking
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
import json
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
from typing import Any, Dict, List, Optional
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class QuotaViewerConfig:
|
| 16 |
+
"""Manages quota viewer configuration including remote proxies."""
|
| 17 |
+
|
| 18 |
+
def __init__(self, config_path: Optional[Path] = None):
|
| 19 |
+
"""
|
| 20 |
+
Initialize the config manager.
|
| 21 |
+
|
| 22 |
+
Args:
|
| 23 |
+
config_path: Path to config file. Defaults to quota_viewer_config.json
|
| 24 |
+
in the current directory or EXE directory.
|
| 25 |
+
"""
|
| 26 |
+
if config_path is None:
|
| 27 |
+
import sys
|
| 28 |
+
|
| 29 |
+
if getattr(sys, "frozen", False):
|
| 30 |
+
base_dir = Path(sys.executable).parent
|
| 31 |
+
else:
|
| 32 |
+
base_dir = Path.cwd()
|
| 33 |
+
config_path = base_dir / "quota_viewer_config.json"
|
| 34 |
+
|
| 35 |
+
self.config_path = config_path
|
| 36 |
+
self.config = self._load()
|
| 37 |
+
|
| 38 |
+
def _load(self) -> Dict[str, Any]:
|
| 39 |
+
"""Load config from file or return defaults."""
|
| 40 |
+
if self.config_path.exists():
|
| 41 |
+
try:
|
| 42 |
+
with open(self.config_path, "r", encoding="utf-8") as f:
|
| 43 |
+
config = json.load(f)
|
| 44 |
+
# Ensure required fields exist
|
| 45 |
+
if "remotes" not in config:
|
| 46 |
+
config["remotes"] = []
|
| 47 |
+
return config
|
| 48 |
+
except (json.JSONDecodeError, IOError):
|
| 49 |
+
pass
|
| 50 |
+
|
| 51 |
+
# Return default config with Local remote
|
| 52 |
+
return {
|
| 53 |
+
"remotes": [
|
| 54 |
+
{
|
| 55 |
+
"name": "Local",
|
| 56 |
+
"host": "127.0.0.1",
|
| 57 |
+
"port": 8000,
|
| 58 |
+
"api_key": None,
|
| 59 |
+
"is_default": True,
|
| 60 |
+
}
|
| 61 |
+
],
|
| 62 |
+
"last_used": "Local",
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
def _save(self) -> bool:
|
| 66 |
+
"""Save config to file. Returns True on success."""
|
| 67 |
+
try:
|
| 68 |
+
with open(self.config_path, "w", encoding="utf-8") as f:
|
| 69 |
+
json.dump(self.config, f, indent=2)
|
| 70 |
+
return True
|
| 71 |
+
except IOError:
|
| 72 |
+
return False
|
| 73 |
+
|
| 74 |
+
def get_remotes(self) -> List[Dict[str, Any]]:
|
| 75 |
+
"""Get list of all configured remotes."""
|
| 76 |
+
return self.config.get("remotes", [])
|
| 77 |
+
|
| 78 |
+
def get_remote_by_name(self, name: str) -> Optional[Dict[str, Any]]:
|
| 79 |
+
"""Get a remote by name."""
|
| 80 |
+
for remote in self.config.get("remotes", []):
|
| 81 |
+
if remote["name"] == name:
|
| 82 |
+
return remote
|
| 83 |
+
return None
|
| 84 |
+
|
| 85 |
+
def get_default_remote(self) -> Optional[Dict[str, Any]]:
|
| 86 |
+
"""Get the default remote."""
|
| 87 |
+
for remote in self.config.get("remotes", []):
|
| 88 |
+
if remote.get("is_default"):
|
| 89 |
+
return remote
|
| 90 |
+
# Fallback to first remote
|
| 91 |
+
remotes = self.config.get("remotes", [])
|
| 92 |
+
return remotes[0] if remotes else None
|
| 93 |
+
|
| 94 |
+
def get_last_used_remote(self) -> Optional[Dict[str, Any]]:
|
| 95 |
+
"""Get the last used remote, or default if not set."""
|
| 96 |
+
last_used_name = self.config.get("last_used")
|
| 97 |
+
if last_used_name:
|
| 98 |
+
remote = self.get_remote_by_name(last_used_name)
|
| 99 |
+
if remote:
|
| 100 |
+
return remote
|
| 101 |
+
return self.get_default_remote()
|
| 102 |
+
|
| 103 |
+
def set_last_used(self, name: str) -> bool:
|
| 104 |
+
"""Set the last used remote name."""
|
| 105 |
+
self.config["last_used"] = name
|
| 106 |
+
return self._save()
|
| 107 |
+
|
| 108 |
+
def add_remote(
|
| 109 |
+
self,
|
| 110 |
+
name: str,
|
| 111 |
+
host: str,
|
| 112 |
+
port: int = 8000,
|
| 113 |
+
api_key: Optional[str] = None,
|
| 114 |
+
is_default: bool = False,
|
| 115 |
+
) -> bool:
|
| 116 |
+
"""
|
| 117 |
+
Add a new remote configuration.
|
| 118 |
+
|
| 119 |
+
Args:
|
| 120 |
+
name: Display name for the remote
|
| 121 |
+
host: Hostname or IP address
|
| 122 |
+
port: Port number (default 8000)
|
| 123 |
+
api_key: Optional API key for authentication
|
| 124 |
+
is_default: Whether this should be the default remote
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
True on success, False if name already exists
|
| 128 |
+
"""
|
| 129 |
+
# Check for duplicate name
|
| 130 |
+
if self.get_remote_by_name(name):
|
| 131 |
+
return False
|
| 132 |
+
|
| 133 |
+
# If setting as default, clear default from others
|
| 134 |
+
if is_default:
|
| 135 |
+
for remote in self.config.get("remotes", []):
|
| 136 |
+
remote["is_default"] = False
|
| 137 |
+
|
| 138 |
+
remote = {
|
| 139 |
+
"name": name,
|
| 140 |
+
"host": host,
|
| 141 |
+
"port": port,
|
| 142 |
+
"api_key": api_key,
|
| 143 |
+
"is_default": is_default,
|
| 144 |
+
}
|
| 145 |
+
self.config.setdefault("remotes", []).append(remote)
|
| 146 |
+
return self._save()
|
| 147 |
+
|
| 148 |
+
def update_remote(self, name: str, **kwargs) -> bool:
|
| 149 |
+
"""
|
| 150 |
+
Update an existing remote configuration.
|
| 151 |
+
|
| 152 |
+
Args:
|
| 153 |
+
name: Name of the remote to update
|
| 154 |
+
**kwargs: Fields to update (host, port, api_key, is_default, new_name)
|
| 155 |
+
|
| 156 |
+
Returns:
|
| 157 |
+
True on success, False if remote not found
|
| 158 |
+
"""
|
| 159 |
+
remote = self.get_remote_by_name(name)
|
| 160 |
+
if not remote:
|
| 161 |
+
return False
|
| 162 |
+
|
| 163 |
+
# Handle rename
|
| 164 |
+
if "new_name" in kwargs:
|
| 165 |
+
new_name = kwargs.pop("new_name")
|
| 166 |
+
if new_name != name and self.get_remote_by_name(new_name):
|
| 167 |
+
return False # New name already exists
|
| 168 |
+
remote["name"] = new_name
|
| 169 |
+
# Update last_used if it was this remote
|
| 170 |
+
if self.config.get("last_used") == name:
|
| 171 |
+
self.config["last_used"] = new_name
|
| 172 |
+
|
| 173 |
+
# If setting as default, clear default from others
|
| 174 |
+
if kwargs.get("is_default"):
|
| 175 |
+
for r in self.config.get("remotes", []):
|
| 176 |
+
r["is_default"] = False
|
| 177 |
+
|
| 178 |
+
# Update other fields
|
| 179 |
+
for key in ("host", "port", "api_key", "is_default"):
|
| 180 |
+
if key in kwargs:
|
| 181 |
+
remote[key] = kwargs[key]
|
| 182 |
+
|
| 183 |
+
return self._save()
|
| 184 |
+
|
| 185 |
+
def delete_remote(self, name: str) -> bool:
|
| 186 |
+
"""
|
| 187 |
+
Delete a remote configuration.
|
| 188 |
+
|
| 189 |
+
Args:
|
| 190 |
+
name: Name of the remote to delete
|
| 191 |
+
|
| 192 |
+
Returns:
|
| 193 |
+
True on success, False if remote not found or is the only one
|
| 194 |
+
"""
|
| 195 |
+
remotes = self.config.get("remotes", [])
|
| 196 |
+
if len(remotes) <= 1:
|
| 197 |
+
return False # Don't delete the last remote
|
| 198 |
+
|
| 199 |
+
for i, remote in enumerate(remotes):
|
| 200 |
+
if remote["name"] == name:
|
| 201 |
+
remotes.pop(i)
|
| 202 |
+
# Update last_used if it was this remote
|
| 203 |
+
if self.config.get("last_used") == name:
|
| 204 |
+
self.config["last_used"] = remotes[0]["name"] if remotes else None
|
| 205 |
+
return self._save()
|
| 206 |
+
return False
|
| 207 |
+
|
| 208 |
+
def set_default_remote(self, name: str) -> bool:
|
| 209 |
+
"""Set a remote as the default."""
|
| 210 |
+
remote = self.get_remote_by_name(name)
|
| 211 |
+
if not remote:
|
| 212 |
+
return False
|
| 213 |
+
|
| 214 |
+
# Clear default from all remotes
|
| 215 |
+
for r in self.config.get("remotes", []):
|
| 216 |
+
r["is_default"] = False
|
| 217 |
+
|
| 218 |
+
# Set new default
|
| 219 |
+
remote["is_default"] = True
|
| 220 |
+
return self._save()
|
| 221 |
+
|
| 222 |
+
def sync_with_launcher_config(self) -> None:
|
| 223 |
+
"""
|
| 224 |
+
Sync the Local remote with launcher_config.json if it exists.
|
| 225 |
+
|
| 226 |
+
This ensures the Local remote always matches the launcher settings.
|
| 227 |
+
"""
|
| 228 |
+
import sys
|
| 229 |
+
|
| 230 |
+
if getattr(sys, "frozen", False):
|
| 231 |
+
base_dir = Path(sys.executable).parent
|
| 232 |
+
else:
|
| 233 |
+
base_dir = Path.cwd()
|
| 234 |
+
|
| 235 |
+
launcher_config_path = base_dir / "launcher_config.json"
|
| 236 |
+
|
| 237 |
+
if launcher_config_path.exists():
|
| 238 |
+
try:
|
| 239 |
+
with open(launcher_config_path, "r", encoding="utf-8") as f:
|
| 240 |
+
launcher_config = json.load(f)
|
| 241 |
+
|
| 242 |
+
host = launcher_config.get("host", "127.0.0.1")
|
| 243 |
+
port = launcher_config.get("port", 8000)
|
| 244 |
+
|
| 245 |
+
# Update Local remote
|
| 246 |
+
local_remote = self.get_remote_by_name("Local")
|
| 247 |
+
if local_remote:
|
| 248 |
+
local_remote["host"] = host
|
| 249 |
+
local_remote["port"] = port
|
| 250 |
+
self._save()
|
| 251 |
+
else:
|
| 252 |
+
# Create Local remote if it doesn't exist
|
| 253 |
+
self.add_remote("Local", host, port, is_default=True)
|
| 254 |
+
|
| 255 |
+
except (json.JSONDecodeError, IOError):
|
| 256 |
+
pass
|
| 257 |
+
|
| 258 |
+
def get_api_key_from_env(self) -> Optional[str]:
|
| 259 |
+
"""
|
| 260 |
+
Get PROXY_API_KEY from .env file for Local remote.
|
| 261 |
+
|
| 262 |
+
Returns:
|
| 263 |
+
API key string or None
|
| 264 |
+
"""
|
| 265 |
+
import sys
|
| 266 |
+
|
| 267 |
+
if getattr(sys, "frozen", False):
|
| 268 |
+
base_dir = Path(sys.executable).parent
|
| 269 |
+
else:
|
| 270 |
+
base_dir = Path.cwd()
|
| 271 |
+
|
| 272 |
+
env_path = base_dir / ".env"
|
| 273 |
+
if not env_path.exists():
|
| 274 |
+
return None
|
| 275 |
+
|
| 276 |
+
try:
|
| 277 |
+
with open(env_path, "r", encoding="utf-8") as f:
|
| 278 |
+
for line in f:
|
| 279 |
+
line = line.strip()
|
| 280 |
+
if line.startswith("PROXY_API_KEY="):
|
| 281 |
+
value = line.split("=", 1)[1].strip()
|
| 282 |
+
# Remove quotes if present
|
| 283 |
+
if value and value[0] in ('"', "'") and value[-1] == value[0]:
|
| 284 |
+
value = value[1:-1]
|
| 285 |
+
return value if value else None
|
| 286 |
+
except IOError:
|
| 287 |
+
pass
|
| 288 |
+
return None
|
|
@@ -2612,3 +2612,280 @@ class RotatingClient:
|
|
| 2612 |
for models in all_provider_models.values():
|
| 2613 |
flat_models.extend(models)
|
| 2614 |
return flat_models
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2612 |
for models in all_provider_models.values():
|
| 2613 |
flat_models.extend(models)
|
| 2614 |
return flat_models
|
| 2615 |
+
|
| 2616 |
+
async def get_quota_stats(
|
| 2617 |
+
self,
|
| 2618 |
+
provider_filter: Optional[str] = None,
|
| 2619 |
+
) -> Dict[str, Any]:
|
| 2620 |
+
"""
|
| 2621 |
+
Get quota and usage stats for all credentials.
|
| 2622 |
+
|
| 2623 |
+
This returns cached/disk data aggregated by provider.
|
| 2624 |
+
For provider-specific quota info (e.g., Antigravity quota groups),
|
| 2625 |
+
it enriches the data from provider plugins.
|
| 2626 |
+
|
| 2627 |
+
Args:
|
| 2628 |
+
provider_filter: If provided, only return stats for this provider
|
| 2629 |
+
|
| 2630 |
+
Returns:
|
| 2631 |
+
Complete stats dict ready for the /v1/quota-stats endpoint
|
| 2632 |
+
"""
|
| 2633 |
+
# Get base stats from usage manager
|
| 2634 |
+
stats = await self.usage_manager.get_stats_for_endpoint(provider_filter)
|
| 2635 |
+
|
| 2636 |
+
# Enrich with provider-specific quota data
|
| 2637 |
+
for provider, prov_stats in stats.get("providers", {}).items():
|
| 2638 |
+
provider_class = self._provider_plugins.get(provider)
|
| 2639 |
+
if not provider_class:
|
| 2640 |
+
continue
|
| 2641 |
+
|
| 2642 |
+
# Get or create provider instance
|
| 2643 |
+
if provider not in self._provider_instances:
|
| 2644 |
+
self._provider_instances[provider] = provider_class()
|
| 2645 |
+
provider_instance = self._provider_instances[provider]
|
| 2646 |
+
|
| 2647 |
+
# Check if provider has quota tracking (like Antigravity)
|
| 2648 |
+
if hasattr(provider_instance, "_get_effective_quota_groups"):
|
| 2649 |
+
# Add quota group summary
|
| 2650 |
+
quota_groups = provider_instance._get_effective_quota_groups()
|
| 2651 |
+
prov_stats["quota_groups"] = {}
|
| 2652 |
+
|
| 2653 |
+
for group_name, group_models in quota_groups.items():
|
| 2654 |
+
group_stats = {
|
| 2655 |
+
"models": group_models,
|
| 2656 |
+
"credentials_total": 0,
|
| 2657 |
+
"credentials_exhausted": 0,
|
| 2658 |
+
"avg_remaining_pct": 0,
|
| 2659 |
+
"total_remaining_pcts": [],
|
| 2660 |
+
}
|
| 2661 |
+
|
| 2662 |
+
# Calculate per-credential quota for this group
|
| 2663 |
+
for cred in prov_stats.get("credentials", []):
|
| 2664 |
+
models_data = cred.get("models", {})
|
| 2665 |
+
group_stats["credentials_total"] += 1
|
| 2666 |
+
|
| 2667 |
+
# Find any model from this group
|
| 2668 |
+
for model in group_models:
|
| 2669 |
+
# Try with and without provider prefix
|
| 2670 |
+
prefixed_model = f"{provider}/{model}"
|
| 2671 |
+
model_stats = models_data.get(
|
| 2672 |
+
prefixed_model
|
| 2673 |
+
) or models_data.get(model)
|
| 2674 |
+
|
| 2675 |
+
if model_stats:
|
| 2676 |
+
baseline = model_stats.get(
|
| 2677 |
+
"baseline_remaining_fraction"
|
| 2678 |
+
)
|
| 2679 |
+
if baseline is not None:
|
| 2680 |
+
remaining_pct = int(baseline * 100)
|
| 2681 |
+
group_stats["total_remaining_pcts"].append(
|
| 2682 |
+
remaining_pct
|
| 2683 |
+
)
|
| 2684 |
+
if baseline <= 0:
|
| 2685 |
+
group_stats["credentials_exhausted"] += 1
|
| 2686 |
+
break
|
| 2687 |
+
|
| 2688 |
+
# Calculate average remaining percentage
|
| 2689 |
+
if group_stats["total_remaining_pcts"]:
|
| 2690 |
+
group_stats["avg_remaining_pct"] = int(
|
| 2691 |
+
sum(group_stats["total_remaining_pcts"])
|
| 2692 |
+
/ len(group_stats["total_remaining_pcts"])
|
| 2693 |
+
)
|
| 2694 |
+
del group_stats["total_remaining_pcts"]
|
| 2695 |
+
|
| 2696 |
+
prov_stats["quota_groups"][group_name] = group_stats
|
| 2697 |
+
|
| 2698 |
+
# Also enrich each credential with formatted quota group info
|
| 2699 |
+
for cred in prov_stats.get("credentials", []):
|
| 2700 |
+
cred["model_groups"] = {}
|
| 2701 |
+
models_data = cred.get("models", {})
|
| 2702 |
+
|
| 2703 |
+
for group_name, group_models in quota_groups.items():
|
| 2704 |
+
# Find representative model from this group
|
| 2705 |
+
for model in group_models:
|
| 2706 |
+
prefixed_model = f"{provider}/{model}"
|
| 2707 |
+
model_stats = models_data.get(
|
| 2708 |
+
prefixed_model
|
| 2709 |
+
) or models_data.get(model)
|
| 2710 |
+
|
| 2711 |
+
if model_stats:
|
| 2712 |
+
baseline = model_stats.get(
|
| 2713 |
+
"baseline_remaining_fraction"
|
| 2714 |
+
)
|
| 2715 |
+
max_req = model_stats.get("quota_max_requests")
|
| 2716 |
+
req_count = model_stats.get("request_count", 0)
|
| 2717 |
+
reset_ts = model_stats.get("quota_reset_ts")
|
| 2718 |
+
|
| 2719 |
+
remaining_pct = (
|
| 2720 |
+
int(baseline * 100)
|
| 2721 |
+
if baseline is not None
|
| 2722 |
+
else None
|
| 2723 |
+
)
|
| 2724 |
+
is_exhausted = baseline is not None and baseline <= 0
|
| 2725 |
+
|
| 2726 |
+
# Format reset time
|
| 2727 |
+
reset_iso = None
|
| 2728 |
+
if reset_ts:
|
| 2729 |
+
try:
|
| 2730 |
+
from datetime import datetime, timezone
|
| 2731 |
+
|
| 2732 |
+
reset_iso = datetime.fromtimestamp(
|
| 2733 |
+
reset_ts, tz=timezone.utc
|
| 2734 |
+
).isoformat()
|
| 2735 |
+
except (ValueError, OSError):
|
| 2736 |
+
pass
|
| 2737 |
+
|
| 2738 |
+
cred["model_groups"][group_name] = {
|
| 2739 |
+
"remaining_pct": remaining_pct,
|
| 2740 |
+
"requests_used": req_count,
|
| 2741 |
+
"requests_max": max_req,
|
| 2742 |
+
"display": f"{req_count}/{max_req}"
|
| 2743 |
+
if max_req
|
| 2744 |
+
else f"{req_count}/?",
|
| 2745 |
+
"is_exhausted": is_exhausted,
|
| 2746 |
+
"reset_time_iso": reset_iso,
|
| 2747 |
+
"models": group_models,
|
| 2748 |
+
"confidence": self._get_baseline_confidence(
|
| 2749 |
+
model_stats
|
| 2750 |
+
),
|
| 2751 |
+
}
|
| 2752 |
+
break
|
| 2753 |
+
|
| 2754 |
+
# Try to get email from provider's cache
|
| 2755 |
+
cred_path = cred.get("full_path", "")
|
| 2756 |
+
if hasattr(provider_instance, "project_tier_cache"):
|
| 2757 |
+
tier = provider_instance.project_tier_cache.get(cred_path)
|
| 2758 |
+
if tier:
|
| 2759 |
+
cred["tier"] = tier
|
| 2760 |
+
|
| 2761 |
+
return stats
|
| 2762 |
+
|
| 2763 |
+
def _get_baseline_confidence(self, model_stats: Dict) -> str:
|
| 2764 |
+
"""
|
| 2765 |
+
Determine confidence level based on baseline age.
|
| 2766 |
+
|
| 2767 |
+
Args:
|
| 2768 |
+
model_stats: Model statistics dict with baseline_fetched_at
|
| 2769 |
+
|
| 2770 |
+
Returns:
|
| 2771 |
+
"high" | "medium" | "low"
|
| 2772 |
+
"""
|
| 2773 |
+
baseline_fetched_at = model_stats.get("baseline_fetched_at")
|
| 2774 |
+
if not baseline_fetched_at:
|
| 2775 |
+
return "low"
|
| 2776 |
+
|
| 2777 |
+
age_seconds = time.time() - baseline_fetched_at
|
| 2778 |
+
if age_seconds < 300: # 5 minutes
|
| 2779 |
+
return "high"
|
| 2780 |
+
elif age_seconds < 1800: # 30 minutes
|
| 2781 |
+
return "medium"
|
| 2782 |
+
return "low"
|
| 2783 |
+
|
| 2784 |
+
async def reload_usage_from_disk(self) -> None:
|
| 2785 |
+
"""
|
| 2786 |
+
Force reload usage data from disk.
|
| 2787 |
+
|
| 2788 |
+
Useful when wanting fresh stats without making external API calls.
|
| 2789 |
+
"""
|
| 2790 |
+
await self.usage_manager.reload_from_disk()
|
| 2791 |
+
|
| 2792 |
+
async def force_refresh_quota(
|
| 2793 |
+
self,
|
| 2794 |
+
provider: Optional[str] = None,
|
| 2795 |
+
credential: Optional[str] = None,
|
| 2796 |
+
) -> Dict[str, Any]:
|
| 2797 |
+
"""
|
| 2798 |
+
Force refresh quota from external API.
|
| 2799 |
+
|
| 2800 |
+
For Antigravity, this fetches live quota data from the API.
|
| 2801 |
+
For other providers, this is a no-op (just reloads from disk).
|
| 2802 |
+
|
| 2803 |
+
Args:
|
| 2804 |
+
provider: If specified, only refresh this provider
|
| 2805 |
+
credential: If specified, only refresh this specific credential
|
| 2806 |
+
|
| 2807 |
+
Returns:
|
| 2808 |
+
Refresh result dict with success/failure info
|
| 2809 |
+
"""
|
| 2810 |
+
result = {
|
| 2811 |
+
"action": "force_refresh",
|
| 2812 |
+
"scope": "credential"
|
| 2813 |
+
if credential
|
| 2814 |
+
else ("provider" if provider else "all"),
|
| 2815 |
+
"provider": provider,
|
| 2816 |
+
"credential": credential,
|
| 2817 |
+
"credentials_refreshed": 0,
|
| 2818 |
+
"success_count": 0,
|
| 2819 |
+
"failed_count": 0,
|
| 2820 |
+
"duration_ms": 0,
|
| 2821 |
+
"errors": [],
|
| 2822 |
+
}
|
| 2823 |
+
|
| 2824 |
+
start_time = time.time()
|
| 2825 |
+
|
| 2826 |
+
# Determine which providers to refresh
|
| 2827 |
+
if provider:
|
| 2828 |
+
providers_to_refresh = (
|
| 2829 |
+
[provider] if provider in self.all_credentials else []
|
| 2830 |
+
)
|
| 2831 |
+
else:
|
| 2832 |
+
providers_to_refresh = list(self.all_credentials.keys())
|
| 2833 |
+
|
| 2834 |
+
for prov in providers_to_refresh:
|
| 2835 |
+
provider_class = self._provider_plugins.get(prov)
|
| 2836 |
+
if not provider_class:
|
| 2837 |
+
continue
|
| 2838 |
+
|
| 2839 |
+
# Get or create provider instance
|
| 2840 |
+
if prov not in self._provider_instances:
|
| 2841 |
+
self._provider_instances[prov] = provider_class()
|
| 2842 |
+
provider_instance = self._provider_instances[prov]
|
| 2843 |
+
|
| 2844 |
+
# Check if provider supports quota refresh (like Antigravity)
|
| 2845 |
+
if hasattr(provider_instance, "fetch_initial_baselines"):
|
| 2846 |
+
# Get credentials to refresh
|
| 2847 |
+
if credential:
|
| 2848 |
+
# Find full path for this credential
|
| 2849 |
+
creds_to_refresh = []
|
| 2850 |
+
for cred_path in self.all_credentials.get(prov, []):
|
| 2851 |
+
if cred_path.endswith(credential) or cred_path == credential:
|
| 2852 |
+
creds_to_refresh.append(cred_path)
|
| 2853 |
+
break
|
| 2854 |
+
else:
|
| 2855 |
+
creds_to_refresh = self.all_credentials.get(prov, [])
|
| 2856 |
+
|
| 2857 |
+
if not creds_to_refresh:
|
| 2858 |
+
continue
|
| 2859 |
+
|
| 2860 |
+
try:
|
| 2861 |
+
# Fetch live quota from API for ALL specified credentials
|
| 2862 |
+
quota_results = await provider_instance.fetch_initial_baselines(
|
| 2863 |
+
creds_to_refresh
|
| 2864 |
+
)
|
| 2865 |
+
|
| 2866 |
+
# Store baselines in usage manager
|
| 2867 |
+
if hasattr(provider_instance, "_store_baselines_to_usage_manager"):
|
| 2868 |
+
stored = (
|
| 2869 |
+
await provider_instance._store_baselines_to_usage_manager(
|
| 2870 |
+
quota_results, self.usage_manager
|
| 2871 |
+
)
|
| 2872 |
+
)
|
| 2873 |
+
result["success_count"] += stored
|
| 2874 |
+
|
| 2875 |
+
result["credentials_refreshed"] += len(creds_to_refresh)
|
| 2876 |
+
|
| 2877 |
+
# Count failures
|
| 2878 |
+
for cred_path, data in quota_results.items():
|
| 2879 |
+
if data.get("status") != "success":
|
| 2880 |
+
result["failed_count"] += 1
|
| 2881 |
+
result["errors"].append(
|
| 2882 |
+
f"{Path(cred_path).name}: {data.get('error', 'Unknown error')}"
|
| 2883 |
+
)
|
| 2884 |
+
|
| 2885 |
+
except Exception as e:
|
| 2886 |
+
lib_logger.error(f"Failed to refresh quota for {prov}: {e}")
|
| 2887 |
+
result["errors"].append(f"{prov}: {str(e)}")
|
| 2888 |
+
result["failed_count"] += len(creds_to_refresh)
|
| 2889 |
+
|
| 2890 |
+
result["duration_ms"] = int((time.time() - start_time) * 1000)
|
| 2891 |
+
return result
|
|
@@ -103,7 +103,7 @@ DEFAULT_MAX_OUTPUT_TOKENS = 64000
|
|
| 103 |
# Empty response retry configuration
|
| 104 |
# When Antigravity returns an empty response (no content, no tool calls),
|
| 105 |
# automatically retry up to this many attempts before giving up (minimum 1)
|
| 106 |
-
EMPTY_RESPONSE_MAX_ATTEMPTS = max(1, _env_int("ANTIGRAVITY_EMPTY_RESPONSE_ATTEMPTS",
|
| 107 |
EMPTY_RESPONSE_RETRY_DELAY = _env_int("ANTIGRAVITY_EMPTY_RESPONSE_RETRY_DELAY", 2)
|
| 108 |
|
| 109 |
# Model alias mappings (internal ↔ public)
|
|
|
|
| 103 |
# Empty response retry configuration
|
| 104 |
# When Antigravity returns an empty response (no content, no tool calls),
|
| 105 |
# automatically retry up to this many attempts before giving up (minimum 1)
|
| 106 |
+
EMPTY_RESPONSE_MAX_ATTEMPTS = max(1, _env_int("ANTIGRAVITY_EMPTY_RESPONSE_ATTEMPTS", 6))
|
| 107 |
EMPTY_RESPONSE_RETRY_DELAY = _env_int("ANTIGRAVITY_EMPTY_RESPONSE_RETRY_DELAY", 2)
|
| 108 |
|
| 109 |
# Model alias mappings (internal ↔ public)
|
|
@@ -1993,3 +1993,479 @@ class UsageManager:
|
|
| 1993 |
"""
|
| 1994 |
# Disabled - see docstring above
|
| 1995 |
pass
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1993 |
"""
|
| 1994 |
# Disabled - see docstring above
|
| 1995 |
pass
|
| 1996 |
+
|
| 1997 |
+
async def get_stats_for_endpoint(
|
| 1998 |
+
self,
|
| 1999 |
+
provider_filter: Optional[str] = None,
|
| 2000 |
+
include_global: bool = True,
|
| 2001 |
+
) -> Dict[str, Any]:
|
| 2002 |
+
"""
|
| 2003 |
+
Get usage stats formatted for the /v1/quota-stats endpoint.
|
| 2004 |
+
|
| 2005 |
+
Aggregates data from key_usage.json grouped by provider.
|
| 2006 |
+
Includes both current period stats and global (lifetime) stats.
|
| 2007 |
+
|
| 2008 |
+
Args:
|
| 2009 |
+
provider_filter: If provided, only return stats for this provider
|
| 2010 |
+
include_global: If True, include global/lifetime stats alongside current
|
| 2011 |
+
|
| 2012 |
+
Returns:
|
| 2013 |
+
{
|
| 2014 |
+
"providers": {
|
| 2015 |
+
"provider_name": {
|
| 2016 |
+
"credential_count": int,
|
| 2017 |
+
"active_count": int,
|
| 2018 |
+
"on_cooldown_count": int,
|
| 2019 |
+
"total_requests": int,
|
| 2020 |
+
"tokens": {
|
| 2021 |
+
"input_cached": int,
|
| 2022 |
+
"input_uncached": int,
|
| 2023 |
+
"input_cache_pct": float,
|
| 2024 |
+
"output": int
|
| 2025 |
+
},
|
| 2026 |
+
"approx_cost": float | None,
|
| 2027 |
+
"credentials": [...],
|
| 2028 |
+
"global": {...} # If include_global is True
|
| 2029 |
+
}
|
| 2030 |
+
},
|
| 2031 |
+
"summary": {...},
|
| 2032 |
+
"global_summary": {...}, # If include_global is True
|
| 2033 |
+
"timestamp": float
|
| 2034 |
+
}
|
| 2035 |
+
"""
|
| 2036 |
+
await self._lazy_init()
|
| 2037 |
+
|
| 2038 |
+
now_ts = time.time()
|
| 2039 |
+
providers: Dict[str, Dict[str, Any]] = {}
|
| 2040 |
+
# Track global stats separately
|
| 2041 |
+
global_providers: Dict[str, Dict[str, Any]] = {}
|
| 2042 |
+
|
| 2043 |
+
async with self._data_lock:
|
| 2044 |
+
if not self._usage_data:
|
| 2045 |
+
return {
|
| 2046 |
+
"providers": {},
|
| 2047 |
+
"summary": {
|
| 2048 |
+
"total_providers": 0,
|
| 2049 |
+
"total_credentials": 0,
|
| 2050 |
+
"active_credentials": 0,
|
| 2051 |
+
"exhausted_credentials": 0,
|
| 2052 |
+
"total_requests": 0,
|
| 2053 |
+
"tokens": {
|
| 2054 |
+
"input_cached": 0,
|
| 2055 |
+
"input_uncached": 0,
|
| 2056 |
+
"input_cache_pct": 0,
|
| 2057 |
+
"output": 0,
|
| 2058 |
+
},
|
| 2059 |
+
"approx_total_cost": 0.0,
|
| 2060 |
+
},
|
| 2061 |
+
"global_summary": {
|
| 2062 |
+
"total_providers": 0,
|
| 2063 |
+
"total_credentials": 0,
|
| 2064 |
+
"total_requests": 0,
|
| 2065 |
+
"tokens": {
|
| 2066 |
+
"input_cached": 0,
|
| 2067 |
+
"input_uncached": 0,
|
| 2068 |
+
"input_cache_pct": 0,
|
| 2069 |
+
"output": 0,
|
| 2070 |
+
},
|
| 2071 |
+
"approx_total_cost": 0.0,
|
| 2072 |
+
},
|
| 2073 |
+
"data_source": "cache",
|
| 2074 |
+
"timestamp": now_ts,
|
| 2075 |
+
}
|
| 2076 |
+
|
| 2077 |
+
for credential, cred_data in self._usage_data.items():
|
| 2078 |
+
# Extract provider from credential path
|
| 2079 |
+
provider = self._get_provider_from_credential(credential)
|
| 2080 |
+
if not provider:
|
| 2081 |
+
continue
|
| 2082 |
+
|
| 2083 |
+
# Apply filter if specified
|
| 2084 |
+
if provider_filter and provider != provider_filter:
|
| 2085 |
+
continue
|
| 2086 |
+
|
| 2087 |
+
# Initialize provider entry
|
| 2088 |
+
if provider not in providers:
|
| 2089 |
+
providers[provider] = {
|
| 2090 |
+
"credential_count": 0,
|
| 2091 |
+
"active_count": 0,
|
| 2092 |
+
"on_cooldown_count": 0,
|
| 2093 |
+
"exhausted_count": 0,
|
| 2094 |
+
"total_requests": 0,
|
| 2095 |
+
"tokens": {
|
| 2096 |
+
"input_cached": 0,
|
| 2097 |
+
"input_uncached": 0,
|
| 2098 |
+
"input_cache_pct": 0,
|
| 2099 |
+
"output": 0,
|
| 2100 |
+
},
|
| 2101 |
+
"approx_cost": 0.0,
|
| 2102 |
+
"credentials": [],
|
| 2103 |
+
}
|
| 2104 |
+
global_providers[provider] = {
|
| 2105 |
+
"total_requests": 0,
|
| 2106 |
+
"tokens": {
|
| 2107 |
+
"input_cached": 0,
|
| 2108 |
+
"input_uncached": 0,
|
| 2109 |
+
"input_cache_pct": 0,
|
| 2110 |
+
"output": 0,
|
| 2111 |
+
},
|
| 2112 |
+
"approx_cost": 0.0,
|
| 2113 |
+
}
|
| 2114 |
+
|
| 2115 |
+
prov_stats = providers[provider]
|
| 2116 |
+
prov_stats["credential_count"] += 1
|
| 2117 |
+
|
| 2118 |
+
# Determine credential status and cooldowns
|
| 2119 |
+
key_cooldown = cred_data.get("key_cooldown_until", 0) or 0
|
| 2120 |
+
model_cooldowns = cred_data.get("model_cooldowns", {})
|
| 2121 |
+
|
| 2122 |
+
# Build active cooldowns with remaining time
|
| 2123 |
+
active_cooldowns = {}
|
| 2124 |
+
for model, cooldown_ts in model_cooldowns.items():
|
| 2125 |
+
if cooldown_ts > now_ts:
|
| 2126 |
+
remaining_seconds = int(cooldown_ts - now_ts)
|
| 2127 |
+
active_cooldowns[model] = {
|
| 2128 |
+
"until_ts": cooldown_ts,
|
| 2129 |
+
"remaining_seconds": remaining_seconds,
|
| 2130 |
+
}
|
| 2131 |
+
|
| 2132 |
+
key_cooldown_remaining = None
|
| 2133 |
+
if key_cooldown > now_ts:
|
| 2134 |
+
key_cooldown_remaining = int(key_cooldown - now_ts)
|
| 2135 |
+
|
| 2136 |
+
has_active_cooldown = key_cooldown > now_ts or len(active_cooldowns) > 0
|
| 2137 |
+
|
| 2138 |
+
# Check if exhausted (all quota groups exhausted for Antigravity)
|
| 2139 |
+
is_exhausted = False
|
| 2140 |
+
models_data = cred_data.get("models", {})
|
| 2141 |
+
if models_data:
|
| 2142 |
+
# Check if any model has remaining quota
|
| 2143 |
+
all_exhausted = True
|
| 2144 |
+
for model_stats in models_data.values():
|
| 2145 |
+
if isinstance(model_stats, dict):
|
| 2146 |
+
baseline = model_stats.get("baseline_remaining_fraction")
|
| 2147 |
+
if baseline is None or baseline > 0:
|
| 2148 |
+
all_exhausted = False
|
| 2149 |
+
break
|
| 2150 |
+
if all_exhausted and len(models_data) > 0:
|
| 2151 |
+
is_exhausted = True
|
| 2152 |
+
|
| 2153 |
+
if is_exhausted:
|
| 2154 |
+
prov_stats["exhausted_count"] += 1
|
| 2155 |
+
status = "exhausted"
|
| 2156 |
+
elif has_active_cooldown:
|
| 2157 |
+
prov_stats["on_cooldown_count"] += 1
|
| 2158 |
+
status = "cooldown"
|
| 2159 |
+
else:
|
| 2160 |
+
prov_stats["active_count"] += 1
|
| 2161 |
+
status = "active"
|
| 2162 |
+
|
| 2163 |
+
# Aggregate token stats (current period)
|
| 2164 |
+
cred_tokens = {
|
| 2165 |
+
"input_cached": 0,
|
| 2166 |
+
"input_uncached": 0,
|
| 2167 |
+
"output": 0,
|
| 2168 |
+
}
|
| 2169 |
+
cred_requests = 0
|
| 2170 |
+
cred_cost = 0.0
|
| 2171 |
+
|
| 2172 |
+
# Aggregate global token stats
|
| 2173 |
+
cred_global_tokens = {
|
| 2174 |
+
"input_cached": 0,
|
| 2175 |
+
"input_uncached": 0,
|
| 2176 |
+
"output": 0,
|
| 2177 |
+
}
|
| 2178 |
+
cred_global_requests = 0
|
| 2179 |
+
cred_global_cost = 0.0
|
| 2180 |
+
|
| 2181 |
+
# Handle per-model structure (current period)
|
| 2182 |
+
if models_data:
|
| 2183 |
+
for model_name, model_stats in models_data.items():
|
| 2184 |
+
if not isinstance(model_stats, dict):
|
| 2185 |
+
continue
|
| 2186 |
+
# Prefer request_count if available and non-zero, else fall back to success+failure
|
| 2187 |
+
req_count = model_stats.get("request_count", 0)
|
| 2188 |
+
if req_count > 0:
|
| 2189 |
+
cred_requests += req_count
|
| 2190 |
+
else:
|
| 2191 |
+
cred_requests += model_stats.get("success_count", 0)
|
| 2192 |
+
cred_requests += model_stats.get("failure_count", 0)
|
| 2193 |
+
# Token stats - track cached separately
|
| 2194 |
+
cred_tokens["input_cached"] += model_stats.get(
|
| 2195 |
+
"prompt_tokens_cached", 0
|
| 2196 |
+
)
|
| 2197 |
+
cred_tokens["input_uncached"] += model_stats.get(
|
| 2198 |
+
"prompt_tokens", 0
|
| 2199 |
+
)
|
| 2200 |
+
cred_tokens["output"] += model_stats.get("completion_tokens", 0)
|
| 2201 |
+
cred_cost += model_stats.get("approx_cost", 0.0)
|
| 2202 |
+
|
| 2203 |
+
# Handle legacy daily structure
|
| 2204 |
+
daily_data = cred_data.get("daily", {})
|
| 2205 |
+
daily_models = daily_data.get("models", {})
|
| 2206 |
+
for model_name, model_stats in daily_models.items():
|
| 2207 |
+
if not isinstance(model_stats, dict):
|
| 2208 |
+
continue
|
| 2209 |
+
cred_requests += model_stats.get("success_count", 0)
|
| 2210 |
+
cred_tokens["input_cached"] += model_stats.get(
|
| 2211 |
+
"prompt_tokens_cached", 0
|
| 2212 |
+
)
|
| 2213 |
+
cred_tokens["input_uncached"] += model_stats.get("prompt_tokens", 0)
|
| 2214 |
+
cred_tokens["output"] += model_stats.get("completion_tokens", 0)
|
| 2215 |
+
cred_cost += model_stats.get("approx_cost", 0.0)
|
| 2216 |
+
|
| 2217 |
+
# Handle global stats
|
| 2218 |
+
global_data = cred_data.get("global", {})
|
| 2219 |
+
global_models = global_data.get("models", {})
|
| 2220 |
+
for model_name, model_stats in global_models.items():
|
| 2221 |
+
if not isinstance(model_stats, dict):
|
| 2222 |
+
continue
|
| 2223 |
+
cred_global_requests += model_stats.get("success_count", 0)
|
| 2224 |
+
cred_global_tokens["input_cached"] += model_stats.get(
|
| 2225 |
+
"prompt_tokens_cached", 0
|
| 2226 |
+
)
|
| 2227 |
+
cred_global_tokens["input_uncached"] += model_stats.get(
|
| 2228 |
+
"prompt_tokens", 0
|
| 2229 |
+
)
|
| 2230 |
+
cred_global_tokens["output"] += model_stats.get(
|
| 2231 |
+
"completion_tokens", 0
|
| 2232 |
+
)
|
| 2233 |
+
cred_global_cost += model_stats.get("approx_cost", 0.0)
|
| 2234 |
+
|
| 2235 |
+
# Add current period stats to global totals
|
| 2236 |
+
cred_global_requests += cred_requests
|
| 2237 |
+
cred_global_tokens["input_cached"] += cred_tokens["input_cached"]
|
| 2238 |
+
cred_global_tokens["input_uncached"] += cred_tokens["input_uncached"]
|
| 2239 |
+
cred_global_tokens["output"] += cred_tokens["output"]
|
| 2240 |
+
cred_global_cost += cred_cost
|
| 2241 |
+
|
| 2242 |
+
# Build credential entry
|
| 2243 |
+
# Mask credential identifier for display
|
| 2244 |
+
if credential.startswith("env://"):
|
| 2245 |
+
identifier = credential
|
| 2246 |
+
else:
|
| 2247 |
+
identifier = Path(credential).name
|
| 2248 |
+
|
| 2249 |
+
cred_entry = {
|
| 2250 |
+
"identifier": identifier,
|
| 2251 |
+
"full_path": credential,
|
| 2252 |
+
"status": status,
|
| 2253 |
+
"last_used_ts": cred_data.get("last_used_ts"),
|
| 2254 |
+
"requests": cred_requests,
|
| 2255 |
+
"tokens": cred_tokens,
|
| 2256 |
+
"approx_cost": cred_cost if cred_cost > 0 else None,
|
| 2257 |
+
}
|
| 2258 |
+
|
| 2259 |
+
# Add cooldown info
|
| 2260 |
+
if key_cooldown_remaining is not None:
|
| 2261 |
+
cred_entry["key_cooldown_remaining"] = key_cooldown_remaining
|
| 2262 |
+
if active_cooldowns:
|
| 2263 |
+
cred_entry["model_cooldowns"] = active_cooldowns
|
| 2264 |
+
|
| 2265 |
+
# Add global stats for this credential
|
| 2266 |
+
if include_global:
|
| 2267 |
+
# Calculate global cache percentage
|
| 2268 |
+
global_total_input = (
|
| 2269 |
+
cred_global_tokens["input_cached"]
|
| 2270 |
+
+ cred_global_tokens["input_uncached"]
|
| 2271 |
+
)
|
| 2272 |
+
global_cache_pct = (
|
| 2273 |
+
round(
|
| 2274 |
+
cred_global_tokens["input_cached"]
|
| 2275 |
+
/ global_total_input
|
| 2276 |
+
* 100,
|
| 2277 |
+
1,
|
| 2278 |
+
)
|
| 2279 |
+
if global_total_input > 0
|
| 2280 |
+
else 0
|
| 2281 |
+
)
|
| 2282 |
+
|
| 2283 |
+
cred_entry["global"] = {
|
| 2284 |
+
"requests": cred_global_requests,
|
| 2285 |
+
"tokens": {
|
| 2286 |
+
"input_cached": cred_global_tokens["input_cached"],
|
| 2287 |
+
"input_uncached": cred_global_tokens["input_uncached"],
|
| 2288 |
+
"input_cache_pct": global_cache_pct,
|
| 2289 |
+
"output": cred_global_tokens["output"],
|
| 2290 |
+
},
|
| 2291 |
+
"approx_cost": cred_global_cost
|
| 2292 |
+
if cred_global_cost > 0
|
| 2293 |
+
else None,
|
| 2294 |
+
}
|
| 2295 |
+
|
| 2296 |
+
# Add model-specific data for providers with per-model tracking
|
| 2297 |
+
if models_data:
|
| 2298 |
+
cred_entry["models"] = {}
|
| 2299 |
+
for model_name, model_stats in models_data.items():
|
| 2300 |
+
if not isinstance(model_stats, dict):
|
| 2301 |
+
continue
|
| 2302 |
+
cred_entry["models"][model_name] = {
|
| 2303 |
+
"requests": model_stats.get("success_count", 0)
|
| 2304 |
+
+ model_stats.get("failure_count", 0),
|
| 2305 |
+
"request_count": model_stats.get("request_count", 0),
|
| 2306 |
+
"success_count": model_stats.get("success_count", 0),
|
| 2307 |
+
"failure_count": model_stats.get("failure_count", 0),
|
| 2308 |
+
"prompt_tokens": model_stats.get("prompt_tokens", 0),
|
| 2309 |
+
"prompt_tokens_cached": model_stats.get(
|
| 2310 |
+
"prompt_tokens_cached", 0
|
| 2311 |
+
),
|
| 2312 |
+
"completion_tokens": model_stats.get(
|
| 2313 |
+
"completion_tokens", 0
|
| 2314 |
+
),
|
| 2315 |
+
"approx_cost": model_stats.get("approx_cost", 0.0),
|
| 2316 |
+
"window_start_ts": model_stats.get("window_start_ts"),
|
| 2317 |
+
"quota_reset_ts": model_stats.get("quota_reset_ts"),
|
| 2318 |
+
# Quota baseline fields (Antigravity-specific)
|
| 2319 |
+
"baseline_remaining_fraction": model_stats.get(
|
| 2320 |
+
"baseline_remaining_fraction"
|
| 2321 |
+
),
|
| 2322 |
+
"baseline_fetched_at": model_stats.get(
|
| 2323 |
+
"baseline_fetched_at"
|
| 2324 |
+
),
|
| 2325 |
+
"quota_max_requests": model_stats.get("quota_max_requests"),
|
| 2326 |
+
"quota_display": model_stats.get("quota_display"),
|
| 2327 |
+
}
|
| 2328 |
+
|
| 2329 |
+
prov_stats["credentials"].append(cred_entry)
|
| 2330 |
+
|
| 2331 |
+
# Aggregate to provider totals (current period)
|
| 2332 |
+
prov_stats["total_requests"] += cred_requests
|
| 2333 |
+
prov_stats["tokens"]["input_cached"] += cred_tokens["input_cached"]
|
| 2334 |
+
prov_stats["tokens"]["input_uncached"] += cred_tokens["input_uncached"]
|
| 2335 |
+
prov_stats["tokens"]["output"] += cred_tokens["output"]
|
| 2336 |
+
if cred_cost > 0:
|
| 2337 |
+
prov_stats["approx_cost"] += cred_cost
|
| 2338 |
+
|
| 2339 |
+
# Aggregate to global provider totals
|
| 2340 |
+
global_providers[provider]["total_requests"] += cred_global_requests
|
| 2341 |
+
global_providers[provider]["tokens"]["input_cached"] += (
|
| 2342 |
+
cred_global_tokens["input_cached"]
|
| 2343 |
+
)
|
| 2344 |
+
global_providers[provider]["tokens"]["input_uncached"] += (
|
| 2345 |
+
cred_global_tokens["input_uncached"]
|
| 2346 |
+
)
|
| 2347 |
+
global_providers[provider]["tokens"]["output"] += cred_global_tokens[
|
| 2348 |
+
"output"
|
| 2349 |
+
]
|
| 2350 |
+
global_providers[provider]["approx_cost"] += cred_global_cost
|
| 2351 |
+
|
| 2352 |
+
# Calculate cache percentages for each provider
|
| 2353 |
+
for provider, prov_stats in providers.items():
|
| 2354 |
+
total_input = (
|
| 2355 |
+
prov_stats["tokens"]["input_cached"]
|
| 2356 |
+
+ prov_stats["tokens"]["input_uncached"]
|
| 2357 |
+
)
|
| 2358 |
+
if total_input > 0:
|
| 2359 |
+
prov_stats["tokens"]["input_cache_pct"] = round(
|
| 2360 |
+
prov_stats["tokens"]["input_cached"] / total_input * 100, 1
|
| 2361 |
+
)
|
| 2362 |
+
# Set cost to None if 0
|
| 2363 |
+
if prov_stats["approx_cost"] == 0:
|
| 2364 |
+
prov_stats["approx_cost"] = None
|
| 2365 |
+
|
| 2366 |
+
# Calculate global cache percentages
|
| 2367 |
+
if include_global and provider in global_providers:
|
| 2368 |
+
gp = global_providers[provider]
|
| 2369 |
+
global_total = (
|
| 2370 |
+
gp["tokens"]["input_cached"] + gp["tokens"]["input_uncached"]
|
| 2371 |
+
)
|
| 2372 |
+
if global_total > 0:
|
| 2373 |
+
gp["tokens"]["input_cache_pct"] = round(
|
| 2374 |
+
gp["tokens"]["input_cached"] / global_total * 100, 1
|
| 2375 |
+
)
|
| 2376 |
+
if gp["approx_cost"] == 0:
|
| 2377 |
+
gp["approx_cost"] = None
|
| 2378 |
+
prov_stats["global"] = gp
|
| 2379 |
+
|
| 2380 |
+
# Build summary (current period)
|
| 2381 |
+
total_creds = sum(p["credential_count"] for p in providers.values())
|
| 2382 |
+
active_creds = sum(p["active_count"] for p in providers.values())
|
| 2383 |
+
exhausted_creds = sum(p["exhausted_count"] for p in providers.values())
|
| 2384 |
+
total_requests = sum(p["total_requests"] for p in providers.values())
|
| 2385 |
+
total_input_cached = sum(
|
| 2386 |
+
p["tokens"]["input_cached"] for p in providers.values()
|
| 2387 |
+
)
|
| 2388 |
+
total_input_uncached = sum(
|
| 2389 |
+
p["tokens"]["input_uncached"] for p in providers.values()
|
| 2390 |
+
)
|
| 2391 |
+
total_output = sum(p["tokens"]["output"] for p in providers.values())
|
| 2392 |
+
total_cost = sum(p["approx_cost"] or 0 for p in providers.values())
|
| 2393 |
+
|
| 2394 |
+
total_input = total_input_cached + total_input_uncached
|
| 2395 |
+
input_cache_pct = (
|
| 2396 |
+
round(total_input_cached / total_input * 100, 1) if total_input > 0 else 0
|
| 2397 |
+
)
|
| 2398 |
+
|
| 2399 |
+
result = {
|
| 2400 |
+
"providers": providers,
|
| 2401 |
+
"summary": {
|
| 2402 |
+
"total_providers": len(providers),
|
| 2403 |
+
"total_credentials": total_creds,
|
| 2404 |
+
"active_credentials": active_creds,
|
| 2405 |
+
"exhausted_credentials": exhausted_creds,
|
| 2406 |
+
"total_requests": total_requests,
|
| 2407 |
+
"tokens": {
|
| 2408 |
+
"input_cached": total_input_cached,
|
| 2409 |
+
"input_uncached": total_input_uncached,
|
| 2410 |
+
"input_cache_pct": input_cache_pct,
|
| 2411 |
+
"output": total_output,
|
| 2412 |
+
},
|
| 2413 |
+
"approx_total_cost": total_cost if total_cost > 0 else None,
|
| 2414 |
+
},
|
| 2415 |
+
"data_source": "cache",
|
| 2416 |
+
"timestamp": now_ts,
|
| 2417 |
+
}
|
| 2418 |
+
|
| 2419 |
+
# Build global summary
|
| 2420 |
+
if include_global:
|
| 2421 |
+
global_total_requests = sum(
|
| 2422 |
+
gp["total_requests"] for gp in global_providers.values()
|
| 2423 |
+
)
|
| 2424 |
+
global_total_input_cached = sum(
|
| 2425 |
+
gp["tokens"]["input_cached"] for gp in global_providers.values()
|
| 2426 |
+
)
|
| 2427 |
+
global_total_input_uncached = sum(
|
| 2428 |
+
gp["tokens"]["input_uncached"] for gp in global_providers.values()
|
| 2429 |
+
)
|
| 2430 |
+
global_total_output = sum(
|
| 2431 |
+
gp["tokens"]["output"] for gp in global_providers.values()
|
| 2432 |
+
)
|
| 2433 |
+
global_total_cost = sum(
|
| 2434 |
+
gp["approx_cost"] or 0 for gp in global_providers.values()
|
| 2435 |
+
)
|
| 2436 |
+
|
| 2437 |
+
global_total_input = global_total_input_cached + global_total_input_uncached
|
| 2438 |
+
global_input_cache_pct = (
|
| 2439 |
+
round(global_total_input_cached / global_total_input * 100, 1)
|
| 2440 |
+
if global_total_input > 0
|
| 2441 |
+
else 0
|
| 2442 |
+
)
|
| 2443 |
+
|
| 2444 |
+
result["global_summary"] = {
|
| 2445 |
+
"total_providers": len(global_providers),
|
| 2446 |
+
"total_credentials": total_creds,
|
| 2447 |
+
"total_requests": global_total_requests,
|
| 2448 |
+
"tokens": {
|
| 2449 |
+
"input_cached": global_total_input_cached,
|
| 2450 |
+
"input_uncached": global_total_input_uncached,
|
| 2451 |
+
"input_cache_pct": global_input_cache_pct,
|
| 2452 |
+
"output": global_total_output,
|
| 2453 |
+
},
|
| 2454 |
+
"approx_total_cost": global_total_cost
|
| 2455 |
+
if global_total_cost > 0
|
| 2456 |
+
else None,
|
| 2457 |
+
}
|
| 2458 |
+
|
| 2459 |
+
return result
|
| 2460 |
+
|
| 2461 |
+
async def reload_from_disk(self) -> None:
|
| 2462 |
+
"""
|
| 2463 |
+
Force reload usage data from disk.
|
| 2464 |
+
|
| 2465 |
+
Useful when another process may have updated the file.
|
| 2466 |
+
"""
|
| 2467 |
+
async with self._init_lock:
|
| 2468 |
+
self._initialized.clear()
|
| 2469 |
+
await self._load_usage()
|
| 2470 |
+
await self._reset_daily_stats_if_needed()
|
| 2471 |
+
self._initialized.set()
|