control-zero 0.2.0__py3-none-any.whl

control_zero/integrations/langchain/graph.py

@@ -0,0 +1,441 @@
+"""
+Control Zero LangGraph Integration.
+
+Provides governance for LangGraph state machines including:
+- Node-level policy enforcement
+- State transition logging
+- Graph execution governance
+"""
+
+from __future__ import annotations
+
+import time
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
+
+try:
+    from langgraph.graph import StateGraph, END
+    from langgraph.graph.state import CompiledStateGraph
+    LANGGRAPH_AVAILABLE = True
+except ImportError:
+    LANGGRAPH_AVAILABLE = False
+    class StateGraph:
+        pass
+    class CompiledStateGraph:
+        pass
+    END = "__end__"
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+
+StateType = TypeVar("StateType", bound=Dict[str, Any])
+
+
+class GovernedNode:
+    """
+    Decorator for governing LangGraph node functions.
+
+    Wraps a node function to add:
+    - Policy checks before execution
+    - State transition logging
+    - Error handling
+
+    Usage:
+        @GovernedNode(client, node_name="process_query")
+        def process_query(state: State) -> State:
+            # Node logic
+            return {"result": "processed"}
+    """
+
+    def __init__(
+        self,
+        client: ControlZeroClient,
+        node_name: Optional[str] = None,
+        log_state: bool = False,
+    ):
+        """
+        Initialize governed node decorator.
+
+        Args:
+            client: Control Zero client
+            node_name: Name for the node (defaults to function name)
+            log_state: Whether to log state content
+        """
+        self._client = client
+        self._node_name = node_name
+        self._log_state = log_state
+
+    def __call__(self, func: Callable) -> Callable:
+        """Wrap the node function."""
+        node_name = self._node_name or func.__name__
+
+        @wraps(func)
+        def governed_node(state: Dict[str, Any], *args, **kwargs) -> Dict[str, Any]:
+            return self._execute_governed(func, node_name, state, *args, **kwargs)
+
+        @wraps(func)
+        async def governed_node_async(state: Dict[str, Any], *args, **kwargs) -> Dict[str, Any]:
+            return await self._execute_governed_async(func, node_name, state, *args, **kwargs)
+
+        # Return async or sync based on function type
+        import asyncio
+        if asyncio.iscoroutinefunction(func):
+            return governed_node_async
+        return governed_node
+
+    def _check_policy(self, node_name: str) -> PolicyDecision:
+        """Check policy for node execution."""
+        if hasattr(self._client, '_policy_cache') and self._client._policy_cache:
+            return self._client._policy_cache.evaluate(f"langgraph:{node_name}", "execute")
+        return PolicyDecision(effect="allow")
+
+    def _execute_governed(
+        self,
+        func: Callable,
+        node_name: str,
+        state: Dict[str, Any],
+        *args,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Execute node with governance (sync)."""
+        decision = self._check_policy(node_name)
+
+        if decision.effect == "deny":
+            self._log(node_name, 0, "denied", decision, state)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = func(state, *args, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log(node_name, latency_ms, "success", decision, state, result)
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log(node_name, latency_ms, "error", decision, state, error=e)
+            raise
+
+    async def _execute_governed_async(
+        self,
+        func: Callable,
+        node_name: str,
+        state: Dict[str, Any],
+        *args,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Execute node with governance (async)."""
+        decision = self._check_policy(node_name)
+
+        if decision.effect == "deny":
+            self._log(node_name, 0, "denied", decision, state)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = await func(state, *args, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log(node_name, latency_ms, "success", decision, state, result)
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log(node_name, latency_ms, "error", decision, state, error=e)
+            raise
+
+    def _log(
+        self,
+        node_name: str,
+        latency_ms: int,
+        status: str,
+        decision: PolicyDecision,
+        input_state: Optional[Dict] = None,
+        output_state: Optional[Dict] = None,
+        error: Optional[Exception] = None,
+    ) -> None:
+        """Log node execution."""
+        log_data = {}
+
+        if self._log_state:
+            if input_state:
+                log_data["input_state"] = {k: str(v)[:100] for k, v in input_state.items()}
+            if output_state:
+                log_data["output_state"] = {k: str(v)[:100] for k, v in output_state.items()}
+
+        self._client._log(
+            tool=f"langgraph:{node_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
+        )
+
+
+def governed_graph(
+    client: ControlZeroClient,
+    graph_name: str = "langgraph",
+) -> Callable:
+    """
+    Decorator to wrap an entire LangGraph with governance.
+
+    Usage:
+        @governed_graph(client, graph_name="my_workflow")
+        def create_graph():
+            graph = StateGraph(State)
+            graph.add_node("process", process_node)
+            graph.add_edge("process", END)
+            return graph.compile()
+    """
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Create the graph
+            compiled_graph = func(*args, **kwargs)
+
+            # Wrap with governance
+            return GovernedStateGraph(
+                graph=compiled_graph,
+                client=client,
+                graph_name=graph_name,
+            )
+        return wrapper
+    return decorator
+
+
+class GovernedStateGraph:
+    """
+    Governance wrapper for compiled LangGraph.
+
+    Provides:
+    - Graph-level policy enforcement
+    - Execution logging
+    - State tracking
+
+    Usage:
+        graph = StateGraph(State)
+        # ... add nodes and edges
+        compiled = graph.compile()
+
+        governed = GovernedStateGraph(
+            graph=compiled,
+            client=client,
+            graph_name="my_workflow"
+        )
+
+        result = governed.invoke({"input": "data"})
+    """
+
+    def __init__(
+        self,
+        graph: Any,  # CompiledStateGraph
+        client: ControlZeroClient,
+        graph_name: str = "langgraph",
+        log_states: bool = False,
+    ):
+        """
+        Initialize governed graph.
+
+        Args:
+            graph: Compiled LangGraph
+            client: Control Zero client
+            graph_name: Name for logging
+            log_states: Whether to log state content
+        """
+        if not LANGGRAPH_AVAILABLE:
+            raise ImportError(
+                "langgraph is required. Install with: pip install langgraph"
+            )
+
+        self._graph = graph
+        self._client = client
+        self._graph_name = graph_name
+        self._log_states = log_states
+
+    def _check_policy(self, method: str) -> PolicyDecision:
+        """Check policy for graph execution."""
+        if hasattr(self._client, '_policy_cache') and self._client._policy_cache:
+            return self._client._policy_cache.evaluate(f"langgraph:{self._graph_name}", method)
+        return PolicyDecision(effect="allow")
+
+    def invoke(
+        self,
+        input: Dict[str, Any],
+        config: Optional[Dict] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Invoke the graph with governance.
+
+        Args:
+            input: Initial state
+            config: Run configuration
+            **kwargs: Additional arguments
+
+        Returns:
+            Final state
+        """
+        decision = self._check_policy("invoke")
+        if decision.effect == "deny":
+            self._log("invoke", 0, "denied", decision, input)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = self._graph.invoke(input, config=config, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("invoke", latency_ms, "success", decision, input, result)
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("invoke", latency_ms, "error", decision, input, error=e)
+            raise
+
+    async def ainvoke(
+        self,
+        input: Dict[str, Any],
+        config: Optional[Dict] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Async invoke with governance."""
+        decision = self._check_policy("invoke")
+        if decision.effect == "deny":
+            self._log("invoke", 0, "denied", decision, input)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = await self._graph.ainvoke(input, config=config, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("invoke", latency_ms, "success", decision, input, result)
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("invoke", latency_ms, "error", decision, input, error=e)
+            raise
+
+    def stream(
+        self,
+        input: Dict[str, Any],
+        config: Optional[Dict] = None,
+        **kwargs,
+    ):
+        """
+        Stream graph execution with governance.
+
+        Args:
+            input: Initial state
+            config: Run configuration
+            **kwargs: Additional arguments
+
+        Yields:
+            State updates
+        """
+        decision = self._check_policy("stream")
+        if decision.effect == "deny":
+            self._log("stream", 0, "denied", decision, input)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        states = []
+
+        try:
+            for state in self._graph.stream(input, config=config, **kwargs):
+                states.append(state)
+                yield state
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("stream", latency_ms, "success", decision, input, metadata={"num_states": len(states)})
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("stream", latency_ms, "error", decision, input, error=e)
+            raise
+
+    async def astream(
+        self,
+        input: Dict[str, Any],
+        config: Optional[Dict] = None,
+        **kwargs,
+    ):
+        """Async stream with governance."""
+        decision = self._check_policy("stream")
+        if decision.effect == "deny":
+            self._log("stream", 0, "denied", decision, input)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        states = []
+
+        try:
+            async for state in self._graph.astream(input, config=config, **kwargs):
+                states.append(state)
+                yield state
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("stream", latency_ms, "success", decision, input, metadata={"num_states": len(states)})
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("stream", latency_ms, "error", decision, input, error=e)
+            raise
+
+    def _log(
+        self,
+        method: str,
+        latency_ms: int,
+        status: str,
+        decision: PolicyDecision,
+        input_state: Optional[Dict] = None,
+        output_state: Optional[Dict] = None,
+        error: Optional[Exception] = None,
+        metadata: Optional[Dict] = None,
+    ) -> None:
+        """Log graph execution."""
+        log_data = metadata or {}
+
+        if self._log_states:
+            if input_state:
+                log_data["input_state"] = {k: str(v)[:100] for k, v in input_state.items()}
+            if output_state:
+                log_data["output_state"] = {k: str(v)[:100] for k, v in output_state.items()}
+
+        self._client._log(
+            tool=f"langgraph:{self._graph_name}",
+            method=method,
+            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 graph(self) -> Any:
+        """Get underlying graph."""
+        return self._graph
+
+    def get_graph(self, **kwargs):
+        """Get graph visualization."""
+        return self._graph.get_graph(**kwargs)

control_zero/integrations/langchain/tool.py

@@ -0,0 +1,271 @@
+"""
+Governed LangChain Tool wrappers.
+
+Provides governance enforcement for LangChain tools including:
+- Policy checks before execution
+- Audit logging
+- Error tracking
+- Secret injection
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Callable, Dict, List, Optional, Type, Union
+from functools import wraps
+
+try:
+    from langchain_core.tools import BaseTool, Tool, StructuredTool
+    from langchain_core.pydantic_v1 import BaseModel, Field
+    from langchain_core.callbacks import CallbackManagerForToolRun
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    # Dummy classes for type hints
+    class BaseTool:
+        pass
+    class Tool:
+        pass
+    class StructuredTool:
+        pass
+    class BaseModel:
+        pass
+    class CallbackManagerForToolRun:
+        pass
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+
+
+class GovernanceTool(BaseTool):
+    """
+    Wraps a LangChain Tool to enforce Control Zero governance policies.
+
+    This is the original wrapper that directly wraps BaseTool instances.
+
+    Usage:
+        client = ControlZeroClient(api_key="...")
+        client.initialize()
+
+        base_tool = MyCustomTool()
+        gov_tool = GovernanceTool(
+            base_tool=base_tool,
+            client=client,
+            tool_name="my_custom_tool"
+        )
+
+        agent = initialize_agent([gov_tool], ...)
+    """
+
+    base_tool: BaseTool
+    client: ControlZeroClient
+    tool_name: str
+
+    # Forward metadata from base tool
+    name: str = ""
+    description: str = ""
+    args_schema: Optional[Type[BaseModel]] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __init__(
+        self,
+        base_tool: BaseTool,
+        client: ControlZeroClient,
+        tool_name: Optional[str] = None,
+        **kwargs
+    ):
+        if not LANGCHAIN_AVAILABLE:
+            raise ImportError(
+                "langchain is required for LangChain integration. "
+                "Install with: pip install langchain-core"
+            )
+
+        name = base_tool.name
+        description = base_tool.description
+        args_schema = getattr(base_tool, 'args_schema', None)
+        tool_name = tool_name or name
+
+        super().__init__(
+            base_tool=base_tool,
+            client=client,
+            tool_name=tool_name,
+            name=name,
+            description=description,
+            args_schema=args_schema,
+            **kwargs
+        )
+
+    def _run(
+        self,
+        *args: Any,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
+        **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, run_manager=run_manager, **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,
+        run_manager: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Async execution with governance."""
+        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
+            kwargs = self._inject_secrets(kwargs)
+
+            # 3. Run Tool
+            result = await self.base_tool._arun(*args, run_manager=run_manager, **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:
+            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(
+            self.tool_name, method, "denied", 0,
+            policy_decision=decision
+        )
+
+    def _log_success(self, method: str, latency_ms: int, decision: PolicyDecision) -> None:
+        """Log successful execution."""
+        self.client._log(
+            self.tool_name, method, "success", 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(
+            self.tool_name, method, "error", latency_ms,
+            policy_decision=decision,
+            error_type=type(error).__name__,
+            error_message=str(error)
+        )
+
+
+# Alias for consistency
+GovernedTool = GovernanceTool
+
+
+def governed_tool(
+    client: ControlZeroClient,
+    tool_name: Optional[str] = None,
+) -> Callable:
+    """
+    Decorator to create a governed tool from a function.
+
+    Usage:
+        @governed_tool(client, tool_name="my_tool")
+        def my_function(query: str) -> str:
+            '''Search for information.'''
+            return search(query)
+
+        # Creates a governed StructuredTool
+    """
+    def decorator(func: Callable) -> StructuredTool:
+        if not LANGCHAIN_AVAILABLE:
+            raise ImportError("langchain is required")
+
+        # Create base tool
+        base_tool = StructuredTool.from_function(
+            func=func,
+            name=tool_name or func.__name__,
+            description=func.__doc__ or "",
+        )
+
+        # Wrap with governance
+        return GovernanceTool(
+            base_tool=base_tool,
+            client=client,
+            tool_name=tool_name or func.__name__,
+        )
+
+    return decorator
+
+
+def wrap_tools(
+    tools: List[BaseTool],
+    client: ControlZeroClient,
+) -> List[GovernanceTool]:
+    """
+    Wrap multiple tools with governance.
+
+    Usage:
+        tools = [tool1, tool2, tool3]
+        governed_tools = wrap_tools(tools, client)
+    """
+    return [
+        GovernanceTool(base_tool=tool, client=client)
+        for tool in tools
+    ]