control-zero 0.2.0__py3-none-any.whl

control_zero/integrations/crewai/agent.py

@@ -0,0 +1,267 @@
+"""
+Governed CrewAI Agent wrapper.
+
+Provides governance enforcement for CrewAI agents including:
+- Role-based access control
+- Tool filtering based on agent role
+- Action logging
+- Execution tracking
+"""
+
+from __future__ import annotations
+
+import time
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional
+
+try:
+    from crewai import Agent
+    from crewai.tools import BaseTool as CrewAIBaseTool
+    CREWAI_AVAILABLE = True
+except ImportError:
+    CREWAI_AVAILABLE = False
+    class Agent:
+        pass
+    class CrewAIBaseTool:
+        pass
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+from control_zero.integrations.crewai.tool import GovernedCrewTool, wrap_crew_tools
+
+
+class GovernedCrewAgent:
+    """
+    Governance wrapper for CrewAI Agent.
+
+    Wraps a CrewAI Agent to provide:
+    - Role-based policy enforcement
+    - Tool-level governance
+    - Action audit logging
+    - Execution tracking
+
+    Usage:
+        from crewai import Agent
+
+        researcher = Agent(
+            role="Senior Researcher",
+            goal="Research topics thoroughly",
+            backstory="You are an expert researcher...",
+            tools=[search_tool, scrape_tool],
+        )
+
+        governed_researcher = GovernedCrewAgent(
+            agent=researcher,
+            client=client,
+            wrap_tools=True,
+        )
+
+        # Use in a crew
+        crew = Crew(
+            agents=[governed_researcher.agent],
+            tasks=[...],
+        )
+    """
+
+    def __init__(
+        self,
+        agent: Agent,
+        client: ControlZeroClient,
+        wrap_tools: bool = True,
+        agent_name: Optional[str] = None,
+        allowed_tools: Optional[List[str]] = None,
+    ):
+        """
+        Initialize governed agent.
+
+        Args:
+            agent: The CrewAI Agent to wrap
+            client: Control Zero client
+            wrap_tools: Whether to wrap agent tools with governance
+            agent_name: Name for logging (defaults to agent.role)
+            allowed_tools: List of allowed tool names (filters tools)
+        """
+        if not CREWAI_AVAILABLE:
+            raise ImportError(
+                "crewai is required. Install with: pip install crewai"
+            )
+
+        self._agent = agent
+        self._client = client
+        self._agent_name = agent_name or agent.role
+        self._allowed_tools = allowed_tools
+
+        # Filter tools if allowlist specified
+        if allowed_tools and agent.tools:
+            agent.tools = [
+                t for t in agent.tools
+                if getattr(t, 'name', str(t)) in allowed_tools
+            ]
+
+        # Wrap tools with governance
+        if wrap_tools and agent.tools:
+            self._wrap_agent_tools()
+
+        # Tracking
+        self._total_actions = 0
+        self._session_start = time.time()
+
+    def _wrap_agent_tools(self) -> None:
+        """Wrap all agent tools with governance."""
+        governed_tools = []
+        for tool in self._agent.tools:
+            if isinstance(tool, GovernedCrewTool):
+                governed_tools.append(tool)
+            elif isinstance(tool, CrewAIBaseTool):
+                governed_tools.append(
+                    GovernedCrewTool(
+                        base_tool=tool,
+                        client=self._client,
+                    )
+                )
+            else:
+                # Callable tools - wrap inline
+                governed_tools.append(tool)
+        self._agent.tools = governed_tools
+
+    def _check_policy(self, action: str) -> PolicyDecision:
+        """Check policy for agent action."""
+        if hasattr(self._client, '_policy_cache') and self._client._policy_cache:
+            return self._client._policy_cache.evaluate(
+                f"crewai:agent:{self._agent_name}",
+                action
+            )
+        return PolicyDecision(effect="allow")
+
+    def execute_task(
+        self,
+        task: Any,
+        context: Optional[str] = None,
+        tools: Optional[List] = None,
+    ) -> str:
+        """
+        Execute a task with governance.
+
+        Args:
+            task: The task to execute
+            context: Additional context
+            tools: Override tools for this execution
+
+        Returns:
+            Task result string
+        """
+        decision = self._check_policy("execute_task")
+
+        if decision.effect == "deny":
+            self._log_action("execute_task", "denied", 0, decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        self._total_actions += 1
+
+        try:
+            # Execute task using agent
+            result = self._agent.execute_task(
+                task=task,
+                context=context,
+                tools=tools or self._agent.tools,
+            )
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_action("execute_task", "success", latency_ms, decision, {
+                "task_description": str(task.description)[:200] if hasattr(task, 'description') else None,
+                "result_length": len(result) if result else 0,
+            })
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_action("execute_task", "error", latency_ms, decision, error=e)
+            raise
+
+    def _log_action(
+        self,
+        action: str,
+        status: str,
+        latency_ms: int,
+        decision: PolicyDecision,
+        metadata: Optional[Dict] = None,
+        error: Optional[Exception] = None,
+    ) -> None:
+        """Log agent action."""
+        log_data = metadata or {}
+        log_data["agent_role"] = self._agent.role
+        log_data["total_actions"] = self._total_actions
+
+        self._client._log(
+            tool=f"crewai:agent:{self._agent_name}",
+            method=action,
+            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 agent(self) -> Agent:
+        """Get underlying CrewAI agent."""
+        return self._agent
+
+    @property
+    def role(self) -> str:
+        """Get agent role."""
+        return self._agent.role
+
+    @property
+    def goal(self) -> str:
+        """Get agent goal."""
+        return self._agent.goal
+
+    @property
+    def tools(self) -> List:
+        """Get agent tools."""
+        return self._agent.tools
+
+    @property
+    def total_actions(self) -> int:
+        """Get total actions executed."""
+        return self._total_actions
+
+
+def governed_agent(
+    client: ControlZeroClient,
+    wrap_tools: bool = True,
+    allowed_tools: Optional[List[str]] = None,
+) -> Callable:
+    """
+    Decorator to create a governed CrewAI agent.
+
+    Usage:
+        @governed_agent(client, wrap_tools=True)
+        def create_researcher():
+            return Agent(
+                role="Researcher",
+                goal="Research topics",
+                backstory="You are an expert researcher",
+                tools=[search_tool],
+            )
+
+        researcher = create_researcher()  # Returns GovernedCrewAgent
+    """
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args, **kwargs) -> GovernedCrewAgent:
+            agent = func(*args, **kwargs)
+            return GovernedCrewAgent(
+                agent=agent,
+                client=client,
+                wrap_tools=wrap_tools,
+                allowed_tools=allowed_tools,
+            )
+        return wrapper
+    return decorator

control_zero/integrations/crewai/crew.py

@@ -0,0 +1,381 @@
+"""
+Governed CrewAI Crew wrapper.
+
+Provides governance enforcement for CrewAI crews including:
+- Crew-level policy enforcement
+- Agent role-based access control
+- Task execution governance
+- Inter-agent communication logging
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, List, Optional
+
+try:
+    from crewai import Crew, Agent, Task, Process
+    CREWAI_AVAILABLE = True
+except ImportError:
+    CREWAI_AVAILABLE = False
+    class Crew:
+        pass
+    class Agent:
+        pass
+    class Task:
+        pass
+    class Process:
+        sequential = "sequential"
+        hierarchical = "hierarchical"
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+from control_zero.integrations.crewai.agent import GovernedCrewAgent
+from control_zero.integrations.crewai.task import GovernedTask
+
+
+class GovernedCrew:
+    """
+    Governance wrapper for CrewAI Crew.
+
+    Wraps a CrewAI Crew to provide:
+    - Crew-level policy enforcement
+    - Automatic agent and task governance wrapping
+    - Execution audit logging
+    - Resource tracking
+
+    Usage:
+        from crewai import Crew, Agent, Task
+
+        researcher = Agent(role="Researcher", ...)
+        writer = Agent(role="Writer", ...)
+
+        research_task = Task(description="...", agent=researcher)
+        writing_task = Task(description="...", agent=writer)
+
+        crew = Crew(
+            agents=[researcher, writer],
+            tasks=[research_task, writing_task],
+        )
+
+        governed_crew = GovernedCrew(
+            crew=crew,
+            client=client,
+            crew_name="content_creation",
+        )
+
+        result = governed_crew.kickoff()
+    """
+
+    def __init__(
+        self,
+        crew: Crew,
+        client: ControlZeroClient,
+        crew_name: str = "crewai_crew",
+        wrap_agents: bool = True,
+        wrap_tasks: bool = True,
+        max_rpm: Optional[int] = None,
+        max_iterations: Optional[int] = None,
+        log_outputs: bool = False,
+    ):
+        """
+        Initialize governed crew.
+
+        Args:
+            crew: The CrewAI Crew to wrap
+            client: Control Zero client
+            crew_name: Name for logging purposes
+            wrap_agents: Whether to wrap agents with governance
+            wrap_tasks: Whether to wrap tasks with governance
+            max_rpm: Maximum requests per minute (rate limit)
+            max_iterations: Maximum iterations for the crew
+            log_outputs: Whether to log output content
+        """
+        if not CREWAI_AVAILABLE:
+            raise ImportError(
+                "crewai is required. Install with: pip install crewai"
+            )
+
+        self._crew = crew
+        self._client = client
+        self._crew_name = crew_name
+        self._log_outputs = log_outputs
+
+        # Apply limits
+        if max_rpm:
+            self._crew.max_rpm = max_rpm
+        if max_iterations:
+            self._crew.max_iterations = max_iterations
+
+        # Wrap agents and tasks with governance
+        if wrap_agents:
+            self._governed_agents = self._wrap_agents()
+        else:
+            self._governed_agents = []
+
+        if wrap_tasks:
+            self._governed_tasks = self._wrap_tasks()
+        else:
+            self._governed_tasks = []
+
+        # Tracking
+        self._kickoff_count = 0
+        self._total_tokens = 0
+        self._session_start = time.time()
+
+    def _wrap_agents(self) -> List[GovernedCrewAgent]:
+        """Wrap crew agents with governance."""
+        governed = []
+        for agent in self._crew.agents:
+            if isinstance(agent, GovernedCrewAgent):
+                governed.append(agent)
+            else:
+                governed.append(
+                    GovernedCrewAgent(
+                        agent=agent,
+                        client=self._client,
+                        wrap_tools=True,
+                    )
+                )
+        return governed
+
+    def _wrap_tasks(self) -> List[GovernedTask]:
+        """Wrap crew tasks with governance."""
+        governed = []
+        for task in self._crew.tasks:
+            if isinstance(task, GovernedTask):
+                governed.append(task)
+            else:
+                governed.append(
+                    GovernedTask(
+                        task=task,
+                        client=self._client,
+                        log_outputs=self._log_outputs,
+                    )
+                )
+        return governed
+
+    def _check_policy(self, action: str) -> PolicyDecision:
+        """Check policy for crew action."""
+        if hasattr(self._client, '_policy_cache') and self._client._policy_cache:
+            return self._client._policy_cache.evaluate(
+                f"crewai:crew:{self._crew_name}",
+                action
+            )
+        return PolicyDecision(effect="allow")
+
+    def kickoff(
+        self,
+        inputs: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """
+        Execute the crew with governance.
+
+        Args:
+            inputs: Input variables for the crew
+
+        Returns:
+            Crew execution result
+        """
+        decision = self._check_policy("kickoff")
+
+        if decision.effect == "deny":
+            self._log_kickoff("denied", 0, decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        self._kickoff_count += 1
+
+        try:
+            # Execute crew
+            result = self._crew.kickoff(inputs=inputs)
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+
+            metadata = {
+                "kickoff_count": self._kickoff_count,
+                "num_agents": len(self._crew.agents),
+                "num_tasks": len(self._crew.tasks),
+                "process_type": str(getattr(self._crew, 'process', 'sequential')),
+            }
+
+            if self._log_outputs and result:
+                metadata["result_preview"] = str(result)[:500]
+
+            self._log_kickoff("success", latency_ms, decision, metadata)
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_kickoff("error", latency_ms, decision, error=e)
+            raise
+
+    async def akickoff(
+        self,
+        inputs: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """
+        Async execute the crew with governance.
+
+        Args:
+            inputs: Input variables for the crew
+
+        Returns:
+            Crew execution result
+        """
+        decision = self._check_policy("kickoff")
+
+        if decision.effect == "deny":
+            self._log_kickoff("denied", 0, decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        self._kickoff_count += 1
+
+        try:
+            # Execute crew async
+            if hasattr(self._crew, 'akickoff'):
+                result = await self._crew.akickoff(inputs=inputs)
+            else:
+                # Fallback to sync
+                result = self._crew.kickoff(inputs=inputs)
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+
+            metadata = {
+                "kickoff_count": self._kickoff_count,
+                "num_agents": len(self._crew.agents),
+                "num_tasks": len(self._crew.tasks),
+                "async": True,
+            }
+
+            if self._log_outputs and result:
+                metadata["result_preview"] = str(result)[:500]
+
+            self._log_kickoff("success", latency_ms, decision, metadata)
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_kickoff("error", latency_ms, decision, error=e)
+            raise
+
+    def kickoff_for_each(
+        self,
+        inputs: List[Dict[str, Any]],
+    ) -> List[Any]:
+        """
+        Execute crew for each input set with governance.
+
+        Args:
+            inputs: List of input dictionaries
+
+        Returns:
+            List of results
+        """
+        decision = self._check_policy("kickoff_for_each")
+
+        if decision.effect == "deny":
+            self._log_kickoff("denied", 0, decision, {"action": "kickoff_for_each"})
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        results = []
+
+        try:
+            for input_set in inputs:
+                result = self._crew.kickoff(inputs=input_set)
+                results.append(result)
+                self._kickoff_count += 1
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+
+            self._log_kickoff("success", latency_ms, decision, {
+                "action": "kickoff_for_each",
+                "batch_size": len(inputs),
+                "successful": len(results),
+            })
+
+            return results
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_kickoff("error", latency_ms, decision, {
+                "action": "kickoff_for_each",
+                "completed": len(results),
+            }, error=e)
+            raise
+
+    def _log_kickoff(
+        self,
+        status: str,
+        latency_ms: int,
+        decision: PolicyDecision,
+        metadata: Optional[Dict] = None,
+        error: Optional[Exception] = None,
+    ) -> None:
+        """Log crew kickoff."""
+        log_data = metadata or {}
+        log_data["crew_name"] = self._crew_name
+
+        self._client._log(
+            tool=f"crewai:crew:{self._crew_name}",
+            method="kickoff",
+            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 crew(self) -> Crew:
+        """Get underlying CrewAI crew."""
+        return self._crew
+
+    @property
+    def agents(self) -> List:
+        """Get crew agents."""
+        return self._crew.agents
+
+    @property
+    def tasks(self) -> List:
+        """Get crew tasks."""
+        return self._crew.tasks
+
+    @property
+    def governed_agents(self) -> List[GovernedCrewAgent]:
+        """Get governed agent wrappers."""
+        return self._governed_agents
+
+    @property
+    def governed_tasks(self) -> List[GovernedTask]:
+        """Get governed task wrappers."""
+        return self._governed_tasks
+
+    @property
+    def kickoff_count(self) -> int:
+        """Get total kickoff count."""
+        return self._kickoff_count
+
+    @property
+    def usage_metrics(self) -> Dict[str, Any]:
+        """Get usage metrics for the crew."""
+        session_duration = time.time() - self._session_start
+        return {
+            "crew_name": self._crew_name,
+            "kickoff_count": self._kickoff_count,
+            "num_agents": len(self._crew.agents),
+            "num_tasks": len(self._crew.tasks),
+            "session_duration_seconds": round(session_duration, 2),
+            "total_tokens": self._total_tokens,
+        }