control-zero 0.2.0__py3-none-any.whl

control_zero/llm/__init__.py

@@ -0,0 +1,77 @@
+"""
+Control Zero LLM Governance Module.
+
+This module provides governance wrappers for vanilla LLM API calls,
+enabling policy enforcement, cost tracking, and audit logging
+without requiring agent frameworks like LangChain or CrewAI.
+
+Supported providers:
+- OpenAI (GPT-4, GPT-3.5, etc.)
+- Anthropic (Claude)
+- Google Gemini
+- Groq
+- Mistral
+- Cohere
+- Ollama (local LLMs)
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.openai import GovernedOpenAI
+    from openai import OpenAI
+
+    # Initialize Control Zero client
+    cz_client = ControlZeroClient(api_key="cz_live_xxx")
+    cz_client.initialize()
+
+    # Wrap OpenAI client with governance
+    openai_client = OpenAI()
+    governed = GovernedOpenAI(client=openai_client, control_zero=cz_client)
+
+    # All calls are now governed
+    response = governed.chat.completions.create(
+        model="gpt-4",
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+"""
+
+from control_zero.llm.base import (
+    GovernedLLM,
+    LLMGovernanceConfig,
+    LLMUsageMetrics,
+    ModelPolicy,
+    CostPolicy,
+    FunctionPolicy,
+    ContentPolicy,
+    GovernanceAction,
+    estimate_cost,
+)
+
+# Provider-specific imports
+from control_zero.llm.openai import GovernedOpenAI
+from control_zero.llm.anthropic import GovernedAnthropic
+from control_zero.llm.gemini import GovernedGemini
+from control_zero.llm.groq import GovernedGroq
+from control_zero.llm.mistral import GovernedMistral
+from control_zero.llm.cohere import GovernedCohere
+from control_zero.llm.ollama import GovernedOllama
+
+__all__ = [
+    # Base classes
+    "GovernedLLM",
+    "LLMGovernanceConfig",
+    "LLMUsageMetrics",
+    "ModelPolicy",
+    "CostPolicy",
+    "FunctionPolicy",
+    "ContentPolicy",
+    "GovernanceAction",
+    "estimate_cost",
+    # Provider wrappers
+    "GovernedOpenAI",
+    "GovernedAnthropic",
+    "GovernedGemini",
+    "GovernedGroq",
+    "GovernedMistral",
+    "GovernedCohere",
+    "GovernedOllama",
+]

control_zero/llm/anthropic/__init__.py

@@ -0,0 +1,35 @@
+"""
+Control Zero Anthropic Governance Wrapper.
+
+This module provides governance wrappers for the Anthropic Python SDK,
+enabling policy enforcement, cost tracking, and audit logging for
+all Claude API calls.
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.anthropic import GovernedAnthropic
+    from anthropic import Anthropic
+
+    # Initialize Control Zero
+    cz_client = ControlZeroClient(api_key="cz_live_xxx")
+    cz_client.initialize()
+
+    # Wrap Anthropic client with governance
+    anthropic_client = Anthropic()
+    governed = GovernedAnthropic(client=anthropic_client, control_zero=cz_client)
+
+    # All calls are now governed
+    response = governed.messages.create(
+        model="claude-3-sonnet-20240229",
+        max_tokens=1024,
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+"""
+
+from control_zero.llm.anthropic.client import GovernedAnthropic
+from control_zero.llm.anthropic.messages import GovernedMessages
+
+__all__ = [
+    "GovernedAnthropic",
+    "GovernedMessages",
+]

control_zero/llm/anthropic/client.py

@@ -0,0 +1,136 @@
+"""
+Governed Anthropic client wrapper.
+
+Provides governance features for the Anthropic Python SDK including:
+- Model access control
+- Cost tracking and limits
+- Tool use governance
+- PII detection and masking
+- Audit logging
+"""
+
+from typing import Any, Dict, Optional
+
+from control_zero.llm.base import (
+    GovernedLLM,
+    LLMGovernanceConfig,
+)
+from control_zero.llm.anthropic.messages import GovernedMessages
+
+
+class GovernedAnthropic(GovernedLLM):
+    """
+    Governed wrapper for the Anthropic Python SDK.
+
+    This class wraps an Anthropic client instance and adds governance
+    features including policy enforcement, cost tracking, and audit logging.
+
+    Example:
+        from control_zero import ControlZeroClient
+        from control_zero.llm.anthropic import GovernedAnthropic
+        from anthropic import Anthropic
+
+        cz = ControlZeroClient(api_key="...")
+        cz.initialize()
+
+        client = Anthropic()
+        governed = GovernedAnthropic(client=client, control_zero=cz)
+
+        # Configure governance
+        from control_zero.llm import LLMGovernanceConfig, ModelPolicy, CostPolicy
+
+        governed = GovernedAnthropic(
+            client=client,
+            control_zero=cz,
+            config=LLMGovernanceConfig(
+                model_policy=ModelPolicy(
+                    allowed_models=["claude-3-sonnet", "claude-3-haiku"],
+                    max_tokens_per_request=4096
+                ),
+                cost_policy=CostPolicy(
+                    max_cost_per_day=10.00,
+                    max_requests_per_day=1000
+                )
+            ),
+            user_context={"user_id": "user_123", "role": "developer"}
+        )
+
+        # Make governed API call
+        response = governed.messages.create(
+            model="claude-3-sonnet-20240229",
+            max_tokens=1024,
+            messages=[{"role": "user", "content": "Hello!"}]
+        )
+    """
+
+    def __init__(
+        self,
+        client: Any,  # anthropic.Anthropic
+        control_zero: Any,  # ControlZeroClient
+        config: Optional[LLMGovernanceConfig] = None,
+        user_context: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize a governed Anthropic client.
+
+        Args:
+            client: An Anthropic client instance
+            control_zero: Control Zero client for policy and logging
+            config: Governance configuration
+            user_context: Context about the current user
+        """
+        super().__init__(client, control_zero, config, user_context)
+
+        # Create governed namespaces
+        self._messages = GovernedMessages(self)
+
+    @property
+    def provider_name(self) -> str:
+        return "anthropic"
+
+    @property
+    def messages(self) -> GovernedMessages:
+        """Access governed messages API."""
+        return self._messages
+
+    @property
+    def completions(self):
+        """Access the underlying completions endpoint (legacy, pass-through)."""
+        return self._client.completions
+
+    def with_user_context(self, user_context: Dict[str, Any]) -> "GovernedAnthropic":
+        """
+        Create a new governed client with updated user context.
+
+        This is useful for per-request user context in multi-tenant applications.
+
+        Args:
+            user_context: New user context (merged with existing)
+
+        Returns:
+            New GovernedAnthropic instance with updated context
+        """
+        merged_context = {**self._user_context, **user_context}
+        return GovernedAnthropic(
+            client=self._client,
+            control_zero=self._cz,
+            config=self._config,
+            user_context=merged_context,
+        )
+
+    def with_config(self, config: LLMGovernanceConfig) -> "GovernedAnthropic":
+        """
+        Create a new governed client with updated configuration.
+
+        Args:
+            config: New governance configuration
+
+        Returns:
+            New GovernedAnthropic instance with updated config
+        """
+        return GovernedAnthropic(
+            client=self._client,
+            control_zero=self._cz,
+            config=config,
+            user_context=self._user_context,
+        )

control_zero/llm/anthropic/messages.py

@@ -0,0 +1,375 @@
+"""
+Governed Anthropic Messages API.
+
+Provides governance wrapper for Anthropic's messages API,
+including support for tool use and streaming.
+"""
+
+import time
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from control_zero.llm.base import (
+    GovernanceAction,
+    GovernedChatMixin,
+    LLMUsageMetrics,
+    estimate_cost,
+)
+from control_zero.policy import PolicyDeniedError
+
+
+class GovernedMessages(GovernedChatMixin):
+    """
+    Governed wrapper for Anthropic messages API.
+
+    Supports:
+    - Standard message creation
+    - Streaming responses
+    - Tool use
+    - System prompts
+    - Multi-modal content (text + images)
+    """
+
+    def __init__(self, governed_client: Any):  # GovernedAnthropic
+        self._governed = governed_client
+        self._client = governed_client._client
+
+    def create(
+        self,
+        *,
+        model: str,
+        max_tokens: int,
+        messages: List[Dict[str, Any]],
+        system: Optional[Union[str, List[Dict[str, Any]]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Dict[str, Any]] = None,
+        stream: bool = False,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        stop_sequences: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, str]] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Create a governed message.
+
+        All parameters match the Anthropic API, with added governance:
+        - Model access is checked against policy
+        - Cost is estimated and checked against limits
+        - Tools are filtered by policy
+        - PII is detected and optionally masked
+        - Request is logged for audit
+
+        Args:
+            model: Model to use (e.g., "claude-3-sonnet-20240229")
+            max_tokens: Maximum tokens in response
+            messages: List of message dicts with role and content
+            system: System prompt (string or structured content)
+            tools: Tool definitions for tool use
+            tool_choice: How to handle tool selection
+            stream: Whether to stream the response
+            temperature: Sampling temperature
+            top_p: Nucleus sampling parameter
+            top_k: Top-k sampling parameter
+            stop_sequences: Stop sequences
+            metadata: Request metadata
+            **kwargs: Additional parameters passed to Anthropic
+
+        Returns:
+            Message response or stream iterator
+
+        Raises:
+            PolicyDeniedError: If request violates governance policy
+        """
+        start_time = time.time()
+
+        # Estimate tokens for governance checks
+        estimated_input_tokens = self._estimate_message_tokens(messages)
+        if system:
+            if isinstance(system, str):
+                estimated_input_tokens += len(system) // 4
+            elif isinstance(system, list):
+                for block in system:
+                    if block.get("type") == "text":
+                        estimated_input_tokens += len(block.get("text", "")) // 4
+
+        # Convert tools to standard format for checking
+        tools_to_check = []
+        if tools:
+            for tool in tools:
+                tools_to_check.append({
+                    "name": tool.get("name", ""),
+                    "type": "function",
+                    "function": tool,
+                })
+
+        # Run pre-request governance checks
+        self._governed._pre_request_checks(
+            model=model,
+            action=GovernanceAction.CHAT_COMPLETION,
+            messages=messages,
+            functions=tools_to_check,
+            estimated_tokens=estimated_input_tokens,
+        )
+
+        # Process messages according to governance policies
+        processed_messages = self._process_messages_for_governance(messages)
+
+        # Handle system prompt
+        processed_system = system
+        if system and self._governed._config.content_policy.enable_pii_detection:
+            pii_action = self._governed._config.content_policy.pii_action
+            if pii_action == "mask":
+                if isinstance(system, str):
+                    processed_system = self._governed._mask_pii(system)
+                elif isinstance(system, list):
+                    processed_system = []
+                    for block in system:
+                        if block.get("type") == "text":
+                            new_block = block.copy()
+                            new_block["text"] = self._governed._mask_pii(block.get("text", ""))
+                            processed_system.append(new_block)
+                        else:
+                            processed_system.append(block)
+
+        # Add required system prompt if configured
+        content_policy = self._governed._config.content_policy
+        if content_policy.required_system_prompt:
+            if processed_system is None:
+                processed_system = content_policy.required_system_prompt
+            elif isinstance(processed_system, str):
+                processed_system = f"{content_policy.required_system_prompt}\n\n{processed_system}"
+
+        # Filter tools according to policy
+        filtered_tools = self._filter_tools_for_governance(tools)
+
+        # Apply max_tokens limit from governance
+        governed_max_tokens = max_tokens
+        if content_policy.max_output_tokens:
+            governed_max_tokens = min(max_tokens, content_policy.max_output_tokens)
+
+        # Build request kwargs
+        request_kwargs = {
+            "model": model,
+            "max_tokens": governed_max_tokens,
+            "messages": processed_messages,
+        }
+
+        # Add optional parameters
+        if processed_system is not None:
+            request_kwargs["system"] = processed_system
+        if filtered_tools is not None:
+            request_kwargs["tools"] = filtered_tools
+        if tool_choice is not None and filtered_tools:
+            request_kwargs["tool_choice"] = tool_choice
+        if temperature is not None:
+            request_kwargs["temperature"] = temperature
+        if top_p is not None:
+            request_kwargs["top_p"] = top_p
+        if top_k is not None:
+            request_kwargs["top_k"] = top_k
+        if stop_sequences is not None:
+            request_kwargs["stop_sequences"] = stop_sequences
+        if metadata is not None:
+            request_kwargs["metadata"] = metadata
+
+        # Add user_id to metadata if available
+        if self._governed._user_context.get("user_id"):
+            if "metadata" not in request_kwargs:
+                request_kwargs["metadata"] = {}
+            request_kwargs["metadata"]["user_id"] = str(self._governed._user_context["user_id"])
+
+        # Add any extra kwargs
+        request_kwargs.update(kwargs)
+
+        # Handle streaming
+        if stream:
+            request_kwargs["stream"] = True
+            return self._create_stream(request_kwargs, start_time, model)
+
+        # Make the API call
+        try:
+            response = self._client.messages.create(**request_kwargs)
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Extract usage metrics
+            usage = getattr(response, "usage", None)
+            input_tokens = usage.input_tokens if usage else estimated_input_tokens
+            output_tokens = usage.output_tokens if usage else 0
+            total_tokens = input_tokens + output_tokens
+
+            # Count tool uses
+            tool_use_count = 0
+            for block in response.content:
+                if getattr(block, "type", None) == "tool_use":
+                    tool_use_count += 1
+
+            # Record metrics
+            metrics = LLMUsageMetrics(
+                provider="anthropic",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                total_tokens=total_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=estimate_cost(model, input_tokens, output_tokens),
+                function_calls=tool_use_count,
+                cached=getattr(usage, "cache_read_input_tokens", 0) > 0 if usage else False,
+            )
+
+            # Update tracking and log
+            self._governed._post_request_update(metrics)
+            self._governed._log_request(model, GovernanceAction.CHAT_COMPLETION, metrics)
+
+            return response
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Log error
+            metrics = LLMUsageMetrics(
+                provider="anthropic",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._governed._log_request(
+                model, GovernanceAction.CHAT_COMPLETION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    def _create_stream(
+        self,
+        request_kwargs: Dict[str, Any],
+        start_time: float,
+        model: str,
+    ) -> Iterator[Any]:
+        """
+        Create a governed streaming response.
+
+        Streams the response while tracking tokens and logging at completion.
+        """
+        input_tokens = 0
+        output_tokens = 0
+        tool_use_count = 0
+
+        try:
+            with self._client.messages.stream(**{k: v for k, v in request_kwargs.items() if k != "stream"}) as stream:
+                for event in stream:
+                    # Track usage from events
+                    if hasattr(event, "type"):
+                        if event.type == "message_start":
+                            if hasattr(event, "message") and hasattr(event.message, "usage"):
+                                input_tokens = event.message.usage.input_tokens
+                        elif event.type == "message_delta":
+                            if hasattr(event, "usage"):
+                                output_tokens = event.usage.output_tokens
+                        elif event.type == "content_block_start":
+                            if hasattr(event, "content_block"):
+                                if getattr(event.content_block, "type", None) == "tool_use":
+                                    tool_use_count += 1
+
+                    yield event
+
+            # Calculate final metrics
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="anthropic",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                total_tokens=input_tokens + output_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=estimate_cost(model, input_tokens, output_tokens),
+                function_calls=tool_use_count,
+            )
+
+            self._governed._post_request_update(metrics)
+            self._governed._log_request(model, GovernanceAction.CHAT_COMPLETION, metrics)
+
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="anthropic",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._governed._log_request(
+                model, GovernanceAction.CHAT_COMPLETION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    def _filter_tools_for_governance(
+        self,
+        tools: Optional[List[Dict[str, Any]]],
+    ) -> Optional[List[Dict[str, Any]]]:
+        """Filter tools according to governance policies."""
+        if not tools:
+            return tools
+
+        config = self._governed._config
+        policy = config.function_policy
+
+        if not policy.allowed_functions and not policy.denied_functions:
+            return tools
+
+        filtered = []
+        for tool in tools:
+            tool_name = tool.get("name", "")
+
+            # Skip denied
+            if policy.denied_functions:
+                denied = False
+                for denied_name in policy.denied_functions:
+                    if denied_name.lower() in tool_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 tool_name.lower():
+                        allowed = True
+                        break
+                if not allowed:
+                    continue
+
+            filtered.append(tool)
+
+        return filtered if filtered else None
+
+    def _estimate_message_tokens(self, messages: List[Dict[str, Any]]) -> int:
+        """
+        Estimate token count for messages.
+
+        This is a rough estimate based on character count.
+        For accurate counts, use Anthropic's token counting API.
+        """
+        total_chars = 0
+        for msg in messages:
+            content = msg.get("content", "")
+            if isinstance(content, str):
+                total_chars += len(content)
+            elif isinstance(content, list):
+                # Multi-modal content
+                for part in content:
+                    if isinstance(part, dict):
+                        if part.get("type") == "text":
+                            total_chars += len(part.get("text", ""))
+                        elif part.get("type") == "image":
+                            # Images cost approximately 1600 tokens
+                            total_chars += 6400
+
+        # Rough estimate: ~4 chars per token
+        return max(1, total_chars // 4)