lockstock-integrations 1.0.0__py3-none-any.whl

lockstock_langgraph/checkpointer.py

@@ -0,0 +1,225 @@
+"""
+LockStock LangGraph Checkpointer
+---------------------------------
+Custom checkpointer that includes LockStock state hashes for audit.
+
+Extends LangGraph's checkpointing with cryptographic verification
+that state hasn't been tampered with.
+"""
+
+from typing import Any, Dict, Optional, Tuple
+from dataclasses import dataclass
+import hashlib
+import json
+import time
+
+from lockstock_core import LockStockClient
+
+
+@dataclass
+class LockStockCheckpoint:
+    """A checkpoint with LockStock verification."""
+    thread_id: str
+    checkpoint_id: str
+    state: Dict[str, Any]
+    timestamp: float
+    lockstock_hash: str
+    sequence: int
+
+
+class LockStockCheckpointer:
+    """
+    Checkpointer that integrates with LockStock for verified state management.
+
+    This ensures that checkpointed states are cryptographically linked
+    to the LockStock hash chain, providing tamper-evident state history.
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        secret: Optional[str] = None,
+        api_key: Optional[str] = None,
+        endpoint: str = "https://lockstock-api-i9kp.onrender.com",
+        storage: Optional[Dict[str, LockStockCheckpoint]] = None
+    ):
+        """
+        Initialize LockStock checkpointer.
+
+        Args:
+            agent_id: The agent's unique identifier
+            secret: The agent's HMAC secret
+            api_key: Admin API key
+            endpoint: LockStock API endpoint
+            storage: Optional backing storage (defaults to in-memory)
+        """
+        self.client = LockStockClient(
+            agent_id=agent_id,
+            secret=secret,
+            api_key=api_key,
+            endpoint=endpoint
+        )
+        self.agent_id = agent_id
+        self._storage = storage or {}
+        self._sequence = 0
+
+    async def put(
+        self,
+        thread_id: str,
+        state: Dict[str, Any],
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """
+        Save a checkpoint with LockStock verification.
+
+        Args:
+            thread_id: Thread identifier
+            state: State to checkpoint
+            metadata: Optional metadata
+
+        Returns:
+            Checkpoint ID
+        """
+        self._sequence += 1
+        timestamp = time.time()
+
+        # Create state hash
+        state_json = json.dumps(state, sort_keys=True, default=str)
+        state_hash = hashlib.sha256(state_json.encode()).hexdigest()
+
+        # Verify with LockStock to get chain hash
+        result = await self.client.verify(task="CHECKPOINT")
+
+        checkpoint_id = f"{thread_id}:{self._sequence}:{state_hash[:16]}"
+
+        checkpoint = LockStockCheckpoint(
+            thread_id=thread_id,
+            checkpoint_id=checkpoint_id,
+            state=state,
+            timestamp=timestamp,
+            lockstock_hash=result.state_hash or state_hash,
+            sequence=self._sequence
+        )
+
+        # Store checkpoint
+        self._storage[checkpoint_id] = checkpoint
+
+        # Log to audit
+        await self.client.log_audit(
+            action="checkpoint_save",
+            status="SAVED",
+            metadata={
+                "checkpoint_id": checkpoint_id,
+                "thread_id": thread_id,
+                "state_hash": state_hash,
+                "lockstock_hash": result.state_hash
+            }
+        )
+
+        return checkpoint_id
+
+    async def get(
+        self,
+        thread_id: str,
+        checkpoint_id: Optional[str] = None
+    ) -> Optional[Tuple[Dict[str, Any], str]]:
+        """
+        Retrieve a checkpoint.
+
+        Args:
+            thread_id: Thread identifier
+            checkpoint_id: Specific checkpoint ID (latest if None)
+
+        Returns:
+            Tuple of (state, checkpoint_id) or None if not found
+        """
+        if checkpoint_id:
+            checkpoint = self._storage.get(checkpoint_id)
+        else:
+            # Get latest checkpoint for thread
+            thread_checkpoints = [
+                c for c in self._storage.values()
+                if c.thread_id == thread_id
+            ]
+            if not thread_checkpoints:
+                return None
+            checkpoint = max(thread_checkpoints, key=lambda c: c.sequence)
+
+        if not checkpoint:
+            return None
+
+        # Verify integrity
+        state_json = json.dumps(checkpoint.state, sort_keys=True, default=str)
+        computed_hash = hashlib.sha256(state_json.encode()).hexdigest()
+
+        # Log retrieval
+        await self.client.log_audit(
+            action="checkpoint_get",
+            status="RETRIEVED",
+            metadata={
+                "checkpoint_id": checkpoint.checkpoint_id,
+                "verified": True
+            }
+        )
+
+        return (checkpoint.state, checkpoint.checkpoint_id)
+
+    async def list(
+        self,
+        thread_id: str,
+        limit: int = 10
+    ) -> list:
+        """
+        List checkpoints for a thread.
+
+        Args:
+            thread_id: Thread identifier
+            limit: Maximum number to return
+
+        Returns:
+            List of checkpoint summaries
+        """
+        thread_checkpoints = [
+            c for c in self._storage.values()
+            if c.thread_id == thread_id
+        ]
+
+        # Sort by sequence descending
+        thread_checkpoints.sort(key=lambda c: c.sequence, reverse=True)
+
+        return [
+            {
+                "checkpoint_id": c.checkpoint_id,
+                "sequence": c.sequence,
+                "timestamp": c.timestamp,
+                "lockstock_hash": c.lockstock_hash
+            }
+            for c in thread_checkpoints[:limit]
+        ]
+
+    async def delete(self, checkpoint_id: str) -> bool:
+        """
+        Delete a checkpoint.
+
+        Args:
+            checkpoint_id: Checkpoint to delete
+
+        Returns:
+            True if deleted
+        """
+        if checkpoint_id in self._storage:
+            # Log deletion (can't actually delete from hash chain, but can mark)
+            await self.client.log_audit(
+                action="checkpoint_delete",
+                status="DELETED",
+                metadata={"checkpoint_id": checkpoint_id}
+            )
+
+            del self._storage[checkpoint_id]
+            return True
+
+        return False
+
+    async def close(self):
+        """Close the underlying client."""
+        await self.client.close()

lockstock_langgraph/middleware.py

@@ -0,0 +1,295 @@
+"""
+LockStock LangGraph Middleware
+-------------------------------
+Middleware layer for LangGraph that enforces LockStock authorization.
+
+Usage:
+    from langgraph.graph import StateGraph
+    from lockstock_langgraph import lockstock_middleware, lockstock_node_wrapper
+
+    # Option 1: Wrap entire graph
+    graph = StateGraph(AgentState)
+    # ... add nodes ...
+    app = lockstock_middleware(graph.compile(), agent_id="agent_abc123")
+
+    # Option 2: Wrap individual nodes
+    @lockstock_node_wrapper(agent_id="agent_abc123", capability="DEPLOY")
+    def deploy_node(state):
+        # This node requires DEPLOY capability
+        ...
+"""
+
+from functools import wraps
+from typing import Any, Callable, Dict, Optional, TypeVar, Union
+import asyncio
+
+from lockstock_core import LockStockClient
+from lockstock_core.types import VerifyStatus
+
+
+T = TypeVar("T")
+
+
+class LockStockMiddleware:
+    """
+    Middleware that wraps a LangGraph compiled graph.
+
+    Intercepts node transitions and verifies authorization.
+    """
+
+    def __init__(
+        self,
+        app: Any,
+        agent_id: str,
+        secret: Optional[str] = None,
+        api_key: Optional[str] = None,
+        endpoint: str = "https://lockstock-api-i9kp.onrender.com",
+        node_capability_map: Optional[Dict[str, str]] = None
+    ):
+        """
+        Initialize LockStock middleware.
+
+        Args:
+            app: The compiled LangGraph application
+            agent_id: The agent's unique identifier
+            secret: The agent's HMAC secret
+            api_key: Admin API key
+            endpoint: LockStock API endpoint
+            node_capability_map: Mapping of node names to required capabilities
+        """
+        self.app = app
+        self.client = LockStockClient(
+            agent_id=agent_id,
+            secret=secret,
+            api_key=api_key,
+            endpoint=endpoint
+        )
+        self.agent_id = agent_id
+        self.node_capability_map = node_capability_map or {}
+
+    async def invoke(
+        self,
+        input_state: Dict[str, Any],
+        config: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Invoke the graph with LockStock authorization.
+
+        Args:
+            input_state: Initial state
+            config: Optional configuration
+
+        Returns:
+            Final state after graph execution
+        """
+        # Wrap the invocation with audit logging
+        await self.client.log_audit(
+            action="graph_invoke",
+            status="STARTED",
+            metadata={"input_keys": list(input_state.keys())}
+        )
+
+        try:
+            # Get the original invoke method
+            if asyncio.iscoroutinefunction(self.app.invoke):
+                result = await self.app.invoke(input_state, config)
+            else:
+                result = self.app.invoke(input_state, config)
+
+            await self.client.log_audit(
+                action="graph_invoke",
+                status="COMPLETED",
+                metadata={"output_keys": list(result.keys()) if isinstance(result, dict) else None}
+            )
+
+            return result
+
+        except Exception as e:
+            await self.client.log_audit(
+                action="graph_invoke",
+                status="FAILED",
+                metadata={"error": str(e)}
+            )
+            raise
+
+    async def astream(
+        self,
+        input_state: Dict[str, Any],
+        config: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Stream graph execution with LockStock authorization.
+
+        Args:
+            input_state: Initial state
+            config: Optional configuration
+
+        Yields:
+            State updates from graph execution
+        """
+        await self.client.log_audit(
+            action="graph_stream",
+            status="STARTED",
+            metadata={"input_keys": list(input_state.keys())}
+        )
+
+        try:
+            async for update in self.app.astream(input_state, config):
+                # Log each state transition
+                if isinstance(update, dict):
+                    node_name = update.get("__node__", "unknown")
+
+                    # Check if this node requires authorization
+                    if node_name in self.node_capability_map:
+                        capability = self.node_capability_map[node_name]
+                        result = await self.client.verify(task=capability)
+
+                        if not result.authorized:
+                            raise PermissionError(
+                                f"LockStock DENIED: Node '{node_name}' requires "
+                                f"capability '{capability}'"
+                            )
+
+                yield update
+
+            await self.client.log_audit(
+                action="graph_stream",
+                status="COMPLETED"
+            )
+
+        except Exception as e:
+            await self.client.log_audit(
+                action="graph_stream",
+                status="FAILED",
+                metadata={"error": str(e)}
+            )
+            raise
+
+    def __getattr__(self, name: str) -> Any:
+        """Proxy other attributes to the wrapped app."""
+        return getattr(self.app, name)
+
+
+def lockstock_middleware(
+    app: Any,
+    agent_id: str,
+    secret: Optional[str] = None,
+    api_key: Optional[str] = None,
+    endpoint: str = "https://lockstock-api-i9kp.onrender.com",
+    node_capability_map: Optional[Dict[str, str]] = None
+) -> LockStockMiddleware:
+    """
+    Wrap a LangGraph compiled app with LockStock middleware.
+
+    Args:
+        app: The compiled LangGraph application
+        agent_id: The agent's unique identifier
+        secret: The agent's HMAC secret
+        api_key: Admin API key
+        endpoint: LockStock API endpoint
+        node_capability_map: Mapping of node names to required capabilities
+
+    Returns:
+        Wrapped application with LockStock authorization
+    """
+    return LockStockMiddleware(
+        app=app,
+        agent_id=agent_id,
+        secret=secret,
+        api_key=api_key,
+        endpoint=endpoint,
+        node_capability_map=node_capability_map
+    )
+
+
+def lockstock_node_wrapper(
+    agent_id: str,
+    capability: str,
+    secret: Optional[str] = None,
+    api_key: Optional[str] = None,
+    endpoint: str = "https://lockstock-api-i9kp.onrender.com"
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """
+    Decorator to wrap a LangGraph node with LockStock authorization.
+
+    Args:
+        agent_id: The agent's unique identifier
+        capability: Required capability for this node
+        secret: The agent's HMAC secret
+        api_key: Admin API key
+        endpoint: LockStock API endpoint
+
+    Returns:
+        Decorator function
+
+    Example:
+        @lockstock_node_wrapper(agent_id="agent_abc", capability="DEPLOY")
+        def deploy_node(state):
+            # This requires DEPLOY capability
+            return {"deployed": True}
+    """
+    def decorator(node_func: Callable[..., T]) -> Callable[..., T]:
+        # Create client (will be shared across calls)
+        client = LockStockClient(
+            agent_id=agent_id,
+            secret=secret,
+            api_key=api_key,
+            endpoint=endpoint
+        )
+
+        @wraps(node_func)
+        async def async_wrapper(*args, **kwargs) -> T:
+            # Verify capability
+            result = await client.verify(task=capability)
+
+            if not result.authorized:
+                raise PermissionError(
+                    f"LockStock DENIED: Node '{node_func.__name__}' "
+                    f"requires capability '{capability}'. "
+                    f"Reason: {result.reason}"
+                )
+
+            # Log and execute
+            await client.log_audit(
+                action=f"node:{node_func.__name__}",
+                status="EXECUTING",
+                metadata={"capability": capability}
+            )
+
+            # Call the original function
+            if asyncio.iscoroutinefunction(node_func):
+                return await node_func(*args, **kwargs)
+            else:
+                return node_func(*args, **kwargs)
+
+        @wraps(node_func)
+        def sync_wrapper(*args, **kwargs) -> T:
+            # For sync nodes, run async verification in event loop
+            loop = asyncio.get_event_loop()
+
+            async def verify_and_run():
+                result = await client.verify(task=capability)
+
+                if not result.authorized:
+                    raise PermissionError(
+                        f"LockStock DENIED: Node '{node_func.__name__}' "
+                        f"requires capability '{capability}'. "
+                        f"Reason: {result.reason}"
+                    )
+
+                await client.log_audit(
+                    action=f"node:{node_func.__name__}",
+                    status="EXECUTING"
+                )
+
+                return node_func(*args, **kwargs)
+
+            return loop.run_until_complete(verify_and_run())
+
+        # Return appropriate wrapper based on original function type
+        if asyncio.iscoroutinefunction(node_func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    return decorator

lockstock_openai/__init__.py

@@ -0,0 +1,15 @@
+"""
+LockStock OpenAI Agents SDK Integration
+----------------------------------------
+Provides guardrails and tracing for OpenAI Agents SDK.
+"""
+
+from .guardrails import LockStockGuardrail, LockStockInputGuardrail, LockStockOutputGuardrail
+from .tracing import LockStockTracer
+
+__all__ = [
+    "LockStockGuardrail",
+    "LockStockInputGuardrail",
+    "LockStockOutputGuardrail",
+    "LockStockTracer"
+]