ai-token-usage 0.1.5__py2.py3-none-any.whl

ai_token_usage-0.1.5.dist-info/METADATA

@@ -0,0 +1,3 @@
+Metadata-Version: 2.4
+Name: ai-token-usage
+Version: 0.1.5

ai_token_usage-0.1.5.dist-info/RECORD

@@ -0,0 +1,14 @@
+token_tracker/__init__.py,sha256=3v6ELg7ohMsCl1Cv3ldXYCyz0qO-0mSPp3TdEwQWMy4,29
+token_tracker/cli.py,sha256=tDQp7ysLiMnlidLo-dqXDprtniGmn4F0EAa1aNMpnUc,2872
+token_tracker/providers/__init__.py,sha256=g5Ixyv5KjMt3JMHmxvW52RYTWORxawWF2u2TiKXlzu4,35
+token_tracker/providers/base.py,sha256=aWz0KnkXzJeiK-x0bCcYBPyHN5aRwgXzumJZe3fy7FQ,248
+token_tracker/providers/codex_direct.py,sha256=ihrVHYhoQIiuWqVeW5ulAqqRjAnguI_SqNA5DTgbl3A,3554
+token_tracker/providers/codexbar.py,sha256=hbiGBRNVEVYDChEFHGjvpNpOB6FLdANgPLAX8u6YbVg,5195
+token_tracker/providers/openrouter.py,sha256=-7beOd0GKOUZePG5u_NHshx6RbR5N1xXz7UJpcDJzag,1124
+token_tracker/stores/__init__.py,sha256=h7ad-GQ1g8LfS8xSljhq1ITJvPh1kpRbVh5Twt3OSoQ,32
+token_tracker/stores/base.py,sha256=sWUuZ0ZT2uzWXy7tYacO8oiEZ_ey9-D2iu3dRdbkwZ8,236
+token_tracker/stores/google_sheets.py,sha256=Fvo1d0YZve33RQlARcbwhXHXBjfjZwhRz1tczzJPW_E,4257
+ai_token_usage-0.1.5.dist-info/METADATA,sha256=SqyEwPOWA9oAryRjdbEw4tsqIhceDTgwkTi15O1iMRc,58
+ai_token_usage-0.1.5.dist-info/WHEEL,sha256=tkmg4JIqwd9H8mL30xA7crRmoStyCtGp0VWshokd1Jc,105
+ai_token_usage-0.1.5.dist-info/entry_points.txt,sha256=c0h8ZPdyOmwYe19N8A0SxbKx5zkuNZBFbNVTrN9wfT0,57
+ai_token_usage-0.1.5.dist-info/RECORD,,

ai_token_usage-0.1.5.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py2-none-any
+Tag: py3-none-any

ai_token_usage-0.1.5.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+token-tracker = token_tracker.cli:main

token_tracker/__init__.py

@@ -0,0 +1 @@
+"""Token Tracker package."""

token_tracker/cli.py

@@ -0,0 +1,79 @@
+import argparse
+import logging
+import sys
+import os
+from pathlib import Path
+from datetime import datetime
+from zoneinfo import ZoneInfo
+
+from .providers.openrouter import OpenRouterProvider
+from .providers.codex_direct import CodexDirectProvider
+from .providers.codexbar import CodexBarProvider
+from .stores.google_sheets import GoogleSheetsStore
+
+def parse_args(argv=None):
+    parser = argparse.ArgumentParser(description="Collect token usage and push to Google Sheets")
+    parser.add_argument("-i", "--init", action="store_true",
+                        help="Overwrite the sheets (uses df_to_sheet). Default: append.")
+    parser.add_argument("--log-level", default="INFO",
+                        help="Logging level (DEBUG, INFO, WARNING, ERROR)")
+    parser.add_argument("--codex", action="store_true", help="Only run direct Codex API fetch")
+    parser.add_argument("--codex-bar", action="store_true", help="Only run codexbar fetch")
+    parser.add_argument("--openrouter", action="store_true", help="Only run OpenRouter fetch")
+    return parser.parse_args(argv)
+
+def main(argv=None) -> None:
+    args = parse_args(argv)
+
+    numeric_level = getattr(logging, args.log_level.upper(), logging.INFO)
+    logging.basicConfig(level=numeric_level, format="%(asctime)s %(levelname)s: %(message)s")
+
+    now_stamp = datetime.now(ZoneInfo("America/New_York")).strftime("%Y-%m-%d %H:%M:%S")
+    had_error = False
+
+    run_all = not (args.codex or args.codex_bar or args.openrouter)
+    run_codex = run_all or args.codex
+    run_codex_bar = run_all or args.codex_bar
+    run_openrouter = run_all or args.openrouter
+
+    try:
+        store = GoogleSheetsStore()
+    except Exception:
+        logging.exception("Failed to initialize GoogleSheetsStore")
+        sys.exit(1)
+
+    if run_codex:
+        try:
+            provider = CodexDirectProvider()
+            df = provider.fetch(now=now_stamp)
+            store.save(df, name="codex", init=args.init)
+        except Exception:
+            logging.exception("Failed to process direct Codex data")
+            had_error = True
+
+    if run_codex_bar:
+        try:
+            provider = CodexBarProvider()
+            df = provider.fetch(now=now_stamp)
+            store.save(df, name="codexbar", init=args.init)
+        except Exception:
+            logging.exception("Failed to process codexbar data")
+            had_error = True
+
+    if run_openrouter:
+        try:
+            provider = OpenRouterProvider()
+            df = provider.fetch(now=now_stamp)
+            store.save(df, name="openrouter", init=args.init)
+        except Exception:
+            logging.exception("Failed to process OpenRouter data")
+            had_error = True
+
+    if had_error:
+        logging.error("Completed with errors")
+        sys.exit(1)
+
+    logging.info("All sources written successfully")
+
+if __name__ == "__main__":
+    main()

token_tracker/providers/__init__.py

@@ -0,0 +1 @@
+"""Providers for Token Tracker."""

token_tracker/providers/base.py

@@ -0,0 +1,9 @@
+from abc import ABC, abstractmethod
+import pandas as pd
+from typing import Optional
+
+class BaseProvider(ABC):
+    @abstractmethod
+    def fetch(self, now: str) -> pd.DataFrame:
+        """Fetch usage data and return as a DataFrame."""
+        pass

token_tracker/providers/codex_direct.py

@@ -0,0 +1,82 @@
+import requests
+import pandas as pd
+import json
+from pathlib import Path
+from datetime import datetime
+from zoneinfo import ZoneInfo
+from typing import Optional
+from .base import BaseProvider
+from project_toolkit import settings
+
+CODEX_DIRECT_COLS = [
+    "time", "email", "plan_type",
+    "rate_limit.p.used_percent", "rate_limit.p.limit_window_seconds",
+    "rate_limit.p.reset_after_seconds", "rate_limit.p.reset_at",
+    "rate_limit.s.used_percent", "rate_limit.s.limit_window_seconds",
+    "rate_limit.s.reset_after_seconds", "rate_limit.s.reset_at",
+    "credits.has_credits", "credits.unlimited",
+    "credits.overage_limit_reached", "credits.balance",
+    "spend_control.reached", "spend_control.individual_limit"
+]
+
+class CodexDirectProvider(BaseProvider):
+    def __init__(self, api_key: Optional[str] = None):
+        self.api_key = api_key or settings.env("CODEX_API_KEY")
+        if not self.api_key:
+            auth_file = Path.home() / ".config" / "codex" / "auth.json"
+            if auth_file.exists():
+                try:
+                    with auth_file.open("r") as f:
+                        auth_data = json.load(f)
+                        self.api_key = auth_data.get("token") or auth_data.get("CODEX_API_KEY")
+                except Exception:
+                    pass
+
+    def fetch(self, now: str) -> pd.DataFrame:
+        if not self.api_key:
+            raise ValueError("CODEX_API_KEY not found in environment or auth file")
+
+        response = requests.get(
+            "https://chatgpt.com/backend-api/wham/usage",
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Accept": "application/json"
+            }
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        def convert_ts(ts):
+            if not ts:
+                return None
+            return datetime.fromtimestamp(ts, tz=ZoneInfo("America/New_York")).isoformat()
+
+        rate_limit = data.get("rate_limit", {}) or {}
+        primary = rate_limit.get("primary_window", {}) or {}
+        secondary = rate_limit.get("secondary_window", {}) or {}
+        credits_data = data.get("credits", {}) or {}
+        spend_control = data.get("spend_control", {}) or {}
+
+        row = {
+            "time": now,
+            "email": data.get("email"),
+            "plan_type": data.get("plan_type"),
+            "rate_limit.p.used_percent": primary.get("used_percent"),
+            "rate_limit.p.limit_window_seconds": primary.get("limit_window_seconds"),
+            "rate_limit.p.reset_after_seconds": primary.get("reset_after_seconds"),
+            "rate_limit.p.reset_at": convert_ts(primary.get("reset_at")),
+            "rate_limit.s.used_percent": secondary.get("used_percent"),
+            "rate_limit.s.limit_window_seconds": secondary.get("limit_window_seconds"),
+            "rate_limit.s.reset_after_seconds": secondary.get("reset_after_seconds"),
+            "rate_limit.s.reset_at": convert_ts(secondary.get("reset_at")),
+            "credits.has_credits": credits_data.get("has_credits"),
+            "credits.unlimited": credits_data.get("unlimited"),
+            "credits.overage_limit_reached": credits_data.get("overage_limit_reached"),
+            "credits.balance": credits_data.get("balance"),
+            "spend_control.reached": spend_control.get("reached"),
+            "spend_control.individual_limit": spend_control.get("individual_limit")
+        }
+
+        df = pd.DataFrame([row], columns=CODEX_DIRECT_COLS)
+        df.attrs["raw_json"] = data
+        return df

token_tracker/providers/codexbar.py

@@ -0,0 +1,119 @@
+import subprocess
+import shlex
+import json
+import re
+import pandas as pd
+from typing import Optional
+from .base import BaseProvider
+
+COL_TIME = "time"
+COL_CODEX_USAGE = "Codex Usage"
+COL_CODEX_RESETS = "Codex Resets"
+COL_CODEX_5HR_USAGE = "Codex 5hr Usage"
+COL_CODEX_5HR_RESET = "Codex 5hr Reset"
+COL_CODEX_WEEKLY_USAGE = "Codex Weekly Usage"
+COL_CODEX_WEEKLY_RESET = "Codex Weekly Reset"
+COL_CODEX_CREDITS = "Codex Credits"
+COL_CODEX_PLAN = "Codex Plan"
+COL_COPILOT_PREMIUM = "Copilot Premium"
+COL_COPILOT_CHAT = "Copilot Chat"
+COL_COPILOT_PLAN = "Copilot Plan"
+COL_GEMINI_PRO = "Gemini Pro"
+COL_GEMINI_PRO_RESETS = "Gemini Pro Resets"
+COL_GEMINI_FLASH = "Gemini Flash"
+COL_GEMINI_FLASH_RESET = "Gemini Flash Reset"
+COL_GEMINI_FLASH_LITE = "Gemini Flash Lite"
+COL_GEMINI_FLASHLITE_RESET = "Gemini Flashlite Reset"
+COL_ANTIGRAVITY_CLAUDE = "Antigravity Claude"
+COL_ANTIGRAVITY_CLAUDE_RESET = "Antigravity Claude Reset"
+COL_ANTIGRAVITY_GEMINI_PRO = "Antigravity Gemini Pro"
+COL_ANTIGRAVITY_GEMINI_PRO_RESET = "Antigravity Gemini Pro Reset"
+COL_ANTIGRAVITY_FLASHLITE = "Antigravity Flashlite"
+COL_ANTIGRAVITY_FLASHLITE_RESETS = "Antigravity Flashlite Resets"
+COL_ANTIGRAVITY_PLAN = "Antigravity Plan"
+
+CODEXBAR_COLS = [
+    COL_TIME, COL_CODEX_USAGE, COL_CODEX_RESETS, COL_CODEX_5HR_USAGE, COL_CODEX_5HR_RESET,
+    COL_CODEX_WEEKLY_USAGE, COL_CODEX_WEEKLY_RESET, COL_CODEX_CREDITS, COL_CODEX_PLAN,
+    COL_COPILOT_PREMIUM, COL_COPILOT_CHAT, COL_COPILOT_PLAN, COL_GEMINI_PRO, COL_GEMINI_PRO_RESETS,
+    COL_GEMINI_FLASH, COL_GEMINI_FLASH_RESET, COL_GEMINI_FLASH_LITE, COL_GEMINI_FLASHLITE_RESET,
+    COL_ANTIGRAVITY_CLAUDE, COL_ANTIGRAVITY_CLAUDE_RESET, COL_ANTIGRAVITY_GEMINI_PRO,
+    COL_ANTIGRAVITY_GEMINI_PRO_RESET, COL_ANTIGRAVITY_FLASHLITE, COL_ANTIGRAVITY_FLASHLITE_RESETS,
+    COL_ANTIGRAVITY_PLAN,
+]
+
+PROVIDER_MAP = {
+    "codex": {
+        "primary": (COL_CODEX_5HR_USAGE, COL_CODEX_5HR_RESET),
+        "secondary": (COL_CODEX_WEEKLY_USAGE, COL_CODEX_WEEKLY_RESET)
+    },
+    "copilot": {
+        "primary": (COL_COPILOT_PREMIUM, None),
+        "secondary": (COL_COPILOT_CHAT, None)
+    },
+    "gemini": {
+        "primary": (COL_GEMINI_PRO, COL_GEMINI_PRO_RESETS),
+        "secondary": (COL_GEMINI_FLASH_LITE, COL_GEMINI_FLASHLITE_RESET),
+        "tertiary": (COL_GEMINI_FLASH, COL_GEMINI_FLASH_RESET)
+    },
+    "antigravity": {
+        "primary": (COL_ANTIGRAVITY_CLAUDE, COL_ANTIGRAVITY_CLAUDE_RESET),
+        "secondary": (COL_ANTIGRAVITY_GEMINI_PRO, COL_ANTIGRAVITY_GEMINI_PRO_RESET),
+        "tertiary": (COL_ANTIGRAVITY_FLASHLITE, COL_ANTIGRAVITY_FLASHLITE_RESETS)
+    },
+}
+
+class CodexBarProvider(BaseProvider):
+    def __init__(self, command: str = "codexbar usage --json"):
+        self.command = command
+
+    def fetch(self, now: str) -> pd.DataFrame:
+        try:
+            proc = subprocess.run(shlex.split(self.command), capture_output=True, text=True, check=True)
+            out_text = proc.stdout
+        except Exception:
+            proc = subprocess.run(self.command, capture_output=True, text=True, shell=True)
+            if proc.returncode != 0:
+                raise RuntimeError(f"Command failed: {proc.stderr or proc.stdout}")
+            out_text = proc.stdout
+
+        m = re.search(r"(\[|\{)", out_text)
+        if not m:
+            raise ValueError("No JSON found in command output")
+        data = json.loads(out_text[m.start():].strip())
+        items = data if isinstance(data, list) else [data]
+
+        out = {c: None for c in CODEXBAR_COLS}
+        out[COL_TIME] = now
+
+        for item in items:
+            prov = (item.get("provider") or item.get("source") or "").lower()
+            usage = item.get("usage") or {}
+
+            plan = usage.get("loginMethod") or (usage.get("identity") or {}).get("loginMethod")
+            if ("antigravity" in prov or "local" in prov) and plan and out.get(COL_ANTIGRAVITY_PLAN) is None:
+                out[COL_ANTIGRAVITY_PLAN] = plan
+
+            if "codex" in prov:
+                out[COL_CODEX_CREDITS] = (item.get("credits") or {}).get("remaining")
+                if plan and out.get(COL_CODEX_PLAN) is None:
+                    out[COL_CODEX_PLAN] = plan
+
+            for key, mapping in PROVIDER_MAP.items():
+                if key in prov:
+                    for slot, (val_col, reset_col) in mapping.items():
+                        slot_data = usage.get(slot) or {}
+                        val = slot_data.get("usedPercent")
+                        rst = slot_data.get("resetDescription") or slot_data.get("resetsAt")
+                        
+                        if val_col and out.get(val_col) is None:
+                            out[val_col] = val
+                        if reset_col and out.get(reset_col) is None:
+                            out[reset_col] = rst
+
+        out[COL_CODEX_USAGE] = out.get(COL_CODEX_5HR_USAGE) if out.get(COL_CODEX_5HR_USAGE) is not None else out.get(COL_CODEX_WEEKLY_USAGE)
+        out[COL_CODEX_RESETS] = out.get(COL_CODEX_5HR_RESET) if out.get(COL_CODEX_5HR_RESET) is not None else out.get(COL_CODEX_WEEKLY_RESET)
+
+        df = pd.DataFrame([[out[c] for c in CODEXBAR_COLS]], columns=CODEXBAR_COLS)
+        df.attrs["raw_json"] = data
+        return df

token_tracker/providers/openrouter.py

@@ -0,0 +1,34 @@
+import requests
+import pandas as pd
+from typing import Optional
+from .base import BaseProvider
+from project_toolkit import settings
+
+OPENROUTER_KEYS = [
+    "limit", "limit_reset", "limit_remaining", "include_byok_in_limit",
+    "usage", "usage_daily", "usage_weekly", "usage_monthly",
+    "byok_usage", "byok_usage_daily", "byok_usage_weekly", "byok_usage_monthly",
+    "is_free_tier", "expires_at"
+]
+
+class OpenRouterProvider(BaseProvider):
+    def __init__(self, api_key: Optional[str] = None):
+        self.api_key = api_key or settings.env("OPENROUTER_API_KEY")
+
+    def fetch(self, now: str) -> pd.DataFrame:
+        if not self.api_key:
+            raise ValueError("OPENROUTER_API_KEY not found")
+
+        response = requests.get(
+            "https://openrouter.ai/api/v1/key",
+            headers={"Authorization": f"Bearer {self.api_key}"},
+        )
+        response.raise_for_status()
+        data = response.json().get("data", {})
+
+        row = {"time": now}
+        row.update({k: data.get(k) for k in OPENROUTER_KEYS})
+
+        df = pd.DataFrame([row])
+        df.attrs["raw_json"] = data
+        return df

token_tracker/stores/__init__.py

@@ -0,0 +1 @@
+"""Stores for Token Tracker."""

token_tracker/stores/base.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+import pandas as pd
+
+class BaseStore(ABC):
+    @abstractmethod
+    def save(self, df: pd.DataFrame, name: str, init: bool = False) -> None:
+        """Save the DataFrame to the store."""
+        pass

token_tracker/stores/google_sheets.py

@@ -0,0 +1,98 @@
+import logging
+import json
+import hashlib
+from pathlib import Path
+import pandas as pd
+from .base import BaseStore
+from project_toolkit import settings
+
+class GoogleSheetsStore(BaseStore):
+    def __init__(self, sheet_id: str = None):
+        self.sheet_id = sheet_id or settings.env("TOKEN_USAGE_SHEET_ID")
+        if not self.sheet_id:
+            raise RuntimeError("Environment variable TOKEN_USAGE_SHEET_ID is not set")
+        
+        from project_toolkit.google.sheet import GoogleSheetsManager
+        self.manager = GoogleSheetsManager()
+
+    def _last_fetch_file(self, name: str) -> Path:
+        return Path.cwd() / f"{name}_last_fetch.json"
+
+    def _load_hash_data(self, name: str) -> dict:
+        p = self._last_fetch_file(name)
+        if not p.exists():
+            return {}
+        try:
+            with p.open("r", encoding="utf-8") as fh:
+                return json.load(fh)
+        except Exception:
+            logging.exception("Failed to load %s", p)
+            return {}
+
+    def _write_hash_data(self, name: str, data: dict) -> None:
+        p = self._last_fetch_file(name)
+        tmp = p.with_suffix(".tmp")
+        with tmp.open("w", encoding="utf-8") as fh:
+            json.dump(data, fh, indent=2, sort_keys=True, ensure_ascii=False, default=str)
+        tmp.replace(p)
+
+    def _snapshot_hash(self, df: pd.DataFrame) -> str:
+        df2 = df.copy()
+        if "time" in getattr(df2, "columns", []):
+            df2 = df2.drop(columns=["time"], errors=True)
+
+        records = df2.where(pd.notnull(df2), None).to_dict(orient="records")
+        s = json.dumps(records, sort_keys=True, default=str, separators=(',', ':'))
+        return hashlib.sha256(s.encode("utf-8")).hexdigest()
+
+    def save(self, df: pd.DataFrame, name: str, init: bool = False) -> None:
+        if df is None or getattr(df, "empty", False):
+            logging.warning("Provided DataFrame is empty; not writing to sheet %s", name)
+            return
+
+        raw = getattr(df, "attrs", {}).get("raw_json")
+        if raw is None:
+            try:
+                raw = df.where(pd.notnull(df), None).to_dict(orient="records")
+            except Exception:
+                raw = df.to_dict(orient="records") if hasattr(df, "to_dict") else (df if isinstance(df, (list, dict)) else None)
+
+        try:
+            curr_hash = self._snapshot_hash(df)
+        except Exception:
+            logging.exception("Failed to compute snapshot hash for %s", name)
+            curr_hash = None
+            if raw is not None:
+                try:
+                    if isinstance(raw, list) and all(isinstance(r, dict) for r in raw):
+                        raw_no_time = [{k: v for k, v in r.items() if k != "time"} for r in raw]
+                    elif isinstance(raw, dict):
+                        raw_no_time = {k: v for k, v in raw.items() if k != "time"}
+                    else:
+                        raw_no_time = raw
+                    s = json.dumps(raw_no_time, sort_keys=True, default=str, separators=(',', ':'))
+                    curr_hash = hashlib.sha256(s.encode('utf-8')).hexdigest()
+                except Exception:
+                    logging.exception("Failed to compute fallback hash from raw for %s", name)
+
+        stored_data = self._load_hash_data(name)
+        stored_hash = stored_data.get("hash")
+
+        if not init and curr_hash is not None and stored_hash == curr_hash:
+            logging.info("No change detected for %s — skipping write", name)
+        else:
+            try:
+                if init:
+                    logging.info("Initializing sheet '%s' (sheet_id=%s)", name, self.sheet_id)
+                    self.manager.df_to_sheet(df, sheet_id=self.sheet_id, sheet_name=name)
+                else:
+                    logging.info("Appending to sheet '%s' (sheet_id=%s)", name, self.sheet_id)
+                    self.manager.append_to_sheet(df, sheet_id=self.sheet_id, sheet_name=name)
+            except Exception:
+                logging.exception("Failed to write DataFrame to Google Sheets")
+                raise
+
+        try:
+            self._write_hash_data(name, {"hash": curr_hash, "raw": raw})
+        except Exception:
+            logging.exception("Failed to update %s_last_fetch.json", name)