Project Spend Caps とアプリ側ソフト上限で、無人パイプラインの費用事故を二重に止める

# pricing.py — 単価は公式の料金ページの最新値で必ず上書きしてください
# 単位: USD / 100万トークン
PRICING = {
    "gemini-flash-latest": {"in": 0.10, "out": 0.40},   # 3.5 Flash（GA）
    "gemini-pro-latest":   {"in": 1.25, "out": 5.00},   # Pro 系
}
 
def estimate_cost(model: str, in_tokens: int, out_tokens: int) -> float:
    p = PRICING[model]
    return (in_tokens * p["in"] + out_tokens * p["out"]) / 1_000_000

# ledger.py
import sqlite3, datetime, os
 
class CostLedger:
    def __init__(self, db="cost.db", monthly_cap_usd=40.0, soft_ratio=0.8):
        self.db = db
        self.monthly_cap = monthly_cap_usd
        self.soft_cap = monthly_cap_usd * soft_ratio
        with sqlite3.connect(self.db) as c:
            c.execute("CREATE TABLE IF NOT EXISTS spend("
                      "ts TEXT, model TEXT, usd REAL, request_id TEXT UNIQUE)")
 
    def _month_key(self):
        # 集計は JST 基準。素の UTC だと日付の境界がずれて当日分を取りこぼします
        now = datetime.datetime.now(datetime.timezone(datetime.timedelta(hours=9)))
        return now.strftime("%Y-%m")
 
    def month_to_date(self) -> float:
        with sqlite3.connect(self.db) as c:
            cur = c.execute("SELECT COALESCE(SUM(usd),0) FROM spend "
                            "WHERE substr(ts,1,7)=?", (self._month_key(),))
            return float(cur.fetchone()[0])
 
    def record(self, model: str, usd: float, request_id: str):
        # request_id を UNIQUE にして、リトライ時の二重計上を防ぎます
        ts = datetime.datetime.now().isoformat()
        with sqlite3.connect(self.db) as c:
            try:
                c.execute("INSERT INTO spend VALUES(?,?,?,?)",
                          (ts, model, usd, request_id))
            except sqlite3.IntegrityError:
                pass  # 同じ request_id は記録済み。黙って無視します
 
    def status(self) -> str:
        mtd = self.month_to_date()
        if mtd >= self.monthly_cap:
            return "HARD"
        if mtd >= self.soft_cap:
            return "SOFT"
        return "OK"

# governor.py
from pricing import estimate_cost
from ledger import CostLedger
 
ledger = CostLedger(monthly_cap_usd=40.0, soft_ratio=0.8)
 
class BudgetExceeded(Exception):
    pass
 
def guarded_generate(client, model, contents, request_id, priority="normal"):
    state = ledger.status()
    if state == "HARD":
        raise BudgetExceeded("月のハード上限に到達。処理を停止します")
    if state == "SOFT" and priority != "high":
        # ソフト上限の内側では、優先度の低いジョブを軽いモデルへ落とします
        model = "gemini-flash-latest"
 
    resp = client.models.generate_content(model=model, contents=contents)
    usage = resp.usage_metadata
    usd = estimate_cost(model, usage.prompt_token_count,
                        usage.candidates_token_count)
    ledger.record(model, usd, request_id)
    return resp