LLM Denial
of Service

import tiktoken

MAX_INPUT_TOKENS = 2000  # Set appropriate for your use case

def validate_prompt(user_message: str) -> str:
    enc = tiktoken.encoding_for_model("gpt-4")
    tokens = enc.encode(user_message)
    
    if len(tokens) > MAX_INPUT_TOKENS:
        raise ValueError(
            f"Input too long: {len(tokens)} tokens. "
            f"Maximum allowed: {MAX_INPUT_TOKENS}"
        )
    return user_message

# Usage
try:
    safe_prompt = validate_prompt(user_input)
    response = llm.chat(safe_prompt)
except ValueError as e:
    return {"error": str(e), "code": 400}  # Reject cleanly

import openai

# ❌ VULNERABLE — no max_tokens limit
response = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": user_prompt}]
)

# ✅ FIXED — hard limit set
response = client.chat.completions.create(
    model="gpt-4",
    messages=[{"role": "user", "content": user_prompt}],
    max_tokens=500,          # Customer service: 500 max
    timeout=30,              # Hard timeout in seconds
)

# For different use cases, use different limits:
# Code review tool → max_tokens=2000
# Summarization   → max_tokens=500
# Q&A chatbot     → max_tokens=300

import redis
from functools import wraps

r = redis.Redis()

def rate_limit(user_id: str, tokens_used: int) -> bool:
    """Returns False if user has exceeded limits."""
    pipe = r.pipeline()
    
    # Requests per minute: max 10
    req_key = f"rate:req:{user_id}:{int(time.time()//60)}"
    pipe.incr(req_key)
    pipe.expire(req_key, 60)
    
    # Tokens per hour: max 50,000
    tok_key = f"rate:tok:{user_id}:{int(time.time()//3600)}"
    pipe.incrby(tok_key, tokens_used)
    pipe.expire(tok_key, 3600)
    
    results = pipe.execute()
    
    if results[0] > 10:  # Request limit
        raise RateLimitError("Too many requests. Try again in 1 minute.")
    if results[2] > 50000:  # Token limit
        raise RateLimitError("Token quota exceeded for this hour.")
    
    return True

DAILY_BUDGET_USD = 50.0
COST_PER_1K_TOKENS = 0.002  # gpt-4o pricing example

def track_cost(tokens_used: int) -> None:
    cost = (tokens_used / 1000) * COST_PER_1K_TOKENS
    
    daily_key = f"cost:daily:{datetime.date.today()}"
    new_total = r.incrbyfloat(daily_key, cost)
    r.expire(daily_key, 86400)
    
    if new_total > DAILY_BUDGET_USD * 0.8:
        alert_team(f"⚠ 80% of daily LLM budget used: ${new_total:.2f}")
    
    if new_total > DAILY_BUDGET_USD:
        # Circuit breaker: stop ALL LLM calls
        r.set("llm:circuit_breaker", "open", ex=3600)
        raise BudgetExceededError("Daily LLM budget exceeded. Service paused.")

def call_llm(prompt: str):
    if r.get("llm:circuit_breaker") == b"open":
        raise ServiceUnavailableError("AI service temporarily paused.")
    # ... make API call

MAX_DOCUMENT_TOKENS = 3000
MAX_DOCUMENTS = 3

def safe_inject_documents(docs: list[str]) -> str:
    """Safely inject user documents into prompt context."""
    
    # Limit number of documents
    docs = docs[:MAX_DOCUMENTS]
    
    safe_parts = []
    total_tokens = 0
    enc = tiktoken.encoding_for_model("gpt-4")
    
    for i, doc in enumerate(docs):
        tokens = enc.encode(doc)
        
        # Truncate if too long
        if len(tokens) > MAX_DOCUMENT_TOKENS:
            tokens = tokens[:MAX_DOCUMENT_TOKENS]
            doc = enc.decode(tokens) + "\n[DOCUMENT TRUNCATED]"
        
        total_tokens += len(tokens)
        safe_parts.append(f"Document {i+1}:\n{doc}")
    
    return "\n\n".join(safe_parts)

def detect_anomaly(user_id: str, tokens: int) -> bool:
    """Returns True if usage looks anomalous."""
    
    # Get rolling average for this user (last 7 days)
    history = get_user_token_history(user_id, days=7)
    
    if len(history) < 5:
        return False  # Not enough data
    
    avg = sum(history) / len(history)
    std = statistics.stdev(history)
    
    # Flag if current usage > 3 standard deviations above average
    z_score = (tokens - avg) / (std + 1)  # +1 to avoid div by zero
    
    if z_score > 3.0:
        log_security_event({
            "type": "llm_dos_suspected",
            "user_id": user_id,
            "tokens": tokens,
            "avg_tokens": avg,
            "z_score": z_score
        })
        return True
    
    return False