Enterprise LLM Deployment: Scaling AI in Production with Security & Performance

# Enterprise LLM Gateway - Multi-Model Orchestration
from fastapi import FastAPI, HTTPException, Depends
from pydantic import BaseModel
import asyncio
import time
from typing import Optional, Dict, List
from enum import Enum
import logging
from dataclasses import dataclass

app = FastAPI(title="Enterprise LLM Gateway")

# Model tier configuration
@dataclass
class ModelConfig:
    name: str
    endpoint: str
    max_tokens: int
    cost_per_token: float
    latency_p95: float  # milliseconds
    availability: float
    context_window: int
    
class ModelTier(Enum):
    PREMIUM = "premium"      # GPT-4, Claude-3 Opus
    STANDARD = "standard"    # GPT-3.5 Turbo, Claude-3 Sonnet  
    EFFICIENT = "efficient"  # Llama-2, Mistral
    SPECIALIZED = "specialized"  # Code, Math, etc.

# Model registry with fallback chains
MODEL_REGISTRY = {
    ModelTier.PREMIUM: [
        ModelConfig("gpt-4-1106-preview", "https://api.openai.com/v1", 
                   128000, 0.03, 2000, 0.999, 128000),
        ModelConfig("claude-3-opus", "https://api.anthropic.com/v1",
                   200000, 0.015, 1500, 0.998, 200000)
    ],
    ModelTier.STANDARD: [
        ModelConfig("gpt-3.5-turbo-1106", "https://api.openai.com/v1",
                   16385, 0.002, 800, 0.999, 16385),
        ModelConfig("claude-3-sonnet", "https://api.anthropic.com/v1",
                   200000, 0.003, 1000, 0.999, 200000)
    ],
    ModelTier.EFFICIENT: [
        ModelConfig("llama-2-70b-chat", "http://internal-llm.company.com/v1",
                   4096, 0.0005, 1200, 0.995, 4096),
        ModelConfig("mistral-7b-instruct", "http://internal-llm.company.com/v1",
                   8192, 0.0002, 600, 0.997, 8192)
    ]
}

class LLMRequest(BaseModel):
    prompt: str
    max_tokens: Optional[int] = 1000
    temperature: Optional[float] = 0.7
    tier_preference: ModelTier = ModelTier.STANDARD
    user_id: str
    organization_id: str
    use_case: Optional[str] = None
    priority: Optional[int] = 5  # 1-10, higher = more important

class LLMGateway:
    def __init__(self):
        self.request_counts = {}  # Rate limiting
        self.model_health = {}    # Health tracking
        self.cost_tracker = {}    # Cost monitoring
        
    async def route_request(self, request: LLMRequest) -> Dict:
        """Intelligent request routing based on requirements and availability"""
        
        # 1. Check rate limits and quotas
        if not await self.check_rate_limits(request.user_id, request.organization_id):
            raise HTTPException(429, "Rate limit exceeded")
            
        # 2. Select optimal model based on request characteristics
        selected_model = await self.select_model(request)
        
        # 3. Execute with fallback chain
        response = await self.execute_with_fallback(request, selected_model)
        
        # 4. Track metrics and costs
        await self.track_request_metrics(request, response, selected_model)
        
        return response
    
    async def select_model(self, request: LLMRequest) -> ModelConfig:
        """Smart model selection based on request requirements"""
        
        # Get models from preferred tier
        candidate_models = MODEL_REGISTRY.get(request.tier_preference, [])
        
        # Filter by context window requirements
        prompt_tokens = len(request.prompt) // 4  # Rough estimation
        candidates = [m for m in candidate_models 
                     if m.context_window >= prompt_tokens + request.max_tokens]
        
        if not candidates:
            # Fallback to larger context models
            for tier in [ModelTier.PREMIUM, ModelTier.STANDARD, ModelTier.EFFICIENT]:
                candidates = [m for m in MODEL_REGISTRY.get(tier, [])
                            if m.context_window >= prompt_tokens + request.max_tokens]
                if candidates:
                    break
                    
        if not candidates:
            raise HTTPException(400, "Request too large for available models")
        
        # Score models based on availability, cost, and latency
        best_model = None
        best_score = -1
        
        for model in candidates:
            health = self.model_health.get(model.name, {"availability": model.availability})
            
            # Scoring algorithm (weights can be configured per organization)
            score = (
                health["availability"] * 0.4 +           # Availability weight
                (1 / (model.cost_per_token * 1000)) * 0.3 +  # Cost efficiency
                (1 / (model.latency_p95 / 1000)) * 0.2 +     # Speed
                (request.priority / 10) * 0.1                # Priority boost
            )
            
            if score > best_score:
                best_score = score
                best_model = model
                
        return best_model
    
    async def execute_with_fallback(self, request: LLMRequest, 
                                  primary_model: ModelConfig) -> Dict:
        """Execute request with automatic fallback on failure"""
        
        # Build fallback chain
        fallback_chain = [primary_model]
        
        # Add models from same tier
        same_tier_models = MODEL_REGISTRY.get(request.tier_preference, [])
        fallback_chain.extend([m for m in same_tier_models if m != primary_model])
        
        # Add models from lower tiers if needed
        if request.tier_preference != ModelTier.EFFICIENT:
            fallback_chain.extend(MODEL_REGISTRY.get(ModelTier.EFFICIENT, []))
            
        for i, model in enumerate(fallback_chain):
            try:
                start_time = time.time()
                response = await self.call_model(model, request)
                latency = (time.time() - start_time) * 1000
                
                return {
                    "response": response,
                    "model_used": model.name,
                    "tier_used": request.tier_preference.value,
                    "latency_ms": latency,
                    "fallback_level": i,
                    "tokens_used": len(response.split()) * 1.3  # Rough estimation
                }
                
            except Exception as e:
                logging.warning(f"Model {model.name} failed: {e}")
                if i == len(fallback_chain) - 1:
                    # All models failed
                    raise HTTPException(503, f"All models unavailable: {e}")
                continue
    
    async def call_model(self, model: ModelConfig, request: LLMRequest) -> str:
        """Call specific model API with timeout and error handling"""
        
        # Implementation would vary by model provider
        # This is a simplified example
        
        payload = {
            "messages": [{"role": "user", "content": request.prompt}],
            "max_tokens": request.max_tokens,
            "temperature": request.temperature
        }
        
        timeout = 30.0  # 30 second timeout
        
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout)) as session:
            headers = self.get_auth_headers(model)
            
            async with session.post(f"{model.endpoint}/chat/completions", 
                                  json=payload, headers=headers) as response:
                if response.status != 200:
                    raise Exception(f"API error: {response.status}")
                    
                result = await response.json()
                return result["choices"][0]["message"]["content"]
    
    async def check_rate_limits(self, user_id: str, org_id: str) -> bool:
        """Check user and organization rate limits"""
        
        current_time = time.time()
        window_size = 3600  # 1 hour window
        
        # Check user limits (e.g., 1000 requests/hour)
        user_key = f"user:{user_id}"
        user_requests = self.request_counts.get(user_key, [])
        user_requests = [t for t in user_requests if current_time - t < window_size]
        
        if len(user_requests) >= 1000:
            return False
            
        # Check org limits (e.g., 50000 requests/hour)
        org_key = f"org:{org_id}"
        org_requests = self.request_counts.get(org_key, [])
        org_requests = [t for t in org_requests if current_time - t < window_size]
        
        if len(org_requests) >= 50000:
            return False
            
        # Update counts
        user_requests.append(current_time)
        org_requests.append(current_time)
        self.request_counts[user_key] = user_requests
        self.request_counts[org_key] = org_requests
        
        return True
    
    async def track_request_metrics(self, request: LLMRequest, 
                                  response: Dict, model: ModelConfig):
        """Track usage metrics for monitoring and billing"""
        
        metrics = {
            "timestamp": time.time(),
            "user_id": request.user_id,
            "organization_id": request.organization_id,
            "model_used": response["model_used"],
            "tokens_used": response["tokens_used"],
            "latency_ms": response["latency_ms"],
            "cost_usd": response["tokens_used"] * model.cost_per_token,
            "use_case": request.use_case,
            "fallback_level": response["fallback_level"]
        }
        
        # Store in metrics database (implementation depends on your setup)
        await self.store_metrics(metrics)

# Initialize gateway
llm_gateway = LLMGateway()

@app.post("/v1/chat/completions")
async def chat_completion(request: LLMRequest):
    """Main endpoint for LLM requests"""
    try:
        response = await llm_gateway.route_request(request)
        return response
    except HTTPException:
        raise
    except Exception as e:
        logging.error(f"Unexpected error: {e}")
        raise HTTPException(500, "Internal server error")

@app.get("/health")
async def health_check():
    """Health check endpoint for load balancers"""
    # Check model health, database connectivity, etc.
    return {"status": "healthy", "models_available": len(MODEL_REGISTRY)}

@app.get("/metrics")
async def get_metrics():
    """Metrics endpoint for monitoring"""
    return {
        "requests_per_hour": len(llm_gateway.request_counts),
        "model_health": llm_gateway.model_health,
        "cost_tracking": llm_gateway.cost_tracker
    }

# Kubernetes deployment for enterprise LLM infrastructure
apiVersion: v1
kind: Namespace
metadata:
  name: llm-platform
  labels:
    security-policy: "high"
    data-classification: "confidential"

---
# ConfigMap for model configurations
apiVersion: v1
kind: ConfigMap
metadata:
  name: llm-config
  namespace: llm-platform
data:
  models.yaml: |
    premium_models:
      - name: "gpt-4-turbo"
        replicas: 3
        resources:
          requests:
            memory: "32Gi"
            cpu: "8"
            nvidia.com/gpu: "1"
          limits:
            memory: "64Gi"
            cpu: "16"
            nvidia.com/gpu: "1"
        node_selector:
          instance-type: "gpu-optimized"
          model-tier: "premium"
          
    standard_models:
      - name: "gpt-3.5-turbo"
        replicas: 5
        resources:
          requests:
            memory: "16Gi"
            cpu: "4"
          limits:
            memory: "32Gi"
            cpu: "8"
        node_selector:
          instance-type: "cpu-optimized"
          model-tier: "standard"

---
# LLM Gateway Deployment
apiVersion: apps/v1
kind: Deployment
metadata:
  name: llm-gateway
  namespace: llm-platform
  labels:
    app: llm-gateway
    tier: api
spec:
  replicas: 6
  strategy:
    type: RollingUpdate
    rollingUpdate:
      maxUnavailable: 1
      maxSurge: 2
  selector:
    matchLabels:
      app: llm-gateway
  template:
    metadata:
      labels:
        app: llm-gateway
        tier: api
    spec:
      securityContext:
        runAsNonRoot: true
        runAsUser: 1000
        fsGroup: 1000
      containers:
      - name: gateway
        image: company.registry.com/llm-gateway:v2.1.0
        ports:
        - containerPort: 8000
          name: api
        - containerPort: 8080
          name: metrics  
        env:
        - name: DATABASE_URL
          valueFrom:
            secretKeyRef:
              name: llm-secrets
              key: database-url
        - name: REDIS_URL
          valueFrom:
            secretKeyRef:
              name: llm-secrets
              key: redis-url
        - name: OPENAI_API_KEY
          valueFrom:
            secretKeyRef:
              name: model-api-keys
              key: openai-key
        resources:
          requests:
            memory: "2Gi"
            cpu: "1"
          limits:
            memory: "4Gi"
            cpu: "2"
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 30
          periodSeconds: 10
        readinessProbe:
          httpGet:
            path: /ready
            port: 8000
          initialDelaySeconds: 5
          periodSeconds: 5
        volumeMounts:
        - name: config
          mountPath: /app/config
          readOnly: true
        - name: cache
          mountPath: /app/cache
      volumes:
      - name: config
        configMap:
          name: llm-config
      - name: cache
        emptyDir:
          sizeLimit: "1Gi"

---
# Horizontal Pod Autoscaler
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: llm-gateway-hpa
  namespace: llm-platform
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: llm-gateway
  minReplicas: 6
  maxReplicas: 50
  metrics:
  - type: Resource
    resource:
      name: cpu
      target:
        type: Utilization
        averageUtilization: 70
  - type: Resource
    resource:
      name: memory
      target:
        type: Utilization
        averageUtilization: 80
  - type: Pods
    pods:
      metric:
        name: requests_per_second
      target:
        type: AverageValue
        averageValue: "100"

---
# Service for LLM Gateway
apiVersion: v1
kind: Service
metadata:
  name: llm-gateway
  namespace: llm-platform
  labels:
    app: llm-gateway
spec:
  selector:
    app: llm-gateway
  ports:
  - name: api
    port: 80
    targetPort: 8000
    protocol: TCP
  - name: metrics
    port: 8080
    targetPort: 8080
    protocol: TCP
  type: ClusterIP

---
# Ingress with SSL termination and rate limiting
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
  name: llm-gateway-ingress
  namespace: llm-platform
  annotations:
    kubernetes.io/ingress.class: "nginx"
    cert-manager.io/cluster-issuer: "letsencrypt-prod"
    nginx.ingress.kubernetes.io/rate-limit: "1000"
    nginx.ingress.kubernetes.io/rate-limit-window: "1m"
    nginx.ingress.kubernetes.io/ssl-redirect: "true"
    nginx.ingress.kubernetes.io/backend-protocol: "HTTP"
    nginx.ingress.kubernetes.io/proxy-connect-timeout: "30"
    nginx.ingress.kubernetes.io/proxy-send-timeout: "120"
    nginx.ingress.kubernetes.io/proxy-read-timeout: "120"
spec:
  tls:
  - hosts:
    - llm-api.company.com
    secretName: llm-gateway-tls
  rules:
  - host: llm-api.company.com
    http:
      paths:
      - path: /
        pathType: Prefix
        backend:
          service:
            name: llm-gateway
            port:
              number: 80

---
# Network Policy for security
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
  name: llm-network-policy
  namespace: llm-platform
spec:
  podSelector:
    matchLabels:
      app: llm-gateway
  policyTypes:
  - Ingress
  - Egress
  ingress:
  - from:
    - namespaceSelector:
        matchLabels:
          name: "ingress-nginx"
    ports:
    - protocol: TCP
      port: 8000
  egress:
  - to: []
    ports:
    - protocol: TCP
      port: 443  # HTTPS to external APIs
    - protocol: TCP
      port: 5432 # PostgreSQL
    - protocol: TCP
      port: 6379 # Redis

# Enterprise AI Security & Compliance Framework
from typing import List, Dict, Optional, Enum
from dataclasses import dataclass
from datetime import datetime
import hashlib
import re
import logging
from functools import wraps

class DataClassification(Enum):
    PUBLIC = "public"
    INTERNAL = "internal"  
    CONFIDENTIAL = "confidential"
    RESTRICTED = "restricted"

class ComplianceFramework(Enum):
    GDPR = "gdpr"
    HIPAA = "hipaa"
    SOX = "sox"
    PCI_DSS = "pci_dss"
    SOC2 = "soc2"

@dataclass
class SecurityContext:
    user_id: str
    organization_id: str
    access_level: str
    data_classification: DataClassification
    compliance_frameworks: List[ComplianceFramework]
    geographic_region: str
    session_id: str

class DataSanitizer:
    """Sanitize and classify data before LLM processing"""
    
    # PII detection patterns
    PII_PATTERNS = {
        'ssn': r'd{3}-?d{2}-?d{4}',
        'email': r'[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+.[A-Z|a-z]{2,}',
        'phone': r'(+d{1,3}[- ]?)?d{10}',
        'credit_card': r'(?:d{4}[-s]?){3}d{4}',
        'ip_address': r'(?:d{1,3}.){3}d{1,3}',
        'api_key': r'[A-Za-z0-9]{32,}'
    }
    
    MEDICAL_PATTERNS = {
        'medical_record': r'MR[N]?s*:?s*d+',
        'diagnosis_code': r'[A-Z]d{2}(.d{1,2})?',
        'medication': r'd+mg|d+s*ml'
    }
    
    def __init__(self):
        self.audit_log = []
    
    def classify_content(self, content: str, context: SecurityContext) -> DataClassification:
        """Automatically classify content based on detected patterns"""
        
        # Check for PII patterns
        pii_detected = []
        for pii_type, pattern in self.PII_PATTERNS.items():
            if re.search(pattern, content, re.IGNORECASE):
                pii_detected.append(pii_type)
        
        # Check for medical information if HIPAA compliance required
        medical_detected = []
        if ComplianceFramework.HIPAA in context.compliance_frameworks:
            for medical_type, pattern in self.MEDICAL_PATTERNS.items():
                if re.search(pattern, content, re.IGNORECASE):
                    medical_detected.append(medical_type)
        
        # Determine classification level
        if medical_detected or 'ssn' in pii_detected or 'credit_card' in pii_detected:
            classification = DataClassification.RESTRICTED
        elif pii_detected:
            classification = DataClassification.CONFIDENTIAL
        elif any(keyword in content.lower() for keyword in 
               ['proprietary', 'confidential', 'internal only']):
            classification = DataClassification.CONFIDENTIAL
        else:
            classification = DataClassification.INTERNAL
            
        # Log classification decision
        self.audit_log.append({
            'timestamp': datetime.utcnow(),
            'user_id': context.user_id,
            'session_id': context.session_id,
            'classification': classification.value,
            'pii_detected': pii_detected,
            'medical_detected': medical_detected,
            'content_hash': hashlib.sha256(content.encode()).hexdigest()[:16]
        })
        
        return classification
    
    def sanitize_for_llm(self, content: str, classification: DataClassification,
                        context: SecurityContext) -> str:
        """Sanitize content based on classification and compliance requirements"""
        
        sanitized = content
        
        if classification in [DataClassification.RESTRICTED, DataClassification.CONFIDENTIAL]:
            # Replace PII with placeholders
            for pii_type, pattern in self.PII_PATTERNS.items():
                sanitized = re.sub(pattern, f'[{pii_type.upper()}_REDACTED]', 
                                 sanitized, flags=re.IGNORECASE)
            
            # Medical data sanitization for HIPAA
            if ComplianceFramework.HIPAA in context.compliance_frameworks:
                for medical_type, pattern in self.MEDICAL_PATTERNS.items():
                    sanitized = re.sub(pattern, f'[{medical_type.upper()}_REDACTED]', 
                                     sanitized, flags=re.IGNORECASE)
        
        return sanitized

class AccessController:
    """Role-based access control for LLM operations"""
    
    def __init__(self):
        self.role_permissions = {
            'admin': ['create', 'read', 'update', 'delete', 'manage_users'],
            'data_scientist': ['create', 'read', 'update', 'analyze'],
            'business_user': ['read', 'limited_create'],
            'readonly': ['read']
        }
        
        self.data_access_matrix = {
            DataClassification.PUBLIC: ['admin', 'data_scientist', 'business_user', 'readonly'],
            DataClassification.INTERNAL: ['admin', 'data_scientist', 'business_user'],
            DataClassification.CONFIDENTIAL: ['admin', 'data_scientist'],
            DataClassification.RESTRICTED: ['admin']
        }
    
    def check_access(self, context: SecurityContext, operation: str, 
                    data_classification: DataClassification) -> bool:
        """Check if user has permission for the operation on data classification"""
        
        user_role = self.get_user_role(context.user_id, context.organization_id)
        
        # Check operation permission
        if operation not in self.role_permissions.get(user_role, []):
            logging.warning(f"Access denied: {context.user_id} lacks {operation} permission")
            return False
            
        # Check data classification access
        if user_role not in self.data_access_matrix.get(data_classification, []):
            logging.warning(f"Access denied: {context.user_id} cannot access {data_classification.value} data")
            return False
            
        # Geographic restrictions (GDPR example)
        if (ComplianceFramework.GDPR in context.compliance_frameworks and 
            context.geographic_region not in ['EU', 'EEA']):
            logging.warning(f"Access denied: GDPR data accessed from {context.geographic_region}")
            return False
            
        return True
    
    def get_user_role(self, user_id: str, org_id: str) -> str:
        """Get user role from identity provider"""
        # Implementation would integrate with your identity provider
        # This is a simplified example
        return "business_user"  # Default role

class ComplianceMonitor:
    """Monitor and ensure compliance across LLM operations"""
    
    def __init__(self):
        self.compliance_rules = {
            ComplianceFramework.GDPR: {
                'data_retention_days': 90,
                'require_consent': True,
                'allow_automated_decisions': False,
                'data_portability': True
            },
            ComplianceFramework.HIPAA: {
                'data_retention_days': 2555,  # 7 years
                'require_encryption': True,
                'audit_trail_required': True,
                'minimum_access_logging': True
            },
            ComplianceFramework.SOC2: {
                'availability_sla': 99.9,
                'security_controls': True,
                'processing_integrity': True,
                'confidentiality': True
            }
        }
    
    def validate_request(self, content: str, context: SecurityContext) -> Dict[str, bool]:
        """Validate request against applicable compliance frameworks"""
        
        violations = {}
        
        for framework in context.compliance_frameworks:
            rules = self.compliance_rules.get(framework, {})
            
            if framework == ComplianceFramework.GDPR:
                # Check for automated decision making
                if ('decision' in content.lower() and 
                    not rules.get('allow_automated_decisions', True)):
                    violations[f"{framework.value}_automated_decision"] = False
                    
            elif framework == ComplianceFramework.HIPAA:
                # Check for PHI handling requirements
                if any(pattern in content.lower() for pattern in 
                      ['patient', 'diagnosis', 'treatment', 'medical']):
                    if not self.verify_encryption_in_transit():
                        violations[f"{framework.value}_encryption"] = False
                        
        return violations
    
    def verify_encryption_in_transit(self) -> bool:
        """Verify that data is encrypted in transit"""
        # Implementation would check TLS configuration
        return True

class AuditLogger:
    """Comprehensive audit logging for compliance"""
    
    def __init__(self):
        self.logger = logging.getLogger('llm_audit')
    
    def log_llm_request(self, context: SecurityContext, content_hash: str,
                       model_used: str, response_hash: str,
                       classification: DataClassification):
        """Log LLM request for audit trail"""
        
        audit_record = {
            'timestamp': datetime.utcnow().isoformat(),
            'event_type': 'llm_request',
            'user_id': context.user_id,
            'organization_id': context.organization_id,
            'session_id': context.session_id,
            'content_hash': content_hash,
            'model_used': model_used,
            'response_hash': response_hash,
            'data_classification': classification.value,
            'compliance_frameworks': [f.value for f in context.compliance_frameworks],
            'geographic_region': context.geographic_region
        }
        
        self.logger.info(f"AUDIT: {audit_record}")
        
        # Store in secure audit database
        self.store_audit_record(audit_record)
    
    def store_audit_record(self, record: Dict):
        """Store audit record in tamper-evident storage"""
        # Implementation would use append-only database or blockchain
        pass

# Security middleware decorator
def require_security_context(func):
    """Decorator to enforce security context on LLM operations"""
    
    @wraps(func)
    async def wrapper(*args, **kwargs):
        # Extract security context from request
        context = kwargs.get('security_context')
        if not context:
            raise ValueError("Security context required")
            
        # Initialize security components
        sanitizer = DataSanitizer()
        access_controller = AccessController()
        compliance_monitor = ComplianceMonitor()
        audit_logger = AuditLogger()
        
        # Extract content from request
        content = kwargs.get('content', '')
        
        # Classify and sanitize content
        classification = sanitizer.classify_content(content, context)
        sanitized_content = sanitizer.sanitize_for_llm(content, classification, context)
        
        # Check access permissions
        if not access_controller.check_access(context, 'create', classification):
            raise PermissionError("Access denied for this operation")
            
        # Validate compliance
        violations = compliance_monitor.validate_request(content, context)
        if violations:
            raise ValueError(f"Compliance violations: {violations}")
            
        # Update kwargs with sanitized content
        kwargs['content'] = sanitized_content
        kwargs['classification'] = classification
        
        # Execute the function
        result = await func(*args, **kwargs)
        
        # Log audit trail
        content_hash = hashlib.sha256(content.encode()).hexdigest()[:16]
        response_hash = hashlib.sha256(str(result).encode()).hexdigest()[:16]
        
        audit_logger.log_llm_request(
            context, content_hash, kwargs.get('model_used', 'unknown'),
            response_hash, classification
        )
        
        return result
    
    return wrapper