aiproxyguard-python-sdk 0.1.0__py3-none-any.whl

aiproxyguard/decorators.py

@@ -0,0 +1,245 @@
+"""AIProxyGuard decorator utilities."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import warnings
+from functools import wraps
+from typing import Any, Callable, TypeVar, overload
+
+from .exceptions import ContentBlockedError
+
+if __import__("typing").TYPE_CHECKING:
+    from .client import AIProxyGuard
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class GuardConfigurationError(ValueError):
+    """Raised when the guard decorator is misconfigured."""
+
+    pass
+
+
+@overload
+def guard(
+    client: AIProxyGuard,
+    *,
+    input_arg: str = "prompt",
+    raise_on_block: bool = True,
+    fail_closed: bool = True,
+) -> Callable[[F], F]: ...
+
+
+@overload
+def guard(
+    client: AIProxyGuard,
+    *,
+    input_arg: int,
+    raise_on_block: bool = True,
+    fail_closed: bool = True,
+) -> Callable[[F], F]: ...
+
+
+def guard(
+    client: AIProxyGuard,
+    *,
+    input_arg: str | int = "prompt",
+    raise_on_block: bool = True,
+    fail_closed: bool = True,
+) -> Callable[[F], F]:
+    """Decorator to guard a function with prompt injection detection.
+
+    Checks the specified input argument before the function executes.
+    If the content is blocked and raise_on_block is True, raises ContentBlockedError.
+
+    Args:
+        client: AIProxyGuard client instance.
+        input_arg: Name or index of the argument to check. Defaults to "prompt".
+                   Can be a string (kwarg name) or int (positional index).
+        raise_on_block: If True, raise ContentBlockedError when content is blocked.
+                        If False, the function is not called but no error is raised.
+        fail_closed: If True (default), raise GuardConfigurationError when input_arg
+                     cannot be resolved. If False, issue a warning and skip checking.
+
+    Returns:
+        Decorated function that checks input before execution.
+
+    Raises:
+        GuardConfigurationError: If input_arg doesn't match any parameter
+            (at decoration time for string args, or at call time if fail_closed).
+
+    Example:
+        >>> client = AIProxyGuard("http://localhost:8080")
+        >>> @guard(client)
+        ... def call_llm(prompt: str) -> str:
+        ...     return "response"
+        ...
+        >>> @guard(client, input_arg="user_input")
+        ... def process(user_input: str, system: str) -> str:
+        ...     return "processed"
+        ...
+        >>> @guard(client, input_arg=0)
+        ... def handle(text: str) -> str:
+        ...     return "handled"
+    """
+
+    def decorator(func: F) -> F:
+        is_async = asyncio.iscoroutinefunction(func)
+
+        # Cache signature introspection at decoration time for performance
+        sig = inspect.signature(func)
+        param_names = list(sig.parameters.keys())
+
+        # Validate string input_arg at decoration time
+        if isinstance(input_arg, str):
+            if input_arg not in param_names:
+                raise GuardConfigurationError(
+                    f"guard(): input_arg '{input_arg}' not found in function "
+                    f"'{func.__name__}' parameters: {param_names}"
+                )
+            # Pre-compute the positional index for this parameter
+            cached_arg_index: int | None = param_names.index(input_arg)
+        else:
+            cached_arg_index = None
+
+        def _extract_text(args: tuple[Any, ...], kwargs: dict[str, Any]) -> str | None:
+            """Extract text to check from function arguments."""
+            text: Any = None
+            resolved = False
+
+            if isinstance(input_arg, int):
+                if len(args) > input_arg:
+                    text = args[input_arg]
+                    resolved = True
+                elif fail_closed:
+                    raise GuardConfigurationError(
+                        f"guard(): input_arg index {input_arg} out of range. "
+                        f"'{func.__name__}' got {len(args)} positional args."
+                    )
+                else:
+                    warnings.warn(
+                        f"guard(): input_arg index {input_arg} out of range for "
+                        f"'{func.__name__}'. Skipping security check.",
+                        RuntimeWarning,
+                        stacklevel=4,
+                    )
+            else:
+                # String input_arg - use cached index
+                if input_arg in kwargs:
+                    text = kwargs[input_arg]
+                    resolved = True
+                elif cached_arg_index is not None and cached_arg_index < len(args):
+                    text = args[cached_arg_index]
+                    resolved = True
+                else:
+                    # This shouldn't happen if validation passed, but handle edge cases
+                    if fail_closed:
+                        raise GuardConfigurationError(
+                            f"guard(): Could not resolve input_arg '{input_arg}' "
+                            f"for function '{func.__name__}'."
+                        )
+                    else:
+                        warnings.warn(
+                            f"guard(): Could not resolve '{input_arg}' for "
+                            f"'{func.__name__}'. Skipping security check.",
+                            RuntimeWarning,
+                            stacklevel=4,
+                        )
+
+            if not resolved:
+                return None
+
+            return str(text) if text is not None else ""
+
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            text = _extract_text(args, kwargs)
+
+            # Check even empty strings (text is not None means we resolved the arg)
+            if text is not None:
+                result = client.check(text)
+                if result.is_blocked:
+                    if raise_on_block:
+                        raise ContentBlockedError(result)
+                    return None
+
+            return func(*args, **kwargs)
+
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            text = _extract_text(args, kwargs)
+
+            # Check even empty strings (text is not None means we resolved the arg)
+            if text is not None:
+                result = await client.check_async(text)
+                if result.is_blocked:
+                    if raise_on_block:
+                        raise ContentBlockedError(result)
+                    return None
+
+            return await func(*args, **kwargs)
+
+        return async_wrapper if is_async else sync_wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def guard_output(
+    client: AIProxyGuard,
+    *,
+    raise_on_block: bool = True,
+) -> Callable[[F], F]:
+    """Decorator to guard a function's output with prompt injection detection.
+
+    Checks the function's return value after execution.
+    Useful for validating LLM responses before returning them.
+
+    Args:
+        client: AIProxyGuard client instance.
+        raise_on_block: If True, raise ContentBlockedError when content is blocked.
+                        If False, returns None instead of the blocked content.
+
+    Returns:
+        Decorated function that checks output after execution.
+
+    Example:
+        >>> client = AIProxyGuard("http://localhost:8080")
+        >>> @guard_output(client)
+        ... def get_llm_response(prompt: str) -> str:
+        ...     return llm.generate(prompt)
+    """
+
+    def decorator(func: F) -> F:
+        is_async = asyncio.iscoroutinefunction(func)
+
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            output = func(*args, **kwargs)
+
+            if output is not None:
+                result = client.check(str(output))
+                if result.is_blocked:
+                    if raise_on_block:
+                        raise ContentBlockedError(result)
+                    return None
+
+            return output
+
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            output = await func(*args, **kwargs)
+
+            if output is not None:
+                result = await client.check_async(str(output))
+                if result.is_blocked:
+                    if raise_on_block:
+                        raise ContentBlockedError(result)
+                    return None
+
+            return output
+
+        return async_wrapper if is_async else sync_wrapper  # type: ignore[return-value]
+
+    return decorator

aiproxyguard/exceptions.py

@@ -0,0 +1,76 @@
+"""AIProxyGuard SDK exceptions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .models import CheckResult
+
+
+class AIProxyGuardError(Exception):
+    """Base exception for AIProxyGuard SDK."""
+
+    def __init__(self, message: str, code: str | None = None) -> None:
+        super().__init__(message)
+        self.message = message
+        self.code = code
+
+    def __repr__(self) -> str:
+        cls = self.__class__.__name__
+        return f"{cls}(message={self.message!r}, code={self.code!r})"
+
+
+class ValidationError(AIProxyGuardError):
+    """Raised when the request is invalid (400 errors)."""
+
+    pass
+
+
+class ConnectionError(AIProxyGuardError):
+    """Raised when connection to the service fails."""
+
+    pass
+
+
+class TimeoutError(AIProxyGuardError):
+    """Raised when a request times out."""
+
+    pass
+
+
+class ServerError(AIProxyGuardError):
+    """Raised when the server returns a 5xx error (retryable)."""
+
+    def __init__(self, message: str, status_code: int) -> None:
+        super().__init__(message, code="server_error")
+        self.status_code = status_code
+
+    def __repr__(self) -> str:
+        return f"ServerError(status_code={self.status_code}, message={self.message!r})"
+
+
+class RateLimitError(AIProxyGuardError):
+    """Raised when rate limited (429 errors)."""
+
+    def __init__(
+        self, message: str = "Rate limited", retry_after: int | None = None
+    ) -> None:
+        super().__init__(message, code="rate_limit")
+        self.retry_after = retry_after
+
+    def __repr__(self) -> str:
+        return f"RateLimitError(message={self.message!r}, retry={self.retry_after})"
+
+
+class ContentBlockedError(AIProxyGuardError):
+    """Raised when content is blocked due to prompt injection detection."""
+
+    def __init__(self, result: CheckResult) -> None:
+        super().__init__(f"Content blocked: {result.category}", code="content_blocked")
+        self.result = result
+
+    def __repr__(self) -> str:
+        cat = self.result.category
+        conf = self.result.confidence
+        return f"ContentBlockedError(category={cat!r}, confidence={conf})"

aiproxyguard/models.py

@@ -0,0 +1,222 @@
+"""AIProxyGuard SDK data models."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+
+class Action(str, Enum):
+    """Action taken by AIProxyGuard on scanned content."""
+
+    ALLOW = "allow"
+    LOG = "log"
+    WARN = "warn"
+    BLOCK = "block"
+
+
+@dataclass(frozen=True)
+class ThreatDetail:
+    """Details about a detected threat (cloud API only).
+
+    Attributes:
+        type: Threat category (e.g., "prompt-injection").
+        confidence: Detection confidence (0.0 to 1.0).
+        rule: Rule/signature ID that triggered the detection.
+    """
+
+    type: str
+    confidence: float
+    rule: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ThreatDetail:
+        """Create ThreatDetail from API response dictionary."""
+        return cls(
+            type=data["type"],
+            confidence=float(data.get("confidence", 0.0)),
+            rule=data.get("rule"),
+        )
+
+
+@dataclass(frozen=True)
+class CheckResult:
+    """Result from scanning text for prompt injection.
+
+    Attributes:
+        action: The action taken (allow, log, warn, or block).
+        category: Category of the detected threat, if any.
+        signature_name: Name of the matching signature, if any.
+        confidence: Confidence score of the detection (0.0 to 1.0).
+    """
+
+    action: Action
+    category: str | None
+    signature_name: str | None
+    confidence: float
+
+    @property
+    def is_safe(self) -> bool:
+        """Returns True if the text was not blocked."""
+        return self.action != Action.BLOCK
+
+    @property
+    def is_blocked(self) -> bool:
+        """Returns True if the text was blocked."""
+        return self.action == Action.BLOCK
+
+    @property
+    def requires_attention(self) -> bool:
+        """Returns True if the text requires attention (warn or block)."""
+        return self.action in (Action.WARN, Action.BLOCK)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> CheckResult:
+        """Create a CheckResult from a proxy API response dictionary."""
+        return cls(
+            action=Action(data["action"]),
+            category=data.get("category"),
+            signature_name=data.get("signature_name"),
+            confidence=float(data.get("confidence", 0.0)),
+        )
+
+    @classmethod
+    def from_cloud_dict(cls, data: dict[str, Any]) -> CheckResult:
+        """Create a CheckResult from a cloud API response dictionary.
+
+        The cloud API returns a different format with threats array.
+        """
+        threats = data.get("threats", [])
+        category = None
+        signature_name = None
+        confidence = 0.0
+
+        if threats:
+            # Use the first threat for backwards compatibility
+            first_threat = threats[0]
+            category = first_threat.get("type")
+            signature_name = first_threat.get("rule")
+            confidence = float(first_threat.get("confidence", 0.0))
+
+        return cls(
+            action=Action(data["action"]),
+            category=category,
+            signature_name=signature_name,
+            confidence=confidence,
+        )
+
+
+@dataclass(frozen=True)
+class CloudCheckResult:
+    """Extended result from the cloud API with additional metadata.
+
+    Attributes:
+        id: Unique check ID.
+        flagged: Whether any threat was detected.
+        action: The action taken (allow, log, warn, or block).
+        threats: List of detected threats.
+        latency_ms: Processing time in milliseconds.
+        cached: Whether result was served from cache.
+    """
+
+    id: str
+    flagged: bool
+    action: Action
+    threats: list[ThreatDetail]
+    latency_ms: float
+    cached: bool
+
+    @property
+    def is_safe(self) -> bool:
+        """Returns True if the text was not blocked."""
+        return self.action != Action.BLOCK
+
+    @property
+    def is_blocked(self) -> bool:
+        """Returns True if the text was blocked."""
+        return self.action == Action.BLOCK
+
+    @property
+    def category(self) -> str | None:
+        """Returns the primary threat category, if any."""
+        return self.threats[0].type if self.threats else None
+
+    @property
+    def confidence(self) -> float:
+        """Returns the primary threat confidence."""
+        return self.threats[0].confidence if self.threats else 0.0
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> CloudCheckResult:
+        """Create CloudCheckResult from API response dictionary."""
+        return cls(
+            id=data["id"],
+            flagged=data["flagged"],
+            action=Action(data["action"]),
+            threats=[ThreatDetail.from_dict(t) for t in data.get("threats", [])],
+            latency_ms=float(data.get("latency_ms", 0.0)),
+            cached=data.get("cached", False),
+        )
+
+
+@dataclass(frozen=True)
+class ServiceInfo:
+    """Service information from the AIProxyGuard API.
+
+    Attributes:
+        service: Service name.
+        version: Service version.
+    """
+
+    service: str
+    version: str
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ServiceInfo:
+        """Create ServiceInfo from an API response dictionary."""
+        return cls(service=data["service"], version=data["version"])
+
+
+@dataclass(frozen=True)
+class HealthStatus:
+    """Health status from the AIProxyGuard API.
+
+    Attributes:
+        status: Health status string (e.g., "healthy").
+        healthy: Boolean indicating if the service is healthy.
+    """
+
+    status: str
+    healthy: bool
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> HealthStatus:
+        """Create HealthStatus from an API response dictionary."""
+        status = data.get("status", "unknown")
+        return cls(status=status, healthy=status == "healthy")
+
+
+@dataclass(frozen=True)
+class ReadyStatus:
+    """Readiness status from the AIProxyGuard API.
+
+    Attributes:
+        status: Readiness status string (e.g., "ready").
+        ready: Boolean indicating if the service is ready.
+        checks: Dictionary of individual check results.
+    """
+
+    status: str
+    ready: bool
+    checks: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ReadyStatus:
+        """Create ReadyStatus from an API response dictionary."""
+        status = data.get("status", "unknown")
+        return cls(
+            status=status,
+            ready=status == "ready",
+            checks=data.get("checks", {}),
+        )