Gemini 2.5 Pro の thinking_budget を制御する — コストを3分の1にしながら推論品質を守る実装パターン

import google.generativeai as genai
import re
from dataclasses import dataclass
from typing import Literal
 
BudgetLevel = Literal["off", "light", "medium", "full"]
 
@dataclass
class TaskConfig:
    level: BudgetLevel
    thinking_budget: int
 
def classify_task(prompt: str) -> TaskConfig:
    """
    プロンプトの内容と長さから thinking budget レベルを自動判定する
    ヒューリスティック分類器。
    """
    word_count = len(prompt.split())
 
    # Level 3: フル思考が必要なキーワード
    full_thinking_keywords = [
        "設計", "アーキテクチャ", "最適化", "比較して分析", "矛盾", "トレードオフ",
        "design", "architecture", "optimize", "analyze", "evaluate", "complex"
    ]
 
    # Level 2: 中量思考が必要なキーワード
    medium_thinking_keywords = [
        "バグ", "エラー", "なぜ", "理由", "計画", "デバッグ",
        "bug", "error", "why", "reason", "plan", "strategy"
    ]
 
    # Level 0: 思考不要な単純タスクのパターン
    simple_task_patterns = [
        r"^翻訳[：:]",
        r"以下のJSONを整形",
        r"^要約[:：].{0,20}$",
        r"^\d+\s*[\+\-\*\/]\s*\d+",
    ]
 
    for pattern in simple_task_patterns:
        if re.search(pattern, prompt):
            return TaskConfig(level="off", thinking_budget=0)
 
    if any(kw in prompt for kw in full_thinking_keywords) or word_count > 300:
        return TaskConfig(level="full", thinking_budget=20480)
 
    if any(kw in prompt for kw in medium_thinking_keywords) or word_count > 100:
        return TaskConfig(level="medium", thinking_budget=8192)
 
    return TaskConfig(level="light", thinking_budget=2048)
 
def generate_with_adaptive_budget(prompt: str, verbose: bool = False) -> str:
    """タスク複雑度に応じた thinking budget で生成する。"""
    genai.configure(api_key="YOUR_GEMINI_API_KEY")
 
    task_config = classify_task(prompt)
 
    if verbose:
        print(f"Task level: {task_config.level} (budget: {task_config.thinking_budget})")
 
    model = genai.GenerativeModel("gemini-2.5-pro")
 
    if task_config.thinking_budget > 0:
        gen_config = genai.GenerationConfig(
            thinking_config=genai.types.ThinkingConfig(
                thinking_budget=task_config.thinking_budget
            )
        )
    else:
        gen_config = genai.GenerationConfig()
 
    response = model.generate_content(prompt, generation_config=gen_config)
 
    if verbose and hasattr(response, "usage_metadata"):
        meta = response.usage_metadata
        thinking = getattr(meta, "thoughts_token_count", 0) or 0
        print(f"Tokens: {meta.prompt_token_count}in / "
              f"{meta.candidates_token_count}out / {thinking}think")
 
    return response.text

test_cases = [
    ("翻訳：Hello, world!", "off"),
    ("PythonのGILがなぜ真の並列性を妨げるのか教えてください", "medium"),
    ("高書き込みのSaaSでイベントソーシングとCQRSのアーキテクチャのトレードオフを比較して分析してください", "full"),
]
 
for prompt, expected_level in test_cases:
    config = classify_task(prompt)
    match = "OK" if config.level == expected_level else "MISMATCH"
    print(f"[{match}] Expected={expected_level}, Got={config.level} | {prompt[:30]}")

import google.generativeai as genai
 
genai.configure(api_key="YOUR_GEMINI_API_KEY")
 
def stream_with_budget(prompt: str, thinking_budget: int = 4096) -> str:
    """
    thinking_budget を指定しながらストリーミング生成する。
    thinking tokens はストリームに含まれない。
    使用量は resolve() 後に取得する。
    """
    model = genai.GenerativeModel("gemini-2.5-pro")
    gen_config = genai.GenerationConfig(
        thinking_config=genai.types.ThinkingConfig(
            thinking_budget=thinking_budget
        )
    )
 
    stream = model.generate_content(prompt, generation_config=gen_config, stream=True)
 
    collected = []
    for chunk in stream:
        if chunk.text:
            print(chunk.text, end="", flush=True)
            collected.append(chunk.text)
    print()
 
    try:
        stream.resolve()
        thinking = getattr(stream.usage_metadata, "thoughts_token_count", 0) or 0
        print(f"Thinking tokens used: {thinking:,}")
    except Exception:
        pass
 
    return "".join(collected)
 
stream_with_budget(
    "セッションストレージにおけるRedisとMemcachedのトレードオフを教えてください",
    thinking_budget=8192
)

import time
from collections import defaultdict
from typing import Optional
import google.generativeai as genai
 
class ThinkingTokenMonitor:
    """
    thinking token の消費を追跡し、異常を検知するモニター。
    個人開発でも回せるよう軽量に保つ。
    """
 
    def __init__(
        self,
        hourly_limit: int = 500_000,
        single_request_limit: int = 20_000,
        alert_callback=None
    ):
        self.hourly_limit = hourly_limit
        self.single_request_limit = single_request_limit
        self.alert_callback = alert_callback or self._log_alert
        self.hourly_counts = defaultdict(int)
        self.request_log = []
 
    def _log_alert(self, message: str):
        # 実際は Slack Webhook やメール等に置き換える
        print(f"ALERT [{time.strftime('%Y-%m-%d %H:%M:%S')}]: {message}")
 
    def record(
        self,
        thinking_tokens: int,
        output_tokens: int,
        budget_level: str,
        prompt_snippet: Optional[str] = None
    ):
        hour_key = int(time.time() // 3600)
        self.hourly_counts[hour_key] += thinking_tokens
        self.request_log.append({
            "ts": time.time(),
            "thinking": thinking_tokens,
            "output": output_tokens,
            "level": budget_level
        })
 
        if self.hourly_counts[hour_key] > self.hourly_limit:
            self.alert_callback(
                f"Hourly thinking token limit exceeded. "
                f"Used: {self.hourly_counts[hour_key]:,} / Limit: {self.hourly_limit:,}. "
                f"Last prompt: {prompt_snippet or 'N/A'}"
            )
 
        if thinking_tokens > self.single_request_limit:
            self.alert_callback(
                f"Single request used {thinking_tokens:,} thinking tokens. "
                f"Prompt: {prompt_snippet or 'N/A'}"
            )
 
    def summary(self) -> dict:
        if not self.request_log:
            return {"status": "no data"}
        total_thinking = sum(r["thinking"] for r in self.request_log)
        total_output = sum(r["output"] for r in self.request_log)
        n = len(self.request_log)
        level_dist = defaultdict(int)
        for r in self.request_log:
            level_dist[r["level"]] += 1
        return {
            "total_requests": n,
            "avg_thinking_tokens": round(total_thinking / n, 1),
            "avg_output_tokens": round(total_output / n, 1),
            "thinking_to_output_ratio": round(total_thinking / max(total_output, 1), 2),
            "level_distribution": dict(level_dist)
        }
 
# 分類器と組み合わせる
monitor = ThinkingTokenMonitor(hourly_limit=200_000)
 
def monitored_generate(prompt: str) -> str:
    task_config = classify_task(prompt)
    genai.configure(api_key="YOUR_GEMINI_API_KEY")
    model = genai.GenerativeModel("gemini-2.5-pro")
 
    gen_config_kwargs = {}
    if task_config.thinking_budget > 0:
        gen_config_kwargs["thinking_config"] = genai.types.ThinkingConfig(
            thinking_budget=task_config.thinking_budget
        )
 
    response = model.generate_content(
        prompt,
        generation_config=genai.GenerationConfig(**gen_config_kwargs) if gen_config_kwargs else None
    )
 
    if hasattr(response, "usage_metadata"):
        meta = response.usage_metadata
        thinking_tokens = getattr(meta, "thoughts_token_count", 0) or 0
        monitor.record(
            thinking_tokens=thinking_tokens,
            output_tokens=meta.candidates_token_count or 0,
            budget_level=task_config.level,
            prompt_snippet=prompt[:80]
        )
 
    return response.text

import json
import time
 
def log_request(prompt: str, task_config, usage_metadata):
    """分類器チューニング用にリクエスト詳細をログ出力する。"""
    thinking_tokens = getattr(usage_metadata, "thoughts_token_count", 0) or 0
    output_tokens = getattr(usage_metadata, "candidates_token_count", 0) or 0
    budget_utilization = thinking_tokens / max(task_config.thinking_budget, 1)
 
    log_entry = {
        "ts": time.time(),
        "classified_level": task_config.level,
        "thinking_budget": task_config.thinking_budget,
        "thinking_tokens_used": thinking_tokens,
        "budget_utilization": round(budget_utilization, 3),
        "prompt_snippet": prompt[:100]
    }
 
    log_path = f"/tmp/budget_log_{time.strftime('%Y-%m-%d')}.jsonl"
    with open(log_path, "a") as f:
        f.write(json.dumps(log_entry, ensure_ascii=False) + "\n")

meta = response.usage_metadata
thinking_tokens = (
    getattr(meta, "thoughts_token_count", None) or
    getattr(meta, "thinking_token_count", None) or  # 古いSDKの命名
    0
)

import json
import time
import glob
from collections import defaultdict
 
def generate_weekly_report(log_dir: str = "/tmp") -> dict:
    """過去7日間の thinking token 使用量を集計する。"""
    seven_days_ago = time.time() - (7 * 24 * 3600)
    all_entries = []
    for log_file in glob.glob(f"{log_dir}/budget_log_*.jsonl"):
        with open(log_file) as f:
            for line in f:
                try:
                    entry = json.loads(line.strip())
                    if entry.get("ts", 0) > seven_days_ago:
                        all_entries.append(entry)
                except json.JSONDecodeError:
                    continue
 
    if not all_entries:
        return {"status": "no data in past 7 days"}
 
    n = len(all_entries)
    total_thinking = sum(e["thinking_tokens_used"] for e in all_entries)
    total_output = sum(e["output_tokens"] for e in all_entries)
 
    level_dist = defaultdict(int)
    for e in all_entries:
        level_dist[e["classified_level"]] += 1
 
    thinking_values = sorted(e["thinking_tokens_used"] for e in all_entries)
    p95_idx = int(0.95 * n)
    p95_thinking = thinking_values[p95_idx] if thinking_values else 0
 
    return {
        "total_requests": n,
        "avg_thinking_tokens": round(total_thinking / n, 1),
        "thinking_to_output_ratio": round(total_thinking / max(total_output, 1), 2),
        "p95_thinking_tokens": p95_thinking,
        "level_distribution": {level: count for level, count in sorted(level_dist.items())}
    }

import random
from typing import List
 
def estimate_batch_cost(
    prompts: List[str],
    sample_rate: float = 0.01,
    cost_per_1k_output: float = 0.012,
    cost_per_1k_thinking: float = 0.012
) -> dict:
    """
    本番投入前にバッチ処理コストを見積もる。
    プロンプトの一部をサンプリングし、分類して外挿する。
    """
    sample_size = max(10, int(len(prompts) * sample_rate))
    sample = random.sample(prompts, min(sample_size, len(prompts)))
 
    level_counts = {"off": 0, "light": 0, "medium": 0, "full": 0}
    for prompt in sample:
        config = classify_task(prompt)
        level_counts[config.level] += 1
 
    # 各レベルの典型的な消費量から thinking tokens を見積もる
    avg_thinking_by_level = {"off": 0, "light": 600, "medium": 3000, "full": 12000}
    avg_output_per_request = 300  # ユースケースに応じて調整
 
    total = len(sample)
    weighted_thinking = sum(
        (count / total) * avg_thinking_by_level[level]
        for level, count in level_counts.items()
    )
 
    total_prompts = len(prompts)
    estimated_thinking_tokens = weighted_thinking * total_prompts
    estimated_output_tokens = avg_output_per_request * total_prompts
 
    estimated_cost = (
        (estimated_thinking_tokens / 1000) * cost_per_1k_thinking +
        (estimated_output_tokens / 1000) * cost_per_1k_output
    )
 
    return {
        "total_prompts": total_prompts,
        "sample_size": len(sample),
        "level_distribution": {
            level: round(count / total * 100, 1) for level, count in level_counts.items()
        },
        "estimated_thinking_tokens": int(estimated_thinking_tokens),
        "estimated_output_tokens": estimated_output_tokens,
        "estimated_cost_usd": round(estimated_cost, 4)
    }
 
# 例: 50,000件のアプリレビューを処理する前に見積もる
reviews = ["素晴らしいアプリ！", "起動時にクラッシュする", "もっと速くなると良い"]  # 実データに置き換える
estimate = estimate_batch_cost(reviews * 1000)  # 3,000件を模擬
print(f"Estimated cost for {estimate['total_prompts']:,} requests: ")
print(f"  Thinking tokens: {estimate['estimated_thinking_tokens']:,}")
print(f"  Cost: ${estimate['estimated_cost_usd']:.4f} USD")