Cost Optimization for AI Workloads

import tiktoken
from dataclasses import dataclass, field
from typing import Optional
import time
import json


@dataclass
class TokenUsage:
    """Track token usage for a single LLM call."""
    prompt_tokens: int
    completion_tokens: int
    model: str
    endpoint: str
    timestamp: float = field(default_factory=time.time)
    cache_hit: bool = False
    estimated_cost_usd: float = 0.0


# Pricing per 1M tokens (input / output) -- update as prices change
MODEL_PRICING = {
    "gpt-4o":          {"input": 2.50,  "output": 10.00},
    "gpt-4o-mini":     {"input": 0.15,  "output": 0.60},
    "claude-sonnet-4-20250514":   {"input": 3.00,  "output": 15.00},
    "claude-haiku-3-5": {"input": 0.80,  "output": 4.00},
}


def estimate_cost(usage: TokenUsage) -> float:
    """Estimate cost in USD for a single LLM call."""
    pricing = MODEL_PRICING.get(usage.model)
    if not pricing:
        return 0.0
    input_cost = (usage.prompt_tokens / 1_000_000) * pricing["input"]
    output_cost = (usage.completion_tokens / 1_000_000) * pricing["output"]
    return round(input_cost + output_cost, 6)


def count_tokens(text: str, model: str = "gpt-4o") -> int:
    """Count tokens for a given text using tiktoken."""
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")
    return len(encoding.encode(text))

def optimize_system_prompt(prompt: str) -> dict:
    """Analyze a system prompt and suggest optimizations.

    Returns a dict with the original token count,
    specific recommendations, and estimated savings.
    """
    token_count = count_tokens(prompt)
    recommendations = []

    # Check for common waste patterns
    lines = prompt.split("\n")

    # 1. Redundant instructions
    seen_instructions = set()
    for i, line in enumerate(lines):
        normalized = line.strip().lower()
        if normalized in seen_instructions and len(normalized) > 20:
            recommendations.append(
                f"Line {i+1}: Duplicate instruction detected"
            )
        seen_instructions.add(normalized)

    # 2. Verbose phrasing
    verbose_patterns = {
        "I want you to act as": "You are",
        "Please make sure to": "",
        "It is important that you": "",
        "You should always remember to": "",
        "Under no circumstances should you ever": "Never",
    }
    for verbose, concise in verbose_patterns.items():
        if verbose.lower() in prompt.lower():
            replacement = f"Replace with '{concise}'" if concise else "Remove"
            recommendations.append(
                f"Verbose phrasing: '{verbose}' -> {replacement}"
            )

    # 3. Few-shot example length
    if prompt.count("Example:") > 3 or prompt.count("###") > 6:
        recommendations.append(
            "Consider reducing few-shot examples to 2-3 "
            "representative cases instead of exhaustive coverage"
        )

    return {
        "original_tokens": token_count,
        "recommendations": recommendations,
        "estimated_savings_pct": min(len(recommendations) * 5, 40),
    }

from typing import List, Dict


def manage_conversation_context(
    messages: List[Dict[str, str]],
    max_context_tokens: int = 4000,
    model: str = "gpt-4o",
) -> List[Dict[str, str]]:
    """Trim conversation history to fit within a token budget.

    Strategy: Keep the system prompt and last N messages,
    summarize older messages if needed.
    """
    if not messages:
        return messages

    system_msgs = [m for m in messages if m["role"] == "system"]
    non_system = [m for m in messages if m["role"] != "system"]

    # Always keep system prompt
    system_tokens = sum(count_tokens(m["content"]) for m in system_msgs)
    remaining_budget = max_context_tokens - system_tokens

    # Work backwards from most recent, keeping messages
    # until we exhaust the budget
    kept_messages = []
    used_tokens = 0

    for msg in reversed(non_system):
        msg_tokens = count_tokens(msg["content"])
        if used_tokens + msg_tokens > remaining_budget:
            break
        kept_messages.insert(0, msg)
        used_tokens += msg_tokens

    return system_msgs + kept_messages


def select_rag_context(
    query: str,
    retrieved_chunks: List[Dict],
    max_context_tokens: int = 3000,
) -> List[Dict]:
    """Select the most relevant RAG chunks within a token budget.

    Chunks should already be sorted by relevance score.
    We greedily add chunks until the budget is exhausted.
    """
    selected = []
    used_tokens = 0

    for chunk in retrieved_chunks:
        chunk_tokens = count_tokens(chunk["text"])
        if used_tokens + chunk_tokens > max_context_tokens:
            continue  # Skip this chunk, try smaller ones
        selected.append(chunk)
        used_tokens += chunk_tokens

    return selected

from enum import Enum
from typing import Optional


class TaskComplexity(Enum):
    SIMPLE = "simple"
    MODERATE = "moderate"
    COMPLEX = "complex"


# Map complexity to model -- adjust based on your quality requirements
ROUTING_TABLE = {
    TaskComplexity.SIMPLE: "gpt-4o-mini",
    TaskComplexity.MODERATE: "gpt-4o-mini",
    TaskComplexity.COMPLEX: "gpt-4o",
}


def classify_task_complexity(
    task_type: str,
    input_tokens: int,
    requires_reasoning: bool = False,
    requires_code_generation: bool = False,
) -> TaskComplexity:
    """Rule-based task complexity classifier.

    Start with rules, then replace with a trained classifier
    once you have labeled data from production traffic.
    """
    # Tasks that always need frontier models
    if requires_reasoning or requires_code_generation:
        return TaskComplexity.COMPLEX

    # Classification, extraction, and formatting are
    # typically simple regardless of input length
    simple_tasks = {
        "classification", "extraction", "formatting",
        "translation", "sentiment", "summarize_short",
    }
    if task_type in simple_tasks:
        return TaskComplexity.SIMPLE

    # Long-form generation and analysis need more capability
    complex_tasks = {
        "analysis", "long_generation", "multi_step",
        "code_review", "architecture",
    }
    if task_type in complex_tasks:
        return TaskComplexity.COMPLEX

    # Default to moderate for unknown task types
    return TaskComplexity.MODERATE


def route_request(
    task_type: str,
    input_text: str,
    **kwargs,
) -> str:
    """Select the appropriate model for a given request."""
    input_tokens = count_tokens(input_text)
    complexity = classify_task_complexity(
        task_type, input_tokens, **kwargs
    )
    return ROUTING_TABLE[complexity]

import hashlib
import numpy as np
from typing import Optional, Tuple
import time


class SemanticCache:
    """LLM response cache using embedding similarity.

    Uses cosine similarity to match semantically equivalent
    queries, avoiding redundant LLM calls.
    """

    def __init__(
        self,
        similarity_threshold: float = 0.95,
        max_entries: int = 10_000,
        ttl_seconds: int = 3600,
    ):
        self.threshold = similarity_threshold
        self.max_entries = max_entries
        self.ttl = ttl_seconds
        self.cache: dict[str, dict] = {}
        self.embeddings: list[Tuple[str, np.ndarray]] = []
        self.stats = {"hits": 0, "misses": 0, "evictions": 0}

    def _cosine_similarity(
        self, a: np.ndarray, b: np.ndarray
    ) -> float:
        return float(
            np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
        )

    def get(
        self, query_embedding: np.ndarray
    ) -> Optional[str]:
        """Look up a semantically similar cached response."""
        best_score = 0.0
        best_key = None

        for key, emb in self.embeddings:
            # Skip expired entries
            entry = self.cache.get(key)
            if not entry:
                continue
            if time.time() - entry["timestamp"] > self.ttl:
                continue

            score = self._cosine_similarity(query_embedding, emb)
            if score > best_score:
                best_score = score
                best_key = key

        if best_score >= self.threshold and best_key:
            self.stats["hits"] += 1
            return self.cache[best_key]["response"]

        self.stats["misses"] += 1
        return None

    def put(
        self,
        query: str,
        query_embedding: np.ndarray,
        response: str,
    ) -> None:
        """Store a query-response pair in the cache."""
        key = hashlib.sha256(query.encode()).hexdigest()

        # Evict oldest if at capacity
        if len(self.cache) >= self.max_entries:
            oldest_key = min(
                self.cache, key=lambda k: self.cache[k]["timestamp"]
            )
            del self.cache[oldest_key]
            self.embeddings = [
                (k, e) for k, e in self.embeddings if k != oldest_key
            ]
            self.stats["evictions"] += 1

        self.cache[key] = {
            "query": query,
            "response": response,
            "timestamp": time.time(),
        }
        self.embeddings.append((key, query_embedding))

    @property
    def hit_rate(self) -> float:
        total = self.stats["hits"] + self.stats["misses"]
        return self.stats["hits"] / total if total > 0 else 0.0

# Anthropic prompt caching: mark cacheable content with
# cache_control to avoid re-processing stable prefixes.
#
# The system prompt and few-shot examples are cached on first
# request. Subsequent requests with the same prefix pay only
# the cache read cost (typically 90% cheaper than processing).

def build_cached_request(user_query: str) -> dict:
    """Build an API request that leverages prompt caching.

    The system prompt and examples are stable across requests,
    so they benefit from caching. Only the user query changes.
    """
    return {
        "model": "claude-sonnet-4-20250514",
        "max_tokens": 1024,
        "system": [
            {
                "type": "text",
                "text": LONG_SYSTEM_PROMPT,  # 2000+ tokens
                "cache_control": {"type": "ephemeral"},
            }
        ],
        "messages": [
            {"role": "user", "content": user_query},
        ],
    }


# Cost comparison for a 3000-token system prompt:
# Without caching: 3000 input tokens charged at full rate
# With caching (first request): 3000 tokens + small cache write fee
# With caching (subsequent): ~300 token equivalent (90% savings)
# Break-even: typically after 2-3 requests with the same prefix

import hashlib
import json
from typing import Optional, List
import redis


class EmbeddingCache:
    """Cache embeddings by content hash to avoid redundant
    embedding API calls.

    Especially valuable when:
    - Rebuilding vector indices after code changes
    - Processing documents that overlap across pipelines
    - Running evaluation suites with fixed test data
    """

    def __init__(self, redis_url: str = "redis://localhost:6379"):
        self.redis = redis.from_url(redis_url)
        self.prefix = "emb_cache:"

    def _content_hash(self, text: str, model: str) -> str:
        """Hash content + model to create cache key."""
        key_input = f"{model}:{text}"
        return hashlib.sha256(key_input.encode()).hexdigest()

    def get(
        self, text: str, model: str
    ) -> Optional[List[float]]:
        """Retrieve cached embedding if available."""
        key = self.prefix + self._content_hash(text, model)
        cached = self.redis.get(key)
        if cached:
            return json.loads(cached)
        return None

    def put(
        self, text: str, model: str, embedding: List[float],
        ttl_seconds: int = 86400 * 30,  # 30 days default
    ) -> None:
        """Cache an embedding with expiration."""
        key = self.prefix + self._content_hash(text, model)
        self.redis.setex(key, ttl_seconds, json.dumps(embedding))

    def get_or_compute(
        self,
        text: str,
        model: str,
        compute_fn,
    ) -> List[float]:
        """Get from cache or compute and cache."""
        cached = self.get(text, model)
        if cached is not None:
            return cached

        embedding = compute_fn(text, model)
        self.put(text, model, embedding)
        return embedding

import asyncio
from dataclasses import dataclass
from typing import List, Callable, Any
from collections import deque
import time


@dataclass
class BatchItem:
    """A single item in the processing queue."""
    payload: dict
    callback: Callable
    enqueued_at: float = 0.0


class BatchProcessor:
    """Collect requests and process them in batches.

    Flushes when batch is full or max_wait_seconds elapses,
    whichever comes first. This lets you use batch API pricing
    while keeping latency bounded.
    """

    def __init__(
        self,
        process_batch_fn: Callable,
        batch_size: int = 20,
        max_wait_seconds: float = 5.0,
    ):
        self.process_fn = process_batch_fn
        self.batch_size = batch_size
        self.max_wait = max_wait_seconds
        self.queue: deque[BatchItem] = deque()
        self.stats = {
            "batches_processed": 0,
            "items_processed": 0,
            "avg_batch_size": 0.0,
        }

    async def enqueue(self, payload: dict) -> Any:
        """Add a request to the batch queue.

        Returns a future that resolves when the batch
        containing this item is processed.
        """
        future = asyncio.get_event_loop().create_future()
        item = BatchItem(
            payload=payload,
            callback=lambda result: future.set_result(result),
            enqueued_at=time.time(),
        )
        self.queue.append(item)

        # Flush if batch is full
        if len(self.queue) >= self.batch_size:
            await self._flush()

        return await future

    async def _flush(self) -> None:
        """Process all queued items as a single batch."""
        if not self.queue:
            return

        batch = []
        while self.queue and len(batch) < self.batch_size:
            batch.append(self.queue.popleft())

        payloads = [item.payload for item in batch]
        results = await self.process_fn(payloads)

        for item, result in zip(batch, results):
            item.callback(result)

        # Update stats
        self.stats["batches_processed"] += 1
        self.stats["items_processed"] += len(batch)
        total = self.stats["items_processed"]
        batches = self.stats["batches_processed"]
        self.stats["avg_batch_size"] = total / batches

# Example auto-scaling configuration for an inference endpoint.
# These values are starting points -- tune based on your
# traffic patterns and latency requirements.

SCALING_CONFIG = {
    "min_replicas": 1,        # Never scale to zero for prod
    "max_replicas": 10,       # Budget cap
    "target_gpu_utilization": 0.7,  # Scale up above 70%

    # Scale-up: react quickly to traffic spikes
    "scale_up_cooldown_seconds": 60,
    "scale_up_threshold_duration_seconds": 30,

    # Scale-down: be more conservative to avoid thrashing
    "scale_down_cooldown_seconds": 300,
    "scale_down_threshold_duration_seconds": 120,

    # Queue-based scaling: scale on pending requests,
    # not just GPU utilization
    "queue_length_per_replica": 5,
    "max_queue_wait_seconds": 10,
}


def calculate_required_replicas(
    current_rps: float,
    avg_inference_time_ms: float,
    target_utilization: float = 0.7,
) -> int:
    """Estimate required replicas for a target utilization.

    Args:
        current_rps: current requests per second
        avg_inference_time_ms: average inference latency
        target_utilization: target GPU utilization (0-1)
    """
    # Each replica can handle (1000 / avg_inference_time_ms) RPS
    # at 100% utilization
    capacity_per_replica = 1000 / avg_inference_time_ms
    effective_capacity = capacity_per_replica * target_utilization

    import math
    replicas = math.ceil(current_rps / effective_capacity)
    return max(
        SCALING_CONFIG["min_replicas"],
        min(replicas, SCALING_CONFIG["max_replicas"]),
    )

from dataclasses import dataclass, field
from typing import Dict, List, Optional
from collections import defaultdict
import time


@dataclass
class CostAlert:
    """A triggered cost alert."""
    alert_type: str
    message: str
    current_value: float
    threshold: float
    timestamp: float


class CostMonitor:
    """Track and alert on AI spending by feature and model.

    Designed for integration with your observability stack.
    Emit metrics to Prometheus/Datadog/CloudWatch as needed.
    """

    def __init__(self):
        self.hourly_costs: Dict[str, float] = defaultdict(float)
        self.daily_costs: Dict[str, float] = defaultdict(float)
        self.feature_costs: Dict[str, float] = defaultdict(float)
        self.alerts: List[CostAlert] = []

        # Configure thresholds per feature
        self.hourly_thresholds: Dict[str, float] = {}
        self.daily_thresholds: Dict[str, float] = {}

    def record_cost(
        self,
        cost_usd: float,
        model: str,
        feature: str,
    ) -> Optional[CostAlert]:
        """Record a cost event and check thresholds."""
        hour_key = f"{feature}:{int(time.time() // 3600)}"
        day_key = f"{feature}:{int(time.time() // 86400)}"

        self.hourly_costs[hour_key] += cost_usd
        self.daily_costs[day_key] += cost_usd
        self.feature_costs[feature] += cost_usd

        # Check hourly threshold
        hourly_threshold = self.hourly_thresholds.get(feature)
        if hourly_threshold:
            current_hourly = self.hourly_costs[hour_key]
            if current_hourly > hourly_threshold:
                alert = CostAlert(
                    alert_type="hourly_budget_exceeded",
                    message=(
                        f"Feature '{feature}' spent "
                        f"${current_hourly:.2f} this hour "
                        f"(threshold: ${hourly_threshold:.2f})"
                    ),
                    current_value=current_hourly,
                    threshold=hourly_threshold,
                    timestamp=time.time(),
                )
                self.alerts.append(alert)
                return alert

        return None

    def get_cost_summary(self) -> Dict[str, float]:
        """Get current cost summary by feature."""
        return dict(self.feature_costs)

    def set_threshold(
        self,
        feature: str,
        hourly: Optional[float] = None,
        daily: Optional[float] = None,
    ) -> None:
        """Set cost alert thresholds for a feature."""
        if hourly is not None:
            self.hourly_thresholds[feature] = hourly
        if daily is not None:
            self.daily_thresholds[feature] = daily