First pass at an OpenAI limit script
This commit is contained in:
@@ -63,7 +63,7 @@
|
||||
"margin-right": 16,
|
||||
"modules-left": ["gamemode", "battery", "temperature", "cpu", "memory", "network"],
|
||||
"modules-center": [],
|
||||
"modules-right": ["mpris", "pulseaudio", "custom/output-device", "backlight", "idle_inhibitor", "clock"],
|
||||
"modules-right": ["mpris", "pulseaudio", "custom/output-device", "custom/openai-rate", "backlight", "idle_inhibitor", "clock"],
|
||||
"clock": {
|
||||
"format": "{:%a %b %d %I:%M %p}",
|
||||
"tooltip": false
|
||||
@@ -155,6 +155,14 @@
|
||||
"exec": "flatpak remote-ls --updates --app | wc -l",
|
||||
"exec-if": "test $(flatpak remote-ls --updates --app | wc -l) -gt 0"
|
||||
},
|
||||
"custom/openai-rate": {
|
||||
"interval": 15,
|
||||
"return-type": "json",
|
||||
"exec": "$HOME/.config/waybar/scripts/openai-rate.py",
|
||||
"exec-if": "test -r ~/.codex/auth.json",
|
||||
"max-length": 90,
|
||||
"format": "{}"
|
||||
},
|
||||
"custom/backup": {
|
||||
"interval": 60,
|
||||
"format": "",
|
||||
|
||||
+195
@@ -0,0 +1,195 @@
|
||||
#!/usr/bin/env python3
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
DEFAULT_BASE_URL = "https://chatgpt.com/backend-api"
|
||||
|
||||
|
||||
def waybar(text, klass="good", tooltip=None):
|
||||
print(json.dumps({"text": text, "class": klass, "tooltip": tooltip or text}))
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def load_auth():
|
||||
auth_path = Path(os.environ.get("OPENAI_AUTH_FILE", "~/.codex/auth.json")).expanduser()
|
||||
auth = {}
|
||||
if auth_path.is_file():
|
||||
try:
|
||||
auth = json.loads(auth_path.read_text())
|
||||
except (OSError, json.JSONDecodeError):
|
||||
waybar("AI auth unreadable", "critical", f"Could not parse {auth_path}")
|
||||
|
||||
tokens = auth.get("tokens") or {}
|
||||
access_token = tokens.get("access_token")
|
||||
account_id = tokens.get("account_id")
|
||||
if access_token:
|
||||
return access_token, account_id, "~/.codex/auth.json:tokens.access_token"
|
||||
|
||||
api_key = auth.get("OPENAI_API_KEY") or os.environ.get("OPENAI_API_KEY")
|
||||
if api_key:
|
||||
source = "~/.codex/auth.json:OPENAI_API_KEY" if auth.get("OPENAI_API_KEY") else "env:OPENAI_API_KEY"
|
||||
return api_key, account_id, source
|
||||
|
||||
waybar(
|
||||
"AI auth missing",
|
||||
"warning",
|
||||
"No Codex access token found in ~/.codex/auth.json",
|
||||
)
|
||||
|
||||
|
||||
def usage_url():
|
||||
base = os.environ.get("CODEX_RATE_BASE_URL") or os.environ.get("OPENAI_BASE_URL") or DEFAULT_BASE_URL
|
||||
base = base.rstrip("/")
|
||||
if base.endswith("/backend-api/codex"):
|
||||
base = base.removesuffix("/codex")
|
||||
if base.startswith(("https://chatgpt.com", "https://chat.openai.com")) and "/backend-api" not in base:
|
||||
base += "/backend-api"
|
||||
return f"{base}/wham/usage" if "/backend-api" in base else f"{base}/api/codex/usage"
|
||||
|
||||
|
||||
def fetch_usage(token, account_id):
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "codex-waybar-rate",
|
||||
}
|
||||
if account_id:
|
||||
headers["ChatGPT-Account-ID"] = account_id
|
||||
|
||||
req = urllib.request.Request(usage_url(), headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=15) as res:
|
||||
return json.load(res)
|
||||
except urllib.error.HTTPError as exc:
|
||||
detail = exc.read().decode("utf-8", "replace")[:240]
|
||||
waybar("AI rate failed", "critical", f"Usage endpoint returned HTTP {exc.code}\n{detail}")
|
||||
except (OSError, TimeoutError) as exc:
|
||||
waybar("AI rate failed", "critical", f"Could not reach Codex usage endpoint\n{exc}")
|
||||
except json.JSONDecodeError as exc:
|
||||
waybar("AI rate unreadable", "critical", f"Usage endpoint returned invalid JSON\n{exc}")
|
||||
|
||||
|
||||
def window_label(seconds):
|
||||
if not isinstance(seconds, int) or seconds <= 0:
|
||||
return "usage"
|
||||
|
||||
minutes = math.ceil(seconds / 60)
|
||||
windows = (
|
||||
(5 * 60, "5h"),
|
||||
(24 * 60, "daily"),
|
||||
(7 * 24 * 60, "weekly"),
|
||||
(30 * 24 * 60, "monthly"),
|
||||
(365 * 24 * 60, "annual"),
|
||||
)
|
||||
for expected, label in windows:
|
||||
if expected * 0.95 <= minutes <= expected * 1.05:
|
||||
return label
|
||||
return "usage"
|
||||
|
||||
|
||||
def reset_text(seconds):
|
||||
if not isinstance(seconds, int) or seconds < 0:
|
||||
return None
|
||||
if seconds < 90:
|
||||
return f"{seconds}s"
|
||||
minutes = round(seconds / 60)
|
||||
if minutes < 90:
|
||||
return f"{minutes}m"
|
||||
hours = round(minutes / 60)
|
||||
if hours < 48:
|
||||
return f"{hours}h"
|
||||
return f"{round(hours / 24)}d"
|
||||
|
||||
|
||||
def format_window(name, window):
|
||||
if not isinstance(window, dict):
|
||||
return None
|
||||
|
||||
used = window.get("used_percent")
|
||||
if not isinstance(used, int):
|
||||
return None
|
||||
|
||||
label = window_label(window.get("limit_window_seconds"))
|
||||
left = max(0, 100 - used)
|
||||
reset = reset_text(window.get("reset_after_seconds"))
|
||||
tooltip = f"{name} {label}: {used}% used"
|
||||
if reset:
|
||||
tooltip += f", resets in {reset}"
|
||||
return {"text": f"{left}% {label}", "used": used, "tooltip": tooltip}
|
||||
|
||||
|
||||
def credit_text(credits):
|
||||
if not isinstance(credits, dict) or not credits.get("has_credits"):
|
||||
return None
|
||||
if credits.get("unlimited"):
|
||||
return "credits unlimited"
|
||||
balance = credits.get("balance")
|
||||
if balance is None:
|
||||
return "credits enabled"
|
||||
try:
|
||||
return f"credits {round(float(balance))}"
|
||||
except (TypeError, ValueError):
|
||||
return f"credits {balance}"
|
||||
|
||||
|
||||
def spend_control_text(spend_control):
|
||||
limit = (spend_control or {}).get("individual_limit")
|
||||
if not isinstance(limit, dict):
|
||||
return None
|
||||
|
||||
remaining = limit.get("remaining_percent")
|
||||
used = limit.get("used")
|
||||
cap = limit.get("limit")
|
||||
reset = reset_text(limit.get("reset_after_seconds"))
|
||||
if not isinstance(remaining, int):
|
||||
return None
|
||||
|
||||
text = f"monthly credits: {100 - remaining}% used"
|
||||
if used and cap:
|
||||
text += f" ({used}/{cap})"
|
||||
if reset:
|
||||
text += f", resets in {reset}"
|
||||
return text
|
||||
|
||||
|
||||
def main():
|
||||
token, account_id, source = load_auth()
|
||||
data = fetch_usage(token, account_id)
|
||||
rate_limit = data.get("rate_limit") or {}
|
||||
|
||||
windows = [
|
||||
format_window("primary", rate_limit.get("primary_window")),
|
||||
format_window("secondary", rate_limit.get("secondary_window")),
|
||||
]
|
||||
windows = [window for window in windows if window]
|
||||
|
||||
parts = [window["text"] for window in windows]
|
||||
credits = credit_text(data.get("credits"))
|
||||
if credits:
|
||||
parts.append(credits)
|
||||
|
||||
if not parts:
|
||||
waybar("AI n/a", "warning", f"No displayable usage data\nsource: {source}")
|
||||
|
||||
worst_used = max(window["used"] for window in windows) if windows else 0
|
||||
klass = "critical" if rate_limit.get("limit_reached") or worst_used >= 95 else "warning" if worst_used >= 90 else "good"
|
||||
|
||||
tooltip = [window["tooltip"] for window in windows]
|
||||
spend = spend_control_text(data.get("spend_control"))
|
||||
if credits:
|
||||
tooltip.append(credits)
|
||||
if spend:
|
||||
tooltip.append(spend)
|
||||
tooltip.append(f"source: {source}")
|
||||
|
||||
waybar(f"AI {' '.join(parts)}", klass, "\n".join(tooltip))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -238,3 +238,13 @@ window#waybar.fullscreen #window {
|
||||
color: #ebdbb2;
|
||||
padding: 0 1em;
|
||||
}
|
||||
#custom-openai-rate {
|
||||
padding: 0 1em;
|
||||
color: #ebdbb2;
|
||||
}
|
||||
#custom-openai-rate.warning {
|
||||
color: #fabd2f;
|
||||
}
|
||||
#custom-openai-rate.critical {
|
||||
color: #fb4934;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user