| 1 | """Provider-reported charges for a scoped watchlist subprocess.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import math |
| 6 | import os |
| 7 | import sqlite3 |
| 8 | import uuid |
| 9 | from contextlib import contextmanager |
| 10 | from pathlib import Path |
| 11 | from urllib.parse import urlsplit |
| 12 | |
| 13 | JOURNAL_ENV = "LAST30DAYS_USAGE_JOURNAL" |
| 14 | |
| 15 | _PAID_HOSTS = { |
| 16 | "openrouter.ai": "openrouter", |
| 17 | "api.perplexity.ai": "perplexity", |
| 18 | "api.openai.com": "openai", |
| 19 | "api.x.ai": "xai", |
| 20 | "generativelanguage.googleapis.com": "gemini", |
| 21 | "api.scrapecreators.com": "scrapecreators", |
| 22 | "api.search.brave.com": "brave", |
| 23 | "api.tavily.com": "tavily", |
| 24 | "api.exa.ai": "exa", |
| 25 | "google.serper.dev": "serper", |
| 26 | "api.parallel.ai": "parallel", |
| 27 | "api.x.com": "x", |
| 28 | "api.twitter.com": "x", |
| 29 | "xquik.com": "xquik", |
| 30 | "api.groq.com": "groq", |
| 31 | } |
| 32 | |
| 33 | |
| 34 | @contextmanager |
| 35 | def _connection(path): |
| 36 | conn = sqlite3.connect(path) |
| 37 | try: |
| 38 | with conn: |
| 39 | yield conn |
| 40 | finally: |
| 41 | conn.close() |
| 42 | |
| 43 | |
| 44 | def create_journal(path: Path) -> None: |
| 45 | fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600) |
| 46 | os.close(fd) |
| 47 | with _connection(path) as conn: |
| 48 | conn.execute( |
| 49 | "CREATE TABLE attempts (id TEXT PRIMARY KEY, provider TEXT NOT NULL, " |
| 50 | "cost REAL, unknown INTEGER NOT NULL DEFAULT 1, " |
| 51 | "prompt_tokens INTEGER DEFAULT 0, completion_tokens INTEGER DEFAULT 0)" |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | def begin(provider: str) -> tuple[str, str, str] | None: |
| 56 | path = os.environ.get(JOURNAL_ENV) |
| 57 | if not path: |
| 58 | return None |
| 59 | attempt_id = uuid.uuid4().hex |
| 60 | with _connection(path) as conn: |
| 61 | conn.execute("INSERT INTO attempts (id, provider) VALUES (?, ?)", (attempt_id, provider)) |
| 62 | return path, attempt_id, provider |
| 63 | |
| 64 | |
| 65 | def begin_http(url: str, method: str) -> tuple[str, str, str] | None: |
| 66 | if not os.environ.get(JOURNAL_ENV): |
| 67 | return None |
| 68 | parts = urlsplit(url) |
| 69 | provider = _PAID_HOSTS.get(parts.hostname) |
| 70 | if provider == "perplexity" and method == "GET": |
| 71 | return None # Retrieving an existing background operation does not start another charge. |
| 72 | if provider is None: |
| 73 | for name in ("OPENAI", "XAI", "OPENROUTER"): |
| 74 | override = os.environ.get(f"{name}_BASE_URL") |
| 75 | if override and urlsplit(override).netloc == parts.netloc: |
| 76 | provider = name.lower() |
| 77 | break |
| 78 | return begin(provider) if provider else None |
| 79 | |
| 80 | |
| 81 | def cancel(attempt: tuple[str, str, str] | None) -> None: |
| 82 | """Remove an attempt when the transport provably never started.""" |
| 83 | if attempt is not None: |
| 84 | path, attempt_id, _ = attempt |
| 85 | with _connection(path) as conn: |
| 86 | conn.execute("DELETE FROM attempts WHERE id = ?", (attempt_id,)) |
| 87 | |
| 88 | |
| 89 | def finish(attempt: tuple[str, str, str] | None, payload: object) -> None: |
| 90 | if attempt is None or not isinstance(payload, dict): |
| 91 | return |
| 92 | path, attempt_id, provider = attempt |
| 93 | usage = payload.get("usage") |
| 94 | if not isinstance(usage, dict): |
| 95 | return |
| 96 | cost = usage.get("cost") |
| 97 | if provider == "perplexity": |
| 98 | if not isinstance(cost, dict) or cost.get("currency") != "USD": |
| 99 | return |
| 100 | cost = cost.get("total_cost") |
| 101 | elif provider != "openrouter": |
| 102 | return |
| 103 | if isinstance(cost, bool) or not isinstance(cost, (int, float)) or not math.isfinite(cost) or cost < 0: |
| 104 | return |
| 105 | # A background response may expose a running subtotal rather than its final bill. |
| 106 | unknown = int(payload.get("status") in {"queued", "in_progress"}) |
| 107 | prompt = usage.get("prompt_tokens", usage.get("input_tokens", 0)) |
| 108 | completion = usage.get("completion_tokens", usage.get("output_tokens", 0)) |
| 109 | prompt = prompt if type(prompt) is int and prompt >= 0 else 0 |
| 110 | completion = completion if type(completion) is int and completion >= 0 else 0 |
| 111 | with _connection(path) as conn: |
| 112 | conn.execute( |
| 113 | "UPDATE attempts SET cost = ?, unknown = ?, prompt_tokens = ?, completion_tokens = ? WHERE id = ?", |
| 114 | (cost, unknown, prompt, completion, attempt_id), |
| 115 | ) |
| 116 | |
| 117 | |
| 118 | def read_journal(path: Path) -> dict: |
| 119 | try: |
| 120 | with _connection(path) as conn: |
| 121 | row = conn.execute( |
| 122 | "SELECT COALESCE(SUM(cost), 0), COALESCE(SUM(unknown), 0), " |
| 123 | "COALESCE(SUM(prompt_tokens), 0), COALESCE(SUM(completion_tokens), 0) FROM attempts" |
| 124 | ).fetchone() |
| 125 | except sqlite3.Error: |
| 126 | return {"token_cost": 0.0, "cost_unknown": 1, "prompt_tokens": 0, "completion_tokens": 0} |
| 127 | return { |
| 128 | "token_cost": row[0], |
| 129 | "cost_unknown": int(row[1] > 0), |
| 130 | "prompt_tokens": row[2], |
| 131 | "completion_tokens": row[3], |
| 132 | } |
| 133 |