countTokens でバッチ投入前にコスト超過を止める — 既定モデルが変わっても崩れない予算ゲートの設計

from google import genai
 
client = genai.Client(api_key="YOUR_GEMINI_API_KEY")
MODEL = "gemini-2.5-flash"  # 実運用では設定から注入します
 
def count_input_tokens(contents: str) -> int:
    """課金されない count_tokens で入力トークン数を返します。"""
    resp = client.models.count_tokens(model=MODEL, contents=contents)
    return resp.total_tokens
 
def sum_batch_input_tokens(items: list[str]) -> int:
    """バッチ全件の入力トークンを実数で積み上げます。"""
    return sum(count_input_tokens(text) for text in items)

def measure_thinking_ratio(sample_prompt: str) -> float:
    """サンプル1件を実行し、出力に占める思考トークンの比率を計測します。
    （この呼び出しだけは課金されますが、係数校正のための投資です）"""
    resp = client.models.generate_content(model=MODEL, contents=sample_prompt)
    um = resp.usage_metadata
    visible = um.candidates_token_count or 0
    thinking = getattr(um, "thoughts_token_count", 0) or 0
    if visible == 0:
        return 1.0
    # 可視出力に対して、思考分をどれだけ上乗せすべきかの係数
    return (visible + thinking) / visible
 
# 例: 可視 120 トークン + 思考 52 トークン → 係数 ≈ 1.43
THINKING_FACTOR = measure_thinking_ratio("分類タスクのサンプルプロンプト")

from dataclasses import dataclass
 
# 試算用の単価（執筆時点の例。実運用では公式の最新単価を設定から注入してください）
PRICE_INPUT_PER_M = 0.30   # 入力 100万トークンあたりのドル単価（例）
PRICE_OUTPUT_PER_M = 2.50  # 出力 100万トークンあたりのドル単価（例）
 
@dataclass
class BudgetEstimate:
    input_tokens: int
    output_tokens: int          # 思考トークン込み
    projected_cost_usd: float
    within_budget: bool
 
def estimate_batch_cost(
    items: list[str],
    median_visible_output: int,
    thinking_factor: float,
    budget_usd: float,
) -> BudgetEstimate:
    input_tokens = sum_batch_input_tokens(items)
    # 可視出力に思考係数を掛けて、課金される出力トークンを見積もります
    output_per_item = round(median_visible_output * thinking_factor)
    output_tokens = output_per_item * len(items)
 
    cost = (
        input_tokens / 1_000_000 * PRICE_INPUT_PER_M
        + output_tokens / 1_000_000 * PRICE_OUTPUT_PER_M
    )
    return BudgetEstimate(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        projected_cost_usd=round(cost, 4),
        within_budget=cost <= budget_usd,
    )
 
def run_batch_with_gate(items, *, median_visible_output, thinking_factor, budget_usd):
    est = estimate_batch_cost(
        items, median_visible_output, thinking_factor, budget_usd
    )
    print(
        f"投入予定 {len(items)}件 / 入力 {est.input_tokens:,}tok / "
        f"出力(思考込み) {est.output_tokens:,}tok / 推定 ${est.projected_cost_usd}"
    )
    if not est.within_budget:
        raise RuntimeError(
            f"予算ゲートで停止しました: 推定 ${est.projected_cost_usd} > 予算 ${budget_usd}"
        )
    # ここで初めて、課金される generate_content をバッチ実行します
    return _execute_batch(items)

# Before: 平均入力 × 件数 × 単価 を一度だけ手計算
# 出力は「だいたい100トークンくらい」と決め打ち、思考トークンは未考慮
monthly_cost = avg_input * item_count * price  # ← モデルが変わると静かにずれる

# After: 投入直前に実数で測り、思考係数を掛け、予算で止める
est = estimate_batch_cost(items, median_visible_output=120,
                          thinking_factor=THINKING_FACTOR, budget_usd=3.0)
if not est.within_budget:
    skip_or_split(items)  # 分割投入や翌日送りに退避させます