controlzero 1.0.0__tar.gz → 1.1.0__tar.gz

{controlzero-1.0.0 → controlzero-1.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: controlzero
-Version: 1.0.0
+Version: 1.1.0
 Summary: AI agent governance: policies, audit, and observability for tool calls. Works locally with no signup.
 Project-URL: Homepage, https://controlzero.ai
 Project-URL: Documentation, https://docs.controlzero.ai

{controlzero-1.0.0 → controlzero-1.1.0}/controlzero/__init__.py

@@ -28,7 +28,7 @@ from controlzero.errors import (
 )
 from controlzero.policy_loader import load_policy
 
-__version__ = "1.0.0"
+__version__ = "1.1.0"
 
 __all__ = [
     "Client",

controlzero-1.1.0/controlzero/integrations/__init__.py

@@ -0,0 +1,20 @@
+"""ControlZero framework integrations.
+
+Each submodule provides a thin wrapper or callback handler for a
+specific AI framework. Import directly from the submodule you need:
+
+    from controlzero.integrations.openai import wrap_openai
+    from controlzero.integrations.anthropic import wrap_anthropic
+    from controlzero.integrations.google import wrap_google
+    from controlzero.integrations.langchain import ControlZeroCallbackHandler
+    from controlzero.integrations.crewai import create_step_callback
+    from controlzero.integrations.litellm import ControlZeroLiteLLMCallback
+    from controlzero.integrations.google_adk import GovernedAgent, GovernedTool
+    from controlzero.integrations.vercel_ai import ControlZeroVercelAIMiddleware
+    from controlzero.integrations.braintrust import BraintrustGovernanceLogger
+    from controlzero.integrations.langfuse import LangfuseGovernanceHandler
+
+No framework dependencies are installed by default. Install the
+framework you need separately (e.g. ``pip install openai``,
+``pip install google-adk``, ``pip install langfuse``).
+"""

controlzero-1.1.0/controlzero/integrations/anthropic.py

@@ -0,0 +1,188 @@
+"""Anthropic SDK integration for ControlZero.
+
+Provides ``wrap_anthropic()`` which transparently intercepts
+``client.messages.create()`` and ``client.messages.stream()``
+to enforce governance policies before the request reaches Anthropic.
+
+All other client methods pass through unchanged.
+
+Usage::
+
+    from controlzero import Client
+    from controlzero.integrations.anthropic import wrap_anthropic
+    import anthropic
+
+    cz = Client(policy={"rules": [{"deny": "delete_*"}]})
+    client = wrap_anthropic(anthropic.Anthropic(), cz)
+
+    # Policy enforcement happens automatically
+    response = client.messages.create(
+        model="claude-sonnet-4-20250514",
+        max_tokens=1024,
+        messages=[{"role": "user", "content": "Hello"}],
+    )
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+try:
+    import anthropic as _anthropic  # noqa: F401
+except ImportError:
+    raise ImportError(
+        "The 'anthropic' package is required for this integration. "
+        "Install it with: pip install anthropic"
+    )
+
+from controlzero.client import Client
+
+
+def _estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
+    """Rough cost estimate fallback (USD per token)."""
+    return (input_tokens + output_tokens) / 1_000_000 * 1.0
+
+
+class _WrappedMessages:
+    """Proxy for ``client.messages`` that enforces policy on create() and stream()."""
+
+    __slots__ = ("_inner", "_cz", "_agent_id")
+
+    def __init__(
+        self,
+        inner: Any,
+        cz: Client,
+        agent_id: str,
+    ) -> None:
+        self._inner = inner
+        self._cz = cz
+        self._agent_id = agent_id
+
+    def _enforce_model(self, kwargs: dict, method: str) -> None:
+        """Enforce policy for a messages API call."""
+        model = kwargs.get("model", "unknown")
+        context: dict[str, Any] = {
+            "agent_id": self._agent_id,
+            "provider": "anthropic",
+            "method": method,
+            "resource": f"model/{model}",
+        }
+        tools = kwargs.get("tools")
+        if tools:
+            context["tool_count"] = len(tools)
+            context["tool_names"] = [
+                t.get("name", "") for t in tools if isinstance(t, dict)
+            ]
+        self._cz.guard(
+            tool="llm:anthropic",
+            method=method,
+            raise_on_deny=True,
+            context=context,
+        )
+
+    def create(self, **kwargs: Any) -> Any:
+        """Enforce policy then delegate to the original create()."""
+        model = kwargs.get("model", "unknown")
+        self._enforce_model(kwargs, "messages.create")
+        start = time.perf_counter()
+        response = self._inner.create(**kwargs)
+        latency_ms = int((time.perf_counter() - start) * 1000)
+
+        # Capture token usage from the response
+        input_tokens = 0
+        output_tokens = 0
+        if hasattr(response, "usage") and response.usage is not None:
+            input_tokens = getattr(response.usage, "input_tokens", 0) or 0
+            output_tokens = getattr(response.usage, "output_tokens", 0) or 0
+
+        # Audit the successful call
+        try:
+            from controlzero._internal.enforcer import PolicyDecision
+            decision = PolicyDecision(
+                effect="allow",
+                reason="guard passed",
+            )
+            self._cz._audit_decision(
+                tool="llm:anthropic",
+                method="messages.create",
+                args={
+                    "model": model,
+                    "input_tokens": input_tokens,
+                    "output_tokens": output_tokens,
+                    "latency_ms": latency_ms,
+                },
+                decision=decision,
+            )
+        except Exception:
+            pass
+
+        return response
+
+    def stream(self, **kwargs: Any) -> Any:
+        """Enforce policy then delegate to the original stream()."""
+        self._enforce_model(kwargs, "messages.stream")
+        return self._inner.stream(**kwargs)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._inner, name)
+
+
+class _WrappedAnthropic:
+    """Proxy for ``anthropic.Anthropic`` that wraps the messages resource."""
+
+    __slots__ = ("_inner", "_messages")
+
+    def __init__(
+        self,
+        inner: Any,
+        cz: Client,
+        agent_id: str,
+    ) -> None:
+        self._inner = inner
+        self._messages = _WrappedMessages(inner.messages, cz, agent_id)
+
+    @property
+    def messages(self) -> _WrappedMessages:
+        return self._messages
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._inner, name)
+
+
+def wrap_anthropic(
+    client: Any,
+    cz: Client,
+    agent_id: str = "",
+) -> Any:
+    """Wrap an Anthropic client with ControlZero policy enforcement.
+
+    Returns a proxy object that intercepts ``messages.create()`` and
+    ``messages.stream()`` to enforce governance policies. All other
+    client methods and properties pass through unchanged.
+
+    Args:
+        client: An ``anthropic.Anthropic()`` instance.
+        cz: A ``controlzero.Client`` instance with a policy loaded.
+        agent_id: Optional identifier for the agent making requests.
+            Included in the enforcement context for audit logging.
+
+    Returns:
+        A wrapped client with transparent policy enforcement.
+
+    Example::
+
+        from controlzero import Client
+        from controlzero.integrations.anthropic import wrap_anthropic
+        import anthropic
+
+        cz = Client(policy={"rules": [{"allow": "*"}]})
+        client = wrap_anthropic(anthropic.Anthropic(), cz, agent_id="my-agent")
+
+        response = client.messages.create(
+            model="claude-sonnet-4-20250514",
+            max_tokens=1024,
+            messages=[{"role": "user", "content": "Hello"}],
+        )
+    """
+    return _WrappedAnthropic(client, cz, agent_id)

controlzero-1.1.0/controlzero/integrations/braintrust.py

@@ -0,0 +1,142 @@
+"""Braintrust integration for ControlZero.
+
+Provides a callback handler that sends governance decisions and audit
+events to Braintrust for evaluation and monitoring.
+
+Usage::
+
+    from controlzero import Client
+    from controlzero.integrations.braintrust import BraintrustGovernanceLogger
+
+    cz = Client(policy={"rules": [{"allow": "*"}]})
+
+    logger = BraintrustGovernanceLogger(project_name="my-agent")
+
+    decision = cz.guard("database", {"sql": "SELECT 1"})
+    logger.log_decision(
+        tool="database",
+        method="query",
+        decision=decision.effect,
+        latency_ms=0.12,
+    )
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Optional
+
+try:
+    import braintrust  # noqa: F401
+
+    _HAS_BRAINTRUST = True
+except ImportError:
+    _HAS_BRAINTRUST = False
+
+
+class BraintrustGovernanceLogger:
+    """Logs ControlZero governance decisions to Braintrust for evaluation.
+
+    If the braintrust package is not installed, operations are no-ops.
+    """
+
+    def __init__(
+        self,
+        project_name: str = "controlzero",
+        experiment_name: Optional[str] = None,
+    ) -> None:
+        self._project_name = project_name
+        self._experiment_name = experiment_name
+        self._logger = None
+
+        if _HAS_BRAINTRUST:
+            self._logger = braintrust.init_logger(project=project_name)
+
+    def log_decision(
+        self,
+        tool: str,
+        method: str,
+        decision: str,
+        latency_ms: float = 0,
+        policy_id: Optional[str] = None,
+        context: Optional[dict] = None,
+        metadata: Optional[dict] = None,
+    ) -> None:
+        """Log a governance decision to Braintrust.
+
+        Args:
+            tool: Tool name
+            method: Method name
+            decision: Policy decision (allow/deny)
+            latency_ms: Policy evaluation latency
+            policy_id: ID of the matched policy
+            context: Request context
+            metadata: Additional metadata
+        """
+        if not self._logger:
+            return
+
+        self._logger.log(
+            input={
+                "tool": tool,
+                "method": method,
+                "action": f"{tool}:{method}",
+                "context": context or {},
+            },
+            output={
+                "decision": decision,
+                "policy_id": policy_id,
+            },
+            metadata={
+                "latency_ms": latency_ms,
+                "source": "controlzero",
+                **(metadata or {}),
+            },
+            scores={
+                "governance_pass": 1.0 if decision == "allow" else 0.0,
+            },
+        )
+
+    def log_tool_call(
+        self,
+        tool: str,
+        method: str,
+        arguments: dict,
+        result: Any = None,
+        error: Optional[str] = None,
+        duration_ms: float = 0,
+    ) -> None:
+        """Log a complete tool call (decision + execution) to Braintrust.
+
+        Args:
+            tool: Tool name
+            method: Method name
+            arguments: Tool call arguments
+            result: Tool call result (if successful)
+            error: Error message (if failed)
+            duration_ms: Total call duration
+        """
+        if not self._logger:
+            return
+
+        self._logger.log(
+            input={
+                "tool": tool,
+                "method": method,
+                "arguments": arguments,
+            },
+            output={
+                "result": str(result)[:1000] if result else None,
+                "error": error,
+            },
+            metadata={
+                "duration_ms": duration_ms,
+                "source": "controlzero",
+                "status": "error" if error else "success",
+            },
+        )
+
+    def flush(self) -> None:
+        """Flush any buffered logs to Braintrust."""
+        if self._logger and hasattr(self._logger, "flush"):
+            self._logger.flush()

controlzero-1.1.0/controlzero/integrations/crewai/__init__.py

@@ -0,0 +1,52 @@
+"""
+ControlZero CrewAI Integration.
+
+Provides governance enforcement for CrewAI including:
+- Crew-level policy enforcement
+- Agent role-based access control
+- Task execution governance
+- Tool-level policy checks
+
+Usage:
+    from controlzero import Client
+    from controlzero.integrations.crewai import (
+        GovernedCrew,
+        GovernedCrewAgent,
+        GovernedTask,
+        GovernedCrewTool,
+    )
+
+    client = Client(policy={"rules": [{"allow": "*"}]})
+
+    # Wrap a crew with governance
+    governed_crew = GovernedCrew(
+        crew=my_crew,
+        client=client,
+    )
+
+    result = governed_crew.kickoff()
+"""
+
+from controlzero.integrations.crewai.crew import GovernedCrew
+from controlzero.integrations.crewai.agent import GovernedCrewAgent, governed_agent
+from controlzero.integrations.crewai.task import GovernedTask, governed_task
+from controlzero.integrations.crewai.tool import (
+    GovernedCrewTool,
+    governed_crew_tool,
+    wrap_crew_tools,
+)
+
+__all__ = [
+    # Crew
+    "GovernedCrew",
+    # Agent
+    "GovernedCrewAgent",
+    "governed_agent",
+    # Task
+    "GovernedTask",
+    "governed_task",
+    # Tool
+    "GovernedCrewTool",
+    "governed_crew_tool",
+    "wrap_crew_tools",
+]

controlzero-1.1.0/controlzero/integrations/crewai/agent.py

@@ -0,0 +1,257 @@
+"""
+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 controlzero.client import Client
+from controlzero.errors import PolicyDeniedError, PolicyDecision
+from controlzero.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: Client,
+        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: ControlZero 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 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
+        """
+        tool_name = f"crewai:agent:{self._agent_name}"
+        decision = self._client.guard(
+            tool=tool_name,
+            method="execute_task",
+        )
+
+        if decision.denied:
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        self._total_actions += 1
+
+        try:
+            result = self._agent.execute_task(
+                task=task,
+                context=context,
+                tools=tools or self._agent.tools,
+            )
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            try:
+                self._client._audit_decision(
+                    tool=tool_name,
+                    method="execute_task",
+                    args={
+                        "status": "success",
+                        "latency_ms": latency_ms,
+                        "agent_role": self._agent.role,
+                        "total_actions": self._total_actions,
+                    },
+                    decision=decision,
+                )
+            except Exception:
+                pass
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            try:
+                err_decision = PolicyDecision(
+                    effect="allow",
+                    reason=f"error: {type(e).__name__}: {e}",
+                )
+                self._client._audit_decision(
+                    tool=tool_name,
+                    method="execute_task",
+                    args={"status": "error", "latency_ms": latency_ms},
+                    decision=err_decision,
+                )
+            except Exception:
+                pass
+            raise
+
+    @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: Client,
+    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