並行して回す実験が共有予算を食い尽くす前に — AI Studio 費用上限をプロジェクト隔離の境界にする設計

import threading
from dataclasses import dataclass, field
 
# Gemini 3.5 Flash の概算単価（USD / 100万トークン）。実値は必ず請求ページで確認してください。
PRICE_PER_MTOK = {"input": 0.075, "output": 0.30}
 
def estimate_cost_usd(input_tokens: int, output_tokens: int) -> float:
    return (input_tokens * PRICE_PER_MTOK["input"]
            + output_tokens * PRICE_PER_MTOK["output"]) / 1_000_000
 
class BudgetExhausted(Exception):
    """ソフト天井に達したことを処理側へ伝える。ハード上限より先に投げる。"""
 
@dataclass
class ProjectBudget:
    project_id: str
    hard_cap_usd: float          # AI Studio 側のハード上限と必ず一致させる
    soft_ratio: float = 0.8      # ハード上限の80%をソフト天井にする
    spent_usd: float = 0.0
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
 
    @property
    def soft_cap_usd(self) -> float:
        return self.hard_cap_usd * self.soft_ratio
 
    def reserve(self, projected_usd: float) -> None:
        with self._lock:
            if self.spent_usd + projected_usd > self.soft_cap_usd:
                raise BudgetExhausted(
                    f"{self.project_id}: soft cap {self.soft_cap_usd:.2f} USD 到達"
                    f"（消費 {self.spent_usd:.2f} + 見込み {projected_usd:.4f}）")
 
    def commit(self, actual_usd: float) -> None:
        with self._lock:
            self.spent_usd += actual_usd

def guarded_generate(client, budget: ProjectBudget, *, model, contents, max_output_tokens):
    # 最悪ケース（出力上限まで使う想定）で予約してから呼ぶ
    input_tokens = client.models.count_tokens(model=model, contents=contents).total_tokens
    projected = estimate_cost_usd(input_tokens, max_output_tokens)
    budget.reserve(projected)  # 天井を超えるなら BudgetExhausted で即座に止まる
 
    resp = client.models.generate_content(
        model=model, contents=contents,
        config={"max_output_tokens": max_output_tokens})
 
    usage = resp.usage_metadata
    actual = estimate_cost_usd(usage.prompt_token_count, usage.candidates_token_count)
    budget.commit(actual)  # 実消費で台帳を更新する
    return resp

# Before: 上限に当たっても「待てば直る」と信じて叩き続ける
for attempt in range(7):
    try:
        return client.models.generate_content(model=model, contents=contents)
    except ResourceExhausted:
        time.sleep(2 ** attempt)  # 予算枯渇のときは何回待っても通らない

try:
    resp = guarded_generate(client, budget, model="gemini-flash-latest",
                            contents=chunk, max_output_tokens=1024)
except BudgetExhausted:
    queue_for_next_month(chunk)  # 止めずに退避し、本番機能の予算は温存する