LLM Cost Optimization in Production: Semantic Caching, Batching, and Smart Model Routing

import asyncio
import hashlib
import json
from dataclasses import dataclass
from enum import Enum
from typing import Optional
import numpy as np
from openai import AsyncOpenAI
from redis.asyncio import Redis


class ModelTier(str, Enum):
    BUDGET = "gpt-4o-mini"          # example 2025 budget-tier model
    STANDARD = "gpt-4o"             # example 2025 standard-tier model
    REASONING = "o1-mini"           # Reserved for explicit reasoning tasks


@dataclass
class RoutedRequest:
    prompt: str
    system: str
    model: ModelTier
    cache_hit: bool
    estimated_cost_usd: float


COMPLEXITY_CLASSIFIER_PROMPT = """Classify the complexity of this task.
Simple: classification, extraction, lookup, format conversion, yes/no questions.
Complex: multi-step reasoning, code generation, detailed analysis, creative generation.
Respond with exactly one word: simple or complex."""

# Illustrative per-token input prices from 2025: replace with your provider's current list prices
PRICING = {
    ModelTier.BUDGET: 0.00000015,    # illustrative
    ModelTier.STANDARD: 0.000005,    # illustrative
    ModelTier.REASONING: 0.000003,   # illustrative
}


class IntelligentModelRouter:
    """
    Routes LLM requests to appropriate model tier based on complexity,
    with semantic caching to avoid redundant API calls.
    """

    CACHE_SIMILARITY_THRESHOLD = 0.95
    CACHE_TTL_SECONDS = 3600  # 1 hour for most enterprise queries

    def __init__(self, redis_url: str = "redis://localhost:6379"):
        self.openai = AsyncOpenAI()
        self._redis_url = redis_url
        self._redis: Optional[Redis] = None
        self._cache_index: list[tuple[list[float], str]] = []  # (embedding, cache_key)
        self._stats = {"total": 0, "cache_hits": 0, "budget_routes": 0, "standard_routes": 0}

    async def connect(self):
        from redis.asyncio import from_url
        self._redis = await from_url(self._redis_url, decode_responses=True)

    async def _embed(self, text: str) -> list[float]:
        resp = await self.openai.embeddings.create(
            input=text[:2000],  # Truncate for cache key generation
            model="text-embedding-3-small",
            dimensions=768
        )
        return resp.data[0].embedding

    async def _check_semantic_cache(
        self, query_embedding: list[float]
    ) -> Optional[str]:
        if not self._cache_index:
            return None
        query_vec = np.array(query_embedding)
        best_score = 0.0
        best_key = None
        for cached_emb, cache_key in self._cache_index[-1000:]:  # Check last 1000
            cached_vec = np.array(cached_emb)
            score = float(np.dot(query_vec, cached_vec) /
                          (np.linalg.norm(query_vec) * np.linalg.norm(cached_vec)))
            if score > best_score:
                best_score = score
                best_key = cache_key
        if best_score >= self.CACHE_SIMILARITY_THRESHOLD and best_key:
            return best_key
        return None

    async def _store_cache(self, query_embedding: list[float],
                           response: str, ttl: int) -> None:
        cache_key = "llm:" + hashlib.sha256(
            str(query_embedding[:10]).encode()
        ).hexdigest()[:16]
        await self._redis.setex(cache_key, ttl, response)
        self._cache_index.append((query_embedding, cache_key))

    async def _classify_complexity(self, prompt: str) -> ModelTier:
        resp = await self.openai.chat.completions.create(
            model="gpt-4o-mini",  # Always use budget model for classification
            messages=[
                {"role": "system", "content": COMPLEXITY_CLASSIFIER_PROMPT},
                {"role": "user", "content": prompt[:500]},  # Truncate for speed
            ],
            max_tokens=5,
            temperature=0,
        )
        classification = resp.choices[0].message.content.strip().lower()
        return ModelTier.STANDARD if classification == "complex" else ModelTier.BUDGET

    def _estimate_cost(self, prompt: str, model: ModelTier) -> float:
        token_estimate = len(prompt.split()) * 1.3  # Rough token estimate
        return token_estimate * PRICING[model]

    async def complete(
        self,
        prompt: str,
        system: str = "",
        force_model: Optional[ModelTier] = None,
        cache_ttl: int = CACHE_TTL_SECONDS,
    ) -> RoutedRequest:
        self._stats["total"] += 1
        full_prompt = f"{system}\n\n{prompt}" if system else prompt

        # Check semantic cache first
        query_embedding = await self._embed(full_prompt)
        cache_key = await self._check_semantic_cache(query_embedding)
        if cache_key:
            cached = await self._redis.get(cache_key)
            if cached:
                self._stats["cache_hits"] += 1
                return RoutedRequest(prompt=prompt, system=system,
                                     model=ModelTier.BUDGET, cache_hit=True,
                                     estimated_cost_usd=0.0)

        # Route to appropriate model
        model = force_model or await self._classify_complexity(prompt)
        if model == ModelTier.BUDGET:
            self._stats["budget_routes"] += 1
        else:
            self._stats["standard_routes"] += 1

        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        resp = await self.openai.chat.completions.create(
            model=model.value, messages=messages, temperature=0.2
        )
        response_text = resp.choices[0].message.content
        await self._store_cache(query_embedding, response_text, cache_ttl)

        return RoutedRequest(
            prompt=prompt, system=system, model=model, cache_hit=False,
            estimated_cost_usd=self._estimate_cost(full_prompt, model)
        )

    def print_stats(self):
        n = self._stats["total"] or 1
        print(f"Total: {n} | Cache hits: {self._stats['cache_hits']/n:.1%} | "
              f"Budget routes: {self._stats['budget_routes']/n:.1%} | "
              f"Standard routes: {self._stats['standard_routes']/n:.1%}")


async def main():
    router = IntelligentModelRouter()
    await router.connect()
    result = await router.complete(
        prompt="What is the refund policy for enterprise customers?",
        system="You are a helpful assistant for Acme Corp."
    )
    print(f"Model: {result.model.value}, Cache hit: {result.cache_hit}, "
          f"Est. cost: ${result.estimated_cost_usd:.6f}")
    router.print_stats()


if __name__ == "__main__":
    asyncio.run(main())