Gemini API Semantic Router: Implementation Notes for Splitting Flash and Pro Smartly

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
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
 
genai.configure(api_key=os.environ["YOUR_GEMINI_API_KEY"])
 
# =====================
# Data Classes
# =====================
 
@dataclass
class RouteResult:
    """Result of a routing decision."""
    model_name: str
    route_method: str  # "rule_based" or "semantic"
    confidence: float
    reasoning: str
    estimated_cost_ratio: float  # Normalized with Pro as 1.0
 
@dataclass
class RouterMetrics:
    """Tracks cost and performance."""
    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: Rule-Based Router
# =====================
 
class RuleBasedRouter:
    """Fast rule-based filter."""
 
    SIMPLE_PATTERNS = [
        r"^(yes|no)[?]?$",
        r"^.{1,30}[?]$",  # Short questions under 30 chars
        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|literature review)",
        r"(optimization algorithm|complexity)",
    ]
 
    def route(self, query: str) -> Optional[RouteResult]:
        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,
                )
 
        # Length heuristic
        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  # Proceed to Stage 2
 
# =====================
# Stage 2: Semantic Router
# =====================
 
class SemanticRouter:
    """Embedding-based semantic router."""
 
    ROUTE_EXAMPLES = {
        "gemini-2.5-flash": [
            "Translate this text into Japanese",
            "Is this review positive or negative?",
            "Summarize this in one sentence",
            "Classify this into the right category",
            "Suggest five email subject lines",
            "Fix the typos in this paragraph",
        ],
        "gemini-2.5-pro": [
            "Find and debug the bug in this code",
            "Draft a marketing strategy outline",
            "Analyze trends in this dataset",
            "Review this API design and propose improvements",
            "Summarize the key points of this contract",
            "Implement a machine learning model in Python",
        ],
        "gemini-2.5-pro-ultra": [
            "Walk me through quantum computing from mathematical foundations to applications",
            "Design a distributed system architecture from scratch",
            "Critically evaluate the methodology of this research paper",
            "Prove the time complexity of this algorithm",
            "Write a full technical design document for a microservices migration",
        ],
    }
 
    def __init__(self):
        self._route_embeddings: dict[str, list] = {}
        self._initialized = False
 
    def _get_embedding(self, text: str) -> list[float]:
        """Vectorize text via the 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):
        """Precompute and cache route embeddings on startup."""
        logger.info("🔧 Initializing semantic router embeddings...")
        for route_name, examples in self.ROUTE_EXAMPLES.items():
            embeddings = [self._get_embedding(ex) for ex in examples]
            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()
 
        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", ""),
            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),
        )
 
# =====================
# Main Router (Orchestrator)
# =====================
 
class GeminiSmartRouter:
    """
    Main two-stage semantic router. Selects the best model
    via rule-based → semantic order, then calls Gemini and
    returns the response.
    """
 
    MODEL_ALIASES = {
        "gemini-2.5-flash": "gemini-2.5-flash",
        "gemini-2.5-pro": "gemini-2.5-pro",
        "gemini-pro": "gemini-2.5-pro",  # fallback alias
    }
 
    def __init__(self):
        self.rule_router = RuleBasedRouter()
        self.semantic_router = SemanticRouter()
        self.metrics = RouterMetrics()
        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_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:
        """Route the query and call Gemini API."""
        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:
            route = self.rule_router.route(prompt)
            if route is None:
                route = self.semantic_router.route(prompt)
 
        model_name = self.MODEL_ALIASES.get(route.model_name, route.model_name)
 
        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:
            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),
            }
 
# =====================
# Usage Example
# =====================
 
if __name__ == "__main__":
    router = GeminiSmartRouter()
 
    test_queries = [
        ("Is this review positive? 'Best product ever!'", "free"),
        ("Implement and explain the Python decorator pattern", "free"),
        ("Design a distributed transaction system from scratch", "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!);
 
// =====================
// Types
// =====================
 
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;
}
 
// =====================
// Rule-based router (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}`,
      };
    }
  }
  if (query.split(/\s+/).length < 8) {
    return {
      modelName: "gemini-2.5-flash",
      routeMethod: "rule_based",
      confidence: 0.72,
      reasoning: "Short query heuristic",
    };
  }
  return null;
}
 
// =====================
// Semantic router (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": [
    "Translate this text into English",
    "Run sentiment analysis on this",
    "Summarize this in one sentence",
    "Classify this into the right category",
  ],
  "gemini-2.5-pro": [
    "Debug this code",
    "Draft a marketing strategy",
    "Analyze the data and tell me the trend",
    "Design an API for me",
  ],
};
 
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)}`,
  };
}
 
// =====================
// Main router function
// =====================
 
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);
  }
 
  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 usage
// =====================
// 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 });
// }

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 days
  });
}

if route.confidence < 0.70 and route.route_method == "semantic":
    route.model_name = "gemini-2.5-pro"  # safe side for low confidence