AI Gateway Pattern: Intelligent LLM Routing

Request
  ↓
┌─────────────────────────────────────┐
│ 1. AUTHENTICATION                   │ Who are you?
│    Validate API Keys                │
├─────────────────────────────────────┤
│ 2. RATE LIMITING                    │ How much can you use?
│    Tokens/min, Requests/min per Key │
├─────────────────────────────────────┤
│ 3. SMART ROUTING                    │ Which model?
│    Cost vs Speed vs Quality         │
├─────────────────────────────────────┤
│ 4. MONITORING & CACHING             │ Did we see this before?
│    Token Counting, Semantic Cache   │
└─────────────────────────────────────┘
  ↓
LLM Backends (vLLM, OpenAI, Anthropic, ...)

import secrets
import hashlib

class APIKeyManager:
    def create_key(self, user_id: str, quota_tokens: int):
        """Create an API key with token quota"""
        key_bytes = secrets.token_bytes(32)
        key_hex = key_bytes.hex()
        key_hash = hashlib.sha256(key_hex.encode()).hexdigest()

        self.db.execute("""
            INSERT INTO api_keys (user_id, hash, quota_tokens, created_at)
            VALUES (?, ?, ?, datetime('now'))
        """, (user_id, key_hash, quota_tokens))

        return key_hex

    def validate_key(self, key_hex: str) -> dict | None:
        """Validate API key and get user limits"""
        key_hash = hashlib.sha256(key_hex.encode()).hexdigest()

        result = self.db.execute("""
            SELECT user_id, quota_tokens, tokens_used
            FROM api_keys WHERE hash = ? AND valid_until > now()
        """, (key_hash,))

        if not result or result[0]['tokens_used'] >= result[0]['quota_tokens']:
            return None

        return {
            'user_id': result[0]['user_id'],
            'tokens_remaining': result[0]['quota_tokens'] - result[0]['tokens_used']
        }

from collections import defaultdict

class TokenRateLimiter:
    def __init__(self, tokens_per_minute=100000):
        self.limit = tokens_per_minute
        self.user_buckets = defaultdict(lambda: {
            "tokens": tokens_per_minute,
            "last_refill": datetime.now()
        })

    def check_and_consume(self, user_id: str, tokens_needed: int) -> bool:
        """Check if user has tokens and consume them"""
        bucket = self.user_buckets[user_id]

        # Token Refill (Leaky Bucket)
        now = datetime.now()
        seconds_since_refill = (now - bucket["last_refill"]).total_seconds()
        refill_rate = self.limit / 60

        bucket["tokens"] = min(
            self.limit,
            bucket["tokens"] + (seconds_since_refill * refill_rate)
        )
        bucket["last_refill"] = now

        if bucket["tokens"] < tokens_needed:
            return False

        bucket["tokens"] -= tokens_needed
        return True

from enum import Enum

class RoutingStrategy(Enum):
    SPEED = "minimize_latency"
    COST = "minimize_cost"
    QUALITY = "maximize_quality"

class GatewayRouter:
    def select_model(
        self,
        request_tokens: int,
        strategy: RoutingStrategy,
        budget_remaining: float
    ) -> str:
        """Choose best model for request"""

        if strategy == RoutingStrategy.SPEED:
            return min(self.endpoints, key=lambda ep: ep.latency_p95)

        elif strategy == RoutingStrategy.COST:
            request_cost = (request_tokens / 1000) * 0.0005  # Example cost
            if request_cost > budget_remaining:
                return min(self.endpoints, key=lambda ep: ep.cost_per_1k_tokens)
            return min(self.endpoints, key=lambda ep: ep.cost_per_1k_tokens)

        else:  # QUALITY
            return max(self.endpoints, key=lambda ep: ep.quality_score)

class CostTracker:
    def log_completion(self, user_id: str, model: str, tokens: int):
        """Log API call with cost"""
        pricing = {
            "gpt-4": {"prompt": 0.03, "completion": 0.06},
            "llama-70b": {"prompt": 0.0005, "completion": 0.0008}
        }

        cost_usd = (tokens / 1000) * pricing.get(model, {}).get("prompt", 0)

        self.db.execute("""
            INSERT INTO cost_events (user_id, model, tokens, cost_usd)
            VALUES (?, ?, ?, ?)
        """, (user_id, model, tokens, cost_usd))

    def get_user_costs(self, user_id: str):
        """Aggregate costs for user"""
        result = self.db.execute("""
            SELECT SUM(cost_usd) as total FROM cost_events
            WHERE user_id = ? AND timestamp > datetime('now', '-30 days')
        """, (user_id,))

        return result[0]['total'] if result else 0

class SemanticCache:
    def __init__(self, embedding_model):
        self.embedding_model = embedding_model
        self.cache = []  # (embedding, prompt, response)

    def get(self, prompt: str, threshold: float = 0.95) -> str | None:
        """Find similar prompts in cache"""
        query_embedding = self.embedding_model.encode(prompt)

        for cached_embedding, cached_prompt, cached_response in self.cache:
            similarity = cosine_similarity(query_embedding, cached_embedding)
            if similarity >= threshold:
                return cached_response

        return None

    def put(self, prompt: str, response: str):
        """Store prompt-response pair"""
        embedding = self.embedding_model.encode(prompt)
        self.cache.append((embedding, prompt, response))

from litellm import completion

litellm.model_list = [
    {
        "model_name": "my-gpt4",
        "litellm_params": {
            "model": "gpt-4",
            "api_key": "$OPENAI_API_KEY"
        }
    },
    {
        "model_name": "my-gpt4",
        "litellm_params": {
            "model": "claude-3-opus",
            "api_key": "$ANTHROPIC_API_KEY"
        }
    }
]

@app.post("/v1/chat/completions")
async def unified_chat(request: dict):
    """One endpoint, multiple backends"""
    response = completion(
        model="my-gpt4",
        messages=request["messages"],
        num_retries=2  # Fallback automatic
    )
    return response