control-zero 0.2.0__py3-none-any.whl

control_zero/llm/gemini/client.py

@@ -0,0 +1,486 @@
+"""
+Governed Google Gemini client wrapper.
+
+Provides governance features for the Google Generative AI 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 GovernedGemini(GovernedLLM, GovernedChatMixin):
+    """
+    Governed wrapper for Google Gemini.
+
+    This class wraps Gemini's GenerativeModel and adds governance
+    features including policy enforcement, cost tracking, and audit logging.
+
+    Example:
+        from control_zero import ControlZeroClient
+        from control_zero.llm.gemini import GovernedGemini
+        import google.generativeai as genai
+
+        genai.configure(api_key="...")
+
+        cz = ControlZeroClient(api_key="...")
+        cz.initialize()
+
+        governed = GovernedGemini(
+            model_name="gemini-1.5-pro",
+            control_zero=cz,
+            config=LLMGovernanceConfig(
+                model_policy=ModelPolicy(
+                    allowed_models=["gemini-1.5-pro", "gemini-1.5-flash"]
+                ),
+                cost_policy=CostPolicy(
+                    max_cost_per_day=10.00
+                )
+            )
+        )
+
+        # Make governed API call
+        response = governed.generate_content("Hello!")
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        control_zero: Any,  # ControlZeroClient
+        config: Optional[LLMGovernanceConfig] = None,
+        user_context: Optional[Dict[str, Any]] = None,
+        generation_config: Optional[Dict[str, Any]] = None,
+        safety_settings: Optional[List[Dict[str, Any]]] = None,
+        tools: Optional[List[Any]] = None,
+        system_instruction: Optional[str] = None,
+    ):
+        """
+        Initialize a governed Gemini model.
+
+        Args:
+            model_name: Name of the Gemini model (e.g., "gemini-1.5-pro")
+            control_zero: Control Zero client for policy and logging
+            config: Governance configuration
+            user_context: Context about the current user
+            generation_config: Gemini generation configuration
+            safety_settings: Gemini safety settings
+            tools: Function/tool definitions
+            system_instruction: System instruction for the model
+        """
+        # Import here to avoid requiring google-generativeai at module load
+        try:
+            import google.generativeai as genai
+        except ImportError:
+            raise ImportError(
+                "google-generativeai is required for Gemini support. "
+                "Install it with: pip install google-generativeai"
+            )
+
+        # Create the underlying model
+        model_kwargs = {"model_name": model_name}
+        if generation_config:
+            model_kwargs["generation_config"] = generation_config
+        if safety_settings:
+            model_kwargs["safety_settings"] = safety_settings
+        if tools:
+            model_kwargs["tools"] = tools
+        if system_instruction:
+            model_kwargs["system_instruction"] = system_instruction
+
+        client = genai.GenerativeModel(**model_kwargs)
+
+        super().__init__(client, control_zero, config, user_context)
+
+        self._model_name = model_name
+        self._tools = tools or []
+        self._system_instruction = system_instruction
+
+    @property
+    def provider_name(self) -> str:
+        return "gemini"
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    def generate_content(
+        self,
+        contents: Union[str, List[Any]],
+        *,
+        generation_config: Optional[Dict[str, Any]] = None,
+        safety_settings: Optional[List[Dict[str, Any]]] = None,
+        tools: Optional[List[Any]] = None,
+        tool_config: Optional[Dict[str, Any]] = None,
+        stream: bool = False,
+        **kwargs,
+    ) -> Any:
+        """
+        Generate content with governance.
+
+        Args:
+            contents: Input content (text or parts)
+            generation_config: Override generation config
+            safety_settings: Override safety settings
+            tools: Override tools
+            tool_config: Tool configuration
+            stream: Whether to stream the response
+            **kwargs: Additional parameters
+
+        Returns:
+            GenerateContentResponse or stream iterator
+
+        Raises:
+            PolicyDeniedError: If request violates governance policy
+        """
+        start_time = time.time()
+
+        # Convert contents to messages format for governance checks
+        messages = self._contents_to_messages(contents)
+
+        # Estimate tokens
+        estimated_tokens = self._estimate_content_tokens(contents)
+
+        # Get tools for policy check
+        tools_to_check = tools or self._tools
+        tools_for_policy = []
+        if tools_to_check:
+            for tool in tools_to_check:
+                if hasattr(tool, "function_declarations"):
+                    for func in tool.function_declarations:
+                        tools_for_policy.append({
+                            "name": getattr(func, "name", ""),
+                            "type": "function",
+                        })
+
+        # Run pre-request governance checks
+        self._pre_request_checks(
+            model=self._model_name,
+            action=GovernanceAction.CHAT_COMPLETION,
+            messages=messages,
+            functions=tools_for_policy,
+            estimated_tokens=estimated_tokens,
+        )
+
+        # Process content for PII if enabled
+        processed_contents = self._process_content_for_governance(contents)
+
+        # Build kwargs
+        request_kwargs = {}
+        if generation_config:
+            request_kwargs["generation_config"] = generation_config
+        if safety_settings:
+            request_kwargs["safety_settings"] = safety_settings
+        if tools:
+            request_kwargs["tools"] = tools
+        if tool_config:
+            request_kwargs["tool_config"] = tool_config
+        request_kwargs.update(kwargs)
+
+        # Handle streaming
+        if stream:
+            return self._generate_stream(processed_contents, request_kwargs, start_time)
+
+        # Make the API call
+        try:
+            response = self._client.generate_content(processed_contents, **request_kwargs)
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Extract usage metrics
+            usage = getattr(response, "usage_metadata", None)
+            input_tokens = getattr(usage, "prompt_token_count", estimated_tokens) if usage else estimated_tokens
+            output_tokens = getattr(usage, "candidates_token_count", 0) if usage else 0
+            total_tokens = getattr(usage, "total_token_count", input_tokens + output_tokens) if usage else input_tokens + output_tokens
+
+            # Count function calls
+            function_call_count = 0
+            for candidate in getattr(response, "candidates", []):
+                content = getattr(candidate, "content", None)
+                if content:
+                    for part in getattr(content, "parts", []):
+                        if hasattr(part, "function_call"):
+                            function_call_count += 1
+
+            # Record metrics
+            metrics = LLMUsageMetrics(
+                provider="gemini",
+                model=self._model_name,
+                action=GovernanceAction.CHAT_COMPLETION,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                total_tokens=total_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=estimate_cost(self._model_name, input_tokens, output_tokens),
+                function_calls=function_call_count,
+            )
+
+            # Update tracking and log
+            self._post_request_update(metrics)
+            self._log_request(self._model_name, GovernanceAction.CHAT_COMPLETION, metrics)
+
+            return response
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="gemini",
+                model=self._model_name,
+                action=GovernanceAction.CHAT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._log_request(
+                self._model_name, GovernanceAction.CHAT_COMPLETION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    def _generate_stream(
+        self,
+        contents: Any,
+        request_kwargs: Dict[str, Any],
+        start_time: float,
+    ) -> Iterator[Any]:
+        """Generate a governed streaming response."""
+        total_tokens = 0
+        function_call_count = 0
+
+        try:
+            stream = self._client.generate_content(contents, stream=True, **request_kwargs)
+
+            for chunk in stream:
+                # Track tokens and function calls
+                usage = getattr(chunk, "usage_metadata", None)
+                if usage:
+                    total_tokens = getattr(usage, "total_token_count", total_tokens)
+
+                for candidate in getattr(chunk, "candidates", []):
+                    content = getattr(candidate, "content", None)
+                    if content:
+                        for part in getattr(content, "parts", []):
+                            if hasattr(part, "function_call"):
+                                function_call_count += 1
+
+                yield chunk
+
+            # Calculate final metrics
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="gemini",
+                model=self._model_name,
+                action=GovernanceAction.CHAT_COMPLETION,
+                total_tokens=total_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=estimate_cost(self._model_name, total_tokens // 2, total_tokens // 2),
+                function_calls=function_call_count,
+            )
+
+            self._post_request_update(metrics)
+            self._log_request(self._model_name, GovernanceAction.CHAT_COMPLETION, metrics)
+
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="gemini",
+                model=self._model_name,
+                action=GovernanceAction.CHAT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._log_request(
+                self._model_name, GovernanceAction.CHAT_COMPLETION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    def start_chat(
+        self,
+        history: Optional[List[Any]] = None,
+        **kwargs,
+    ) -> "GovernedChat":
+        """
+        Start a governed chat session.
+
+        Args:
+            history: Initial chat history
+            **kwargs: Additional parameters
+
+        Returns:
+            GovernedChat instance
+        """
+        chat = self._client.start_chat(history=history, **kwargs)
+        return GovernedChat(chat, self)
+
+    def _contents_to_messages(self, contents: Any) -> List[Dict[str, Any]]:
+        """Convert Gemini contents to standard message format."""
+        if isinstance(contents, str):
+            return [{"role": "user", "content": contents}]
+        elif isinstance(contents, list):
+            messages = []
+            for item in contents:
+                if isinstance(item, str):
+                    messages.append({"role": "user", "content": item})
+                elif hasattr(item, "text"):
+                    messages.append({"role": "user", "content": item.text})
+            return messages
+        return [{"role": "user", "content": str(contents)}]
+
+    def _estimate_content_tokens(self, contents: Any) -> int:
+        """Estimate token count for content."""
+        if isinstance(contents, str):
+            return len(contents) // 4
+        elif isinstance(contents, list):
+            total = 0
+            for item in contents:
+                if isinstance(item, str):
+                    total += len(item) // 4
+                elif hasattr(item, "text"):
+                    total += len(item.text) // 4
+            return total
+        return len(str(contents)) // 4
+
+    def _process_content_for_governance(self, contents: Any) -> Any:
+        """Process content for PII masking if enabled."""
+        if not self._config.content_policy.enable_pii_detection:
+            return contents
+
+        if self._config.content_policy.pii_action != "mask":
+            return contents
+
+        if isinstance(contents, str):
+            return self._mask_pii(contents)
+        elif isinstance(contents, list):
+            processed = []
+            for item in contents:
+                if isinstance(item, str):
+                    processed.append(self._mask_pii(item))
+                else:
+                    processed.append(item)
+            return processed
+        return contents
+
+    def with_user_context(self, user_context: Dict[str, Any]) -> "GovernedGemini":
+        """Create a new governed client with updated user context."""
+        merged_context = {**self._user_context, **user_context}
+        return GovernedGemini(
+            model_name=self._model_name,
+            control_zero=self._cz,
+            config=self._config,
+            user_context=merged_context,
+        )
+
+
+class GovernedChat:
+    """Governed wrapper for Gemini chat sessions."""
+
+    def __init__(self, chat: Any, governed_model: GovernedGemini):
+        self._chat = chat
+        self._governed = governed_model
+
+    def send_message(
+        self,
+        content: Union[str, List[Any]],
+        *,
+        generation_config: Optional[Dict[str, Any]] = None,
+        safety_settings: Optional[List[Dict[str, Any]]] = None,
+        stream: bool = False,
+        **kwargs,
+    ) -> Any:
+        """
+        Send a governed message in the chat.
+
+        Args:
+            content: Message content
+            generation_config: Override generation config
+            safety_settings: Override safety settings
+            stream: Whether to stream the response
+            **kwargs: Additional parameters
+
+        Returns:
+            GenerateContentResponse or stream iterator
+        """
+        start_time = time.time()
+
+        # Governance checks
+        messages = self._governed._contents_to_messages(content)
+        estimated_tokens = self._governed._estimate_content_tokens(content)
+
+        self._governed._pre_request_checks(
+            model=self._governed._model_name,
+            action=GovernanceAction.CHAT_COMPLETION,
+            messages=messages,
+            estimated_tokens=estimated_tokens,
+        )
+
+        # Process content
+        processed_content = self._governed._process_content_for_governance(content)
+
+        # Build kwargs
+        request_kwargs = {}
+        if generation_config:
+            request_kwargs["generation_config"] = generation_config
+        if safety_settings:
+            request_kwargs["safety_settings"] = safety_settings
+        request_kwargs.update(kwargs)
+
+        try:
+            if stream:
+                response = self._chat.send_message(processed_content, stream=True, **request_kwargs)
+            else:
+                response = self._chat.send_message(processed_content, **request_kwargs)
+
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            # Extract metrics
+            usage = getattr(response, "usage_metadata", None)
+            input_tokens = getattr(usage, "prompt_token_count", estimated_tokens) if usage else estimated_tokens
+            output_tokens = getattr(usage, "candidates_token_count", 0) if usage else 0
+
+            metrics = LLMUsageMetrics(
+                provider="gemini",
+                model=self._governed._model_name,
+                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(self._governed._model_name, input_tokens, output_tokens),
+            )
+
+            self._governed._post_request_update(metrics)
+            self._governed._log_request(self._governed._model_name, GovernanceAction.CHAT_COMPLETION, metrics)
+
+            return response
+
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="gemini",
+                model=self._governed._model_name,
+                action=GovernanceAction.CHAT_COMPLETION,
+                latency_ms=latency_ms,
+            )
+            self._governed._log_request(
+                self._governed._model_name, GovernanceAction.CHAT_COMPLETION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    @property
+    def history(self) -> List[Any]:
+        """Get the chat history."""
+        return self._chat.history

control_zero/llm/groq/__init__.py

@@ -0,0 +1,32 @@
+"""
+Control Zero Groq Governance Wrapper.
+
+This module provides governance wrappers for the Groq Python SDK,
+enabling policy enforcement, cost tracking, and audit logging for
+fast inference API calls.
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.groq import GovernedGroq
+    from groq import Groq
+
+    # Initialize Control Zero
+    cz_client = ControlZeroClient(api_key="cz_live_xxx")
+    cz_client.initialize()
+
+    # Wrap Groq client with governance
+    groq_client = Groq()
+    governed = GovernedGroq(client=groq_client, control_zero=cz_client)
+
+    # All calls are now governed
+    response = governed.chat.completions.create(
+        model="llama-3.1-70b-versatile",
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+"""
+
+from control_zero.llm.groq.client import GovernedGroq
+
+__all__ = [
+    "GovernedGroq",
+]