control-zero 0.2.0__py3-none-any.whl

control_zero/integrations/crewai/task.py

@@ -0,0 +1,291 @@
+"""
+Governed CrewAI Task wrapper.
+
+Provides governance enforcement for CrewAI tasks including:
+- Task-level policy checks
+- Delegation governance
+- Output validation
+- Execution logging
+"""
+
+from __future__ import annotations
+
+import time
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional
+
+try:
+    from crewai import Task, Agent
+    CREWAI_AVAILABLE = True
+except ImportError:
+    CREWAI_AVAILABLE = False
+    class Task:
+        pass
+    class Agent:
+        pass
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+
+
+class GovernedTask:
+    """
+    Governance wrapper for CrewAI Task.
+
+    Wraps a CrewAI Task to provide:
+    - Task execution policy enforcement
+    - Delegation governance
+    - Output validation
+    - Execution audit logging
+
+    Usage:
+        from crewai import Task
+
+        research_task = Task(
+            description="Research the topic...",
+            agent=researcher,
+            expected_output="A detailed report...",
+        )
+
+        governed_task = GovernedTask(
+            task=research_task,
+            client=client,
+            task_name="research_task",
+        )
+    """
+
+    def __init__(
+        self,
+        task: Task,
+        client: ControlZeroClient,
+        task_name: Optional[str] = None,
+        allow_delegation: bool = True,
+        max_delegation_depth: int = 3,
+        log_outputs: bool = False,
+    ):
+        """
+        Initialize governed task.
+
+        Args:
+            task: The CrewAI Task to wrap
+            client: Control Zero client
+            task_name: Name for logging purposes
+            allow_delegation: Whether to allow task delegation
+            max_delegation_depth: Maximum delegation depth
+            log_outputs: Whether to log output content
+        """
+        if not CREWAI_AVAILABLE:
+            raise ImportError(
+                "crewai is required. Install with: pip install crewai"
+            )
+
+        self._task = task
+        self._client = client
+        self._task_name = task_name or self._extract_task_name(task)
+        self._allow_delegation = allow_delegation
+        self._max_delegation_depth = max_delegation_depth
+        self._log_outputs = log_outputs
+
+        # Tracking
+        self._delegation_depth = 0
+        self._execution_count = 0
+
+    def _extract_task_name(self, task: Task) -> str:
+        """Extract a name from task description."""
+        if hasattr(task, 'description'):
+            # Take first 30 chars of description
+            desc = str(task.description)[:30]
+            return desc.replace(" ", "_").lower()
+        return "unnamed_task"
+
+    def _check_policy(self, action: str) -> PolicyDecision:
+        """Check policy for task action."""
+        if hasattr(self._client, '_policy_cache') and self._client._policy_cache:
+            return self._client._policy_cache.evaluate(
+                f"crewai:task:{self._task_name}",
+                action
+            )
+        return PolicyDecision(effect="allow")
+
+    def execute(
+        self,
+        agent: Optional[Agent] = None,
+        context: Optional[str] = None,
+        tools: Optional[List] = None,
+    ) -> Any:
+        """
+        Execute the task with governance.
+
+        Args:
+            agent: Agent to execute (overrides task.agent)
+            context: Additional context
+            tools: Override tools for execution
+
+        Returns:
+            Task output
+        """
+        decision = self._check_policy("execute")
+
+        if decision.effect == "deny":
+            self._log_execution("denied", 0, decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        self._execution_count += 1
+
+        try:
+            # Check delegation policy
+            if self._delegation_depth >= self._max_delegation_depth:
+                raise PolicyDeniedError(PolicyDecision(
+                    effect="deny",
+                    reason=f"Maximum delegation depth ({self._max_delegation_depth}) exceeded"
+                ))
+
+            # Execute task
+            executing_agent = agent or self._task.agent
+            result = self._task.execute(
+                agent=executing_agent,
+                context=context,
+                tools=tools,
+            )
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+
+            metadata = {
+                "execution_count": self._execution_count,
+                "delegation_depth": self._delegation_depth,
+            }
+
+            if self._log_outputs and result:
+                metadata["output_preview"] = str(result)[:200]
+
+            self._log_execution("success", latency_ms, decision, metadata)
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_execution("error", latency_ms, decision, error=e)
+            raise
+
+    def delegate(
+        self,
+        to_agent: Agent,
+        context: Optional[str] = None,
+    ) -> Any:
+        """
+        Delegate task to another agent with governance.
+
+        Args:
+            to_agent: Agent to delegate to
+            context: Context for delegation
+
+        Returns:
+            Delegation result
+        """
+        if not self._allow_delegation:
+            decision = PolicyDecision(
+                effect="deny",
+                reason="Task delegation is not allowed"
+            )
+            self._log_execution("denied", 0, decision, {"reason": "delegation_disabled"})
+            raise PolicyDeniedError(decision)
+
+        decision = self._check_policy("delegate")
+
+        if decision.effect == "deny":
+            self._log_execution("denied", 0, decision, {"action": "delegate"})
+            raise PolicyDeniedError(decision)
+
+        # Increment delegation depth
+        self._delegation_depth += 1
+
+        try:
+            return self.execute(agent=to_agent, context=context)
+        finally:
+            self._delegation_depth -= 1
+
+    def _log_execution(
+        self,
+        status: str,
+        latency_ms: int,
+        decision: PolicyDecision,
+        metadata: Optional[Dict] = None,
+        error: Optional[Exception] = None,
+    ) -> None:
+        """Log task execution."""
+        log_data = metadata or {}
+        log_data["task_name"] = self._task_name
+
+        if hasattr(self._task, 'description'):
+            log_data["description_preview"] = str(self._task.description)[:100]
+
+        self._client._log(
+            tool=f"crewai:task:{self._task_name}",
+            method="execute",
+            status=status,
+            latency_ms=latency_ms,
+            policy_decision=decision,
+            error_type=type(error).__name__ if error else None,
+            error_message=str(error) if error else None,
+            **log_data
+        )
+
+    @property
+    def task(self) -> Task:
+        """Get underlying CrewAI task."""
+        return self._task
+
+    @property
+    def description(self) -> str:
+        """Get task description."""
+        return getattr(self._task, 'description', '')
+
+    @property
+    def agent(self) -> Optional[Agent]:
+        """Get assigned agent."""
+        return getattr(self._task, 'agent', None)
+
+    @property
+    def execution_count(self) -> int:
+        """Get total execution count."""
+        return self._execution_count
+
+
+def governed_task(
+    client: ControlZeroClient,
+    task_name: Optional[str] = None,
+    allow_delegation: bool = True,
+    max_delegation_depth: int = 3,
+    log_outputs: bool = False,
+) -> Callable:
+    """
+    Decorator to create a governed CrewAI task.
+
+    Usage:
+        @governed_task(client, task_name="research")
+        def create_research_task(agent):
+            return Task(
+                description="Research the topic...",
+                agent=agent,
+                expected_output="A detailed report",
+            )
+
+        task = create_research_task(researcher)  # Returns GovernedTask
+    """
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args, **kwargs) -> GovernedTask:
+            task = func(*args, **kwargs)
+            return GovernedTask(
+                task=task,
+                client=client,
+                task_name=task_name,
+                allow_delegation=allow_delegation,
+                max_delegation_depth=max_delegation_depth,
+                log_outputs=log_outputs,
+            )
+        return wrapper
+    return decorator

control_zero/integrations/crewai/tool.py

@@ -0,0 +1,299 @@
+"""
+Governed CrewAI Tool wrappers.
+
+Provides governance enforcement for CrewAI tools including:
+- Policy checks before execution
+- Audit logging
+- Error tracking
+- Secret injection
+"""
+
+from __future__ import annotations
+
+import time
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, Type
+
+try:
+    from crewai.tools import BaseTool as CrewAIBaseTool
+    from pydantic import BaseModel, Field
+    CREWAI_AVAILABLE = True
+except ImportError:
+    CREWAI_AVAILABLE = False
+    class CrewAIBaseTool:
+        pass
+    class BaseModel:
+        pass
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+
+
+class GovernedCrewTool(CrewAIBaseTool):
+    """
+    Wraps a CrewAI Tool to enforce Control Zero governance policies.
+
+    Usage:
+        from crewai.tools import BaseTool
+
+        class MyTool(BaseTool):
+            name: str = "my_tool"
+            description: str = "Does something"
+
+            def _run(self, query: str) -> str:
+                return "result"
+
+        governed_tool = GovernedCrewTool(
+            base_tool=MyTool(),
+            client=client,
+        )
+    """
+
+    name: str = ""
+    description: str = ""
+    base_tool: Any = None
+    client: Any = None
+    tool_name: str = ""
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __init__(
+        self,
+        base_tool: CrewAIBaseTool,
+        client: ControlZeroClient,
+        tool_name: Optional[str] = None,
+        **kwargs
+    ):
+        if not CREWAI_AVAILABLE:
+            raise ImportError(
+                "crewai is required for CrewAI integration. "
+                "Install with: pip install crewai"
+            )
+
+        name = base_tool.name
+        description = base_tool.description
+        tool_name = tool_name or name
+
+        super().__init__(
+            name=name,
+            description=description,
+            base_tool=base_tool,
+            client=client,
+            tool_name=tool_name,
+            **kwargs
+        )
+
+    def _run(self, *args: Any, **kwargs: Any) -> Any:
+        """Execute the tool with governance checks."""
+        method = "run"
+
+        # 1. Check Policy
+        decision = self._check_policy(method)
+
+        if decision.effect == "deny":
+            self._log_denied(method, decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        try:
+            # 2. Inject secrets if needed
+            kwargs = self._inject_secrets(kwargs)
+
+            # 3. Run Tool
+            result = self.base_tool._run(*args, **kwargs)
+
+            # 4. Log Success
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_success(method, latency_ms, decision)
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            # 5. Log Error
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_error(method, latency_ms, decision, e)
+            raise
+
+    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
+        """Async execution with governance."""
+        method = "run"
+
+        decision = self._check_policy(method)
+
+        if decision.effect == "deny":
+            self._log_denied(method, decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        try:
+            kwargs = self._inject_secrets(kwargs)
+
+            # Check if base tool has async support
+            if hasattr(self.base_tool, '_arun'):
+                result = await self.base_tool._arun(*args, **kwargs)
+            else:
+                result = self.base_tool._run(*args, **kwargs)
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_success(method, latency_ms, decision)
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_error(method, latency_ms, decision, e)
+            raise
+
+    def _check_policy(self, method: str) -> PolicyDecision:
+        """Check policy for tool execution."""
+        if hasattr(self.client, '_policy_cache') and self.client._policy_cache:
+            return self.client._policy_cache.evaluate(self.tool_name, method)
+        return PolicyDecision(effect="allow")
+
+    def _inject_secrets(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        """Inject secrets into kwargs if configured."""
+        if hasattr(self.client, '_secret_manager') and self.client._secret_manager:
+            return self.client._secret_manager.inject(self.tool_name, kwargs)
+        return kwargs
+
+    def _log_denied(self, method: str, decision: PolicyDecision) -> None:
+        """Log a denied execution."""
+        self.client._log(
+            tool=self.tool_name,
+            method=method,
+            status="denied",
+            latency_ms=0,
+            policy_decision=decision
+        )
+
+    def _log_success(self, method: str, latency_ms: int, decision: PolicyDecision) -> None:
+        """Log successful execution."""
+        self.client._log(
+            tool=self.tool_name,
+            method=method,
+            status="success",
+            latency_ms=latency_ms,
+            policy_decision=decision
+        )
+
+    def _log_error(
+        self,
+        method: str,
+        latency_ms: int,
+        decision: PolicyDecision,
+        error: Exception
+    ) -> None:
+        """Log execution error."""
+        self.client._log(
+            tool=self.tool_name,
+            method=method,
+            status="error",
+            latency_ms=latency_ms,
+            policy_decision=decision,
+            error_type=type(error).__name__,
+            error_message=str(error)
+        )
+
+
+def governed_crew_tool(
+    client: ControlZeroClient,
+    tool_name: Optional[str] = None,
+) -> Callable:
+    """
+    Decorator to create a governed CrewAI tool from a function.
+
+    Usage:
+        @governed_crew_tool(client, tool_name="search")
+        def search_web(query: str) -> str:
+            '''Search the web for information.'''
+            return search(query)
+    """
+    def decorator(func: Callable) -> Callable:
+        if not CREWAI_AVAILABLE:
+            raise ImportError("crewai is required")
+
+        name = tool_name or func.__name__
+        description = func.__doc__ or ""
+
+        @wraps(func)
+        def governed_func(*args, **kwargs):
+            method = "run"
+            decision = PolicyDecision(effect="allow")
+
+            if hasattr(client, '_policy_cache') and client._policy_cache:
+                decision = client._policy_cache.evaluate(name, method)
+
+            if decision.effect == "deny":
+                client._log(
+                    tool=name,
+                    method=method,
+                    status="denied",
+                    latency_ms=0,
+                    policy_decision=decision
+                )
+                raise PolicyDeniedError(decision)
+
+            start = time.perf_counter()
+            try:
+                # Inject secrets
+                if hasattr(client, '_secret_manager') and client._secret_manager:
+                    kwargs = client._secret_manager.inject(name, kwargs)
+
+                result = func(*args, **kwargs)
+
+                latency_ms = int((time.perf_counter() - start) * 1000)
+                client._log(
+                    tool=name,
+                    method=method,
+                    status="success",
+                    latency_ms=latency_ms,
+                    policy_decision=decision
+                )
+
+                return result
+
+            except PolicyDeniedError:
+                raise
+            except Exception as e:
+                latency_ms = int((time.perf_counter() - start) * 1000)
+                client._log(
+                    tool=name,
+                    method=method,
+                    status="error",
+                    latency_ms=latency_ms,
+                    policy_decision=decision,
+                    error_type=type(e).__name__,
+                    error_message=str(e)
+                )
+                raise
+
+        # Return as a simple callable with metadata
+        governed_func.name = name
+        governed_func.description = description
+        governed_func.is_governed = True
+
+        return governed_func
+
+    return decorator
+
+
+def wrap_crew_tools(
+    tools: List[CrewAIBaseTool],
+    client: ControlZeroClient,
+) -> List[GovernedCrewTool]:
+    """
+    Wrap multiple CrewAI tools with governance.
+
+    Usage:
+        tools = [tool1, tool2, tool3]
+        governed_tools = wrap_crew_tools(tools, client)
+    """
+    return [
+        GovernedCrewTool(base_tool=tool, client=client)
+        for tool in tools
+    ]

control_zero/integrations/langchain/__init__.py

@@ -0,0 +1,58 @@
+"""
+Control Zero LangChain Integration.
+
+This module provides governance wrappers for LangChain components including:
+- Tools (GovernedTool, GovernanceTool)
+- Agents (GovernedAgent)
+- Chains (GovernedChain)
+- LangGraph (GovernedNode, governed_graph)
+- Callbacks (ControlZeroCallbackHandler)
+
+Usage:
+    from control_zero import ControlZeroClient
+    from control_zero.integrations.langchain import (
+        GovernedTool,
+        GovernedAgent,
+        ControlZeroCallbackHandler,
+    )
+
+    # Initialize Control Zero
+    client = ControlZeroClient(api_key="...")
+    client.initialize()
+
+    # Wrap tools with governance
+    governed_tool = GovernedTool(
+        tool=my_tool,
+        client=client,
+    )
+
+    # Use callback handler for comprehensive logging
+    callback = ControlZeroCallbackHandler(client)
+    agent.run("query", callbacks=[callback])
+"""
+
+from control_zero.integrations.langchain.tool import GovernanceTool, GovernedTool
+from control_zero.integrations.langchain.agent import GovernedAgent
+from control_zero.integrations.langchain.chain import GovernedChain
+from control_zero.integrations.langchain.callbacks import ControlZeroCallbackHandler
+from control_zero.integrations.langchain.graph import (
+    GovernedNode,
+    governed_graph,
+    GovernedStateGraph,
+)
+
+__all__ = [
+    # Tools
+    "GovernanceTool",
+    "GovernedTool",
+    # Agents
+    "GovernedAgent",
+    # Chains
+    "GovernedChain",
+    # Callbacks
+    "ControlZeroCallbackHandler",
+    # LangGraph
+    "GovernedNode",
+    "governed_graph",
+    "GovernedStateGraph",
+]