control-zero 0.2.0__py3-none-any.whl

control_zero/integrations/langchain/agent.py

@@ -0,0 +1,311 @@
+"""
+Governed LangChain Agent wrapper.
+
+Provides governance enforcement at the agent level including:
+- Per-action policy checks
+- Cost and token tracking
+- Conversation audit logging
+- Tool filtering
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+try:
+    from langchain.agents import AgentExecutor
+    from langchain_core.agents import AgentAction, AgentFinish
+    from langchain_core.tools import BaseTool
+    from langchain_core.callbacks import BaseCallbackHandler
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    class AgentExecutor:
+        pass
+    class AgentAction:
+        pass
+    class AgentFinish:
+        pass
+    class BaseTool:
+        pass
+    class BaseCallbackHandler:
+        pass
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+from control_zero.integrations.langchain.tool import GovernanceTool
+
+
+class GovernedAgent:
+    """
+    Governance wrapper for LangChain agents.
+
+    Wraps an AgentExecutor to provide:
+    - Tool-level policy enforcement
+    - Agent action logging
+    - Cost tracking
+    - Iteration limits
+
+    Usage:
+        from langchain.agents import initialize_agent
+
+        agent = initialize_agent(tools, llm, agent=AgentType.OPENAI_FUNCTIONS)
+
+        governed_agent = GovernedAgent(
+            agent=agent,
+            client=client,
+            max_iterations=10,
+        )
+
+        result = governed_agent.run("What's the weather?")
+    """
+
+    def __init__(
+        self,
+        agent: AgentExecutor,
+        client: ControlZeroClient,
+        max_iterations: Optional[int] = None,
+        max_execution_time: Optional[float] = None,
+        wrap_tools: bool = True,
+        agent_name: str = "langchain_agent",
+    ):
+        """
+        Initialize governed agent.
+
+        Args:
+            agent: The LangChain AgentExecutor to wrap
+            client: Control Zero client
+            max_iterations: Maximum agent iterations (overrides agent setting)
+            max_execution_time: Maximum execution time in seconds
+            wrap_tools: Whether to wrap agent tools with governance
+            agent_name: Name for logging purposes
+        """
+        if not LANGCHAIN_AVAILABLE:
+            raise ImportError(
+                "langchain is required. Install with: pip install langchain"
+            )
+
+        self._agent = agent
+        self._client = client
+        self._agent_name = agent_name
+        self._max_iterations = max_iterations
+        self._max_execution_time = max_execution_time
+
+        # Apply iteration limit
+        if max_iterations:
+            self._agent.max_iterations = max_iterations
+        if max_execution_time:
+            self._agent.max_execution_time = max_execution_time
+
+        # Wrap tools with governance if requested
+        if wrap_tools:
+            self._wrap_agent_tools()
+
+        # Tracking
+        self._total_iterations = 0
+        self._total_tool_calls = 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, GovernanceTool):
+                governed_tools.append(tool)
+            else:
+                governed_tools.append(
+                    GovernanceTool(
+                        base_tool=tool,
+                        client=self._client,
+                    )
+                )
+        self._agent.tools = governed_tools
+
+    def _check_agent_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(self._agent_name, action)
+        return PolicyDecision(effect="allow")
+
+    def run(
+        self,
+        input: Union[str, Dict[str, Any]],
+        callbacks: Optional[List[BaseCallbackHandler]] = None,
+        **kwargs,
+    ) -> str:
+        """
+        Run the agent with governance.
+
+        Args:
+            input: Agent input (string or dict)
+            callbacks: Additional callbacks
+            **kwargs: Additional arguments passed to agent
+
+        Returns:
+            Agent response string
+
+        Raises:
+            PolicyDeniedError: If agent execution is denied
+        """
+        # Check agent-level policy
+        decision = self._check_agent_policy("run")
+        if decision.effect == "deny":
+            self._log_denied("run", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            # Run agent
+            result = self._agent.run(input, callbacks=callbacks, **kwargs)
+
+            # Log success
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_success("run", latency_ms, decision, result)
+
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_error("run", latency_ms, decision, e)
+            raise
+
+    async def arun(
+        self,
+        input: Union[str, Dict[str, Any]],
+        callbacks: Optional[List[BaseCallbackHandler]] = None,
+        **kwargs,
+    ) -> str:
+        """Async run with governance."""
+        decision = self._check_agent_policy("run")
+        if decision.effect == "deny":
+            self._log_denied("run", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = await self._agent.arun(input, callbacks=callbacks, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_success("run", latency_ms, decision, result)
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_error("run", latency_ms, decision, e)
+            raise
+
+    def invoke(
+        self,
+        input: Union[str, Dict[str, Any]],
+        config: Optional[Dict] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Invoke the agent (LCEL interface) with governance.
+
+        Args:
+            input: Agent input
+            config: Run configuration
+            **kwargs: Additional arguments
+
+        Returns:
+            Agent output dict
+        """
+        decision = self._check_agent_policy("invoke")
+        if decision.effect == "deny":
+            self._log_denied("invoke", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = self._agent.invoke(input, config=config, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_success("invoke", latency_ms, decision, str(result))
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_error("invoke", latency_ms, decision, e)
+            raise
+
+    async def ainvoke(
+        self,
+        input: Union[str, Dict[str, Any]],
+        config: Optional[Dict] = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Async invoke with governance."""
+        decision = self._check_agent_policy("invoke")
+        if decision.effect == "deny":
+            self._log_denied("invoke", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = await self._agent.ainvoke(input, config=config, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_success("invoke", latency_ms, decision, str(result))
+            return result
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log_error("invoke", latency_ms, decision, e)
+            raise
+
+    def _log_denied(self, method: str, decision: PolicyDecision) -> None:
+        """Log denied execution."""
+        self._client._log(
+            self._agent_name, method, "denied", 0,
+            policy_decision=decision
+        )
+
+    def _log_success(
+        self,
+        method: str,
+        latency_ms: int,
+        decision: PolicyDecision,
+        result: str,
+    ) -> None:
+        """Log successful execution."""
+        self._client._log(
+            self._agent_name, method, "success", latency_ms,
+            policy_decision=decision,
+            metadata={"result_length": len(result)}
+        )
+
+    def _log_error(
+        self,
+        method: str,
+        latency_ms: int,
+        decision: PolicyDecision,
+        error: Exception,
+    ) -> None:
+        """Log execution error."""
+        self._client._log(
+            self._agent_name, method, "error", latency_ms,
+            policy_decision=decision,
+            error_type=type(error).__name__,
+            error_message=str(error)
+        )
+
+    @property
+    def tools(self) -> List[BaseTool]:
+        """Get agent tools."""
+        return self._agent.tools
+
+    @property
+    def agent(self) -> AgentExecutor:
+        """Get underlying agent executor."""
+        return self._agent