refactor
This commit is contained in:
parent
7cef56de15
commit
ecb5f68e32
17 changed files with 760 additions and 1626 deletions
|
|
@ -10,20 +10,9 @@ from urllib.parse import parse_qs, urlparse
|
|||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
|
||||
|
||||
from gibby.account_ops import (
|
||||
PermanentAccountFailure,
|
||||
failed_identifier,
|
||||
refresh_account_usage,
|
||||
window_used_percent,
|
||||
)
|
||||
from gibby.client import OpenAIAPIError, OpenAIClient
|
||||
from gibby.models import AccountRecord, format_reset_in
|
||||
from gibby.oauth import (
|
||||
build_authorize_url,
|
||||
generate_pkce_pair,
|
||||
generate_state,
|
||||
make_account_id,
|
||||
)
|
||||
from gibby.models import AccountRecord, format_reset_in, parse_usage_payload
|
||||
from gibby.oauth import build_authorize_url, generate_pkce_pair, generate_state
|
||||
from gibby.settings import Settings
|
||||
from gibby.store import JsonStateStore
|
||||
|
||||
|
|
@ -55,12 +44,10 @@ async def wait_for_callback(
|
|||
code, state = parse_redirect_url(parts[1])
|
||||
if not result.done():
|
||||
result.set_result((code, state))
|
||||
|
||||
while True:
|
||||
line = await reader.readline()
|
||||
if not line or line == b"\r\n":
|
||||
break
|
||||
|
||||
writer.write(
|
||||
b"HTTP/1.1 200 OK\r\n"
|
||||
b"Content-Type: text/plain; charset=utf-8\r\n"
|
||||
|
|
@ -91,35 +78,34 @@ async def exchange_and_store_account(
|
|||
verifier: str,
|
||||
set_active: bool,
|
||||
) -> AccountRecord:
|
||||
access_token, refresh_token, expires_at = await client.exchange_code(code, verifier)
|
||||
access_token, refresh_token, token_refresh_at = await client.exchange_code(code, verifier)
|
||||
account = AccountRecord(
|
||||
id=make_account_id(),
|
||||
email="",
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
expires_at=expires_at,
|
||||
token_refresh_at=token_refresh_at,
|
||||
)
|
||||
try:
|
||||
usage = await refresh_account_usage(
|
||||
account,
|
||||
client,
|
||||
client.settings.exhausted_usage_threshold,
|
||||
)
|
||||
except PermanentAccountFailure:
|
||||
store.append_failed_identifier(failed_identifier(account))
|
||||
raise
|
||||
except OpenAIAPIError as exc:
|
||||
account.last_error = str(exc)
|
||||
payload = await client.fetch_usage_payload(account.access_token)
|
||||
email = payload.get("email")
|
||||
if isinstance(email, str) and email:
|
||||
account.email = email
|
||||
account.usage = parse_usage_payload(payload)
|
||||
account.usage_checked_at = account.usage.checked_at
|
||||
except OpenAIAPIError:
|
||||
if not account.email:
|
||||
raise
|
||||
store.upsert_account(account, set_active=set_active)
|
||||
print("Usage fetch failed, stored account without usage snapshot.")
|
||||
return account
|
||||
|
||||
store.upsert_account(account, set_active=set_active)
|
||||
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)}"
|
||||
f"token ready for {account.email}, "
|
||||
f"primary {account.usage.primary_window.used_percent if account.usage and 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)}"
|
||||
)
|
||||
return account
|
||||
|
||||
|
|
@ -168,14 +154,12 @@ async def run(
|
|||
store = JsonStateStore(settings.accounts_file, settings.failed_file)
|
||||
client = OpenAIClient(settings)
|
||||
try:
|
||||
account = await exchange_and_store_account(
|
||||
store, client, code, verifier, set_active
|
||||
)
|
||||
account = await exchange_and_store_account(store, client, code, verifier, set_active)
|
||||
finally:
|
||||
await client.aclose()
|
||||
|
||||
print(f"Stored account: {account.id}")
|
||||
print(f"Access token expires at: {account.expires_at}")
|
||||
print(f"Stored account: {account.email}")
|
||||
print(f"Access token refresh at: {account.token_refresh_at}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
|
@ -192,7 +176,7 @@ def main() -> None:
|
|||
print("Timed out waiting for OAuth callback.")
|
||||
except ValueError as exc:
|
||||
print(str(exc))
|
||||
except PermanentAccountFailure as exc:
|
||||
except OpenAIAPIError as exc:
|
||||
print(str(exc))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -7,14 +7,8 @@ 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.client import OpenAIAPIError, OpenAIClient
|
||||
from gibby.models import format_reset_in, parse_usage_payload
|
||||
from gibby.settings import Settings
|
||||
from gibby.store import JsonStateStore
|
||||
|
||||
|
|
@ -33,29 +27,23 @@ async def run(data_dir: Path | None = None) -> None:
|
|||
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)
|
||||
payload = await client.fetch_usage_payload(account.access_token)
|
||||
account.usage = parse_usage_payload(payload)
|
||||
account.usage_checked_at = account.usage.checked_at
|
||||
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)}"
|
||||
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 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}")
|
||||
except OpenAIAPIError as exc:
|
||||
if exc.permanent:
|
||||
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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue