control-zero 0.2.0__py3-none-any.whl

control_zero/integrations/langchain/chain.py

@@ -0,0 +1,319 @@
+"""
+Governed LangChain Chain wrapper.
+
+Provides governance enforcement for LangChain chains including:
+- Chain-level policy checks
+- Input/output logging
+- Cost tracking
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Dict, List, Optional, Union
+
+try:
+    from langchain_core.runnables import Runnable, RunnableConfig
+    from langchain.chains.base import Chain
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    LANGCHAIN_AVAILABLE = False
+    class Runnable:
+        pass
+    class RunnableConfig:
+        pass
+    class Chain:
+        pass
+
+from control_zero.client import ControlZeroClient
+from control_zero.policy import PolicyDecision, PolicyDeniedError
+
+
+class GovernedChain:
+    """
+    Governance wrapper for LangChain chains.
+
+    Wraps any LangChain chain (including LCEL runnables) with:
+    - Policy enforcement
+    - Input/output audit logging
+    - Error tracking
+
+    Usage:
+        from langchain_core.prompts import ChatPromptTemplate
+        from langchain_openai import ChatOpenAI
+
+        prompt = ChatPromptTemplate.from_messages([...])
+        llm = ChatOpenAI()
+        chain = prompt | llm
+
+        governed = GovernedChain(
+            chain=chain,
+            client=client,
+            chain_name="my_chain"
+        )
+
+        result = governed.invoke({"question": "Hello"})
+    """
+
+    def __init__(
+        self,
+        chain: Union[Runnable, Chain],
+        client: ControlZeroClient,
+        chain_name: str = "langchain_chain",
+        log_inputs: bool = False,
+        log_outputs: bool = False,
+    ):
+        """
+        Initialize governed chain.
+
+        Args:
+            chain: The LangChain chain or runnable to wrap
+            client: Control Zero client
+            chain_name: Name for logging purposes
+            log_inputs: Whether to log input content
+            log_outputs: Whether to log output content
+        """
+        if not LANGCHAIN_AVAILABLE:
+            raise ImportError(
+                "langchain is required. Install with: pip install langchain-core"
+            )
+
+        self._chain = chain
+        self._client = client
+        self._chain_name = chain_name
+        self._log_inputs = log_inputs
+        self._log_outputs = log_outputs
+
+    def _check_policy(self, method: str) -> PolicyDecision:
+        """Check policy for chain execution."""
+        if hasattr(self._client, '_policy_cache') and self._client._policy_cache:
+            return self._client._policy_cache.evaluate(self._chain_name, method)
+        return PolicyDecision(effect="allow")
+
+    def invoke(
+        self,
+        input: Union[str, Dict[str, Any]],
+        config: Optional[RunnableConfig] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Invoke the chain with governance.
+
+        Args:
+            input: Chain input
+            config: Runnable configuration
+            **kwargs: Additional arguments
+
+        Returns:
+            Chain output
+        """
+        decision = self._check_policy("invoke")
+        if decision.effect == "deny":
+            self._log("invoke", 0, "denied", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = self._chain.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: Union[str, Dict[str, Any]],
+        config: Optional[RunnableConfig] = None,
+        **kwargs,
+    ) -> Any:
+        """Async invoke with governance."""
+        decision = self._check_policy("invoke")
+        if decision.effect == "deny":
+            self._log("invoke", 0, "denied", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            result = await self._chain.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: Union[str, Dict[str, Any]],
+        config: Optional[RunnableConfig] = None,
+        **kwargs,
+    ):
+        """
+        Stream chain output with governance.
+
+        Args:
+            input: Chain input
+            config: Runnable configuration
+            **kwargs: Additional arguments
+
+        Yields:
+            Chain output chunks
+        """
+        decision = self._check_policy("stream")
+        if decision.effect == "deny":
+            self._log("stream", 0, "denied", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        chunks = []
+
+        try:
+            for chunk in self._chain.stream(input, config=config, **kwargs):
+                chunks.append(chunk)
+                yield chunk
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("stream", latency_ms, "success", decision, input, chunks)
+
+        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: Union[str, Dict[str, Any]],
+        config: Optional[RunnableConfig] = None,
+        **kwargs,
+    ):
+        """Async stream with governance."""
+        decision = self._check_policy("stream")
+        if decision.effect == "deny":
+            self._log("stream", 0, "denied", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+        chunks = []
+
+        try:
+            async for chunk in self._chain.astream(input, config=config, **kwargs):
+                chunks.append(chunk)
+                yield chunk
+
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("stream", latency_ms, "success", decision, input, chunks)
+
+        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 batch(
+        self,
+        inputs: List[Union[str, Dict[str, Any]]],
+        config: Optional[RunnableConfig] = None,
+        **kwargs,
+    ) -> List[Any]:
+        """
+        Batch invoke with governance.
+
+        Args:
+            inputs: List of inputs
+            config: Runnable configuration
+            **kwargs: Additional arguments
+
+        Returns:
+            List of outputs
+        """
+        decision = self._check_policy("batch")
+        if decision.effect == "deny":
+            self._log("batch", 0, "denied", decision)
+            raise PolicyDeniedError(decision)
+
+        start = time.perf_counter()
+
+        try:
+            results = self._chain.batch(inputs, config=config, **kwargs)
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("batch", latency_ms, "success", decision, metadata={"batch_size": len(inputs)})
+            return results
+
+        except PolicyDeniedError:
+            raise
+        except Exception as e:
+            latency_ms = int((time.perf_counter() - start) * 1000)
+            self._log("batch", latency_ms, "error", decision, error=e)
+            raise
+
+    def _log(
+        self,
+        method: str,
+        latency_ms: int,
+        status: str,
+        decision: PolicyDecision,
+        input: Any = None,
+        output: Any = None,
+        error: Optional[Exception] = None,
+        metadata: Optional[Dict] = None,
+    ) -> None:
+        """Log chain execution."""
+        log_metadata = metadata or {}
+
+        if self._log_inputs and input:
+            log_metadata["input"] = str(input)[:500]
+        if self._log_outputs and output:
+            log_metadata["output"] = str(output)[:500]
+
+        self._client._log(
+            tool=self._chain_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_metadata
+        )
+
+    @property
+    def chain(self) -> Union[Runnable, Chain]:
+        """Get underlying chain."""
+        return self._chain
+
+    def __or__(self, other: Any) -> "GovernedChain":
+        """Support pipe operator for LCEL composition."""
+        new_chain = self._chain | other
+        return GovernedChain(
+            chain=new_chain,
+            client=self._client,
+            chain_name=self._chain_name,
+            log_inputs=self._log_inputs,
+            log_outputs=self._log_outputs,
+        )
+
+    def __ror__(self, other: Any) -> "GovernedChain":
+        """Support reverse pipe operator."""
+        new_chain = other | self._chain
+        return GovernedChain(
+            chain=new_chain,
+            client=self._client,
+            chain_name=self._chain_name,
+            log_inputs=self._log_inputs,
+            log_outputs=self._log_outputs,
+        )