Networking and AI API Gateway Design#

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from typing import Optional
import asyncio
import time

# Educational example — not for clinical use

class ModelTier(Enum):
    FRONTIER = "frontier"        # claude-opus-4-8, gpt-4o — highest capability, highest cost
    STANDARD = "standard"        # claude-sonnet-4-6, gpt-4o-mini — balanced
    EFFICIENT = "efficient"      # claude-haiku-4-5, gpt-3.5-turbo — lowest cost
    SELF_HOSTED = "self_hosted"  # On-premises GPU inference


@dataclass
class TeamConfig:
    """Per-team rate limits and model access control."""
    team_id: str
    monthly_token_budget: int           # Total tokens per month across all models
    tokens_per_minute_limit: int        # Rate limit (token-level, not request-level)
    allowed_model_tiers: list[ModelTier]
    allowed_use_cases: list[str]        # e.g., ["clinical_rag", "admin_summarization"]
    phi_access: bool = False            # Whether this team can process PHI


@dataclass
class CircuitBreakerState:
    provider: str
    is_open: bool = False              # True = circuit is broken (blocking requests)
    failure_count: int = 0
    last_failure_time: Optional[datetime] = None
    half_open_until: Optional[datetime] = None


class AIAPIGateway:
    """
    Enterprise AI API gateway with authentication, rate limiting,
    intelligent routing, circuit breaking, and PHI-safe logging.
    
    Educational example — not for clinical use.
    """
    
    def __init__(
        self,
        team_configs: dict[str, TeamConfig],
        circuit_breaker_threshold: int = 5,    # failures before opening circuit
        circuit_breaker_timeout: int = 60,     # seconds to wait before half-open
    ):
        self.team_configs = team_configs
        self.circuit_breakers: dict[str, CircuitBreakerState] = {
            "anthropic": CircuitBreakerState("anthropic"),
            "azure_openai": CircuitBreakerState("azure_openai"),
            "vertex": CircuitBreakerState("vertex"),
            "self_hosted": CircuitBreakerState("self_hosted"),
        }
        self.cb_threshold = circuit_breaker_threshold
        self.cb_timeout = circuit_breaker_timeout
        self._token_usage: dict[str, list[tuple[datetime, int]]] = {}  # team_id → [(timestamp, tokens)]
    
    async def route_request(
        self,
        team_id: str,
        messages: list[dict],
        requested_model_tier: ModelTier,
        use_case: str,
        max_tokens: int = 1024,
        phi_context: bool = False,     # Caller must declare if request may contain PHI
    ) -> dict:
        """
        Route an AI request through the gateway with full policy enforcement.
        Returns the LLM response dict.
        """
        
        # 1. Authentication and authorization
        team_config = self._authenticate_team(team_id)
        self._authorize_request(team_config, requested_model_tier, use_case, phi_context)
        
        # 2. Rate limiting (token-level)
        self._check_rate_limit(team_id, team_config, max_tokens)
        
        # 3. Intelligent model selection
        provider, model = self._select_model(requested_model_tier, team_config)
        
        # 4. Circuit breaker check
        self._check_circuit_breaker(provider)
        
        # 5. Execute request with circuit breaker tracking
        try:
            response = await self._execute_request(provider, model, messages, max_tokens)
            self._record_success(provider)
        except Exception as e:
            self._record_failure(provider)
            raise
        
        # 6. Record token usage for budget tracking
        total_tokens = response.get("usage", {}).get("total_tokens", max_tokens)
        self._record_token_usage(team_id, total_tokens)
        
        # 7. Emit metrics
        await self._emit_metrics(team_id, provider, model, use_case, response)
        
        return response
    
    def _authenticate_team(self, team_id: str) -> TeamConfig:
        """Validate team API key and return team configuration."""
        if team_id not in self.team_configs:
            raise PermissionError(f"Unknown team: {team_id}")
        return self.team_configs[team_id]
    
    def _authorize_request(
        self,
        team_config: TeamConfig,
        model_tier: ModelTier,
        use_case: str,
        phi_context: bool,
    ) -> None:
        """Enforce model and use-case authorization."""
        if model_tier not in team_config.allowed_model_tiers:
            raise PermissionError(
                f"Team {team_config.team_id} is not authorized for {model_tier.value} models"
            )
        if use_case not in team_config.allowed_use_cases:
            raise PermissionError(
                f"Team {team_config.team_id} is not authorized for use case: {use_case}"
            )
        if phi_context and not team_config.phi_access:
            raise PermissionError(
                f"Team {team_config.team_id} is not authorized to process PHI. "
                "Contact AI Platform team to request PHI access with appropriate compliance review."
            )
    
    def _check_rate_limit(
        self,
        team_id: str,
        team_config: TeamConfig,
        requested_tokens: int,
    ) -> None:
        """Enforce token-per-minute rate limit using sliding window."""
        now = datetime.utcnow()
        window_start = now - timedelta(minutes=1)
        
        if team_id not in self._token_usage:
            self._token_usage[team_id] = []
        
        # Remove entries outside the sliding window
        self._token_usage[team_id] = [
            (ts, tokens) for ts, tokens in self._token_usage[team_id]
            if ts >= window_start
        ]
        
        current_window_tokens = sum(tokens for _, tokens in self._token_usage[team_id])
        
        if current_window_tokens + requested_tokens > team_config.tokens_per_minute_limit:
            raise Exception(
                f"Rate limit exceeded for team {team_id}. "
                f"Used {current_window_tokens}/{team_config.tokens_per_minute_limit} tokens/min."
            )
    
    def _check_circuit_breaker(self, provider: str) -> None:
        """Check circuit breaker state for the target provider."""
        cb = self.circuit_breakers[provider]
        
        if not cb.is_open:
            return
        
        # Check if half-open window has passed
        if cb.half_open_until and datetime.utcnow() >= cb.half_open_until:
            # Allow one test request through (half-open state)
            return
        
        raise Exception(
            f"Circuit breaker is open for provider {provider}. "
            f"Provider is experiencing failures. "
            f"Retry after {cb.half_open_until.isoformat() if cb.half_open_until else 'N/A'}."
        )
    
    def _record_success(self, provider: str) -> None:
        """Reset circuit breaker on success."""
        cb = self.circuit_breakers[provider]
        cb.is_open = False
        cb.failure_count = 0
        cb.half_open_until = None
    
    def _record_failure(self, provider: str) -> None:
        """Record failure; open circuit if threshold exceeded."""
        cb = self.circuit_breakers[provider]
        cb.failure_count += 1
        cb.last_failure_time = datetime.utcnow()
        
        if cb.failure_count >= self.cb_threshold:
            cb.is_open = True
            cb.half_open_until = datetime.utcnow() + timedelta(seconds=self.cb_timeout)
    
    def _select_model(
        self,
        requested_tier: ModelTier,
        team_config: TeamConfig,
    ) -> tuple[str, str]:
        """
        Select provider and model based on tier, provider health, and team access.
        
        Model ID mapping is illustrative — verify current model IDs in official docs.
        """
        tier_to_models = {
            ModelTier.FRONTIER: [
                ("anthropic", "claude-opus-4-8"),
                ("azure_openai", "gpt-4o"),
            ],
            ModelTier.STANDARD: [
                ("anthropic", "claude-sonnet-4-6"),
                ("azure_openai", "gpt-4o-mini"),
            ],
            ModelTier.EFFICIENT: [
                ("anthropic", "claude-haiku-4-5-20251001"),
            ],
            ModelTier.SELF_HOSTED: [
                ("self_hosted", "clinical-llm"),
            ],
        }
        
        candidates = tier_to_models.get(requested_tier, [])
        
        for provider, model in candidates:
            cb = self.circuit_breakers.get(provider)
            if cb and cb.is_open:
                continue  # Skip providers with open circuit breakers
            return provider, model
        
        raise Exception(
            f"No healthy provider available for tier {requested_tier.value}. "
            f"All providers may be circuit-broken."
        )
    
    async def _execute_request(
        self,
        provider: str,
        model: str,
        messages: list[dict],
        max_tokens: int,
    ) -> dict:
        """Execute the request against the selected provider."""
        if provider == "anthropic":
            return await self._call_anthropic(model, messages, max_tokens)
        elif provider == "azure_openai":
            return await self._call_azure_openai(model, messages, max_tokens)
        elif provider == "self_hosted":
            return await self._call_self_hosted(model, messages, max_tokens)
        raise ValueError(f"Unknown provider: {provider}")
    
    def _record_token_usage(self, team_id: str, tokens: int) -> None:
        """Record token consumption for rate limiting and budget tracking."""
        self._token_usage.setdefault(team_id, []).append((datetime.utcnow(), tokens))
    
    async def _emit_metrics(
        self,
        team_id: str,
        provider: str,
        model: str,
        use_case: str,
        response: dict,
    ) -> None:
        """Emit metrics for observability. PHI must not appear in metric labels."""
        # Implementation: increment Prometheus counters
        # Labels: team_id, provider, model, use_case (never include request content)
        pass
    
    # Provider-specific callers omitted for brevity — see cloud-ai-platforms.md
    async def _call_anthropic(self, model, messages, max_tokens): ...
    async def _call_azure_openai(self, model, messages, max_tokens): ...
    async def _call_self_hosted(self, model, messages, max_tokens): ...

import re
from typing import Any

# HIPAA Safe Harbor redaction patterns (18 identifiers)
# This is a simplified illustrative example — consult a healthcare data expert
# for production PHI de-identification.

PHI_REDACTION_PATTERNS = [
    # Names (very rough — context-dependent; NLP-based detection is more reliable)
    (r'\bDr\.\s+[A-Z][a-z]+\s+[A-Z][a-z]+\b', '[PROVIDER_NAME]'),
    # MRN patterns
    (r'\bMRN[:\s#]*\d{6,10}\b', '[MRN]'),
    # Date of birth
    (r'\bDOB[:\s]*\d{1,2}[/-]\d{1,2}[/-]\d{2,4}\b', '[DOB]'),
    # Phone numbers
    (r'\b\d{3}[-.\s]?\d{3}[-.\s]?\d{4}\b', '[PHONE]'),
    # SSN
    (r'\b\d{3}-\d{2}-\d{4}\b', '[SSN]'),
    # Email
    (r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b', '[EMAIL]'),
    # ZIP codes
    (r'\b\d{5}(?:-\d{4})?\b', '[ZIP]'),
]


def redact_phi_from_log_entry(content: str) -> str:
    """
    Apply PHI redaction patterns to a string before logging.
    
    This is a simplified pattern-based redactor. Production healthcare
    deployments should use a trained NER model or dedicated PHI detection
    service for more accurate de-identification.
    
    Educational Example — Not for clinical use.
    """
    redacted = content
    for pattern, replacement in PHI_REDACTION_PATTERNS:
        redacted = re.sub(pattern, replacement, redacted, flags=re.IGNORECASE)
    return redacted


def build_phi_safe_log_entry(
    team_id: str,
    provider: str,
    model: str,
    use_case: str,
    prompt_tokens: int,
    completion_tokens: int,
    latency_ms: float,
    phi_context: bool,
) -> dict:
    """
    Build a log entry that is safe to write to the logging pipeline.
    Never include request or response content in logs for PHI-handling teams.
    """
    return {
        "timestamp": datetime.utcnow().isoformat(),
        "team_id": team_id,
        "provider": provider,
        "model": model,
        "use_case": use_case,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
        "latency_ms": round(latency_ms, 1),
        "phi_context": phi_context,
        # NEVER log: messages, response content, system prompt content (may contain PHI)
    }