84 lines
3.7 KiB
Python
84 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
|
|
|
|
from gibby.client import OpenAIAPIError, OpenAIClient
|
|
from gibby.models import format_reset_in, parse_usage_payload
|
|
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):
|
|
try:
|
|
if account.token_refresh_at <= int(time.time()) + settings.token_refresh_buffer_seconds:
|
|
access_token, refresh_token, refresh_at = await client.refresh_access_token(
|
|
account.refresh_token
|
|
)
|
|
account.access_token = access_token
|
|
account.refresh_token = refresh_token
|
|
account.token_refresh_at = refresh_at
|
|
payload = await client.fetch_usage_payload(account.access_token)
|
|
email = payload.get("email")
|
|
if isinstance(email, str) and email:
|
|
previous_email = account.email
|
|
account.email = email
|
|
store.update_active_account(state, previous_email, email)
|
|
account.usage = parse_usage_payload(payload)
|
|
account.usage_checked_at = account.usage.checked_at
|
|
print(
|
|
f"token ready for {account.email}, "
|
|
f"primary {account.usage.primary_window.used_percent if account.usage.primary_window else 0}% "
|
|
f"reset in {format_reset_in(account.usage.primary_window.reset_at if account.usage and account.usage.primary_window else None)}, "
|
|
f"secondary {account.usage.secondary_window.used_percent if account.usage and account.usage.secondary_window else 0}% "
|
|
f"reset in {format_reset_in(account.usage.secondary_window.reset_at if account.usage and account.usage.secondary_window else None)}"
|
|
)
|
|
except OpenAIAPIError as exc:
|
|
if exc.permanent:
|
|
usage = account.usage
|
|
print(
|
|
f"moving account to failed.json: email={account.email} reason=usage refresh auth failure: {exc} "
|
|
f"primary={usage.primary_window.used_percent if usage and usage.primary_window else 0}% "
|
|
f"primary reset in {format_reset_in(usage.primary_window.reset_at if usage and usage.primary_window else None)} "
|
|
f"secondary={usage.secondary_window.used_percent if usage and usage.secondary_window else 0}% "
|
|
f"secondary reset in {format_reset_in(usage.secondary_window.reset_at if usage and usage.secondary_window else None)}"
|
|
)
|
|
store.move_to_failed(state, account.email)
|
|
print(f"{account.email}: removed={exc}")
|
|
else:
|
|
print(f"{account.email}: 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()
|