scripts/refresh_limits.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+
+from gibby.account_ops import (
+    PermanentAccountFailure,
+    handle_failed_account,
+    refresh_account_usage,
+    window_used_percent,
+)
+from gibby.client import OpenAIClient
+from gibby.models import format_reset_in
+from gibby.settings import Settings
+from gibby.store import JsonStateStore
+
+
+async def run(data_dir: Path | None = None) -> None:
+    settings = Settings()
+    if data_dir is not None:
+        settings.data_dir = data_dir
+
+    store = JsonStateStore(settings.accounts_file, settings.failed_file)
+    state = store.load()
+    if not state.accounts:
+        print("No accounts configured.")
+        return
+
+    client = OpenAIClient(settings)
+    try:
+        for account in list(state.accounts):
+            previous_id = account.id
+            try:
+                usage = await refresh_account_usage(
+                    account,
+                    client,
+                    settings.exhausted_usage_threshold,
+                )
+                store.update_active_account_id(state, previous_id, account.id)
+                print(
+                    f"token ready for {account.id}, "
+                    f"primary {window_used_percent(usage.primary_window)}% "
+                    f"reset in {format_reset_in(usage.primary_window.reset_at if usage.primary_window else None)}, "
+                    f"secondary {window_used_percent(usage.secondary_window)}% "
+                    f"reset in {format_reset_in(usage.secondary_window.reset_at if usage.secondary_window else None)}"
+                )
+            except PermanentAccountFailure as exc:
+                account.last_error = str(exc)
+                handle_failed_account(store, account)
+                store.remove_account(state, account.id)
+                print(f"{account.id}: removed={exc}")
+            except Exception as exc:
+                account.last_error = str(exc)
+                print(f"{account.id}: error={exc}")
+        store.save(state)
+    finally:
+        await client.aclose()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-d", "--data-dir", type=Path)
+    args = parser.parse_args()
+    try:
+        asyncio.run(run(args.data_dir))
+    except KeyboardInterrupt:
+        print("\nCancelled.")
+
+
+if __name__ == "__main__":
+    main()