Gemini API セマンティックルーター実装ノート ― Flash と Pro を賢く振り分けて API 費用を抑える

pip install google-generativeai numpy scikit-learn python-dotenv

import os
import re
import time
import json
import logging
from dataclasses import dataclass, field
from typing import Optional
import numpy as np
import google.generativeai as genai
from dotenv import load_dotenv
 
load_dotenv()
 
# ロギング設定
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
 
genai.configure(api_key=os.environ["YOUR_GEMINI_API_KEY"])
 
# =====================
# データクラス定義
# =====================
 
@dataclass
class RouteResult:
    """ルーティング結果"""
    model_name: str
    route_method: str  # "rule_based" or "semantic"
    confidence: float
    reasoning: str
    estimated_cost_ratio: float  # Pro を 1.0 として正規化
 
@dataclass
class RouterMetrics:
    """コスト・パフォーマンス追跡"""
    total_requests: int = 0
    flash_count: int = 0
    pro_count: int = 0
    ultra_count: int = 0
    total_tokens_in: int = 0
    total_tokens_out: int = 0
    estimated_cost_usd: float = 0.0
    estimated_savings_usd: float = 0.0
 
    def log_summary(self):
        total = self.total_requests or 1
        print(f"\n📊 Router Metrics Summary")
        print(f"  Total requests : {self.total_requests}")
        print(f"  Flash ratio    : {self.flash_count/total:.1%} ({self.flash_count} 件)")
        print(f"  Pro ratio      : {self.pro_count/total:.1%} ({self.pro_count} 件)")
        print(f"  Ultra ratio    : {self.ultra_count/total:.1%} ({self.ultra_count} 件)")
        print(f"  Est. Cost      : ${self.estimated_cost_usd:.4f}")
        print(f"  Est. Savings   : ${self.estimated_savings_usd:.4f}")
 
# =====================
# Stage 1: ルールベースルーター
# =====================
 
class RuleBasedRouter:
    """高速ルールベースフィルター"""
 
    # シンプルタスクのパターン
    SIMPLE_PATTERNS = [
        r"^(はい|いいえ|yes|no)[?？]?$",
        r"^.{1,30}[?？]$",  # 30 文字以下の短い質問
        r"(翻訳して|translate|要約して|summarize in one sentence)",
        r"(分類して|classify|ラベルを付け|label this)",
        r"(感情分析|sentiment|ポジティブ|ネガティブ|positive|negative)",
        r"(今日の日付|current date|今何時|what time)",
    ]
 
    # 複雑タスクのパターン
    COMPLEX_PATTERNS = [
        r"(アーキテクチャ|architecture|設計書|design document)",
        r"(証明して|prove that|数学的に|mathematically)",
        r"(包括的な|comprehensive|詳細な分析|detailed analysis)",
        r"(コードをゼロから|implement from scratch|完全実装|full implementation)",
        r"(論文|research paper|学術|academic|文献調査)",
        r"(最適化アルゴリズム|optimization algorithm|複雑度|complexity)",
    ]
 
    def route(self, query: str) -> Optional[RouteResult]:
        query_lower = query.lower()
 
        # シンプルパターンチェック
        for pattern in self.SIMPLE_PATTERNS:
            if re.search(pattern, query, re.IGNORECASE):
                return RouteResult(
                    model_name="gemini-2.5-flash",
                    route_method="rule_based",
                    confidence=0.95,
                    reasoning=f"Simple pattern matched: {pattern}",
                    estimated_cost_ratio=0.05,
                )
 
        # 複雑パターンチェック
        for pattern in self.COMPLEX_PATTERNS:
            if re.search(pattern, query, re.IGNORECASE):
                return RouteResult(
                    model_name="gemini-2.5-pro",
                    route_method="rule_based",
                    confidence=0.90,
                    reasoning=f"Complex pattern matched: {pattern}",
                    estimated_cost_ratio=1.0,
                )
 
        # クエリ長による簡易判定
        word_count = len(query.split())
        if word_count < 10:
            return RouteResult(
                model_name="gemini-2.5-flash",
                route_method="rule_based",
                confidence=0.75,
                reasoning=f"Short query ({word_count} words)",
                estimated_cost_ratio=0.05,
            )
 
        return None  # Stage 2 へ
 
# =====================
# Stage 2: セマンティックルーター
# =====================
 
class SemanticRouter:
    """Embedding ベースのセマンティックルーター"""
 
    # 各ルートのサンプルクエリ（事前に定義しておく）
    ROUTE_EXAMPLES = {
        "gemini-2.5-flash": [
            "今日の天気を教えて",
            "このテキストを英語に翻訳して",
            "このレビューはポジティブですか？",
            "一文で要約してください",
            "郵便番号から都道府県を教えて",
            "このカテゴリに分類して",
            "メールの件名を 5 つ考えて",
            "この文章の誤字を直して",
        ],
        "gemini-2.5-pro": [
            "このコードのバグを見つけてデバッグして",
            "マーケティング戦略のアウトラインを作成して",
            "このデータセットのトレンドを分析して",
            "API の設計をレビューして改善点を提案して",
            "この契約書の重要なポイントをまとめて",
            "Python で機械学習モデルを実装して",
        ],
        "gemini-2.5-pro-ultra": [
            "量子コンピューターの数学的基礎から応用まで包括的に解説して",
            "ゼロから分散システムのアーキテクチャを設計して",
            "この研究論文の手法を批判的に評価して",
            "複雑なアルゴリズムの時間複雑度を証明して",
            "マイクロサービス移行の完全な技術設計書を作成して",
        ],
    }
 
    def __init__(self):
        self._route_embeddings: dict[str, list] = {}
        self._initialized = False
 
    def _get_embedding(self, text: str) -> list[float]:
        """Gemini Embedding API でテキストをベクトル化"""
        result = genai.embed_content(
            model="models/text-embedding-004",
            content=text,
            task_type="SEMANTIC_SIMILARITY",
        )
        return result["embedding"]
 
    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
        """コサイン類似度を計算"""
        a_np = np.array(a)
        b_np = np.array(b)
        return float(np.dot(a_np, b_np) / (np.linalg.norm(a_np) * np.linalg.norm(b_np)))
 
    def initialize(self):
        """起動時にルート embedding を事前計算してキャッシュ"""
        logger.info("🔧 Initializing semantic router embeddings...")
        for route_name, examples in self.ROUTE_EXAMPLES.items():
            embeddings = [self._get_embedding(ex) for ex in examples]
            # 各ルートの embedding を平均化（centroid）
            self._route_embeddings[route_name] = np.mean(embeddings, axis=0).tolist()
        self._initialized = True
        logger.info("✅ Semantic router initialized")
 
    def route(self, query: str) -> RouteResult:
        if not self._initialized:
            self.initialize()
 
        # クエリの embedding を取得
        query_embedding = self._get_embedding(query)
 
        # 各ルートとの類似度を計算
        similarities = {}
        for route_name, route_embedding in self._route_embeddings.items():
            sim = self._cosine_similarity(query_embedding, route_embedding)
            similarities[route_name] = sim
 
        # 最も類似度が高いルートを選択
        best_route = max(similarities, key=similarities.get)
        best_score = similarities[best_route]
 
        # コスト比率マッピング
        cost_ratios = {
            "gemini-2.5-flash": 0.05,
            "gemini-2.5-pro": 0.40,
            "gemini-2.5-pro-ultra": 1.0,
        }
 
        return RouteResult(
            model_name=best_route.replace("-ultra", ""),  # API モデル名に変換
            route_method="semantic",
            confidence=best_score,
            reasoning=f"Similarity scores: {json.dumps({k: round(v, 3) for k, v in similarities.items()})}",
            estimated_cost_ratio=cost_ratios.get(best_route, 0.5),
        )
 
# =====================
# メインルーター（オーケストレーター）
# =====================
 
class GeminiSmartRouter:
    """
    2 段階セマンティックルーターのメインクラス。
    ルールベース → セマンティック の順で最適モデルを選択し、
    実際に Gemini API を呼び出して結果を返す。
    """
 
    # モデル名のエイリアス解決
    MODEL_ALIASES = {
        "gemini-2.5-flash": "gemini-2.5-flash",
        "gemini-2.5-pro": "gemini-2.5-pro",
        "gemini-pro": "gemini-2.5-pro",  # フォールバック
    }
 
    def __init__(self):
        self.rule_router = RuleBasedRouter()
        self.semantic_router = SemanticRouter()
        self.metrics = RouterMetrics()
        # コスト定数（USD/1M tokens）
        self._costs = {
            "gemini-2.5-flash": {"in": 0.15, "out": 0.60},
            "gemini-2.5-pro": {"in": 1.25, "out": 10.00},
        }
 
    def _update_metrics(self, model: str, usage):
        """使用量とコストをメトリクスに記録"""
        self.metrics.total_requests += 1
        if "flash" in model:
            self.metrics.flash_count += 1
        elif "pro" in model:
            self.metrics.pro_count += 1
 
        if usage:
            in_tok = getattr(usage, "prompt_token_count", 0) or 0
            out_tok = getattr(usage, "candidates_token_count", 0) or 0
            self.metrics.total_tokens_in += in_tok
            self.metrics.total_tokens_out += out_tok
 
            # 実際コスト計算
            cost_cfg = self._costs.get(model, self._costs["gemini-2.5-pro"])
            actual_cost = (in_tok * cost_cfg["in"] + out_tok * cost_cfg["out"]) / 1_000_000
            self.metrics.estimated_cost_usd += actual_cost
 
            # 仮に全部 Pro で処理した場合のコスト（節約額計算用）
            pro_cfg = self._costs["gemini-2.5-pro"]
            pro_cost = (in_tok * pro_cfg["in"] + out_tok * pro_cfg["out"]) / 1_000_000
            self.metrics.estimated_savings_usd += max(0, pro_cost - actual_cost)
 
    def query(
        self,
        prompt: str,
        system_instruction: Optional[str] = None,
        force_model: Optional[str] = None,
        user_tier: str = "free",  # "free" | "pro" | "premium"
    ) -> dict:
        """
        クエリをルーティングして Gemini API を呼び出す。
 
        Args:
            prompt: ユーザーのクエリテキスト
            system_instruction: システムプロンプト（オプション）
            force_model: モデルを強制指定（デバッグ用）
            user_tier: ユーザープラン（プレミアムは常に高品質モデルを使用）
 
        Returns:
            {"text": str, "model": str, "route": RouteResult, "latency_ms": float}
        """
        start_time = time.perf_counter()
 
        # プレミアムユーザーは常に最高モデルを使用
        if user_tier == "premium":
            route = RouteResult(
                model_name="gemini-2.5-pro",
                route_method="user_tier",
                confidence=1.0,
                reasoning="Premium user always gets top model",
                estimated_cost_ratio=1.0,
            )
        elif force_model:
            route = RouteResult(
                model_name=force_model,
                route_method="forced",
                confidence=1.0,
                reasoning="Model explicitly forced",
                estimated_cost_ratio=0.5,
            )
        else:
            # Stage 1: ルールベースルーティング
            route = self.rule_router.route(prompt)
 
            # Stage 2: セマンティックルーティング（Stage 1 が決定できなかった場合）
            if route is None:
                route = self.semantic_router.route(prompt)
 
        model_name = self.MODEL_ALIASES.get(route.model_name, route.model_name)
 
        # Gemini API 呼び出し
        try:
            gen_config = {
                "temperature": 0.7,
                "max_output_tokens": 4096,
            }
 
            if system_instruction:
                model = genai.GenerativeModel(
                    model_name,
                    system_instruction=system_instruction,
                    generation_config=gen_config,
                )
            else:
                model = genai.GenerativeModel(model_name, generation_config=gen_config)
 
            response = model.generate_content(prompt)
            latency_ms = (time.perf_counter() - start_time) * 1000
 
            # メトリクス更新
            self._update_metrics(model_name, response.usage_metadata)
 
            return {
                "text": response.text,
                "model": model_name,
                "route": route,
                "latency_ms": round(latency_ms, 2),
            }
 
        except Exception as e:
            # フォールバック: エラー時は Pro にアップグレードして再試行
            logger.warning(f"⚠️ {model_name} failed: {e}. Falling back to gemini-2.5-pro")
            model = genai.GenerativeModel("gemini-2.5-pro")
            response = model.generate_content(prompt)
            latency_ms = (time.perf_counter() - start_time) * 1000
            self._update_metrics("gemini-2.5-pro", response.usage_metadata)
 
            return {
                "text": response.text,
                "model": "gemini-2.5-pro (fallback)",
                "route": route,
                "latency_ms": round(latency_ms, 2),
            }
 
# =====================
# 使用例
# =====================
 
if __name__ == "__main__":
    router = GeminiSmartRouter()
 
    test_queries = [
        ("このレビューはポジティブですか？「最高の商品でした！」", "free"),
        ("Python のデコレーターパターンを実装して説明してください", "free"),
        ("ゼロから分散トランザクションシステムを設計してください", "premium"),
    ]
 
    for query, tier in test_queries:
        result = router.query(query, user_tier=tier)
        print(f"\n📝 Query : {query[:50]}...")
        print(f"   Model  : {result['model']}")
        print(f"   Method : {result['route'].route_method}")
        print(f"   Conf.  : {result['route'].confidence:.2f}")
        print(f"   Time   : {result['latency_ms']}ms")
        print(f"   Answer : {result['text'][:100]}...")
 
    # メトリクスサマリーを表示
    router.metrics.log_summary()

import { GoogleGenerativeAI, EmbedContentRequest } from "@google/generative-ai";
 
const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY!);
 
// =====================
// 型定義
// =====================
 
interface RouteResult {
  modelName: string;
  routeMethod: "rule_based" | "semantic" | "user_tier" | "forced";
  confidence: number;
  reasoning: string;
}
 
interface QueryResult {
  text: string;
  model: string;
  route: RouteResult;
  latencyMs: number;
}
 
// =====================
// ルールベースルーター（TypeScript 版）
// =====================
 
const SIMPLE_PATTERNS = [
  /^.{1,40}[?？]$/,
  /(翻訳|translate|要約|summarize in one|感情分析|sentiment|classify)/i,
  /^(yes|no|はい|いいえ)/i,
];
 
const COMPLEX_PATTERNS = [
  /(アーキテクチャ|architecture|設計書|design document)/i,
  /(証明|prove that|数学的|mathematically)/i,
  /(包括的|comprehensive|詳細な分析|detailed analysis)/i,
  /(ゼロから|from scratch|完全実装|full implementation)/i,
];
 
function ruleBasedRoute(query: string): RouteResult | null {
  for (const pattern of SIMPLE_PATTERNS) {
    if (pattern.test(query)) {
      return {
        modelName: "gemini-2.5-flash",
        routeMethod: "rule_based",
        confidence: 0.92,
        reasoning: `Matched simple pattern: ${pattern}`,
      };
    }
  }
  for (const pattern of COMPLEX_PATTERNS) {
    if (pattern.test(query)) {
      return {
        modelName: "gemini-2.5-pro",
        routeMethod: "rule_based",
        confidence: 0.90,
        reasoning: `Matched complex pattern: ${pattern}`,
      };
    }
  }
  // 短いクエリは Flash
  if (query.split(/\s+/).length < 8) {
    return {
      modelName: "gemini-2.5-flash",
      routeMethod: "rule_based",
      confidence: 0.72,
      reasoning: "Short query heuristic",
    };
  }
  return null;
}
 
// =====================
// セマンティックルーター（TypeScript 版）
// =====================
 
async function getEmbedding(text: string): Promise<number[]> {
  const embeddingModel = genAI.getGenerativeModel({ model: "text-embedding-004" });
  const result = await embeddingModel.embedContent({
    content: { parts: [{ text }], role: "user" },
    taskType: "SEMANTIC_SIMILARITY" as any,
  } as EmbedContentRequest);
  return result.embedding.values;
}
 
function cosineSimilarity(a: number[], b: number[]): number {
  let dot = 0, normA = 0, normB = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }
  return dot / (Math.sqrt(normA) * Math.sqrt(normB));
}
 
// キャッシュ（起動時に初期化）
let routeEmbeddingsCache: Record<string, number[]> | null = null;
 
const ROUTE_CENTROIDS: Record<string, string[]> = {
  "gemini-2.5-flash": [
    "このテキストを英語に翻訳して",
    "感情分析してください",
    "一文で要約して",
    "カテゴリを分類して",
  ],
  "gemini-2.5-pro": [
    "このコードをデバッグして",
    "マーケティング戦略を作成して",
    "データを分析してトレンドを教えて",
    "API を設計してください",
  ],
};
 
async function initSemanticRouter(): Promise<Record<string, number[]>> {
  if (routeEmbeddingsCache) return routeEmbeddingsCache;
 
  const cache: Record<string, number[]> = {};
  for (const [route, examples] of Object.entries(ROUTE_CENTROIDS)) {
    const embeddings = await Promise.all(examples.map(getEmbedding));
    // 平均埋め込みベクトルを計算
    const centroid = embeddings[0].map((_, i) =>
      embeddings.reduce((sum, emb) => sum + emb[i], 0) / embeddings.length
    );
    cache[route] = centroid;
  }
  routeEmbeddingsCache = cache;
  return cache;
}
 
async function semanticRoute(query: string): Promise<RouteResult> {
  const routeEmbeddings = await initSemanticRouter();
  const queryEmb = await getEmbedding(query);
 
  let bestRoute = "gemini-2.5-pro";
  let bestScore = -1;
 
  for (const [route, routeEmb] of Object.entries(routeEmbeddings)) {
    const sim = cosineSimilarity(queryEmb, routeEmb);
    if (sim > bestScore) {
      bestScore = sim;
      bestRoute = route;
    }
  }
 
  return {
    modelName: bestRoute,
    routeMethod: "semantic",
    confidence: bestScore,
    reasoning: `Highest semantic similarity: ${bestScore.toFixed(3)}`,
  };
}
 
// =====================
// メインルーター関数
// =====================
 
export async function smartQuery(
  prompt: string,
  options: {
    systemInstruction?: string;
    userTier?: "free" | "pro" | "premium";
    forceModel?: string;
  } = {}
): Promise<QueryResult> {
  const start = performance.now();
  const { systemInstruction, userTier = "free", forceModel } = options;
 
  // ルート決定
  let route: RouteResult;
  if (userTier === "premium") {
    route = { modelName: "gemini-2.5-pro", routeMethod: "user_tier", confidence: 1.0, reasoning: "Premium tier" };
  } else if (forceModel) {
    route = { modelName: forceModel, routeMethod: "forced", confidence: 1.0, reasoning: "Forced" };
  } else {
    route = ruleBasedRoute(prompt) ?? await semanticRoute(prompt);
  }
 
  // API 呼び出し
  try {
    const model = genAI.getGenerativeModel({
      model: route.modelName,
      systemInstruction: systemInstruction,
      generationConfig: { temperature: 0.7, maxOutputTokens: 4096 },
    });
 
    const response = await model.generateContent(prompt);
    const latencyMs = performance.now() - start;
 
    return {
      text: response.response.text(),
      model: route.modelName,
      route,
      latencyMs: Math.round(latencyMs),
    };
  } catch (error) {
    // フォールバック
    console.warn(`[SmartRouter] ${route.modelName} failed, falling back to gemini-2.5-pro`);
    const fallbackModel = genAI.getGenerativeModel({ model: "gemini-2.5-pro" });
    const response = await fallbackModel.generateContent(prompt);
    return {
      text: response.response.text(),
      model: "gemini-2.5-pro (fallback)",
      route,
      latencyMs: Math.round(performance.now() - start),
    };
  }
}
 
// =====================
// Next.js API Route での利用例
// =====================
// app/api/chat/route.ts
//
// import { smartQuery } from "@/lib/smart-router";
//
// export async function POST(req: Request) {
//   const { message, userTier } = await req.json();
//   const result = await smartQuery(message, { userTier });
//   return Response.json({ text: result.text, model: result.model });
// }

// Edge 環境でのコスト追跡（Cloudflare Workers KV）
async function trackUsage(
  env: { COST_KV: KVNamespace },
  model: string,
  inputTokens: number,
  outputTokens: number
): Promise<void> {
  const today = new Date().toISOString().split("T")[0]; // YYYY-MM-DD
  const key = `usage:${today}`;
 
  const existing = await env.COST_KV.get(key, "json") as any || {
    flash: { in: 0, out: 0 },
    pro: { in: 0, out: 0 },
    total_usd: 0,
  };
 
  const modelKey = model.includes("flash") ? "flash" : "pro";
  existing[modelKey].in += inputTokens;
  existing[modelKey].out += outputTokens;
 
  // コスト計算
  const rates = {
    flash: { in: 0.15, out: 0.60 },
    pro: { in: 1.25, out: 10.00 },
  };
  existing.total_usd = (
    existing.flash.in * rates.flash.in +
    existing.flash.out * rates.flash.out +
    existing.pro.in * rates.pro.in +
    existing.pro.out * rates.pro.out
  ) / 1_000_000;
 
  await env.COST_KV.put(key, JSON.stringify(existing), {
    expirationTtl: 60 * 60 * 24 * 30, // 30 日保持
  });
}

if route.confidence < 0.70 and route.route_method == "semantic":
    route.model_name = "gemini-2.5-pro"  # 低信頼時は安全側へ