control-zero 0.2.0__py3-none-any.whl

control_zero/llm/cohere/client.py

@@ -0,0 +1,402 @@
+"""
+Governed Cohere client wrapper.
+
+Provides governance features for the Cohere Python SDK including:
+- Model access control
+- Cost tracking and limits
+- Tool use governance
+- PII detection and masking
+- Audit logging
+
+Supports chat, RAG (retrieval-augmented generation), and reranking.
+"""
+
+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 GovernedCohere(GovernedLLM, GovernedChatMixin):
+    """
+    Governed wrapper for the Cohere Python SDK.
+
+    Supports:
+    - Chat API (command models)
+    - RAG with connectors
+    - Reranking
+    - Embeddings
+
+    Example:
+        from control_zero import ControlZeroClient
+        from control_zero.llm.cohere import GovernedCohere
+        import cohere
+
+        cz = ControlZeroClient(api_key="...")
+        cz.initialize()
+
+        client = cohere.ClientV2()
+        governed = GovernedCohere(client=client, control_zero=cz)
+
+        response = governed.chat(
+            model="command-r-plus",
+            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)
+
+    @property
+    def provider_name(self) -> str:
+        return "cohere"
+
+    def chat(
+        self,
+        *,
+        model: str,
+        messages: List[Dict[str, Any]],
+        tools: Optional[List[Dict[str, Any]]] = None,
+        documents: Optional[List[Dict[str, Any]]] = None,
+        connectors: Optional[List[Dict[str, Any]]] = None,
+        stream: bool = False,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        p: Optional[float] = None,
+        k: Optional[int] = None,
+        stop_sequences: Optional[List[str]] = None,
+        seed: Optional[int] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Create a governed chat completion.
+
+        Args:
+            model: Model to use (e.g., "command-r-plus")
+            messages: List of message dicts
+            tools: Tool definitions
+            documents: Documents for RAG
+            connectors: Connector configurations
+            stream: Whether to stream
+            max_tokens: Maximum output tokens
+            temperature: Sampling temperature
+            p: Nucleus sampling
+            k: Top-k sampling
+            stop_sequences: Stop sequences
+            seed: Random seed
+            **kwargs: Additional parameters
+
+        Returns:
+            ChatResponse or stream iterator
+
+        Raises:
+            PolicyDeniedError: If request violates policy
+        """
+        start_time = time.time()
+
+        # Estimate tokens
+        estimated_input_tokens = self._estimate_message_tokens(messages)
+        if documents:
+            for doc in documents:
+                estimated_input_tokens += len(str(doc.get("text", ""))) // 4
+
+        # Prepare tools for policy check
+        tools_to_check = []
+        if tools:
+            tools_to_check = [{"name": t.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_input_tokens,
+        )
+
+        # Process messages for governance
+        processed_messages = self._process_messages_for_governance(messages)
+
+        # Filter tools
+        filtered_tools = self._filter_tools_for_governance(tools)
+
+        # Apply max_tokens limit
+        governed_max_tokens = max_tokens
+        if self._config.content_policy.max_output_tokens:
+            if max_tokens:
+                governed_max_tokens = min(max_tokens, self._config.content_policy.max_output_tokens)
+            else:
+                governed_max_tokens = self._config.content_policy.max_output_tokens
+
+        # Build request
+        request_kwargs = {
+            "model": model,
+            "messages": processed_messages,
+        }
+
+        if filtered_tools:
+            request_kwargs["tools"] = filtered_tools
+        if documents:
+            request_kwargs["documents"] = documents
+        if connectors:
+            request_kwargs["connectors"] = connectors
+        if governed_max_tokens:
+            request_kwargs["max_tokens"] = governed_max_tokens
+        if temperature is not None:
+            request_kwargs["temperature"] = temperature
+        if p is not None:
+            request_kwargs["p"] = p
+        if k is not None:
+            request_kwargs["k"] = k
+        if stop_sequences:
+            request_kwargs["stop_sequences"] = stop_sequences
+        if seed is not None:
+            request_kwargs["seed"] = seed
+
+        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
+            usage = getattr(response, "usage", None)
+            input_tokens = getattr(usage, "billed_units", {}).get("input_tokens", estimated_input_tokens) if usage else estimated_input_tokens
+            output_tokens = getattr(usage, "billed_units", {}).get("output_tokens", 0) if usage else 0
+            total_tokens = input_tokens + output_tokens
+
+            # Count tool calls
+            tool_call_count = 0
+            if hasattr(response, "tool_calls") and response.tool_calls:
+                tool_call_count = len(response.tool_calls)
+
+            metrics = LLMUsageMetrics(
+                provider="cohere",
+                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._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="cohere",
+                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."""
+        total_tokens = 0
+
+        try:
+            stream = self._client.chat_stream(**request_kwargs)
+
+            for event in stream:
+                total_tokens += 1
+                yield event
+
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="cohere",
+                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),
+            )
+
+            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="cohere",
+                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 rerank(
+        self,
+        *,
+        model: str,
+        query: str,
+        documents: List[Union[str, Dict[str, Any]]],
+        top_n: Optional[int] = None,
+        return_documents: bool = True,
+        **kwargs,
+    ) -> Any:
+        """
+        Rerank documents with governance.
+
+        Args:
+            model: Rerank model (e.g., "rerank-v3.5")
+            query: Query to rerank against
+            documents: Documents to rerank
+            top_n: Number of results to return
+            return_documents: Whether to return document content
+            **kwargs: Additional parameters
+
+        Returns:
+            RerankResponse
+        """
+        start_time = time.time()
+
+        # Governance checks
+        estimated_tokens = len(query) // 4 + sum(len(str(d)) // 4 for d in documents)
+
+        self._pre_request_checks(
+            model=model,
+            action=GovernanceAction.MODERATION,
+            estimated_tokens=estimated_tokens,
+        )
+
+        try:
+            response = self._client.rerank(
+                model=model,
+                query=query,
+                documents=documents,
+                top_n=top_n,
+                return_documents=return_documents,
+                **kwargs,
+            )
+
+            latency_ms = int((time.time() - start_time) * 1000)
+
+            metrics = LLMUsageMetrics(
+                provider="cohere",
+                model=model,
+                action=GovernanceAction.MODERATION,
+                input_tokens=estimated_tokens,
+                latency_ms=latency_ms,
+                estimated_cost=0.0,  # Rerank pricing is different
+            )
+
+            self._post_request_update(metrics)
+            self._log_request(model, GovernanceAction.MODERATION, metrics)
+
+            return response
+
+        except Exception as e:
+            latency_ms = int((time.time() - start_time) * 1000)
+            metrics = LLMUsageMetrics(
+                provider="cohere",
+                model=model,
+                action=GovernanceAction.MODERATION,
+                latency_ms=latency_ms,
+            )
+            self._log_request(
+                model, GovernanceAction.MODERATION, metrics,
+                status="error", error=str(e)
+            )
+            raise
+
+    @property
+    def embed(self):
+        """Access the underlying embed endpoint (pass-through)."""
+        return self._client.embed
+
+    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("name", "")
+
+            # Skip denied
+            if policy.denied_functions:
+                denied = any(d.lower() in tool_name.lower() for d in policy.denied_functions)
+                if denied:
+                    continue
+
+            # Check allowed
+            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]) -> "GovernedCohere":
+        merged_context = {**self._user_context, **user_context}
+        return GovernedCohere(
+            client=self._client,
+            control_zero=self._cz,
+            config=self._config,
+            user_context=merged_context,
+        )
+
+    def with_config(self, config: LLMGovernanceConfig) -> "GovernedCohere":
+        return GovernedCohere(
+            client=self._client,
+            control_zero=self._cz,
+            config=config,
+            user_context=self._user_context,
+        )

control_zero/llm/gemini/__init__.py

@@ -0,0 +1,34 @@
+"""
+Control Zero Google Gemini Governance Wrapper.
+
+This module provides governance wrappers for the Google Generative AI SDK,
+enabling policy enforcement, cost tracking, and audit logging for
+all Gemini API calls.
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.llm.gemini import GovernedGemini
+    import google.generativeai as genai
+
+    # Initialize Control Zero
+    cz_client = ControlZeroClient(api_key="cz_live_xxx")
+    cz_client.initialize()
+
+    # Configure Gemini
+    genai.configure(api_key="your-gemini-api-key")
+
+    # Create governed model
+    governed = GovernedGemini(
+        model_name="gemini-1.5-pro",
+        control_zero=cz_client
+    )
+
+    # All calls are now governed
+    response = governed.generate_content("Hello, how are you?")
+"""
+
+from control_zero.llm.gemini.client import GovernedGemini
+
+__all__ = [
+    "GovernedGemini",
+]