nanobot-runtime/skills/usage/scripts/ollama_usage.py

#!/usr/bin/env python3
"""Ollama Cloud usage — GET https://ollama.com/api/usage with Bearer key.

Also the shared base for the poller and the report: DB location, the sample
row shape, and the session-window arithmetic described below.
"""

from __future__ import annotations

import itertools
import json
import os
import sqlite3
import sys
import urllib.error
import urllib.request
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import NamedTuple

API_URL = "https://ollama.com/api/usage"

WORKSPACE = Path(__file__).resolve().parents[3]
DB_PATH = WORKSPACE / "db" / "ollama_usage.sqlite"

SAMPLE_COLUMNS = "ts, session_usage, weekly_usage, session_models, weekly_models"

# Legacy plans bill in 5-hour session windows; see ollama.com/blog/transparent-pricing
# ("no 5-hour or weekly limits" is what the *new* plans dropped).
SESSION_BLOCK = timedelta(hours=5)


class Sample(NamedTuple):
    ts: str
    session_usage: float
    weekly_usage: float
    session_models: str
    weekly_models: str


def load_api_key() -> str:
    env = os.environ.get("OLLAMA_API_KEY")
    if env:
        return env.strip()
    env_file = WORKSPACE / ".env"
    if env_file.is_file():
        for line in env_file.read_text().splitlines():
            line = line.strip()
            if line.startswith("OLLAMA_API_KEY="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    sys.exit("OLLAMA_API_KEY not found (neither env nor workspace/.env)")


def until_next_monday(now: datetime) -> timedelta:
    # days until Monday (weekday(): Mon == 0)
    days = (7 - now.weekday()) % 7 or 7
    nxt = now.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=days)
    return nxt - now


def fmt_delta_short(td: timedelta) -> str:
    """Human countdown, dashboard style: '1 hour', '55 minutes', '6 days'."""
    total = td.total_seconds()
    if total >= 86400:
        return f"{int(total // 86400)} days"
    if total >= 3600:
        return f"{round(total / 3600)} hours"
    return f"{max(1, round(total / 60))} minutes"


def fmt_countdown(td: timedelta) -> str:
    """Exact countdown: '3 h 55 min'. The dashboard rounds; we do not."""
    hours, minutes = divmod(int(td.total_seconds() // 60), 60)
    return f"{hours} h {minutes} min" if hours else f"{minutes} min"


def is_window_reset(previous: Sample, current: Sample) -> bool:
    """Session window rolled over: its usage or its request total went down."""
    previous_total = sum(json.loads(previous.session_models).values())
    current_total = sum(json.loads(current.session_models).values())
    return current.session_usage < previous.session_usage or current_total < previous_total


def load_samples(conn: sqlite3.Connection, since: str | None = None, until: str | None = None) -> list[Sample]:
    if since is None and until is None:
        rows = conn.execute(f"SELECT {SAMPLE_COLUMNS} FROM samples ORDER BY ts")
    else:
        rows = conn.execute(
            f"SELECT {SAMPLE_COLUMNS} FROM samples WHERE ts >= ? AND ts <= ? ORDER BY ts",
            (since, until),
        )
    return [Sample(*row) for row in rows]


def window_rollovers(samples: list[Sample]) -> list[datetime]:
    """Every sample where the session window rolled over, oldest first.

    The window is not on a fixed grid: it starts with the first request after
    the previous one ran out. On 2026-09-15 usage sat unchanged through 05:00
    UTC and only reset once a request arrived at 06:00 — so a rollover sample
    marks the *start of a new window*, not a boundary that was due anyway.
    """
    return [
        datetime.fromisoformat(current.ts)
        for previous, current in itertools.pairwise(samples)
        if is_window_reset(previous, current)
    ]


def session_window_end() -> datetime | None:
    """End of the window in progress, or None when the history cannot show it."""
    if not DB_PATH.is_file():
        return None
    conn = sqlite3.connect(DB_PATH)
    try:
        rollovers = window_rollovers(load_samples(conn))
    finally:
        conn.close()
    return rollovers[-1] + SESSION_BLOCK if rollovers else None


def fetch_usage() -> dict:
    request = urllib.request.Request(
        API_URL,
        headers={
            "Authorization": f"Bearer {load_api_key()}",
            "Accept": "application/json",
        },
    )
    with urllib.request.urlopen(request, timeout=15) as response:
        return json.load(response)


def print_session(usage: float, now: datetime) -> None:
    percent = f"Session: {usage * 100:.1f} %"
    end = session_window_end()

    if end is None:
        # No rollover recorded yet — the window's start is simply unknown.
        print(percent)
    elif end <= now:
        # The window ran out; the next one only begins with the next request,
        # so there is no time to count down to.
        print(f"{percent}, window expired — the next one starts with the next request")
    else:
        local = end.astimezone()
        print(f"{percent}, resets {local:%H:%M %Z} (in {fmt_countdown(end - now)})")


def main() -> None:
    try:
        data = fetch_usage()
    except urllib.error.HTTPError as e:
        sys.exit(f"ollama.com API error: HTTP {e.code}")
    except urllib.error.URLError as e:
        sys.exit(f"ollama.com unreachable: {e.reason}")

    limits = data.get("limits", {})
    session = limits.get("session", {})
    weekly = limits.get("weekly", {})
    now = datetime.now(UTC)

    print("Ollama Cloud usage")
    print_session(session.get("usage", 0.0), now)
    print(f"Weekly:  {weekly.get('usage', 0) * 100:.1f} %, resets in {fmt_delta_short(until_next_monday(now))}")

    models = weekly.get("models", [])
    if models:
        print("Models (request count, weekly window):")
        for m in models:
            print(f"  {m['name']}: {m['request_count']}")


if __name__ == "__main__":
    main()