control-zero 0.2.0__py3-none-any.whl

control_zero/llm/openai/chat.py

@@ -0,0 +1,331 @@
+"""
+Governed OpenAI Chat Completions.
+
+Provides governance wrapper for OpenAI's chat completions API,
+including support for function calling, streaming, and tool use.
+"""
+
+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 GovernedChatCompletions(GovernedChatMixin):
+    """
+    Governed wrapper for OpenAI chat completions.
+
+    Supports:
+    - Standard chat completions
+    - Streaming responses
+    - Function calling
+    - Tool use
+    - Response format (JSON mode)
+    """
+
+    def __init__(self, governed_client: Any):  # GovernedOpenAI
+        self._governed = governed_client
+        self._client = governed_client._client
+
+    def create(
+        self,
+        *,
+        model: str,
+        messages: List[Dict[str, Any]],
+        functions: Optional[List[Dict[str, Any]]] = None,
+        function_call: Optional[Union[str, Dict[str, str]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
+        stream: bool = False,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        n: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        presence_penalty: Optional[float] = None,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[Dict[str, float]] = None,
+        user: Optional[str] = None,
+        response_format: Optional[Dict[str, str]] = None,
+        seed: Optional[int] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Create a governed chat completion.
+
+        All parameters match the OpenAI API, with added governance:
+        - Model access is checked against policy
+        - Cost is estimated and checked against limits
+        - Functions/tools are filtered by policy
+        - PII is detected and optionally masked
+        - Request is logged for audit
+
+        Args:
+            model: Model to use (e.g., "gpt-4o")
+            messages: List of message dicts with role and content
+            functions: Legacy function definitions (deprecated)
+            function_call: How to handle function calls (deprecated)
+            tools: Tool definitions for function calling
+            tool_choice: How to handle tool selection
+            stream: Whether to stream the response
+            max_tokens: Maximum tokens in response
+            temperature: Sampling temperature
+            top_p: Nucleus sampling parameter
+            n: Number of completions to generate
+            stop: Stop sequences
+            presence_penalty: Presence penalty
+            frequency_penalty: Frequency penalty
+            logit_bias: Token biases
+            user: End-user ID for abuse tracking
+            response_format: Response format (e.g., {"type": "json_object"})
+            seed: Random seed for deterministic outputs
+            **kwargs: Additional parameters passed to OpenAI
+
+        Returns:
+            ChatCompletion 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)
+
+        # Determine which functions/tools to check
+        tools_to_check = tools or []
+        if functions:
+            # Convert legacy functions to tools format for checking
+            tools_to_check = [
+                {"type": "function", "function": f} for f in functions
+            ]
+
+        # 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)
+
+        # Filter functions/tools according to policy
+        filtered_tools = self._filter_functions_for_governance(tools)
+        filtered_functions = None
+        if functions:
+            # Handle legacy functions
+            filtered_functions = self._filter_functions_for_governance(
+                [{"name": f["name"], **f} for f in functions]
+            )
+            if filtered_functions:
+                filtered_functions = [
+                    {k: v for k, v in f.items() if k != "type"}
+                    for f in filtered_functions
+                ]
+
+        # Build request kwargs
+        request_kwargs = {
+            "model": model,
+            "messages": processed_messages,
+        }
+
+        # Add optional parameters
+        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 filtered_functions is not None:
+            request_kwargs["functions"] = filtered_functions
+        if function_call is not None and filtered_functions:
+            request_kwargs["function_call"] = function_call
+        if max_tokens is not None:
+            # Apply governance limit
+            max_output = self._governed._config.content_policy.max_output_tokens
+            if max_output:
+                max_tokens = min(max_tokens, max_output)
+            request_kwargs["max_tokens"] = max_tokens
+        if temperature is not None:
+            request_kwargs["temperature"] = temperature
+        if top_p is not None:
+            request_kwargs["top_p"] = top_p
+        if n is not None:
+            request_kwargs["n"] = n
+        if stop is not None:
+            request_kwargs["stop"] = stop
+        if presence_penalty is not None:
+            request_kwargs["presence_penalty"] = presence_penalty
+        if frequency_penalty is not None:
+            request_kwargs["frequency_penalty"] = frequency_penalty
+        if logit_bias is not None:
+            request_kwargs["logit_bias"] = logit_bias
+        if user is not None:
+            request_kwargs["user"] = user
+        elif self._governed._user_context.get("user_id"):
+            request_kwargs["user"] = str(self._governed._user_context["user_id"])
+        if response_format is not None:
+            request_kwargs["response_format"] = response_format
+        if seed is not None:
+            request_kwargs["seed"] = seed
+
+        # 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.chat.completions.create(**request_kwargs)
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Extract usage metrics
+            usage = getattr(response, "usage", None)
+            input_tokens = usage.prompt_tokens if usage else estimated_input_tokens
+            output_tokens = usage.completion_tokens if usage else 0
+            total_tokens = usage.total_tokens if usage else input_tokens + output_tokens
+
+            # Count function calls
+            function_call_count = 0
+            for choice in response.choices:
+                if hasattr(choice.message, "tool_calls") and choice.message.tool_calls:
+                    function_call_count += len(choice.message.tool_calls)
+                if hasattr(choice.message, "function_call") and choice.message.function_call:
+                    function_call_count += 1
+
+            # Record metrics
+            metrics = LLMUsageMetrics(
+                provider="openai",
+                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=function_call_count,
+                cached=getattr(usage, "cached_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="openai",
+                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.
+        """
+        total_tokens = 0
+        output_content = []
+        function_call_count = 0
+
+        try:
+            stream = self._client.chat.completions.create(**request_kwargs)
+
+            for chunk in stream:
+                total_tokens += 1  # Approximate token counting for streams
+
+                # Track content
+                if chunk.choices:
+                    delta = chunk.choices[0].delta
+                    if hasattr(delta, "content") and delta.content:
+                        output_content.append(delta.content)
+                    if hasattr(delta, "tool_calls") and delta.tool_calls:
+                        function_call_count += len(delta.tool_calls)
+
+                yield chunk
+
+            # Calculate final metrics
+            latency_ms = int((time.time() - start_time) * 1000)
+            estimated_output_tokens = len("".join(output_content).split()) * 1.3
+            estimated_input_tokens = request_kwargs.get("messages", [])
+            estimated_input_tokens = self._estimate_message_tokens(estimated_input_tokens)
+
+            metrics = LLMUsageMetrics(
+                provider="openai",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                input_tokens=int(estimated_input_tokens),
+                output_tokens=int(estimated_output_tokens),
+                total_tokens=int(estimated_input_tokens + estimated_output_tokens),
+                latency_ms=latency_ms,
+                estimated_cost=estimate_cost(
+                    model, int(estimated_input_tokens), int(estimated_output_tokens)
+                ),
+                function_calls=function_call_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="openai",
+                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 _estimate_message_tokens(self, messages: List[Dict[str, Any]]) -> int:
+        """
+        Estimate token count for messages.
+
+        This is a rough estimate based on word count.
+        For accurate counts, use tiktoken.
+        """
+        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) and part.get("type") == "text":
+                        total_chars += len(part.get("text", ""))
+
+        # Rough estimate: ~4 chars per token
+        return max(1, total_chars // 4)

control_zero/llm/openai/client.py

@@ -0,0 +1,182 @@
+"""
+Governed OpenAI client wrapper.
+
+Provides governance features for the OpenAI Python SDK including:
+- Model access control
+- Cost tracking and limits
+- Function calling 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.openai.chat import GovernedChatCompletions
+
+
+class GovernedChat:
+    """Governed wrapper for OpenAI chat namespace."""
+
+    def __init__(self, governed_client: "GovernedOpenAI"):
+        self._governed_client = governed_client
+        self._completions = GovernedChatCompletions(governed_client)
+
+    @property
+    def completions(self) -> GovernedChatCompletions:
+        """Access governed chat completions."""
+        return self._completions
+
+
+class GovernedOpenAI(GovernedLLM):
+    """
+    Governed wrapper for the OpenAI Python SDK.
+
+    This class wraps an OpenAI client instance and adds governance
+    features including policy enforcement, cost tracking, and audit logging.
+
+    Example:
+        from control_zero import ControlZeroClient
+        from control_zero.llm.openai import GovernedOpenAI
+        from openai import OpenAI
+
+        cz = ControlZeroClient(api_key="...")
+        cz.initialize()
+
+        client = OpenAI()
+        governed = GovernedOpenAI(client=client, control_zero=cz)
+
+        # Configure governance
+        from control_zero.llm import LLMGovernanceConfig, ModelPolicy, CostPolicy
+
+        governed = GovernedOpenAI(
+            client=client,
+            control_zero=cz,
+            config=LLMGovernanceConfig(
+                model_policy=ModelPolicy(
+                    allowed_models=["gpt-4o", "gpt-4o-mini"],
+                    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.chat.completions.create(
+            model="gpt-4o",
+            messages=[{"role": "user", "content": "Hello!"}]
+        )
+    """
+
+    def __init__(
+        self,
+        client: Any,  # openai.OpenAI
+        control_zero: Any,  # ControlZeroClient
+        config: Optional[LLMGovernanceConfig] = None,
+        user_context: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize a governed OpenAI client.
+
+        Args:
+            client: An OpenAI 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._chat = GovernedChat(self)
+
+    @property
+    def provider_name(self) -> str:
+        return "openai"
+
+    @property
+    def chat(self) -> GovernedChat:
+        """Access governed chat completions."""
+        return self._chat
+
+    @property
+    def models(self):
+        """Access the underlying models endpoint (pass-through)."""
+        return self._client.models
+
+    @property
+    def embeddings(self):
+        """Access the underlying embeddings endpoint.
+
+        TODO: Add governance wrapper for embeddings.
+        """
+        return self._client.embeddings
+
+    @property
+    def files(self):
+        """Access the underlying files endpoint (pass-through)."""
+        return self._client.files
+
+    @property
+    def images(self):
+        """Access the underlying images endpoint.
+
+        TODO: Add governance wrapper for image generation.
+        """
+        return self._client.images
+
+    @property
+    def audio(self):
+        """Access the underlying audio endpoint.
+
+        TODO: Add governance wrapper for audio.
+        """
+        return self._client.audio
+
+    @property
+    def moderations(self):
+        """Access the underlying moderations endpoint (pass-through)."""
+        return self._client.moderations
+
+    def with_user_context(self, user_context: Dict[str, Any]) -> "GovernedOpenAI":
+        """
+        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 GovernedOpenAI instance with updated context
+        """
+        merged_context = {**self._user_context, **user_context}
+        return GovernedOpenAI(
+            client=self._client,
+            control_zero=self._cz,
+            config=self._config,
+            user_context=merged_context,
+        )
+
+    def with_config(self, config: LLMGovernanceConfig) -> "GovernedOpenAI":
+        """
+        Create a new governed client with updated configuration.
+
+        Args:
+            config: New governance configuration
+
+        Returns:
+            New GovernedOpenAI instance with updated config
+        """
+        return GovernedOpenAI(
+            client=self._client,
+            control_zero=self._cz,
+            config=config,
+            user_context=self._user_context,
+        )

control_zero/logging/__init__.py

@@ -0,0 +1,5 @@
+"""Async logging for Control Zero SDK."""
+
+from control_zero.logging.async_logger import AsyncLogger
+
+__all__ = ["AsyncLogger"]

control_zero/logging/async_logger.py

@@ -0,0 +1,65 @@
+"""Async batched logger for audit logs."""
+
+import asyncio
+from typing import Optional
+
+from control_zero.types import AuditLogEntry
+
+
+class AsyncLogger:
+    """
+    Async logger that batches and sends audit logs.
+
+    Logs are queued locally and flushed periodically to avoid
+    impacting the latency of tool calls.
+    """
+
+    def __init__(
+        self,
+        batch_size: int = 100,
+        flush_interval: float = 5.0,
+    ):
+        self._batch_size = batch_size
+        self._flush_interval = flush_interval
+        self._queue: asyncio.Queue[AuditLogEntry] = asyncio.Queue()
+        self._task: Optional[asyncio.Task[None]] = None
+
+    async def start(self, send_fn) -> None:
+        """Start the background flush task."""
+        self._send_fn = send_fn
+        self._task = asyncio.create_task(self._flush_loop())
+
+    async def _flush_loop(self) -> None:
+        """Background loop that flushes logs periodically."""
+        while True:
+            await asyncio.sleep(self._flush_interval)
+            await self.flush()
+
+    async def log(self, entry: AuditLogEntry) -> None:
+        """Queue a log entry."""
+        await self._queue.put(entry)
+
+        # Flush immediately if batch is full
+        if self._queue.qsize() >= self._batch_size:
+            await self.flush()
+
+    async def flush(self) -> None:
+        """Flush queued logs to the server."""
+        logs: list[AuditLogEntry] = []
+
+        while len(logs) < self._batch_size and not self._queue.empty():
+            logs.append(await self._queue.get())
+
+        if logs and hasattr(self, "_send_fn"):
+            try:
+                await self._send_fn(logs)
+            except Exception:
+                # Re-queue on failure (best effort)
+                for log in logs:
+                    await self._queue.put(log)
+
+    async def stop(self) -> None:
+        """Stop the logger and flush remaining logs."""
+        if self._task:
+            self._task.cancel()
+        await self.flush()

control_zero/mcp/__init__.py

@@ -0,0 +1,5 @@
+"""MCP middleware for Control Zero SDK."""
+
+from control_zero.mcp.middleware import MCPMiddleware
+
+__all__ = ["MCPMiddleware"]