control-zero 0.2.0__py3-none-any.whl

control_zero/llm/groq/client.py

@@ -0,0 +1,330 @@
+"""
+Governed Groq client wrapper.
+
+Provides governance features for the Groq Python SDK including:
+- Model access control
+- Cost tracking and limits
+- Function calling governance
+- PII detection and masking
+- Audit logging
+
+Groq provides fast inference for open-source models like Llama, Mixtral, etc.
+"""
+
+import time
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from control_zero.llm.base import (
+    GovernanceAction,
+    GovernedLLM,
+    GovernedChatMixin,
+    LLMGovernanceConfig,
+    LLMUsageMetrics,
+    estimate_cost,
+)
+from control_zero.policy import PolicyDeniedError
+
+
+class GovernedChatCompletions(GovernedChatMixin):
+    """Governed wrapper for Groq chat completions."""
+
+    def __init__(self, governed_client: "GovernedGroq"):
+        self._governed = governed_client
+        self._client = governed_client._client
+
+    def create(
+        self,
+        *,
+        model: str,
+        messages: List[Dict[str, Any]],
+        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,
+        user: Optional[str] = None,
+        response_format: Optional[Dict[str, str]] = None,
+        seed: Optional[int] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Create a governed chat completion.
+
+        Args:
+            model: Model to use (e.g., "llama-3.1-70b-versatile")
+            messages: List of message dicts
+            tools: Tool definitions
+            tool_choice: Tool selection mode
+            stream: Whether to stream
+            max_tokens: Maximum output tokens
+            temperature: Sampling temperature
+            top_p: Nucleus sampling
+            n: Number of completions
+            stop: Stop sequences
+            presence_penalty: Presence penalty
+            frequency_penalty: Frequency penalty
+            user: End-user ID
+            response_format: Response format
+            seed: Random seed
+            **kwargs: Additional parameters
+
+        Returns:
+            ChatCompletion response or stream iterator
+
+        Raises:
+            PolicyDeniedError: If request violates policy
+        """
+        start_time = time.time()
+
+        # Estimate tokens
+        estimated_input_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._governed._pre_request_checks(
+            model=model,
+            action=GovernanceAction.CHAT_COMPLETION,
+            messages=messages,
+            functions=tools_to_check,
+            estimated_tokens=estimated_input_tokens,
+        )
+
+        # Process messages for governance
+        processed_messages = self._process_messages_for_governance(messages)
+
+        # Filter tools
+        filtered_tools = self._filter_functions_for_governance(tools)
+
+        # Apply max_tokens limit
+        governed_max_tokens = max_tokens
+        if self._governed._config.content_policy.max_output_tokens:
+            if max_tokens:
+                governed_max_tokens = min(max_tokens, self._governed._config.content_policy.max_output_tokens)
+            else:
+                governed_max_tokens = self._governed._config.content_policy.max_output_tokens
+
+        # Build request
+        request_kwargs = {
+            "model": model,
+            "messages": processed_messages,
+        }
+
+        if filtered_tools:
+            request_kwargs["tools"] = filtered_tools
+        if tool_choice and filtered_tools:
+            request_kwargs["tool_choice"] = tool_choice
+        if governed_max_tokens:
+            request_kwargs["max_tokens"] = governed_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 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
+
+        request_kwargs.update(kwargs)
+
+        # Handle streaming
+        if stream:
+            request_kwargs["stream"] = True
+            return self._create_stream(request_kwargs, start_time, model)
+
+        # Make API call
+        try:
+            response = self._client.chat.completions.create(**request_kwargs)
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Extract 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 tool calls
+            tool_call_count = 0
+            for choice in response.choices:
+                if hasattr(choice.message, "tool_calls") and choice.message.tool_calls:
+                    tool_call_count += len(choice.message.tool_calls)
+
+            metrics = LLMUsageMetrics(
+                provider="groq",
+                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_call_count,
+            )
+
+            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)
+            metrics = LLMUsageMetrics(
+                provider="groq",
+                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."""
+        total_tokens = 0
+        tool_call_count = 0
+
+        try:
+            stream = self._client.chat.completions.create(**request_kwargs)
+
+            for chunk in stream:
+                total_tokens += 1
+
+                if chunk.choices:
+                    delta = chunk.choices[0].delta
+                    if hasattr(delta, "tool_calls") and delta.tool_calls:
+                        tool_call_count += len(delta.tool_calls)
+
+                yield chunk
+
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="groq",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                total_tokens=total_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=estimate_cost(model, total_tokens // 2, total_tokens // 2),
+                function_calls=tool_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="groq",
+                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."""
+        total_chars = sum(len(str(m.get("content", ""))) for m in messages)
+        return max(1, total_chars // 4)
+
+
+class GovernedChat:
+    """Governed wrapper for Groq chat namespace."""
+
+    def __init__(self, governed_client: "GovernedGroq"):
+        self._completions = GovernedChatCompletions(governed_client)
+
+    @property
+    def completions(self) -> GovernedChatCompletions:
+        return self._completions
+
+
+class GovernedGroq(GovernedLLM):
+    """
+    Governed wrapper for the Groq Python SDK.
+
+    Groq provides fast inference for open-source models like Llama and Mixtral.
+
+    Example:
+        from control_zero import ControlZeroClient
+        from control_zero.llm.groq import GovernedGroq
+        from groq import Groq
+
+        cz = ControlZeroClient(api_key="...")
+        cz.initialize()
+
+        client = Groq()
+        governed = GovernedGroq(client=client, control_zero=cz)
+
+        response = governed.chat.completions.create(
+            model="llama-3.1-70b-versatile",
+            messages=[{"role": "user", "content": "Hello!"}]
+        )
+    """
+
+    def __init__(
+        self,
+        client: Any,
+        control_zero: Any,
+        config: Optional[LLMGovernanceConfig] = None,
+        user_context: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(client, control_zero, config, user_context)
+        self._chat = GovernedChat(self)
+
+    @property
+    def provider_name(self) -> str:
+        return "groq"
+
+    @property
+    def chat(self) -> GovernedChat:
+        return self._chat
+
+    def with_user_context(self, user_context: Dict[str, Any]) -> "GovernedGroq":
+        merged_context = {**self._user_context, **user_context}
+        return GovernedGroq(
+            client=self._client,
+            control_zero=self._cz,
+            config=self._config,
+            user_context=merged_context,
+        )
+
+    def with_config(self, config: LLMGovernanceConfig) -> "GovernedGroq":
+        return GovernedGroq(
+            client=self._client,
+            control_zero=self._cz,
+            config=config,
+            user_context=self._user_context,
+        )

control_zero/llm/mistral/__init__.py

@@ -0,0 +1,32 @@
+"""
+Control Zero Mistral Governance Wrapper.
+
+This module provides governance wrappers for the Mistral Python SDK,
+enabling policy enforcement, cost tracking, and audit logging for
+all Mistral API calls.
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.mistral import GovernedMistral
+    from mistralai import Mistral
+
+    # Initialize Control Zero
+    cz_client = ControlZeroClient(api_key="cz_live_xxx")
+    cz_client.initialize()
+
+    # Wrap Mistral client with governance
+    mistral_client = Mistral(api_key="your-mistral-key")
+    governed = GovernedMistral(client=mistral_client, control_zero=cz_client)
+
+    # All calls are now governed
+    response = governed.chat.complete(
+        model="mistral-large-latest",
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+"""
+
+from control_zero.llm.mistral.client import GovernedMistral
+
+__all__ = [
+    "GovernedMistral",
+]

control_zero/llm/mistral/client.py

@@ -0,0 +1,319 @@
+"""
+Governed Mistral client wrapper.
+
+Provides governance features for the Mistral Python SDK including:
+- Model access control
+- Cost tracking and limits
+- Function calling governance
+- PII detection and masking
+- Audit logging
+"""
+
+import time
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from control_zero.llm.base import (
+    GovernanceAction,
+    GovernedLLM,
+    GovernedChatMixin,
+    LLMGovernanceConfig,
+    LLMUsageMetrics,
+    estimate_cost,
+)
+from control_zero.policy import PolicyDeniedError
+
+
+class GovernedChatCompletions(GovernedChatMixin):
+    """Governed wrapper for Mistral chat completions."""
+
+    def __init__(self, governed_client: "GovernedMistral"):
+        self._governed = governed_client
+        self._client = governed_client._client
+
+    def complete(
+        self,
+        *,
+        model: str,
+        messages: List[Dict[str, Any]],
+        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,
+        random_seed: Optional[int] = None,
+        safe_prompt: bool = False,
+        response_format: Optional[Dict[str, str]] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Create a governed chat completion.
+
+        Args:
+            model: Model to use (e.g., "mistral-large-latest")
+            messages: List of message dicts
+            tools: Tool definitions
+            tool_choice: Tool selection mode
+            stream: Whether to stream
+            max_tokens: Maximum output tokens
+            temperature: Sampling temperature
+            top_p: Nucleus sampling
+            random_seed: Random seed
+            safe_prompt: Enable safe mode
+            response_format: Response format
+            **kwargs: Additional parameters
+
+        Returns:
+            ChatCompletionResponse or stream iterator
+
+        Raises:
+            PolicyDeniedError: If request violates policy
+        """
+        start_time = time.time()
+
+        # Estimate tokens
+        estimated_input_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._governed._pre_request_checks(
+            model=model,
+            action=GovernanceAction.CHAT_COMPLETION,
+            messages=messages,
+            functions=tools_to_check,
+            estimated_tokens=estimated_input_tokens,
+        )
+
+        # Process messages for governance
+        processed_messages = self._process_messages_for_governance(messages)
+
+        # Filter tools
+        filtered_tools = self._filter_functions_for_governance(tools)
+
+        # Apply max_tokens limit
+        governed_max_tokens = max_tokens
+        if self._governed._config.content_policy.max_output_tokens:
+            if max_tokens:
+                governed_max_tokens = min(max_tokens, self._governed._config.content_policy.max_output_tokens)
+            else:
+                governed_max_tokens = self._governed._config.content_policy.max_output_tokens
+
+        # Build request
+        request_kwargs = {
+            "model": model,
+            "messages": processed_messages,
+        }
+
+        if filtered_tools:
+            request_kwargs["tools"] = filtered_tools
+        if tool_choice and filtered_tools:
+            request_kwargs["tool_choice"] = tool_choice
+        if governed_max_tokens:
+            request_kwargs["max_tokens"] = governed_max_tokens
+        if temperature is not None:
+            request_kwargs["temperature"] = temperature
+        if top_p is not None:
+            request_kwargs["top_p"] = top_p
+        if random_seed is not None:
+            request_kwargs["random_seed"] = random_seed
+        if safe_prompt:
+            request_kwargs["safe_prompt"] = safe_prompt
+        if response_format is not None:
+            request_kwargs["response_format"] = response_format
+
+        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.complete(**request_kwargs)
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Extract 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 tool calls
+            tool_call_count = 0
+            for choice in response.choices:
+                if hasattr(choice.message, "tool_calls") and choice.message.tool_calls:
+                    tool_call_count += len(choice.message.tool_calls)
+
+            metrics = LLMUsageMetrics(
+                provider="mistral",
+                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_call_count,
+            )
+
+            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)
+            metrics = LLMUsageMetrics(
+                provider="mistral",
+                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."""
+        total_tokens = 0
+        tool_call_count = 0
+
+        try:
+            stream = self._client.chat.stream(**request_kwargs)
+
+            for chunk in stream:
+                total_tokens += 1
+
+                data = getattr(chunk, "data", chunk)
+                if hasattr(data, "choices"):
+                    for choice in data.choices:
+                        if hasattr(choice, "delta") and hasattr(choice.delta, "tool_calls"):
+                            if choice.delta.tool_calls:
+                                tool_call_count += len(choice.delta.tool_calls)
+
+                yield chunk
+
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="mistral",
+                model=model,
+                action=GovernanceAction.CHAT_COMPLETION,
+                total_tokens=total_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=estimate_cost(model, total_tokens // 2, total_tokens // 2),
+                function_calls=tool_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="mistral",
+                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."""
+        total_chars = sum(len(str(m.get("content", ""))) for m in messages)
+        return max(1, total_chars // 4)
+
+
+class GovernedChat:
+    """Governed wrapper for Mistral chat namespace."""
+
+    def __init__(self, governed_client: "GovernedMistral"):
+        self._completions = GovernedChatCompletions(governed_client)
+
+    def complete(self, **kwargs) -> Any:
+        """Create a chat completion."""
+        return self._completions.complete(**kwargs)
+
+    def stream(self, **kwargs) -> Iterator[Any]:
+        """Create a streaming chat completion."""
+        kwargs["stream"] = True
+        return self._completions.complete(**kwargs)
+
+
+class GovernedMistral(GovernedLLM):
+    """
+    Governed wrapper for the Mistral Python SDK.
+
+    Example:
+        from control_zero import ControlZeroClient
+        from control_zero.llm.mistral import GovernedMistral
+        from mistralai import Mistral
+
+        cz = ControlZeroClient(api_key="...")
+        cz.initialize()
+
+        client = Mistral(api_key="...")
+        governed = GovernedMistral(client=client, control_zero=cz)
+
+        response = governed.chat.complete(
+            model="mistral-large-latest",
+            messages=[{"role": "user", "content": "Hello!"}]
+        )
+    """
+
+    def __init__(
+        self,
+        client: Any,
+        control_zero: Any,
+        config: Optional[LLMGovernanceConfig] = None,
+        user_context: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(client, control_zero, config, user_context)
+        self._chat = GovernedChat(self)
+
+    @property
+    def provider_name(self) -> str:
+        return "mistral"
+
+    @property
+    def chat(self) -> GovernedChat:
+        return self._chat
+
+    @property
+    def embeddings(self):
+        """Access the underlying embeddings endpoint (pass-through)."""
+        return self._client.embeddings
+
+    def with_user_context(self, user_context: Dict[str, Any]) -> "GovernedMistral":
+        merged_context = {**self._user_context, **user_context}
+        return GovernedMistral(
+            client=self._client,
+            control_zero=self._cz,
+            config=self._config,
+            user_context=merged_context,
+        )
+
+    def with_config(self, config: LLMGovernanceConfig) -> "GovernedMistral":
+        return GovernedMistral(
+            client=self._client,
+            control_zero=self._cz,
+            config=config,
+            user_context=self._user_context,
+        )