用 SQLite 为 OpenRouter 批处理任务设置本地软预算

控制台可以给 Key 设额度；批处理任务还需要进程内的“软预算”，避免循环调用烧完余额。 这篇做 budget.py：用 SQLite 记 usage，按你传入的 $/M 单价估算费用，超预算拒绝继续调用。单价不写死在代码里。

先看清成品

• dry-usage：100+100 tokens @ $10/M → 花费 $0.002；下一笔 1M tokens 预估 $10 被拒绝（退出码 3）。
• 3 项 pytest 通过。未对 OpenRouter 发起计费请求。

合规提醒：使用自己的 OpenRouter 账号与 API Key，不共享；按各模型页面的许可与数据政策选用模型。价格以 openrouter.ai/models 为准，本文不写、不编造价格；本地预算脚本的单价需你自行传入。

运行说明：没有可用的 OpenRouter 付费 Key。真实运行过：客户端对本地 mock 的 fallback / provider 路由、预算脚本的 dry-usage，以及对 https://openrouter.ai 使用无效密钥得到 HTTP 401（2026-10-23，UTC+8）。未产生计费调用。

做完之后你会得到什么？

需要准备哪些环境？

关键事实

文件结构

93-openrouter-C-src/
├── budget.py
└── test_93c.py

第 1 步：预算脚本

#!/usr/bin/env python3
"""Track OpenRouter-style usage and refuse calls when a local budget is exceeded.

OpenRouter keys can have credit limits in the dashboard; this script adds a
local soft budget so a batch job stops before surprising cost. Prices are NOT
hard-coded — pass --prompt-price / --completion-price yourself from the model
page, or leave usage-only mode.
"""
from __future__ import annotations

import argparse
import json
import os
import sqlite3
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path


SCHEMA = """
CREATE TABLE IF NOT EXISTS usage (
  id INTEGER PRIMARY KEY,
  ts REAL,
  model TEXT,
  prompt_tokens INTEGER,
  completion_tokens INTEGER,
  total_tokens INTEGER,
  est_cost_usd REAL
);
"""


def db(path: Path) -> sqlite3.Connection:
    conn = sqlite3.connect(path)
    conn.executescript(SCHEMA)
    return conn


def spent(conn: sqlite3.Connection) -> float:
    row = conn.execute("SELECT COALESCE(SUM(est_cost_usd),0) FROM usage").fetchone()
    return float(row[0])


def estimate(prompt_t: int, completion_t: int, prompt_price: float, completion_price: float) -> float:
    # prices are USD per 1M tokens (OpenRouter model page convention)
    return prompt_t * prompt_price / 1_000_000 + completion_t * completion_price / 1_000_000


def chat_once(*, base_url, api_key, model, message) -> dict:
    url = f"{base_url.rstrip('/')}/api/v1/chat/completions"
    payload = {"model": model, "messages": [{"role": "user", "content": message}]}
    req = urllib.request.Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read().decode())


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("message", nargs="?", default="ping")
    p.add_argument("--base", default=os.environ.get("OPENROUTER_BASE_URL", "https://openrouter.ai"))
    p.add_argument("--key", default=os.environ.get("OPENROUTER_API_KEY", ""))
    p.add_argument("--model", default="meta-llama/llama-3.1-8b-instruct")
    p.add_argument("--db", type=Path, default=Path("usage.db"))
    p.add_argument("--budget-usd", type=float, default=1.0)
    p.add_argument("--prompt-price", type=float, default=0.0, help="USD per 1M prompt tokens from model page")
    p.add_argument("--completion-price", type=float, default=0.0, help="USD per 1M completion tokens")
    p.add_argument("--dry-usage", type=json.loads, default=None, help='JSON usage override for tests')
    args = p.parse_args(argv)

    conn = db(args.db)
    already = spent(conn)
    if already >= args.budget_usd:
        print(f"✗ 本地预算已用尽: ${already:.6f} / ${args.budget_usd:.2f}")
        return 3

    if args.dry_usage is not None:
        usage = args.dry_usage
        model = args.model
        content = "(dry)"
    else:
        if not args.key:
            print("设置 OPENROUTER_API_KEY", file=sys.stderr)
            return 2
        try:
            r = chat_once(base_url=args.base, api_key=args.key, model=args.model, message=args.message)
        except urllib.error.HTTPError as e:
            print(f"HTTP {e.code}: {e.read().decode(errors='replace')}", file=sys.stderr)
            return 1
        usage = r.get("usage") or {}
        model = r.get("model", args.model)
        content = r["choices"][0]["message"]["content"]

    pt = int(usage.get("prompt_tokens") or 0)
    ct = int(usage.get("completion_tokens") or 0)
    tt = int(usage.get("total_tokens") or (pt + ct))
    cost = estimate(pt, ct, args.prompt_price, args.completion_price)
    if already + cost > args.budget_usd and cost > 0:
        print(f"✗ 本次预估 ${cost:.6f} 会超过预算（已用 ${already:.6f}）")
        return 3
    conn.execute(
        "INSERT INTO usage(ts,model,prompt_tokens,completion_tokens,total_tokens,est_cost_usd) VALUES(?,?,?,?,?,?)",
        (time.time(), model, pt, ct, tt, cost),
    )
    conn.commit()
    print(f"model={model} tokens={tt} est_cost=${cost:.6f} spent=${spent(conn):.6f}/{args.budget_usd:.2f}")
    if args.dry_usage is None:
        print(content)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

第 2 步：测试

import json
from pathlib import Path
from budget import main, db, spent, estimate

def test_estimate():
    # 1M prompt tokens at $1 -> $1
    assert abs(estimate(1_000_000, 0, 1.0, 2.0) - 1.0) < 1e-9
    assert abs(estimate(0, 500_000, 1.0, 2.0) - 1.0) < 1e-9

def test_budget_blocks(tmp_path: Path):
    dbpath = tmp_path / "u.db"
    # first call within budget
    assert main(["--db", str(dbpath), "--budget-usd", "0.01", "--prompt-price", "10",
                 "--completion-price", "10", "--dry-usage",
                 json.dumps({"prompt_tokens": 100, "completion_tokens": 100, "total_tokens": 200})]) == 0
    # second would exceed
    rc = main(["--db", str(dbpath), "--budget-usd", "0.01", "--prompt-price", "10",
               "--completion-price", "10", "--dry-usage",
               json.dumps({"prompt_tokens": 1_000_000, "completion_tokens": 0, "total_tokens": 1_000_000})])
    assert rc == 3

def test_spent(tmp_path: Path):
    conn = db(tmp_path / "u.db")
    assert spent(conn) == 0.0

第 3 步：运行

# 2026-10-23 12:37 UTC+8
$ python budget.py --db /tmp/u.db --budget-usd 0.01 --prompt-price 10 --completion-price 10 \
  --dry-usage '{"prompt_tokens":100,"completion_tokens":100,"total_tokens":200}'
model=meta-llama/llama-3.1-8b-instruct tokens=200 est_cost=$0.002000 spent=$0.002000/0.01
$ python budget.py ... --dry-usage '{"prompt_tokens":1000000,...}'
✗ 本次预估 $10.000000 会超过预算（已用 $0.002000）
$ python -m pytest -q
3 passed

怎么检查结果？

• 超预算必须拒绝（非 0 退出码）。
• 估算用你提供的单价；换模型就改参数。
• 控制台 Key 额度与本地软预算可同时用。

调试示例

不要把价目表写进仓库。 模型价格会变；脚本只接受命令行/配置传入的数字。测试里用 10 只是方便算整数，不代表任何真实模型价格。

常见问题

模板规则

• 单价外部传入，不写死。
• 超预算直接停。
• 与控制台额度双保险。

项目资源

软预算让批任务在“还差一点”时停住，而不是刷爆账单。