control-zero 0.2.0__py3-none-any.whl

control_zero/llm/base.py

@@ -0,0 +1,551 @@
+"""
+Base classes for LLM governance wrappers.
+
+This module provides the foundational classes that all LLM provider
+wrappers inherit from, ensuring consistent governance behavior across
+OpenAI, Anthropic, Gemini, and other providers.
+"""
+
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Generic
+from enum import Enum
+
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+
+
+class GovernanceAction(str, Enum):
+    """Actions that can be governed for LLM calls."""
+    CHAT_COMPLETION = "chat_completion"
+    TEXT_COMPLETION = "text_completion"
+    EMBEDDING = "embedding"
+    FUNCTION_CALL = "function_call"
+    TOOL_USE = "tool_use"
+    IMAGE_GENERATION = "image_generation"
+    AUDIO_TRANSCRIPTION = "audio_transcription"
+    MODERATION = "moderation"
+
+
+@dataclass
+class ModelPolicy:
+    """Policy for model access control."""
+    allowed_models: List[str] = field(default_factory=list)
+    denied_models: List[str] = field(default_factory=list)
+    default_model: Optional[str] = None
+    max_tokens_per_request: Optional[int] = None
+    max_context_length: Optional[int] = None
+
+
+@dataclass
+class CostPolicy:
+    """Policy for cost control."""
+    max_cost_per_request: Optional[float] = None
+    max_cost_per_day: Optional[float] = None
+    max_cost_per_month: Optional[float] = None
+    max_tokens_per_day: Optional[int] = None
+    max_requests_per_minute: Optional[int] = None
+    max_requests_per_day: Optional[int] = None
+
+
+@dataclass
+class FunctionPolicy:
+    """Policy for function/tool calling."""
+    allowed_functions: List[str] = field(default_factory=list)
+    denied_functions: List[str] = field(default_factory=list)
+    require_approval: List[str] = field(default_factory=list)
+    max_function_calls_per_request: Optional[int] = None
+
+
+@dataclass
+class ContentPolicy:
+    """Policy for content filtering."""
+    enable_pii_detection: bool = False
+    pii_action: str = "mask"  # mask, block, warn
+    blocked_topics: List[str] = field(default_factory=list)
+    required_system_prompt: Optional[str] = None
+    max_output_tokens: Optional[int] = None
+
+
+@dataclass
+class LLMGovernanceConfig:
+    """Complete governance configuration for LLM calls."""
+    model_policy: ModelPolicy = field(default_factory=ModelPolicy)
+    cost_policy: CostPolicy = field(default_factory=CostPolicy)
+    function_policy: FunctionPolicy = field(default_factory=FunctionPolicy)
+    content_policy: ContentPolicy = field(default_factory=ContentPolicy)
+    enable_audit_logging: bool = True
+    enable_streaming_governance: bool = True
+
+
+@dataclass
+class LLMUsageMetrics:
+    """Metrics for a single LLM call."""
+    provider: str
+    model: str
+    action: GovernanceAction
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    latency_ms: int = 0
+    estimated_cost: float = 0.0
+    function_calls: int = 0
+    cached: bool = False
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "provider": self.provider,
+            "model": self.model,
+            "action": self.action.value,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "total_tokens": self.total_tokens,
+            "latency_ms": self.latency_ms,
+            "estimated_cost": self.estimated_cost,
+            "function_calls": self.function_calls,
+            "cached": self.cached,
+        }
+
+
+# Model pricing (approximate, per 1M tokens)
+MODEL_PRICING = {
+    # OpenAI
+    "gpt-4-turbo": {"input": 10.0, "output": 30.0},
+    "gpt-4": {"input": 30.0, "output": 60.0},
+    "gpt-4o": {"input": 2.50, "output": 10.0},
+    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
+    "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
+    "o1-preview": {"input": 15.0, "output": 60.0},
+    "o1-mini": {"input": 3.0, "output": 12.0},
+    # Anthropic
+    "claude-3-opus": {"input": 15.0, "output": 75.0},
+    "claude-3-sonnet": {"input": 3.0, "output": 15.0},
+    "claude-3-haiku": {"input": 0.25, "output": 1.25},
+    "claude-3.5-sonnet": {"input": 3.0, "output": 15.0},
+    "claude-opus-4": {"input": 15.0, "output": 75.0},
+    # Google
+    "gemini-1.5-pro": {"input": 3.50, "output": 10.50},
+    "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
+    "gemini-2.0-flash": {"input": 0.10, "output": 0.40},
+    # Groq
+    "llama-3.1-70b": {"input": 0.59, "output": 0.79},
+    "llama-3.1-8b": {"input": 0.05, "output": 0.08},
+    "mixtral-8x7b": {"input": 0.24, "output": 0.24},
+    # Mistral
+    "mistral-large": {"input": 2.0, "output": 6.0},
+    "mistral-small": {"input": 0.2, "output": 0.6},
+    "codestral": {"input": 0.2, "output": 0.6},
+}
+
+
+def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
+    """Estimate cost for a model call."""
+    # Normalize model name
+    model_lower = model.lower()
+
+    # Find matching pricing
+    for model_key, pricing in MODEL_PRICING.items():
+        if model_key in model_lower:
+            input_cost = (input_tokens / 1_000_000) * pricing["input"]
+            output_cost = (output_tokens / 1_000_000) * pricing["output"]
+            return input_cost + output_cost
+
+    # Default estimate if model not found
+    return (input_tokens + output_tokens) / 1_000_000 * 1.0
+
+
+T = TypeVar("T")  # Original client type
+
+
+class GovernedLLM(ABC, Generic[T]):
+    """
+    Base class for governed LLM clients.
+
+    This class wraps any LLM provider's client and adds:
+    - Policy enforcement (model access, function calling, content)
+    - Cost tracking and limits
+    - Audit logging
+    - PII detection and masking
+
+    Subclasses implement provider-specific wrapping logic.
+    """
+
+    def __init__(
+        self,
+        client: T,
+        control_zero: Any,  # ControlZeroClient
+        config: Optional[LLMGovernanceConfig] = None,
+        user_context: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize a governed LLM client.
+
+        Args:
+            client: The original LLM provider client (OpenAI, Anthropic, etc.)
+            control_zero: Control Zero client for policy and logging
+            config: Governance configuration (optional, uses defaults if not provided)
+            user_context: Context about the current user (user_id, role, etc.)
+        """
+        self._client = client
+        self._cz = control_zero
+        self._config = config or LLMGovernanceConfig()
+        self._user_context = user_context or {}
+
+        # Usage tracking
+        self._session_metrics: List[LLMUsageMetrics] = []
+        self._daily_cost: float = 0.0
+        self._daily_tokens: int = 0
+        self._daily_requests: int = 0
+        self._last_reset: float = time.time()
+
+    @property
+    @abstractmethod
+    def provider_name(self) -> str:
+        """Return the provider name (e.g., 'openai', 'anthropic')."""
+        pass
+
+    @property
+    def client(self) -> T:
+        """Access the underlying client."""
+        return self._client
+
+    def _check_daily_reset(self) -> None:
+        """Reset daily counters if a new day has started."""
+        current_time = time.time()
+        # Reset if more than 24 hours since last reset
+        if current_time - self._last_reset > 86400:
+            self._daily_cost = 0.0
+            self._daily_tokens = 0
+            self._daily_requests = 0
+            self._last_reset = current_time
+
+    def _check_model_policy(self, model: str) -> PolicyDecision:
+        """Check if the model is allowed."""
+        policy = self._config.model_policy
+
+        # Check denied list first
+        if policy.denied_models:
+            for denied in policy.denied_models:
+                if denied.lower() in model.lower():
+                    return PolicyDecision(
+                        effect="deny",
+                        reason=f"Model '{model}' is not allowed by policy",
+                        policy_id="model_deny_list"
+                    )
+
+        # Check allowed list if specified
+        if policy.allowed_models:
+            allowed = False
+            for allowed_model in policy.allowed_models:
+                if allowed_model.lower() in model.lower():
+                    allowed = True
+                    break
+            if not allowed:
+                return PolicyDecision(
+                    effect="deny",
+                    reason=f"Model '{model}' is not in the allowed models list",
+                    policy_id="model_allow_list"
+                )
+
+        return PolicyDecision(effect="allow")
+
+    def _check_cost_policy(self, estimated_cost: float) -> PolicyDecision:
+        """Check if the request is within cost limits."""
+        self._check_daily_reset()
+        policy = self._config.cost_policy
+
+        # Check per-request limit
+        if policy.max_cost_per_request:
+            if estimated_cost > policy.max_cost_per_request:
+                return PolicyDecision(
+                    effect="deny",
+                    reason=f"Estimated cost ${estimated_cost:.4f} exceeds per-request limit ${policy.max_cost_per_request:.4f}",
+                    policy_id="cost_per_request"
+                )
+
+        # Check daily limit
+        if policy.max_cost_per_day:
+            if self._daily_cost + estimated_cost > policy.max_cost_per_day:
+                return PolicyDecision(
+                    effect="deny",
+                    reason=f"Daily cost limit ${policy.max_cost_per_day:.2f} would be exceeded",
+                    policy_id="cost_per_day"
+                )
+
+        # Check daily request limit
+        if policy.max_requests_per_day:
+            if self._daily_requests >= policy.max_requests_per_day:
+                return PolicyDecision(
+                    effect="deny",
+                    reason=f"Daily request limit of {policy.max_requests_per_day} reached",
+                    policy_id="requests_per_day"
+                )
+
+        return PolicyDecision(effect="allow")
+
+    def _check_function_policy(self, functions: List[Dict[str, Any]]) -> PolicyDecision:
+        """Check if function calls are allowed."""
+        policy = self._config.function_policy
+
+        for func in functions:
+            func_name = func.get("name", func.get("function", {}).get("name", ""))
+
+            # Check denied list
+            if policy.denied_functions:
+                for denied in policy.denied_functions:
+                    if denied.lower() in func_name.lower():
+                        return PolicyDecision(
+                            effect="deny",
+                            reason=f"Function '{func_name}' is not allowed",
+                            policy_id="function_deny_list"
+                        )
+
+            # Check allowed list if specified
+            if policy.allowed_functions:
+                allowed = False
+                for allowed_func in policy.allowed_functions:
+                    if allowed_func.lower() in func_name.lower():
+                        allowed = True
+                        break
+                if not allowed:
+                    return PolicyDecision(
+                        effect="deny",
+                        reason=f"Function '{func_name}' is not in the allowed list",
+                        policy_id="function_allow_list"
+                    )
+
+        return PolicyDecision(effect="allow")
+
+    def _detect_pii(self, text: str) -> List[Dict[str, Any]]:
+        """
+        Simple PII detection.
+        In production, use a more sophisticated solution.
+        """
+        import re
+
+        patterns = {
+            "email": r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
+            "phone": r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b',
+            "ssn": r'\b\d{3}[-]?\d{2}[-]?\d{4}\b',
+            "credit_card": r'\b(?:\d{4}[-\s]?){3}\d{4}\b',
+            "ip_address": r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b',
+        }
+
+        findings = []
+        for pii_type, pattern in patterns.items():
+            matches = re.findall(pattern, text)
+            for match in matches:
+                findings.append({"type": pii_type, "value": match})
+
+        return findings
+
+    def _mask_pii(self, text: str) -> str:
+        """Mask detected PII in text."""
+        findings = self._detect_pii(text)
+        masked = text
+        for finding in findings:
+            pii_type = finding["type"]
+            value = finding["value"]
+            mask = f"[{pii_type.upper()}_MASKED]"
+            masked = masked.replace(value, mask)
+        return masked
+
+    def _pre_request_checks(
+        self,
+        model: str,
+        action: GovernanceAction,
+        messages: Optional[List[Dict[str, Any]]] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        estimated_tokens: int = 0,
+    ) -> None:
+        """
+        Run all pre-request governance checks.
+        Raises PolicyDeniedError if any check fails.
+        """
+        # Check model policy
+        decision = self._check_model_policy(model)
+        if decision.effect == "deny":
+            self._log_denied(model, action, decision)
+            raise PolicyDeniedError(decision)
+
+        # Estimate cost and check cost policy
+        estimated_cost = estimate_cost(model, estimated_tokens, estimated_tokens // 2)
+        decision = self._check_cost_policy(estimated_cost)
+        if decision.effect == "deny":
+            self._log_denied(model, action, decision)
+            raise PolicyDeniedError(decision)
+
+        # Check function policy if functions provided
+        if functions:
+            decision = self._check_function_policy(functions)
+            if decision.effect == "deny":
+                self._log_denied(model, action, decision)
+                raise PolicyDeniedError(decision)
+
+        # Check content policy - PII detection
+        if self._config.content_policy.enable_pii_detection and messages:
+            for msg in messages:
+                content = msg.get("content", "")
+                if isinstance(content, str):
+                    findings = self._detect_pii(content)
+                    if findings:
+                        pii_action = self._config.content_policy.pii_action
+                        if pii_action == "block":
+                            decision = PolicyDecision(
+                                effect="deny",
+                                reason=f"PII detected in request: {[f['type'] for f in findings]}",
+                                policy_id="pii_detection"
+                            )
+                            self._log_denied(model, action, decision)
+                            raise PolicyDeniedError(decision)
+
+    def _post_request_update(self, metrics: LLMUsageMetrics) -> None:
+        """Update tracking after a successful request."""
+        self._session_metrics.append(metrics)
+        self._daily_cost += metrics.estimated_cost
+        self._daily_tokens += metrics.total_tokens
+        self._daily_requests += 1
+
+    def _log_request(
+        self,
+        model: str,
+        action: GovernanceAction,
+        metrics: LLMUsageMetrics,
+        status: str = "success",
+        error: Optional[str] = None,
+    ) -> None:
+        """Log request to Control Zero."""
+        if not self._config.enable_audit_logging:
+            return
+
+        try:
+            # Use Control Zero's logging
+            self._cz._log(
+                tool=f"llm:{self.provider_name}",
+                method=action.value,
+                status=status,
+                latency_ms=metrics.latency_ms,
+                error_type="LLMError" if error else None,
+                error_message=error,
+            )
+        except Exception:
+            pass  # Don't fail on logging errors
+
+    def _log_denied(
+        self,
+        model: str,
+        action: GovernanceAction,
+        decision: PolicyDecision,
+    ) -> None:
+        """Log a denied request."""
+        if not self._config.enable_audit_logging:
+            return
+
+        try:
+            self._cz._log(
+                tool=f"llm:{self.provider_name}",
+                method=action.value,
+                status="denied",
+                latency_ms=0,
+                policy_decision=decision,
+            )
+        except Exception:
+            pass
+
+    def get_session_metrics(self) -> Dict[str, Any]:
+        """Get aggregated metrics for the current session."""
+        total_cost = sum(m.estimated_cost for m in self._session_metrics)
+        total_tokens = sum(m.total_tokens for m in self._session_metrics)
+        total_requests = len(self._session_metrics)
+
+        return {
+            "provider": self.provider_name,
+            "total_cost": total_cost,
+            "total_tokens": total_tokens,
+            "total_requests": total_requests,
+            "daily_cost": self._daily_cost,
+            "daily_tokens": self._daily_tokens,
+            "daily_requests": self._daily_requests,
+        }
+
+    def reset_session_metrics(self) -> None:
+        """Reset session metrics."""
+        self._session_metrics = []
+
+
+class GovernedChatMixin:
+    """Mixin providing governed chat completion functionality."""
+
+    def _process_messages_for_governance(
+        self,
+        messages: List[Dict[str, Any]],
+    ) -> List[Dict[str, Any]]:
+        """Process messages according to governance policies."""
+        processed = []
+
+        # Add required system prompt if configured
+        content_policy = getattr(self, '_config', LLMGovernanceConfig()).content_policy
+        if content_policy.required_system_prompt:
+            has_system = any(m.get("role") == "system" for m in messages)
+            if not has_system:
+                processed.append({
+                    "role": "system",
+                    "content": content_policy.required_system_prompt
+                })
+
+        # Process each message
+        for msg in messages:
+            new_msg = msg.copy()
+            content = msg.get("content", "")
+
+            # Mask PII if enabled
+            if content_policy.enable_pii_detection and content_policy.pii_action == "mask":
+                if isinstance(content, str):
+                    governed_self = getattr(self, '_mask_pii', lambda x: x)
+                    new_msg["content"] = governed_self(content)
+
+            processed.append(new_msg)
+
+        return processed
+
+    def _filter_functions_for_governance(
+        self,
+        functions: Optional[List[Dict[str, Any]]],
+    ) -> Optional[List[Dict[str, Any]]]:
+        """Filter functions/tools according to governance policies."""
+        if not functions:
+            return functions
+
+        config = getattr(self, '_config', LLMGovernanceConfig())
+        policy = config.function_policy
+
+        if not policy.allowed_functions and not policy.denied_functions:
+            return functions
+
+        filtered = []
+        for func in functions:
+            func_name = func.get("name", func.get("function", {}).get("name", ""))
+
+            # Skip denied
+            if policy.denied_functions:
+                denied = False
+                for denied_name in policy.denied_functions:
+                    if denied_name.lower() in func_name.lower():
+                        denied = True
+                        break
+                if denied:
+                    continue
+
+            # Check allowed
+            if policy.allowed_functions:
+                allowed = False
+                for allowed_name in policy.allowed_functions:
+                    if allowed_name.lower() in func_name.lower():
+                        allowed = True
+                        break
+                if not allowed:
+                    continue
+
+            filtered.append(func)
+
+        return filtered if filtered else None

control_zero/llm/cohere/__init__.py

@@ -0,0 +1,32 @@
+"""
+Control Zero Cohere Governance Wrapper.
+
+This module provides governance wrappers for the Cohere Python SDK,
+enabling policy enforcement, cost tracking, and audit logging for
+Cohere API calls including chat, RAG, and reranking.
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.cohere import GovernedCohere
+    import cohere
+
+    # Initialize Control Zero
+    cz_client = ControlZeroClient(api_key="cz_live_xxx")
+    cz_client.initialize()
+
+    # Wrap Cohere client with governance
+    co_client = cohere.ClientV2()
+    governed = GovernedCohere(client=co_client, control_zero=cz_client)
+
+    # All calls are now governed
+    response = governed.chat(
+        model="command-r-plus",
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+"""
+
+from control_zero.llm.cohere.client import GovernedCohere
+
+__all__ = [
+    "GovernedCohere",
+]