Ch 2 — Enterprise AI Architecture¶

"""
gateway/auth.py — JWT authentication middleware for the LLM platform gateway.

Requirements:
    pip install fastapi python-jose[cryptography]
"""
from __future__ import annotations

import os
from typing import Annotated

from fastapi import Depends, HTTPException, Security, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError, jwt

security_scheme = HTTPBearer()
SECRET_KEY: str = os.environ["GATEWAY_JWT_SECRET"]
ALGORITHM = "HS256"


def get_token_claims(
    credentials: Annotated[HTTPAuthorizationCredentials, Security(security_scheme)],
) -> dict:
    """
    Validate a JWT Bearer token and return its claims.

    Claims must include: tenant_id, api_key_id, rate_limit_rpm, tier.
    Raises HTTP 401 on invalid or expired tokens.
    """
    try:
        claims = jwt.decode(credentials.credentials, SECRET_KEY, algorithms=[ALGORITHM])
        required = {"tenant_id", "api_key_id", "rate_limit_rpm"}
        if not required.issubset(claims.keys()):
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=f"Token missing required claims: {required - claims.keys()}",
            )
        return claims
    except JWTError as exc:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired token",
            headers={"WWW-Authenticate": "Bearer"},
        ) from exc

"""
gateway/rate_limit.py — Per-tenant sliding window rate limiter backed by Redis.

Requirements:
    pip install redis
"""
from __future__ import annotations

import time

import redis

_redis: redis.Redis = redis.Redis(host="redis", port=6379, decode_responses=True)


class RateLimitExceeded(Exception):
    pass


def enforce_rate_limit(tenant_id: str, rpm_limit: int) -> None:
    """
    Enforce a per-tenant requests-per-minute (RPM) limit using a sliding window.

    The window key rotates every 60 seconds. Two keys are maintained (current + previous)
    to avoid hard resets at window boundaries.

    Raises:
        RateLimitExceeded: if the tenant has exceeded their RPM limit.
    """
    now = int(time.time())
    window = now // 60
    key = f"ratelimit:{tenant_id}:{window}"
    pipe = _redis.pipeline()
    pipe.incr(key)
    pipe.expire(key, 120)       # keep for 2 windows
    count = pipe.execute()[0]

    if count > rpm_limit:
        raise RateLimitExceeded(
            f"Tenant '{tenant_id}' exceeded {rpm_limit} RPM "
            f"(current window: {count} requests)"
        )

"""
gateway/cost_tracker.py — Per-request cost attribution and audit logging.
"""
from __future__ import annotations

import datetime
import sqlite3
from dataclasses import dataclass, field


@dataclass
class UsageRecord:
    tenant_id: str
    api_key_id: str
    model: str
    input_tokens: int
    output_tokens: int
    cost_usd: float
    latency_ms: float
    request_id: str
    use_case: str = ""
    timestamp: str = field(
        default_factory=lambda: datetime.datetime.utcnow().isoformat()
    )


# Model pricing (USD per million tokens) — update from provider pricing pages
MODEL_PRICING: dict[str, dict[str, float]] = {
    "claude-haiku-4-5":    {"input": 0.80,  "output": 4.00},
    "claude-sonnet-4-5":   {"input": 3.00,  "output": 15.00},
    "claude-opus-4-5":     {"input": 15.00, "output": 75.00},
    "gpt-4o-mini":         {"input": 0.15,  "output": 0.60},
    "gpt-4o":              {"input": 2.50,  "output": 10.00},
}


def compute_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """Calculate USD cost for an LLM call."""
    pricing = MODEL_PRICING.get(model, {"input": 0.0, "output": 0.0})
    return (
        input_tokens  * pricing["input"]  / 1_000_000
        + output_tokens * pricing["output"] / 1_000_000
    )

"""
gateway/prompt_guard.py — Heuristic prompt injection detection.
"""
from __future__ import annotations

import re
from typing import Final

INJECTION_PATTERNS: Final[list[re.Pattern]] = [
    re.compile(r"ignore (all |your )?(previous|prior) instructions?", re.I),
    re.compile(r"you are now", re.I),
    re.compile(r"disregard (the )?system prompt", re.I),
    re.compile(r"(reveal|print|repeat|output|show) (your |the )?(system |)prompt", re.I),
    re.compile(r"DAN (mode|activated)", re.I),
    re.compile(r"new persona", re.I),
]
MAX_MESSAGE_CHARS: Final[int] = 16_000


def screen_message(text: str) -> str:
    """
    Screen a user message for prompt injection patterns.

    Returns:
        The original text if clean.

    Raises:
        ValueError: with a user-safe explanation if a pattern is matched.
    """
    if len(text) > MAX_MESSAGE_CHARS:
        raise ValueError(f"Message too long ({len(text)} chars; limit {MAX_MESSAGE_CHARS}).")

    for pattern in INJECTION_PATTERNS:
        if pattern.search(text):
            raise ValueError(
                "Request blocked: the message contains patterns that are not permitted. "
                "Please rephrase your request."
            )
    return text

"""
tenancy/isolation.py — Namespace-based multi-tenant isolation in the vector store.

Requirements:
    pip install qdrant-client
"""
from __future__ import annotations

from qdrant_client import QdrantClient, models

_qdrant = QdrantClient(url="http://qdrant:6333")


def tenant_collection(tenant_id: str) -> str:
    """Derive a deterministic, safe collection name from a tenant ID."""
    safe = tenant_id.replace("-", "_").replace(".", "_").lower()
    return f"docs_{safe}"


def upsert_document(
    tenant_id: str,
    doc_id: str,
    embedding: list[float],
    payload: dict,
) -> None:
    """Store a document in the tenant's isolated vector collection."""
    _qdrant.upsert(
        collection_name=tenant_collection(tenant_id),
        points=[
            models.PointStruct(
                id=doc_id,
                vector=embedding,
                payload={"tenant_id": tenant_id, **payload},
            )
        ],
    )


def search_documents(
    tenant_id: str,
    query_embedding: list[float],
    top_k: int = 5,
) -> list[dict]:
    """Search only within the tenant's namespace."""
    results = _qdrant.search(
        collection_name=tenant_collection(tenant_id),
        query_vector=query_embedding,
        limit=top_k,
        with_payload=True,
    )
    return [
        {"id": r.id, "score": r.score, **r.payload}
        for r in results
    ]

from string import Template


TENANT_PROMPT_TEMPLATES: dict[str, str] = {
    "default": (
        "You are a helpful AI assistant for ${company_name}. "
        "Answer questions accurately and concisely. "
        "If you do not know the answer, say so."
    ),
    "financial_services": (
        "You are a financial information assistant for ${company_name}. "
        "You provide general information only and do not give personalised financial advice. "
        "Always recommend consulting a qualified financial adviser for personal decisions. "
        "Do not discuss competitor products or make price predictions."
    ),
    "healthcare": (
        "You are a healthcare information assistant for ${company_name}. "
        "Provide general health information only. Always advise users to consult a "
        "qualified healthcare professional for medical advice. Do not diagnose conditions "
        "or recommend specific treatments or medications."
    ),
}


def build_system_prompt(
    tenant_id: str,
    company_name: str,
    template_override: str | None = None,
) -> str:
    """
    Build the system prompt for a tenant from their template.

    Args:
        tenant_id:         The tenant's identifier (used to select template).
        company_name:      The company name to interpolate into the template.
        template_override: If provided, use this template instead of the registry.
    """
    template_str = template_override or TENANT_PROMPT_TEMPLATES.get(
        tenant_id, TENANT_PROMPT_TEMPLATES["default"]
    )
    return Template(template_str).safe_substitute(company_name=company_name)

"""
routing/cost_router.py — Cost-optimised LLM routing by query complexity.

Requirements:
    pip install anthropic
"""
from __future__ import annotations

import anthropic

client = anthropic.Anthropic()

# Model tiers: (model_id, cost_per_million_output_tokens)
TIER_CHEAP  = ("claude-haiku-4-5",  4.00)
TIER_MID    = ("claude-sonnet-4-5", 15.00)
TIER_STRONG = ("claude-opus-4-5",  75.00)


def classify_query_complexity(query: str) -> str:
    """
    Classify a query as 'simple', 'moderate', or 'complex' using a lightweight model.

    Simple:   FAQ, definition, short factual question (< 50 words)
    Moderate: Multi-step reasoning, comparison, summarisation (50–200 words)
    Complex:  Code generation, long-form analysis, multi-document synthesis
    """
    response = client.messages.create(
        model="claude-haiku-4-5",    # classifier itself uses the cheap model
        max_tokens=10,
        system=(
            "Classify the following query as exactly one of: simple, moderate, complex. "
            "Reply with only that single word."
        ),
        messages=[{"role": "user", "content": query}],
    )
    label = response.content[0].text.strip().lower()
    return label if label in {"simple", "moderate", "complex"} else "moderate"


def route_and_complete(query: str, system_prompt: str) -> str:
    """
    Route query to the cheapest model that can handle its complexity.

    Returns:
        The model's text response.
    """
    complexity = classify_query_complexity(query)
    model, _ = {
        "simple":   TIER_CHEAP,
        "moderate": TIER_MID,
        "complex":  TIER_STRONG,
    }[complexity]

    response = client.messages.create(
        model=model,
        max_tokens=1024,
        system=system_prompt,
        messages=[{"role": "user", "content": query}],
    )
    return response.content[0].text

from __future__ import annotations
import anthropic
import logging

client = anthropic.Anthropic()
logger = logging.getLogger(__name__)

FALLBACK_CHAIN: list[str] = [
    "claude-sonnet-4-5",     # primary
    "claude-haiku-4-5",      # fallback 1: cheaper, available
    "claude-opus-4-5",       # fallback 2: most capable (for emergency)
]


def complete_with_fallback(
    messages: list[dict],
    system: str,
    max_tokens: int = 1024,
) -> tuple[str, str]:
    """
    Attempt completion with each model in the fallback chain.

    Returns:
        (response_text, model_used)

    Raises:
        RuntimeError: if all models in the chain fail.
    """
    last_error: Exception | None = None

    for model in FALLBACK_CHAIN:
        try:
            response = client.messages.create(
                model=model,
                max_tokens=max_tokens,
                system=system,
                messages=messages,
            )
            if model != FALLBACK_CHAIN[0]:
                logger.warning("Used fallback model %s (primary failed)", model)
            return response.content[0].text, model
        except (anthropic.APIStatusError, anthropic.APIConnectionError) as exc:
            logger.error("Model %s failed: %s", model, exc)
            last_error = exc

    raise RuntimeError(f"All models in fallback chain failed. Last error: {last_error}")

"""
cache/semantic_cache.py — Embedding-based semantic cache using Qdrant.

Requirements:
    pip install qdrant-client anthropic
"""
from __future__ import annotations

import hashlib
import json
import time

import anthropic
from qdrant_client import QdrantClient, models

_qdrant = QdrantClient(url="http://qdrant:6333")
_anthropic = anthropic.Anthropic()
COLLECTION = "semantic_cache"
SIMILARITY_THRESHOLD = 0.97     # tune: higher = fewer hits, lower = more false positives
CACHE_TTL_SECONDS = 86_400      # 24 hours


def _embed(text: str) -> list[float]:
    """Embed a query string using Voyage's embedding API via Anthropic SDK."""
    # In production: use voyage-3-lite for speed and cost at embedding time
    response = _anthropic.messages.create(
        model="claude-haiku-4-5",
        max_tokens=10,
        system="Compute a semantic embedding. (Placeholder — use a real embedding API.)",
        messages=[{"role": "user", "content": text}],
    )
    # ↑ Placeholder: replace with actual embedding call (Voyage, OpenAI, etc.)
    # Real implementation: voyageai.Client().embed([text], model="voyage-3-lite").embeddings[0]
    raise NotImplementedError("Replace with real embedding call")


def cache_lookup(
    query: str,
    tenant_id: str,
) -> str | None:
    """
    Look up a semantically similar cached response for the query.

    Cache hit rate vs staleness tradeoff:
    - Higher SIMILARITY_THRESHOLD → fewer hits, lower false-positive rate
    - Lower TTL → fresher responses, lower hit rate
    - Solution: different TTLs per query category (FAQs: 7d; current events: 1h)

    Returns:
        Cached response string if a hit above threshold is found, else None.
    """
    embedding = _embed(query)
    results = _qdrant.search(
        collection_name=COLLECTION,
        query_vector=embedding,
        query_filter=models.Filter(
            must=[models.FieldCondition(key="tenant_id", match=models.MatchValue(value=tenant_id))]
        ),
        limit=1,
        score_threshold=SIMILARITY_THRESHOLD,
        with_payload=True,
    )
    if not results:
        return None
    cached = results[0].payload
    if time.time() - cached.get("cached_at", 0) > CACHE_TTL_SECONDS:
        return None                    # expired — treat as miss
    return cached.get("response")


def cache_store(query: str, response: str, tenant_id: str) -> None:
    """Store a query-response pair in the semantic cache."""
    embedding = _embed(query)
    point_id = hashlib.md5(f"{tenant_id}:{query}".encode()).hexdigest()
    _qdrant.upsert(
        collection_name=COLLECTION,
        points=[
            models.PointStruct(
                id=point_id,
                vector=embedding,
                payload={
                    "query": query,
                    "response": response,
                    "tenant_id": tenant_id,
                    "cached_at": time.time(),
                },
            )
        ],
    )

"""
prompts/registry.py — Centralised prompt registry with versioning and A/B routing.
"""
from __future__ import annotations

import random
from dataclasses import dataclass, field
import datetime


@dataclass
class PromptVersion:
    name: str
    version: int
    template: str
    description: str
    created_by: str
    tags: list[str] = field(default_factory=list)
    is_production: bool = False
    created_at: str = field(
        default_factory=lambda: datetime.datetime.utcnow().isoformat()
    )


class PromptRegistry:
    """Centralised prompt registry with versioning, promotion, and A/B routing."""

    def __init__(self) -> None:
        self._store: dict[str, dict[int, PromptVersion]] = {}

    def register(self, prompt: PromptVersion) -> None:
        """Register a new prompt version."""
        self._store.setdefault(prompt.name, {})[prompt.version] = prompt

    def get(self, name: str, version: int | None = None) -> PromptVersion:
        """Get a prompt by name; returns the production version if version is None."""
        versions = self._store.get(name, {})
        if not versions:
            raise KeyError(f"Prompt '{name}' not found in registry.")
        if version is not None:
            return versions[version]
        prod = [v for v in versions.values() if v.is_production]
        if prod:
            return prod[-1]
        return max(versions.values(), key=lambda v: v.version)

    def promote(self, name: str, version: int) -> None:
        """Set a version as production; demote all others."""
        for v in self._store.get(name, {}).values():
            v.is_production = False
        self._store[name][version].is_production = True

    def rollback(self, name: str) -> PromptVersion | None:
        """Revert to the second-latest version (the previous production version)."""
        versions = sorted(self._store.get(name, {}).values(), key=lambda v: v.version)
        if len(versions) < 2:
            return None
        self.promote(name, versions[-2].version)
        return versions[-2]

    def ab_route(
        self,
        name: str,
        challenger_version: int,
        challenger_fraction: float = 0.1,
    ) -> tuple[PromptVersion, str]:
        """
        Route traffic between production and challenger prompt.

        Returns:
            (PromptVersion, cohort_label) where cohort_label is 'control' or 'challenger'.
        """
        if random.random() < challenger_fraction:
            return self.get(name, version=challenger_version), "challenger"
        return self.get(name), "control"

"""
eval/pipeline.py — CI/CD evaluation runner for prompt regression testing.

Requirements:
    pip install anthropic pytest
"""
from __future__ import annotations

from dataclasses import dataclass

import anthropic

client = anthropic.Anthropic()


@dataclass
class EvalCase:
    id: str
    user_message: str
    required_keywords: list[str]       # at least one must appear in response
    forbidden_keywords: list[str]      # none of these may appear


@dataclass
class EvalResult:
    case_id: str
    passed: bool
    response: str
    failure_reason: str | None = None


def run_eval(
    system_prompt: str,
    cases: list[EvalCase],
    model: str = "claude-haiku-4-5",
) -> dict:
    """
    Run evaluation cases against a system prompt; return aggregate results.

    Returns:
        dict with pass_rate, total, passed, and failure details.
    """
    results: list[EvalResult] = []

    for case in cases:
        response = client.messages.create(
            model=model,
            max_tokens=512,
            system=system_prompt,
            messages=[{"role": "user", "content": case.user_message}],
        )
        text = response.content[0].text.lower()

        has_required = any(kw.lower() in text for kw in case.required_keywords)
        has_forbidden = any(kw.lower() in text for kw in case.forbidden_keywords)
        passed = has_required and not has_forbidden

        results.append(
            EvalResult(
                case_id=case.id,
                passed=passed,
                response=text,
                failure_reason=(
                    f"Missing: {case.required_keywords}" if not has_required
                    else f"Contains forbidden: {case.forbidden_keywords}" if has_forbidden
                    else None
                ),
            )
        )

    pass_rate = sum(1 for r in results if r.passed) / len(results)
    return {
        "pass_rate": pass_rate,
        "passed": sum(1 for r in results if r.passed),
        "total": len(results),
        "failures": [vars(r) for r in results if not r.passed],
    }


def ci_gate(pass_rate: float, threshold: float = 0.90) -> None:
    """Exit with code 1 (fail CI) if pass_rate falls below threshold."""
    if pass_rate < threshold:
        print(f"EVAL FAILED: {pass_rate:.1%} < {threshold:.1%} threshold")
        raise SystemExit(1)
    print(f"EVAL PASSED: {pass_rate:.1%}")

"""
optimize/prompt_compress.py — Prompt compression using LLMLingua.

Requirements:
    pip install llmlingua
"""
from llmlingua import PromptCompressor

compressor = PromptCompressor(
    model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
    use_llmlingua2=True,
)


def compress_long_context(
    instruction: str,
    context: str,
    question: str,
    target_token_ratio: float = 0.5,
) -> str:
    """
    Compress a long RAG context to approximately target_token_ratio of its original length.

    Args:
        instruction:          The task instruction (not compressed).
        context:              The retrieved context to compress.
        question:             The user query (not compressed).
        target_token_ratio:   Target fraction of original tokens (0.5 = 50% reduction).

    Returns:
        The compressed prompt string ready to send to the LLM.
    """
    result = compressor.compress_prompt(
        context.split("\n"),
        instruction=instruction,
        question=question,
        target_token=int(len(context.split()) * target_token_ratio),
        rank_method="longllmlingua",
    )
    return result["compressed_prompt"]