98 lines
2.7 KiB
Python
98 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Sample Ollama Cloud usage into db/ollama_usage.sqlite. Run from cron every minute.
|
|
|
|
Writes a `samples` row only when something changed; `meta` records every poll so a
|
|
gap in `samples` can be told apart from a poll that failed or never ran.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
import sys
|
|
import urllib.error
|
|
from datetime import UTC, datetime
|
|
|
|
from ollama_usage import DB_PATH, fetch_usage
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS samples (
|
|
ts TEXT PRIMARY KEY,
|
|
session_usage REAL NOT NULL,
|
|
weekly_usage REAL NOT NULL,
|
|
session_models TEXT NOT NULL,
|
|
weekly_models TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS meta (
|
|
id INTEGER PRIMARY KEY CHECK (id = 1),
|
|
last_ts TEXT NOT NULL,
|
|
last_status TEXT NOT NULL
|
|
);
|
|
"""
|
|
|
|
|
|
def canonical_models(models: list[dict]) -> str:
|
|
"""Stable JSON for change detection — the API returns an unordered list."""
|
|
return json.dumps(
|
|
{m["name"]: m["request_count"] for m in models},
|
|
sort_keys=True,
|
|
)
|
|
|
|
|
|
def connect() -> sqlite3.Connection:
|
|
DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = sqlite3.connect(DB_PATH)
|
|
conn.executescript(SCHEMA)
|
|
return conn
|
|
|
|
|
|
def record_poll(conn: sqlite3.Connection, ts: str, status: str) -> None:
|
|
conn.execute(
|
|
"INSERT INTO meta (id, last_ts, last_status) VALUES (1, ?, ?) "
|
|
"ON CONFLICT(id) DO UPDATE SET last_ts = excluded.last_ts, "
|
|
"last_status = excluded.last_status",
|
|
(ts, status),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def latest_sample(conn: sqlite3.Connection) -> tuple | None:
|
|
row = conn.execute(
|
|
"SELECT session_usage, weekly_usage, session_models, weekly_models FROM samples ORDER BY ts DESC LIMIT 1"
|
|
).fetchone()
|
|
return row
|
|
|
|
|
|
def main() -> None:
|
|
now = datetime.now(UTC).replace(microsecond=0).isoformat()
|
|
conn = connect()
|
|
|
|
try:
|
|
data = fetch_usage()
|
|
except urllib.error.HTTPError as e:
|
|
record_poll(conn, now, f"http_{e.code}")
|
|
print(f"{now} ollama.com API error: HTTP {e.code}", file=sys.stderr)
|
|
return
|
|
except urllib.error.URLError as e:
|
|
record_poll(conn, now, "unreachable")
|
|
print(f"{now} ollama.com unreachable: {e.reason}", file=sys.stderr)
|
|
return
|
|
|
|
limits = data.get("limits", {})
|
|
session = limits.get("session", {})
|
|
weekly = limits.get("weekly", {})
|
|
sample = (
|
|
session.get("usage", 0.0),
|
|
weekly.get("usage", 0.0),
|
|
canonical_models(session.get("models", [])),
|
|
canonical_models(weekly.get("models", [])),
|
|
)
|
|
|
|
if sample != latest_sample(conn):
|
|
conn.execute("INSERT INTO samples VALUES (?, ?, ?, ?, ?)", (now, *sample))
|
|
record_poll(conn, now, "ok")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|