zetro-sentinel-sdk 0.3.0__py3-none-any.whl

zetro_sentinel_sdk/client.py

@@ -0,0 +1,890 @@
+"""AI Sentinel SDK Client."""
+
+from typing import Any, Dict, List, Optional
+
+import httpx
+
+from zetro_sentinel_sdk.exceptions import (
+    AuthenticationError,
+    NetworkError,
+    RateLimitError,
+    SentinelError,
+    ValidationError,
+)
+from zetro_sentinel_sdk.models import (
+    ActionSourceResult,
+    AuthorizeResult,
+    HierarchyResult,
+    Incident,
+    IncidentList,
+    Policy,
+    RateLimitResult,
+    ScanResult,
+    ToolExecution,
+    ToolExecutionList,
+    ToolResultScanResult,
+)
+
+
+class Sentinel:
+    """
+    Synchronous client for AI Sentinel API.
+
+    Usage:
+        sentinel = Sentinel(api_key="your-api-key")
+
+        # Scan user input
+        result = sentinel.scan_input("Hello world", agent_id="my-agent")
+        if not result.allowed:
+            print(f"Blocked: {result.reason}")
+
+        # Authorize tool call
+        auth = sentinel.authorize_tool(
+            agent_id="my-agent",
+            tool_name="send_email",
+            user_role="USER",
+            user_id="user-123"
+        )
+    """
+
+    DEFAULT_BASE_URL = "https://api.aisentinel.io"
+    DEFAULT_TIMEOUT = 30.0
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = None,
+        timeout: float = None,
+    ) -> None:
+        """
+        Initialize the Sentinel client.
+
+        Args:
+            api_key: Your AI Sentinel API key
+            base_url: API base URL (defaults to production)
+            timeout: Request timeout in seconds
+        """
+        self.api_key = api_key
+        self.base_url = (base_url or self.DEFAULT_BASE_URL).rstrip("/")
+        self.timeout = timeout or self.DEFAULT_TIMEOUT
+
+        self._client = httpx.Client(
+            base_url=self.base_url,
+            headers={
+                "X-API-Key": api_key,
+                "Content-Type": "application/json",
+                "User-Agent": "zetro-sentinel-sdk/0.3.0",
+            },
+            timeout=self.timeout,
+        )
+
+    def _handle_response(self, response: httpx.Response) -> Dict[str, Any]:
+        """Handle API response and raise appropriate exceptions."""
+        if response.status_code == 401:
+            raise AuthenticationError(
+                "Invalid API key",
+                status_code=401,
+                response=response.json() if response.content else None,
+            )
+        elif response.status_code == 429:
+            retry_after = response.headers.get("Retry-After")
+            raise RateLimitError(
+                "Rate limit exceeded",
+                retry_after=int(retry_after) if retry_after else None,
+                response=response.json() if response.content else None,
+            )
+        elif response.status_code == 422:
+            data = response.json() if response.content else {}
+            raise ValidationError(
+                "Validation error",
+                errors=data.get("detail", []),
+                response=data,
+            )
+        elif response.status_code >= 400:
+            raise SentinelError(
+                f"API error: {response.status_code}",
+                status_code=response.status_code,
+                response=response.json() if response.content else None,
+            )
+
+        return response.json()
+
+    def _request(self, method: str, path: str, **kwargs) -> Dict[str, Any]:
+        """Make an API request."""
+        try:
+            response = self._client.request(method, f"/v1{path}", **kwargs)
+            return self._handle_response(response)
+        except httpx.RequestError as e:
+            raise NetworkError(f"Network error: {str(e)}")
+
+    # =========================================================================
+    # SCANNING
+    # =========================================================================
+
+    def scan_input(
+        self,
+        text: str,
+        agent_id: str = "default",
+        session_id: str = None,
+    ) -> ScanResult:
+        """
+        Scan user input for prompt injection and policy violations.
+
+        Args:
+            text: User input text to scan
+            agent_id: Agent identifier (for agent-specific policies)
+            session_id: Optional session ID for logging
+
+        Returns:
+            ScanResult with allowed status and detection details
+        """
+        data = self._request(
+            "POST",
+            "/scan/input",
+            json={
+                "text": text,
+                "agent_id": agent_id,
+                "session_id": session_id,
+            },
+        )
+        return ScanResult(**data)
+
+    def scan_output(
+        self,
+        text: str,
+        agent_id: str = "default",
+        session_id: str = None,
+    ) -> ScanResult:
+        """
+        Scan agent output for sensitive data leaks.
+
+        Args:
+            text: Output text to scan
+            agent_id: Agent identifier
+            session_id: Optional session ID for logging
+
+        Returns:
+            ScanResult with allowed status and any detected patterns
+        """
+        data = self._request(
+            "POST",
+            "/scan/output",
+            json={
+                "text": text,
+                "agent_id": agent_id,
+                "session_id": session_id,
+            },
+        )
+        return ScanResult(**data)
+
+    def scan_tool_result(
+        self,
+        text: str,
+        tool_name: str,
+        agent_id: str = "default",
+        session_id: str = None,
+    ) -> ToolResultScanResult:
+        """
+        Scan tool result for indirect injection patterns.
+
+        Use this after executing external tool calls to detect
+        instruction patterns embedded in the response.
+
+        Args:
+            text: Tool result text to scan
+            tool_name: Name of the tool that produced the result
+            agent_id: Agent identifier
+            session_id: Optional session ID for logging
+
+        Returns:
+            ToolResultScanResult with detection details
+        """
+        data = self._request(
+            "POST",
+            "/scan/tool-result",
+            json={
+                "text": text,
+                "tool_name": tool_name,
+                "agent_id": agent_id,
+                "session_id": session_id,
+            },
+        )
+        return ToolResultScanResult(**data)
+
+    # =========================================================================
+    # AUTHORIZATION
+    # =========================================================================
+
+    def authorize_tool(
+        self,
+        agent_id: str,
+        tool_name: str,
+        user_role: str,
+        user_id: str,
+        is_resource_owner: bool = True,
+        arguments: Dict[str, Any] = None,
+        session_id: str = None,
+    ) -> AuthorizeResult:
+        """
+        Authorize a tool call based on policy.
+
+        Args:
+            agent_id: Agent identifier
+            tool_name: Tool to authorize
+            user_role: User's role (e.g., "USER", "ADMIN")
+            user_id: User identifier
+            is_resource_owner: Whether user owns the resource
+            arguments: Tool arguments (for HITL hash verification)
+            session_id: Optional session ID
+
+        Returns:
+            AuthorizeResult with allowed status and approval details
+        """
+        data = self._request(
+            "POST",
+            "/authorize/tool",
+            json={
+                "agent_id": agent_id,
+                "tool_name": tool_name,
+                "user_role": user_role,
+                "user_id": user_id,
+                "is_resource_owner": is_resource_owner,
+                "arguments": arguments,
+                "session_id": session_id,
+            },
+        )
+        return AuthorizeResult(**data)
+
+    def check_rate_limit(
+        self,
+        agent_id: str,
+        tool_name: str,
+        user_id: str,
+    ) -> RateLimitResult:
+        """
+        Check if a tool call is within rate limits.
+
+        Args:
+            agent_id: Agent identifier
+            tool_name: Tool being called
+            user_id: User identifier
+
+        Returns:
+            RateLimitResult with current counts and limits
+        """
+        data = self._request(
+            "POST",
+            "/authorize/rate-limit",
+            json={
+                "agent_id": agent_id,
+                "tool_name": tool_name,
+                "user_id": user_id,
+            },
+        )
+        return RateLimitResult(**data)
+
+    def evaluate_action_source(
+        self,
+        agent_id: str,
+        user_message: str,
+        tool_name: str,
+        tool_arguments: Dict[str, Any],
+        tool_results: List[Dict] = None,
+    ) -> ActionSourceResult:
+        """
+        Evaluate whether an action was directly requested or data-derived.
+
+        Use this for indirect injection defense to determine if
+        a proposed action came from the user or external data.
+
+        Args:
+            agent_id: Agent identifier
+            user_message: Original user request
+            tool_name: Tool being called
+            tool_arguments: Arguments to the tool
+            tool_results: Previous tool results with provenance
+
+        Returns:
+            ActionSourceResult with source classification
+        """
+        data = self._request(
+            "POST",
+            "/authorize/action-source",
+            json={
+                "agent_id": agent_id,
+                "user_message": user_message,
+                "tool_name": tool_name,
+                "tool_arguments": tool_arguments,
+                "tool_results": tool_results or [],
+            },
+        )
+        return ActionSourceResult(**data)
+
+    def check_hierarchy(
+        self,
+        agent_id: str,
+        user_message: str,
+        proposed_action: Dict[str, Any],
+        tool_results_with_instructions: List[Dict],
+    ) -> HierarchyResult:
+        """
+        Check if proposed action respects instruction hierarchy.
+
+        Verifies that external data is not overriding user instructions.
+
+        Args:
+            agent_id: Agent identifier
+            user_message: Original user request
+            proposed_action: Action LLM wants to take
+            tool_results_with_instructions: Tool results flagged with instructions
+
+        Returns:
+            HierarchyResult with violation details if any
+        """
+        data = self._request(
+            "POST",
+            "/authorize/hierarchy",
+            json={
+                "agent_id": agent_id,
+                "user_message": user_message,
+                "proposed_action": proposed_action,
+                "tool_results_with_instructions": tool_results_with_instructions,
+            },
+        )
+        return HierarchyResult(**data)
+
+    # =========================================================================
+    # INCIDENTS
+    # =========================================================================
+
+    def list_incidents(
+        self,
+        page: int = 1,
+        page_size: int = 20,
+        severity: str = None,
+        category: str = None,
+        agent_id: str = None,
+        resolved: bool = None,
+    ) -> IncidentList:
+        """
+        List security incidents with filtering.
+
+        Args:
+            page: Page number (1-indexed)
+            page_size: Items per page (max 100)
+            severity: Filter by severity
+            category: Filter by category
+            agent_id: Filter by agent
+            resolved: Filter by resolution status
+
+        Returns:
+            IncidentList with paginated incidents
+        """
+        params = {"page": page, "page_size": page_size}
+        if severity:
+            params["severity"] = severity
+        if category:
+            params["category"] = category
+        if agent_id:
+            params["agent_id"] = agent_id
+        if resolved is not None:
+            params["resolved"] = resolved
+
+        data = self._request("GET", "/incidents", params=params)
+        return IncidentList(**data)
+
+    def get_incident(self, incident_id: str) -> Incident:
+        """
+        Get a specific incident by ID.
+
+        Args:
+            incident_id: Incident ID
+
+        Returns:
+            Incident details
+        """
+        data = self._request("GET", f"/incidents/{incident_id}")
+        return Incident(**data)
+
+    # =========================================================================
+    # POLICIES
+    # =========================================================================
+
+    def get_policy(self) -> Policy:
+        """
+        Get the current security policy.
+
+        Returns:
+            Policy configuration
+        """
+        data = self._request("GET", "/policies")
+        return Policy(**data)
+
+    def toggle_agent(self, agent_id: str, enabled: bool, reason: str = None) -> Dict:
+        """
+        Toggle the kill switch for an agent.
+
+        Args:
+            agent_id: Agent identifier
+            enabled: Whether to enable or disable
+            reason: Reason for the change
+
+        Returns:
+            Confirmation dict
+        """
+        return self._request(
+            "POST",
+            f"/policies/kill-switch/agent/{agent_id}",
+            json={"enabled": enabled, "reason": reason},
+        )
+
+    def toggle_tool(
+        self, agent_id: str, tool_name: str, enabled: bool, reason: str = None
+    ) -> Dict:
+        """
+        Toggle the kill switch for a specific tool.
+
+        Args:
+            agent_id: Agent identifier
+            tool_name: Tool name
+            enabled: Whether to enable or disable
+            reason: Reason for the change
+
+        Returns:
+            Confirmation dict
+        """
+        return self._request(
+            "POST",
+            f"/policies/kill-switch/tool/{agent_id}/{tool_name}",
+            json={"enabled": enabled, "reason": reason},
+        )
+
+    # =========================================================================
+    # TOOL EXECUTIONS
+    # =========================================================================
+
+    def create_execution(
+        self,
+        agent_id: str,
+        tool_name: str,
+        user_id: str = None,
+        session_id: str = None,
+        tool_arguments: Dict[str, Any] = None,
+        argument_hash: str = None,
+        action_source: str = None,
+    ) -> ToolExecution:
+        """
+        Create a tool execution record to track a tool call.
+
+        Call this when starting a tool execution to begin tracking.
+        Use complete_execution() when the tool call finishes.
+
+        Args:
+            agent_id: Agent making the tool call
+            tool_name: Name of the tool being called
+            user_id: Optional user identifier
+            session_id: Optional session identifier
+            tool_arguments: Arguments passed to the tool
+            argument_hash: Hash of arguments (for verification)
+            action_source: DIRECT_REQUEST, DATA_DERIVED, or HYBRID
+
+        Returns:
+            ToolExecution with ID to use for completion
+
+        Example:
+            execution = sentinel.create_execution(
+                agent_id="my-agent",
+                tool_name="send_email",
+                user_id="user-123",
+                tool_arguments={"to": "user@example.com"}
+            )
+            try:
+                result = execute_tool(...)
+                sentinel.complete_execution(execution.id, "SUCCESS", result=result)
+            except Exception as e:
+                sentinel.complete_execution(execution.id, "FAILED", error=str(e))
+        """
+        data = self._request(
+            "POST",
+            "/tool-executions",
+            json={
+                "agent_id": agent_id,
+                "tool_name": tool_name,
+                "user_id": user_id,
+                "session_id": session_id,
+                "tool_arguments": tool_arguments,
+                "argument_hash": argument_hash,
+                "action_source": action_source,
+            },
+        )
+        return ToolExecution(**data)
+
+    def complete_execution(
+        self,
+        execution_id: str,
+        status: str,
+        result: Dict[str, Any] = None,
+        error: str = None,
+        error_type: str = None,
+    ) -> ToolExecution:
+        """
+        Complete a tool execution record.
+
+        Call this when a tool call finishes to record the outcome.
+
+        Args:
+            execution_id: Execution ID from create_execution()
+            status: Final status - one of:
+                - SUCCESS: Tool completed successfully
+                - FAILED: Tool encountered an error
+                - DENIED: Tool was denied by policy
+                - TIMEOUT: Tool execution timed out
+                - CANCELLED: Tool execution was cancelled
+            result: Optional result data (for SUCCESS)
+            error: Error message (for FAILED/TIMEOUT)
+            error_type: Type of error (e.g., "ValidationError")
+
+        Returns:
+            Updated ToolExecution with completion details
+        """
+        data = self._request(
+            "POST",
+            f"/tool-executions/{execution_id}/complete",
+            json={
+                "status": status,
+                "result": result,
+                "error": error,
+                "error_type": error_type,
+            },
+        )
+        return ToolExecution(**data)
+
+    def list_executions(
+        self,
+        page: int = 1,
+        page_size: int = 20,
+        tool_name: str = None,
+        status: str = None,
+        agent_id: str = None,
+        user_id: str = None,
+        session_id: str = None,
+    ) -> ToolExecutionList:
+        """
+        List tool executions with filtering.
+
+        Args:
+            page: Page number (1-indexed)
+            page_size: Items per page (max 100)
+            tool_name: Filter by tool name
+            status: Filter by status
+            agent_id: Filter by agent
+            user_id: Filter by user
+            session_id: Filter by session
+
+        Returns:
+            ToolExecutionList with paginated executions
+        """
+        params = {"page": page, "page_size": page_size}
+        if tool_name:
+            params["tool_name"] = tool_name
+        if status:
+            params["status"] = status
+        if agent_id:
+            params["agent_id"] = agent_id
+        if user_id:
+            params["user_id"] = user_id
+        if session_id:
+            params["session_id"] = session_id
+
+        data = self._request("GET", "/tool-executions", params=params)
+        return ToolExecutionList(**data)
+
+    def get_execution(self, execution_id: str) -> ToolExecution:
+        """
+        Get a specific tool execution by ID.
+
+        Args:
+            execution_id: Execution ID
+
+        Returns:
+            ToolExecution details
+        """
+        data = self._request("GET", f"/tool-executions/{execution_id}")
+        return ToolExecution(**data)
+
+    def close(self) -> None:
+        """Close the HTTP client."""
+        self._client.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
+class AsyncSentinel:
+    """
+    Asynchronous client for AI Sentinel API.
+
+    Usage:
+        async with AsyncSentinel(api_key="your-api-key") as sentinel:
+            result = await sentinel.scan_input("Hello world")
+    """
+
+    DEFAULT_BASE_URL = "https://api.aisentinel.io"
+    DEFAULT_TIMEOUT = 30.0
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = None,
+        timeout: float = None,
+    ) -> None:
+        self.api_key = api_key
+        self.base_url = (base_url or self.DEFAULT_BASE_URL).rstrip("/")
+        self.timeout = timeout or self.DEFAULT_TIMEOUT
+
+        self._client = httpx.AsyncClient(
+            base_url=self.base_url,
+            headers={
+                "X-API-Key": api_key,
+                "Content-Type": "application/json",
+                "User-Agent": "zetro-sentinel-sdk/0.3.0",
+            },
+            timeout=self.timeout,
+        )
+
+    async def _handle_response(self, response: httpx.Response) -> Dict[str, Any]:
+        """Handle API response and raise appropriate exceptions."""
+        if response.status_code == 401:
+            raise AuthenticationError("Invalid API key", status_code=401)
+        elif response.status_code == 429:
+            retry_after = response.headers.get("Retry-After")
+            raise RateLimitError(
+                "Rate limit exceeded",
+                retry_after=int(retry_after) if retry_after else None,
+            )
+        elif response.status_code == 422:
+            data = response.json() if response.content else {}
+            raise ValidationError("Validation error", errors=data.get("detail", []))
+        elif response.status_code >= 400:
+            raise SentinelError(f"API error: {response.status_code}", status_code=response.status_code)
+
+        return response.json()
+
+    async def _request(self, method: str, path: str, **kwargs) -> Dict[str, Any]:
+        """Make an async API request."""
+        try:
+            response = await self._client.request(method, f"/v1{path}", **kwargs)
+            return await self._handle_response(response)
+        except httpx.RequestError as e:
+            raise NetworkError(f"Network error: {str(e)}")
+
+    async def scan_input(
+        self,
+        text: str,
+        agent_id: str = "default",
+        session_id: str = None,
+    ) -> ScanResult:
+        """Async version of scan_input."""
+        data = await self._request(
+            "POST",
+            "/scan/input",
+            json={"text": text, "agent_id": agent_id, "session_id": session_id},
+        )
+        return ScanResult(**data)
+
+    async def scan_output(
+        self,
+        text: str,
+        agent_id: str = "default",
+        session_id: str = None,
+    ) -> ScanResult:
+        """Async version of scan_output."""
+        data = await self._request(
+            "POST",
+            "/scan/output",
+            json={"text": text, "agent_id": agent_id, "session_id": session_id},
+        )
+        return ScanResult(**data)
+
+    async def scan_tool_result(
+        self,
+        text: str,
+        tool_name: str,
+        agent_id: str = "default",
+        session_id: str = None,
+    ) -> ToolResultScanResult:
+        """Async version of scan_tool_result."""
+        data = await self._request(
+            "POST",
+            "/scan/tool-result",
+            json={
+                "text": text,
+                "tool_name": tool_name,
+                "agent_id": agent_id,
+                "session_id": session_id,
+            },
+        )
+        return ToolResultScanResult(**data)
+
+    async def authorize_tool(
+        self,
+        agent_id: str,
+        tool_name: str,
+        user_role: str,
+        user_id: str,
+        is_resource_owner: bool = True,
+        arguments: Dict[str, Any] = None,
+        session_id: str = None,
+    ) -> AuthorizeResult:
+        """Async version of authorize_tool."""
+        data = await self._request(
+            "POST",
+            "/authorize/tool",
+            json={
+                "agent_id": agent_id,
+                "tool_name": tool_name,
+                "user_role": user_role,
+                "user_id": user_id,
+                "is_resource_owner": is_resource_owner,
+                "arguments": arguments,
+                "session_id": session_id,
+            },
+        )
+        return AuthorizeResult(**data)
+
+    async def evaluate_action_source(
+        self,
+        agent_id: str,
+        user_message: str,
+        tool_name: str,
+        tool_arguments: Dict[str, Any],
+        tool_results: List[Dict] = None,
+    ) -> ActionSourceResult:
+        """Async version of evaluate_action_source."""
+        data = await self._request(
+            "POST",
+            "/authorize/action-source",
+            json={
+                "agent_id": agent_id,
+                "user_message": user_message,
+                "tool_name": tool_name,
+                "tool_arguments": tool_arguments,
+                "tool_results": tool_results or [],
+            },
+        )
+        return ActionSourceResult(**data)
+
+    async def check_hierarchy(
+        self,
+        agent_id: str,
+        user_message: str,
+        proposed_action: Dict[str, Any],
+        tool_results_with_instructions: List[Dict],
+    ) -> HierarchyResult:
+        """Async version of check_hierarchy."""
+        data = await self._request(
+            "POST",
+            "/authorize/hierarchy",
+            json={
+                "agent_id": agent_id,
+                "user_message": user_message,
+                "proposed_action": proposed_action,
+                "tool_results_with_instructions": tool_results_with_instructions,
+            },
+        )
+        return HierarchyResult(**data)
+
+    # =========================================================================
+    # TOOL EXECUTIONS
+    # =========================================================================
+
+    async def create_execution(
+        self,
+        agent_id: str,
+        tool_name: str,
+        user_id: str = None,
+        session_id: str = None,
+        tool_arguments: Dict[str, Any] = None,
+        argument_hash: str = None,
+        action_source: str = None,
+    ) -> ToolExecution:
+        """Async version of create_execution."""
+        data = await self._request(
+            "POST",
+            "/tool-executions",
+            json={
+                "agent_id": agent_id,
+                "tool_name": tool_name,
+                "user_id": user_id,
+                "session_id": session_id,
+                "tool_arguments": tool_arguments,
+                "argument_hash": argument_hash,
+                "action_source": action_source,
+            },
+        )
+        return ToolExecution(**data)
+
+    async def complete_execution(
+        self,
+        execution_id: str,
+        status: str,
+        result: Dict[str, Any] = None,
+        error: str = None,
+        error_type: str = None,
+    ) -> ToolExecution:
+        """Async version of complete_execution."""
+        data = await self._request(
+            "POST",
+            f"/tool-executions/{execution_id}/complete",
+            json={
+                "status": status,
+                "result": result,
+                "error": error,
+                "error_type": error_type,
+            },
+        )
+        return ToolExecution(**data)
+
+    async def list_executions(
+        self,
+        page: int = 1,
+        page_size: int = 20,
+        tool_name: str = None,
+        status: str = None,
+        agent_id: str = None,
+        user_id: str = None,
+        session_id: str = None,
+    ) -> ToolExecutionList:
+        """Async version of list_executions."""
+        params = {"page": page, "page_size": page_size}
+        if tool_name:
+            params["tool_name"] = tool_name
+        if status:
+            params["status"] = status
+        if agent_id:
+            params["agent_id"] = agent_id
+        if user_id:
+            params["user_id"] = user_id
+        if session_id:
+            params["session_id"] = session_id
+
+        data = await self._request("GET", "/tool-executions", params=params)
+        return ToolExecutionList(**data)
+
+    async def get_execution(self, execution_id: str) -> ToolExecution:
+        """Async version of get_execution."""
+        data = await self._request("GET", f"/tool-executions/{execution_id}")
+        return ToolExecution(**data)
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        await self._client.aclose()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        await self.close()