control-zero 0.2.0__py3-none-any.whl

control_zero/llm/ollama/__init__.py

@@ -0,0 +1,31 @@
+"""
+Control Zero Ollama Governance Wrapper.
+
+This module provides governance wrappers for Ollama local LLMs,
+enabling policy enforcement, usage tracking, and audit logging for
+locally-hosted open-source models.
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.ollama import GovernedOllama
+    import ollama
+
+    # Initialize Control Zero
+    cz_client = ControlZeroClient(api_key="cz_live_xxx")
+    cz_client.initialize()
+
+    # Create governed Ollama client
+    governed = GovernedOllama(control_zero=cz_client)
+
+    # All calls are now governed
+    response = governed.chat(
+        model="llama3.2",
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+"""
+
+from control_zero.llm.ollama.client import GovernedOllama
+
+__all__ = [
+    "GovernedOllama",
+]

control_zero/llm/ollama/client.py

@@ -0,0 +1,439 @@
+"""
+Governed Ollama client wrapper.
+
+Provides governance features for Ollama local LLMs including:
+- Model access control
+- Usage tracking
+- PII detection and masking
+- Audit logging
+
+Ollama runs models locally, so there's no cost tracking,
+but governance policies still apply for:
+- Model allowlisting
+- Request limits
+- PII detection
+- Audit logging
+"""
+
+import time
+from typing import Any, Dict, Iterator, List, Optional
+
+from control_zero.llm.base import (
+    GovernanceAction,
+    GovernedLLM,
+    GovernedChatMixin,
+    LLMGovernanceConfig,
+    LLMUsageMetrics,
+)
+from control_zero.policy import PolicyDeniedError
+
+
+class GovernedOllama(GovernedLLM, GovernedChatMixin):
+    """
+    Governed wrapper for Ollama local LLMs.
+
+    Ollama runs models locally, providing privacy and cost benefits.
+    This wrapper adds governance for:
+    - Model access control (which local models can be used)
+    - Request rate limiting
+    - PII detection and masking
+    - Audit logging
+
+    Example:
+        from control_zero import ControlZeroClient
+        from control_zero.llm.ollama import GovernedOllama
+
+        cz = ControlZeroClient(api_key="...")
+        cz.initialize()
+
+        governed = GovernedOllama(
+            control_zero=cz,
+            host="http://localhost:11434",
+            config=LLMGovernanceConfig(
+                model_policy=ModelPolicy(
+                    allowed_models=["llama3.2", "mistral"]
+                )
+            )
+        )
+
+        response = governed.chat(
+            model="llama3.2",
+            messages=[{"role": "user", "content": "Hello!"}]
+        )
+    """
+
+    def __init__(
+        self,
+        control_zero: Any,
+        config: Optional[LLMGovernanceConfig] = None,
+        user_context: Optional[Dict[str, Any]] = None,
+        host: Optional[str] = None,
+        timeout: Optional[float] = None,
+    ):
+        """
+        Initialize a governed Ollama client.
+
+        Args:
+            control_zero: Control Zero client for policy and logging
+            config: Governance configuration
+            user_context: Context about the current user
+            host: Ollama server host (default: http://localhost:11434)
+            timeout: Request timeout in seconds
+        """
+        try:
+            import ollama
+        except ImportError:
+            raise ImportError(
+                "ollama is required for Ollama support. "
+                "Install it with: pip install ollama"
+            )
+
+        # Create client with optional parameters
+        client_kwargs = {}
+        if host:
+            client_kwargs["host"] = host
+        if timeout:
+            client_kwargs["timeout"] = timeout
+
+        client = ollama.Client(**client_kwargs) if client_kwargs else ollama
+
+        super().__init__(client, control_zero, config, user_context)
+
+        self._host = host or "http://localhost:11434"
+
+    @property
+    def provider_name(self) -> str:
+        return "ollama"
+
+    def chat(
+        self,
+        *,
+        model: str,
+        messages: List[Dict[str, Any]],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        stream: bool = False,
+        format: Optional[str] = None,
+        options: Optional[Dict[str, Any]] = None,
+        keep_alive: Optional[str] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Create a governed chat completion.
+
+        Args:
+            model: Model to use (e.g., "llama3.2", "mistral")
+            messages: List of message dicts
+            tools: Tool definitions (if model supports it)
+            stream: Whether to stream
+            format: Response format (e.g., "json")
+            options: Model options (temperature, top_p, etc.)
+            keep_alive: How long to keep model loaded
+            **kwargs: Additional parameters
+
+        Returns:
+            Chat response or stream iterator
+
+        Raises:
+            PolicyDeniedError: If request violates policy
+        """
+        start_time = time.time()
+
+        # Estimate tokens
+        estimated_tokens = self._estimate_message_tokens(messages)
+
+        # Prepare tools for policy check
+        tools_to_check = []
+        if tools:
+            tools_to_check = [{"name": t.get("function", {}).get("name", ""), "type": "function"} for t in tools]
+
+        # Run governance checks
+        self._pre_request_checks(
+            model=model,
+            action=GovernanceAction.CHAT_COMPLETION,
+            messages=messages,
+            functions=tools_to_check,
+            estimated_tokens=estimated_tokens,
+        )
+
+        # Process messages for governance
+        processed_messages = self._process_messages_for_governance(messages)
+
+        # Filter tools
+        filtered_tools = self._filter_tools_for_governance(tools)
+
+        # Build request
+        request_kwargs = {
+            "model": model,
+            "messages": processed_messages,
+        }
+
+        if filtered_tools:
+            request_kwargs["tools"] = filtered_tools
+        if format:
+            request_kwargs["format"] = format
+        if options:
+            request_kwargs["options"] = options
+        if keep_alive:
+            request_kwargs["keep_alive"] = keep_alive
+
+        request_kwargs.update(kwargs)
+
+        # Handle streaming
+        if stream:
+            return self._create_stream(request_kwargs, start_time, model)
+
+        # Make API call
+        try:
+            response = self._client.chat(**request_kwargs)
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Extract metrics (Ollama provides some usage info)
+            prompt_tokens = response.get("prompt_eval_count", estimated_tokens)
+            output_tokens = response.get("eval_count", 0)
+            total_tokens = prompt_tokens + output_tokens
+
+            # Count tool calls
+            tool_call_count = 0
+            message = response.get("message", {})
+            if message.get("tool_calls"):
+                tool_call_count = len(message["tool_calls"])
+
+            metrics = LLMUsageMetrics(
+                provider="ollama",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                input_tokens=prompt_tokens,
+                output_tokens=output_tokens,
+                total_tokens=total_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=0.0,  # Local models have no API cost
+                function_calls=tool_call_count,
+            )
+
+            self._post_request_update(metrics)
+            self._log_request(model, GovernanceAction.CHAT_COMPLETION, metrics)
+
+            return response
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="ollama",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._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."""
+        request_kwargs["stream"] = True
+        total_tokens = 0
+
+        try:
+            stream = self._client.chat(**request_kwargs)
+
+            for chunk in stream:
+                total_tokens += 1
+                yield chunk
+
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="ollama",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                total_tokens=total_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=0.0,
+            )
+
+            self._post_request_update(metrics)
+            self._log_request(model, GovernanceAction.CHAT_COMPLETION, metrics)
+
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="ollama",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._log_request(
+                model, GovernanceAction.CHAT_COMPLETION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    def generate(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        stream: bool = False,
+        format: Optional[str] = None,
+        options: Optional[Dict[str, Any]] = None,
+        system: Optional[str] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Generate text with governance (legacy API).
+
+        Args:
+            model: Model to use
+            prompt: Input prompt
+            stream: Whether to stream
+            format: Response format
+            options: Model options
+            system: System prompt
+            **kwargs: Additional parameters
+
+        Returns:
+            Generate response or stream iterator
+        """
+        start_time = time.time()
+
+        # Governance checks
+        messages = [{"role": "user", "content": prompt}]
+        if system:
+            messages.insert(0, {"role": "system", "content": system})
+
+        estimated_tokens = len(prompt) // 4
+
+        self._pre_request_checks(
+            model=model,
+            action=GovernanceAction.TEXT_COMPLETION,
+            messages=messages,
+            estimated_tokens=estimated_tokens,
+        )
+
+        # Process prompt for PII
+        processed_prompt = prompt
+        if self._config.content_policy.enable_pii_detection:
+            if self._config.content_policy.pii_action == "mask":
+                processed_prompt = self._mask_pii(prompt)
+
+        request_kwargs = {
+            "model": model,
+            "prompt": processed_prompt,
+        }
+
+        if format:
+            request_kwargs["format"] = format
+        if options:
+            request_kwargs["options"] = options
+        if system:
+            request_kwargs["system"] = system
+        if stream:
+            request_kwargs["stream"] = True
+
+        request_kwargs.update(kwargs)
+
+        try:
+            response = self._client.generate(**request_kwargs)
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="ollama",
+                model=model,
+                action=GovernanceAction.TEXT_COMPLETION,
+                input_tokens=response.get("prompt_eval_count", estimated_tokens),
+                output_tokens=response.get("eval_count", 0),
+                latency_ms=latency_ms,
+                estimated_cost=0.0,
+            )
+
+            self._post_request_update(metrics)
+            self._log_request(model, GovernanceAction.TEXT_COMPLETION, metrics)
+
+            return response
+
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="ollama",
+                model=model,
+                action=GovernanceAction.TEXT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._log_request(
+                model, GovernanceAction.TEXT_COMPLETION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    def list(self) -> Any:
+        """List available models (pass-through)."""
+        return self._client.list()
+
+    def show(self, model: str) -> Any:
+        """Show model info (pass-through)."""
+        return self._client.show(model)
+
+    def pull(self, model: str, **kwargs) -> Any:
+        """Pull a model (pass-through)."""
+        return self._client.pull(model, **kwargs)
+
+    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
+
+        policy = self._config.function_policy
+
+        if not policy.allowed_functions and not policy.denied_functions:
+            return tools
+
+        filtered = []
+        for tool in tools:
+            tool_name = tool.get("function", {}).get("name", "")
+
+            if policy.denied_functions:
+                denied = any(d.lower() in tool_name.lower() for d in policy.denied_functions)
+                if denied:
+                    continue
+
+            if policy.allowed_functions:
+                allowed = any(a.lower() in tool_name.lower() for a in policy.allowed_functions)
+                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."""
+        total_chars = sum(len(str(m.get("content", ""))) for m in messages)
+        return max(1, total_chars // 4)
+
+    def with_user_context(self, user_context: Dict[str, Any]) -> "GovernedOllama":
+        merged_context = {**self._user_context, **user_context}
+        return GovernedOllama(
+            control_zero=self._cz,
+            config=self._config,
+            user_context=merged_context,
+            host=self._host,
+        )
+
+    def with_config(self, config: LLMGovernanceConfig) -> "GovernedOllama":
+        return GovernedOllama(
+            control_zero=self._cz,
+            config=config,
+            user_context=self._user_context,
+            host=self._host,
+        )

control_zero/llm/openai/__init__.py

@@ -0,0 +1,34 @@
+"""
+Control Zero OpenAI Governance Wrapper.
+
+This module provides governance wrappers for the OpenAI Python SDK,
+enabling policy enforcement, cost tracking, and audit logging for
+all OpenAI API calls.
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.openai import GovernedOpenAI
+    from openai import OpenAI
+
+    # Initialize Control Zero
+    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.openai.client import GovernedOpenAI
+from control_zero.llm.openai.chat import GovernedChatCompletions
+
+__all__ = [
+    "GovernedOpenAI",
+    "GovernedChatCompletions",
+]