Enterprise AI Governance: Building the Framework Before You Desperately Need It

import hashlib
import json
import os
import time
import uuid
from dataclasses import asdict, dataclass
from datetime import datetime
from typing import Any, Optional

import boto3  # For S3 Object Lock append-only audit storage
from openai import OpenAI
from anthropic import Anthropic


# --- Audit event schema ---
@dataclass
class AuditEvent:
    event_id: str
    timestamp: str
    user_id: str
    session_id: str
    system_id: str           # Which AI system/workflow generated this call
    provider: str
    model: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    cost_usd: float
    prompt_hash: str         # SHA-256 of full prompt — stored separately for PII handling
    policy_violations: list[str]
    tool_calls: list[dict]
    risk_tier: str           # "high", "limited", "minimal"


# --- Token cost table (update quarterly) ---
TOKEN_COSTS_PER_1M: dict[str, dict[str, float]] = {
    "gpt-4o-2024-11-20": {"input": 2.50, "output": 10.00},
    "gpt-4o-mini-2024-07-18": {"input": 0.15, "output": 0.60},
    "claude-sonnet-4-5-20251101": {"input": 3.00, "output": 15.00},
}


def compute_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    costs = TOKEN_COSTS_PER_1M.get(model, {"input": 0.0, "output": 0.0})
    return (input_tokens * costs["input"] + output_tokens * costs["output"]) / 1_000_000


# --- Policy violation detection ---
POLICY_PATTERNS = [
    ("pii_request", ["social security", "ssn", "credit card number", "passport number"]),
    ("jailbreak_attempt", ["ignore previous instructions", "disregard your system prompt", "act as dan"]),
    ("sensitive_data", ["internal salary", "confidential", "top secret", "eyes only"]),
]


def detect_policy_violations(prompt: str) -> list[str]:
    violations = []
    prompt_lower = prompt.lower()
    for violation_type, patterns in POLICY_PATTERNS:
        if any(p in prompt_lower for p in patterns):
            violations.append(violation_type)
    return violations


# --- Audit writer (S3 Object Lock for tamper-evidence) ---
class S3AuditWriter:
    def __init__(self, bucket: str, prefix: str = "ai-audit-logs/"):
        self.s3 = boto3.client("s3")
        self.bucket = bucket
        self.prefix = prefix

    def write(self, event: AuditEvent) -> None:
        key = f"{self.prefix}{event.timestamp[:10]}/{event.event_id}.json"
        self.s3.put_object(
            Bucket=self.bucket,
            Key=key,
            Body=json.dumps(asdict(event), default=str).encode(),
            ContentType="application/json",
            # Object Lock COMPLIANCE mode: immutable for retention_days
        )


# --- Governance middleware ---
class GovernanceLLMClient:
    """
    Drop-in wrapper around OpenAI/Anthropic clients that enforces:
    - Full audit logging with user attribution
    - Cost tracking per call
    - Policy violation detection
    - Risk tier tagging
    """

    def __init__(
        self,
        audit_writer: S3AuditWriter,
        system_id: str,
        risk_tier: str = "limited",  # "high", "limited", "minimal"
    ):
        self.openai = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
        self.anthropic = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
        self.audit_writer = audit_writer
        self.system_id = system_id
        self.risk_tier = risk_tier

    def complete(
        self,
        messages: list[dict],
        model: str,
        user_id: str,
        session_id: str,
        max_tokens: int = 2048,
    ) -> dict:
        full_prompt = json.dumps(messages)
        violations = detect_policy_violations(full_prompt)

        if self.risk_tier == "high" and violations:
            # Hard block for high-risk systems with policy violations
            raise PermissionError(f"Policy violation detected: {violations}. Request blocked for high-risk system.")

        start_time = time.time()
        if model.startswith("gpt") or model.startswith("o3"):
            response = self.openai.chat.completions.create(
                model=model, messages=messages, max_tokens=max_tokens
            )
            content = response.choices[0].message.content
            tool_calls = [
                {"name": tc.function.name, "arguments": tc.function.arguments}
                for tc in (response.choices[0].message.tool_calls or [])
            ]
            input_tokens = response.usage.prompt_tokens
            output_tokens = response.usage.completion_tokens
        else:
            # Anthropic path
            system_msg = next((m["content"] for m in messages if m["role"] == "system"), "")
            user_messages = [m for m in messages if m["role"] != "system"]
            response = self.anthropic.messages.create(
                model=model, max_tokens=max_tokens,
                system=system_msg, messages=user_messages,
            )
            content = response.content[0].text
            tool_calls = []
            input_tokens = response.usage.input_tokens
            output_tokens = response.usage.output_tokens

        latency_ms = (time.time() - start_time) * 1000
        cost = compute_cost(model, input_tokens, output_tokens)

        event = AuditEvent(
            event_id=str(uuid.uuid4()),
            timestamp=datetime.utcnow().isoformat(),
            user_id=user_id,
            session_id=session_id,
            system_id=self.system_id,
            provider="openai" if model.startswith(("gpt", "o3")) else "anthropic",
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            latency_ms=round(latency_ms, 2),
            cost_usd=round(cost, 6),
            prompt_hash=hashlib.sha256(full_prompt.encode()).hexdigest(),
            policy_violations=violations,
            tool_calls=tool_calls,
            risk_tier=self.risk_tier,
        )
        self.audit_writer.write(event)

        if violations:
            # Log violation but allow for non-high-risk systems — alert downstream
            print(f"[GOVERNANCE] Policy violations detected for user {user_id}: {violations}")

        return {"content": content, "cost_usd": cost, "event_id": event.event_id}