raxe 0.4.6__py3-none-any.whl

raxe/sdk/suppression_context.py

@@ -0,0 +1,215 @@
+"""Scoped suppression context manager for SDK.
+
+Provides a context manager that temporarily applies suppressions
+to all scans within its scope.
+
+Example:
+    with raxe.suppressed("pi-*", reason="Testing auth flow"):
+        result = raxe.scan(text)
+        # All scans within this block have pi-* suppressed
+
+Clean Architecture:
+    - SDK layer convenience that wraps domain suppression logic
+    - Thread-safe using contextvars
+"""
+from __future__ import annotations
+
+import contextvars
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from raxe.domain.suppression import Suppression, SuppressionAction
+
+if TYPE_CHECKING:
+    from raxe.sdk.client import Raxe
+
+# Thread-local storage for scoped suppressions
+# Note: We use a factory function to avoid mutable default issues
+_scoped_suppressions: contextvars.ContextVar[list[Suppression]] = contextvars.ContextVar(
+    "scoped_suppressions"
+)
+
+
+def _get_or_init_scoped_suppressions() -> list[Suppression]:
+    """Get scoped suppressions, initializing with empty list if needed."""
+    try:
+        return _scoped_suppressions.get()
+    except LookupError:
+        # First access in this context - initialize with empty list
+        empty_list: list[Suppression] = []
+        _scoped_suppressions.set(empty_list)
+        return empty_list
+
+
+@dataclass(frozen=True)
+class ScopedSuppressionConfig:
+    """Configuration for scoped suppression.
+
+    Attributes:
+        pattern: Rule ID pattern to suppress
+        action: Action to take (default: SUPPRESS)
+        reason: Reason for suppression
+    """
+
+    pattern: str
+    action: SuppressionAction = SuppressionAction.SUPPRESS
+    reason: str = "Scoped suppression"
+
+
+def get_scoped_suppressions() -> list[Suppression]:
+    """Get current thread's scoped suppressions.
+
+    Returns:
+        List of suppressions active in the current context
+    """
+    return _get_or_init_scoped_suppressions()
+
+
+def _create_scoped_suppression(
+    pattern: str,
+    *,
+    action: SuppressionAction = SuppressionAction.SUPPRESS,
+    reason: str = "Scoped suppression",
+) -> Suppression:
+    """Create a scoped suppression.
+
+    Args:
+        pattern: Rule ID pattern
+        action: Action to take
+        reason: Reason for suppression
+
+    Returns:
+        Suppression object
+    """
+    return Suppression(
+        pattern=pattern,
+        reason=reason,
+        action=action,
+        created_at=datetime.now(timezone.utc).isoformat(),
+        created_by="scoped",
+    )
+
+
+@contextmanager
+def suppression_scope(
+    *patterns: str,
+    action: SuppressionAction | str = SuppressionAction.SUPPRESS,
+    reason: str = "Scoped suppression",
+) -> Iterator[None]:
+    """Context manager for scoped suppression.
+
+    All scans within the context will have the specified patterns suppressed.
+    This is thread-safe using contextvars.
+
+    Args:
+        *patterns: One or more rule ID patterns to suppress
+        action: Action to take (SUPPRESS, FLAG, or LOG)
+        reason: Reason for suppression
+
+    Yields:
+        None
+
+    Example:
+        with suppression_scope("pi-*", "jb-*", reason="Testing"):
+            result = pipeline.scan(text)
+            # pi-* and jb-* patterns are suppressed
+    """
+    # Convert string action to enum if needed
+    if isinstance(action, str):
+        action = SuppressionAction(action.upper())
+
+    # Create suppressions for all patterns
+    new_suppressions = [
+        _create_scoped_suppression(pattern, action=action, reason=reason)
+        for pattern in patterns
+    ]
+
+    # Get current suppressions and add new ones
+    current = _get_or_init_scoped_suppressions()
+    combined = list(current) + new_suppressions
+
+    # Set new suppressions
+    token = _scoped_suppressions.set(combined)
+
+    try:
+        yield
+    finally:
+        # Restore previous suppressions
+        _scoped_suppressions.reset(token)
+
+
+class SuppressedContext:
+    """Context manager bound to a specific Raxe client.
+
+    This is returned by Raxe.suppressed() and provides a more
+    convenient API that's bound to the client instance.
+
+    Example:
+        raxe = Raxe()
+        with raxe.suppressed("pi-*", reason="Testing"):
+            result = raxe.scan(text)
+    """
+
+    def __init__(
+        self,
+        client: Raxe,
+        *patterns: str,
+        action: SuppressionAction | str = SuppressionAction.SUPPRESS,
+        reason: str = "Scoped suppression",
+    ) -> None:
+        """Initialize suppressed context.
+
+        Args:
+            client: The Raxe client this context is bound to
+            *patterns: Rule ID patterns to suppress
+            action: Action to take
+            reason: Reason for suppression
+        """
+        self._client = client
+        self._patterns = patterns
+        self._action = action
+        self._reason = reason
+        self._token: contextvars.Token[list[Suppression]] | None = None
+
+    def __enter__(self) -> SuppressedContext:
+        """Enter the suppression context."""
+        # Convert string action to enum if needed
+        action = self._action
+        if isinstance(action, str):
+            action = SuppressionAction(action.upper())
+
+        # Create suppressions for all patterns
+        new_suppressions = [
+            _create_scoped_suppression(pattern, action=action, reason=self._reason)
+            for pattern in self._patterns
+        ]
+
+        # Get current suppressions and add new ones
+        current = _get_or_init_scoped_suppressions()
+        combined = list(current) + new_suppressions
+
+        # Set new suppressions
+        self._token = _scoped_suppressions.set(combined)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: Any,
+    ) -> None:
+        """Exit the suppression context."""
+        if self._token is not None:
+            _scoped_suppressions.reset(self._token)
+            self._token = None
+
+
+__all__ = [
+    "ScopedSuppressionConfig",
+    "SuppressedContext",
+    "get_scoped_suppressions",
+    "suppression_scope",
+]

raxe/sdk/wrappers/__init__.py

@@ -0,0 +1,163 @@
+"""LLM Client Wrappers.
+
+Drop-in replacements for popular LLM clients that add RAXE scanning.
+
+This module provides wrappers for:
+    - OpenAI client wrapper (RaxeOpenAI)
+    - Anthropic client wrapper (RaxeAnthropic)
+    - Google Vertex AI wrapper (RaxeVertexAI)
+    - wrap_client() helper for runtime wrapping
+
+Integrations are available via raxe.sdk.integrations:
+    - LangChain callback handler (RaxeCallbackHandler)
+    - Hugging Face pipeline wrapper (RaxePipeline)
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from raxe.sdk.client import Raxe
+    from raxe.sdk.wrappers.anthropic import RaxeAnthropic
+    from raxe.sdk.wrappers.openai import RaxeOpenAI
+    from raxe.sdk.wrappers.vertexai import RaxeVertexAI
+
+__all__ = [
+    "RaxeAnthropic",
+    "RaxeOpenAI",
+    "RaxeVertexAI",
+    "wrap_client",
+]
+
+# Availability flags
+_OPENAI_AVAILABLE = False
+_ANTHROPIC_AVAILABLE = False
+_VERTEXAI_AVAILABLE = False
+
+# Lazy imports to avoid requiring all dependencies
+RaxeOpenAI = None
+RaxeAnthropic = None
+RaxeVertexAI = None
+
+
+def __getattr__(name: str):
+    """Lazy import wrappers to avoid requiring all dependencies.
+
+    Args:
+        name: Name of the wrapper to import
+
+    Returns:
+        The requested wrapper class
+
+    Raises:
+        ImportError: If the wrapper's dependencies are not installed
+        AttributeError: If the wrapper name is not recognized
+    """
+    global RaxeOpenAI, RaxeAnthropic, RaxeVertexAI
+    global _OPENAI_AVAILABLE, _ANTHROPIC_AVAILABLE, _VERTEXAI_AVAILABLE
+
+    if name == "RaxeOpenAI":
+        if RaxeOpenAI is None:
+            try:
+                from raxe.sdk.wrappers.openai import RaxeOpenAI as _RaxeOpenAI
+                RaxeOpenAI = _RaxeOpenAI
+                _OPENAI_AVAILABLE = True
+            except ImportError as e:
+                raise ImportError(
+                    "OpenAI wrapper requires openai package. "
+                    "Install with: pip install openai"
+                ) from e
+        return RaxeOpenAI
+
+    elif name == "RaxeAnthropic":
+        if RaxeAnthropic is None:
+            try:
+                from raxe.sdk.wrappers.anthropic import RaxeAnthropic as _RaxeAnthropic
+                RaxeAnthropic = _RaxeAnthropic
+                _ANTHROPIC_AVAILABLE = True
+            except ImportError as e:
+                raise ImportError(
+                    "Anthropic wrapper requires anthropic package. "
+                    "Install with: pip install anthropic"
+                ) from e
+        return RaxeAnthropic
+
+    elif name == "RaxeVertexAI":
+        if RaxeVertexAI is None:
+            try:
+                from raxe.sdk.wrappers.vertexai import RaxeVertexAI as _RaxeVertexAI
+                RaxeVertexAI = _RaxeVertexAI
+                _VERTEXAI_AVAILABLE = True
+            except ImportError as e:
+                raise ImportError(
+                    "Vertex AI wrapper requires google-cloud-aiplatform package. "
+                    "Install with: pip install google-cloud-aiplatform"
+                ) from e
+        return RaxeVertexAI
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def wrap_client(raxe_client: Raxe, client: Any) -> Any:
+    """Wrap an LLM client with RAXE scanning.
+
+    This helper function wraps an existing LLM client instance
+    with RAXE scanning capabilities. It detects the client type
+    and returns the appropriate wrapper.
+
+    Args:
+        raxe_client: Raxe instance to use for scanning
+        client: LLM client to wrap (OpenAI, Anthropic, etc.)
+
+    Returns:
+        Wrapped client with automatic scanning
+
+    Raises:
+        ImportError: If required wrapper package not installed
+        NotImplementedError: If client type not supported
+
+    Usage:
+        >>> from raxe import Raxe
+        >>> from openai import OpenAI
+        >>> raxe = Raxe()
+        >>> client = raxe.wrap(OpenAI(api_key="sk-..."))
+        >>> # All calls automatically scanned
+
+    Supported Clients:
+        - openai.OpenAI -> RaxeOpenAI
+        - anthropic.Anthropic -> RaxeAnthropic
+        - More coming soon...
+    """
+    client_type = type(client).__name__
+
+    if client_type == "OpenAI":
+        # Wrap OpenAI client - import directly from submodule to avoid lazy loading issues
+        from raxe.sdk.wrappers.openai import RaxeOpenAI
+
+        # Copy API key from original if available
+        api_key = getattr(client, "api_key", None)
+
+        # Create wrapped version with same config
+        # RaxeOpenAI inherits from OpenAI, so api_key is passed to the constructor
+        wrapped = RaxeOpenAI(api_key=api_key, raxe=raxe_client)
+
+        return wrapped
+
+    elif client_type == "Anthropic":
+        # Wrap Anthropic client - import directly from submodule to avoid lazy loading issues
+        from raxe.sdk.wrappers.anthropic import RaxeAnthropic
+
+        # Create wrapped version with same config
+        wrapped = RaxeAnthropic(raxe=raxe_client)
+
+        # Copy API key from original if available
+        if hasattr(client, "api_key"):
+            wrapped._anthropic_client.api_key = client.api_key
+
+        return wrapped
+
+    else:
+        raise NotImplementedError(
+            f"Wrapper for {client_type} not implemented yet. "
+            f"Supported: OpenAI, Anthropic"
+        )

raxe/sdk/wrappers/anthropic.py

@@ -0,0 +1,310 @@
+"""Anthropic client wrapper with automatic RAXE scanning.
+
+Drop-in replacement for anthropic.Anthropic that scans all prompts and responses.
+
+This wrapper allows users to replace:
+    from anthropic import Anthropic
+    client = Anthropic()
+
+With:
+    from raxe.sdk.wrappers import RaxeAnthropic
+    client = RaxeAnthropic()  # All calls automatically scanned
+
+The wrapper intercepts messages.create calls, scans user messages
+before sending to Claude, and optionally scans responses.
+
+Default behavior is LOG-ONLY (safe to add to production without breaking flows).
+Enable blocking with `raxe_block_on_threat=True` for strict mode.
+"""
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
+
+from raxe.sdk.agent_scanner import (
+    AgentScannerConfig,
+    create_agent_scanner,
+)
+
+if TYPE_CHECKING:
+    from raxe.sdk.client import Raxe
+
+logger = logging.getLogger(__name__)
+
+
+class RaxeAnthropic:
+    """Drop-in replacement for anthropic.Anthropic with automatic scanning.
+
+    This wrapper wraps the Anthropic client and intercepts all
+    messages.create calls to scan prompts and responses.
+
+    Usage:
+        # Instead of:
+        from anthropic import Anthropic
+        client = Anthropic(api_key="sk-ant-...")
+
+        # Use:
+        from raxe.sdk.wrappers import RaxeAnthropic
+        client = RaxeAnthropic(api_key="sk-ant-...")
+
+        # All messages.create calls are automatically scanned
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            messages=[{"role": "user", "content": "Hello"}]
+        )
+
+    Attributes:
+        raxe: Raxe client instance for scanning
+        raxe_block_on_threat: Whether to block requests on threat detection
+        raxe_scan_responses: Whether to scan Claude responses
+    """
+
+    def __init__(
+        self,
+        *args,
+        raxe: Raxe | None = None,
+        raxe_block_on_threat: bool = False,
+        raxe_scan_responses: bool = True,
+        **kwargs
+    ):
+        """Initialize RaxeAnthropic client.
+
+        Args:
+            *args: Passed to Anthropic.__init__
+            raxe: Optional Raxe client (creates default if not provided)
+            raxe_block_on_threat: Block requests on threat detection.
+                Default is False (log-only mode, safe for production).
+            raxe_scan_responses: Also scan Claude responses
+            **kwargs: Passed to Anthropic.__init__
+
+        Example:
+            # With default Raxe client (log-only mode)
+            client = RaxeAnthropic(api_key="sk-ant-...")
+
+            # With custom Raxe client
+            raxe = Raxe(telemetry=False)
+            client = RaxeAnthropic(api_key="sk-ant-...", raxe=raxe)
+
+            # Enable blocking (strict mode)
+            client = RaxeAnthropic(
+                api_key="sk-ant-...",
+                raxe_block_on_threat=True
+            )
+
+        Raises:
+            ImportError: If anthropic package not installed
+        """
+        # Try to import Anthropic
+        try:
+            from anthropic import Anthropic
+        except ImportError as e:
+            raise ImportError(
+                "anthropic package is required for RaxeAnthropic. "
+                "Install with: pip install anthropic"
+            ) from e
+
+        # Initialize parent Anthropic client
+        self._anthropic_client = Anthropic(*args, **kwargs)
+
+        # Create or use provided Raxe client
+        if raxe is None:
+            from raxe.sdk.client import Raxe
+            raxe = Raxe()
+
+        self.raxe = raxe
+        self.raxe_block_on_threat = raxe_block_on_threat
+        self.raxe_scan_responses = raxe_scan_responses
+
+        # Create AgentScanner for unified scanning
+        config = AgentScannerConfig(
+            scan_prompts=True,
+            scan_responses=raxe_scan_responses,
+            on_threat="block" if raxe_block_on_threat else "log",
+        )
+        self._scanner = create_agent_scanner(raxe, config, integration_type="anthropic")
+
+        # Wrap the messages resource
+        self._wrap_messages()
+
+        logger.debug(
+            f"RaxeAnthropic initialized: block={raxe_block_on_threat}, "
+            f"scan_responses={raxe_scan_responses}"
+        )
+
+    def _wrap_messages(self):
+        """Wrap messages.create method with scanning."""
+        # Get original create method
+        original_create = self._anthropic_client.messages.create
+
+        def wrapped_create(*args, **kwargs):
+            """Wrapped messages.create with RAXE scanning."""
+            # Extract messages
+            messages = kwargs.get("messages", [])
+
+            # Scan each user message
+            for message in messages:
+                if isinstance(message, dict):
+                    # Handle dict-style messages
+                    if message.get("role") == "user":
+                        content = message.get("content", "")
+                        if content:
+                            self._scan_message(content)
+                elif hasattr(message, "role") and hasattr(message, "content"):
+                    # Handle object-style messages
+                    if message.role == "user":
+                        self._scan_message(message.content)
+
+            # Check if streaming is requested
+            stream = kwargs.get("stream", False)
+
+            # Call original Anthropic
+            response = original_create(*args, **kwargs)
+
+            # Handle streaming and non-streaming differently
+            if stream:
+                # For streaming, wrap the iterator
+                return self._wrap_streaming_response(response)
+            else:
+                # For non-streaming, scan the complete response
+                if self.raxe_scan_responses:
+                    self._scan_response(response)
+                return response
+
+        # Replace the method
+        self._anthropic_client.messages.create = wrapped_create
+
+    def _scan_message(self, content: Any):
+        """Scan a user message for threats.
+
+        Args:
+            content: Message content to scan (str or list of content blocks)
+
+        Raises:
+            ThreatDetectedError: If threat detected and blocking enabled
+        """
+        # Extract text from content (may be string or list of blocks)
+        text = self._extract_text_from_content(content)
+
+        if not text:
+            return
+
+        # Use AgentScanner for unified scanning with integration telemetry
+        result = self._scanner.scan_prompt(text)
+
+        # AgentScanner handles blocking if configured
+        # Log for monitoring
+        if result.has_threats:
+            logger.warning(
+                f"Threat detected in user message: {result.severity} "
+                f"(action={result.action_taken})"
+            )
+
+    def _scan_response(self, response: Any):
+        """Scan Claude response for threats.
+
+        Args:
+            response: Anthropic response object
+        """
+        try:
+            # Extract text from response content
+            if hasattr(response, "content"):
+                for content_block in response.content:
+                    if hasattr(content_block, "text"):
+                        text = content_block.text
+                        if text:
+                            # Scan response (AgentScanner handles blocking based on config)
+                            result = self._scanner.scan_response(text)
+                            if result.has_threats:
+                                logger.info(
+                                    f"Threat detected in Claude response: {result.severity}"
+                                )
+        except Exception as e:
+            # Don't fail on response scanning
+            logger.error(f"Failed to scan response: {e}")
+
+    def _wrap_streaming_response(self, stream: Iterator) -> Iterator:
+        """Wrap streaming response to scan chunks.
+
+        Args:
+            stream: Original streaming response iterator
+
+        Yields:
+            Response chunks with scanning
+        """
+        accumulated_text = ""
+
+        for chunk in stream:
+            # Accumulate text from chunks
+            if hasattr(chunk, "delta") and hasattr(chunk.delta, "text"):
+                accumulated_text += chunk.delta.text
+
+            yield chunk
+
+        # Scan accumulated text after stream completes
+        if self.raxe_scan_responses and accumulated_text:
+            try:
+                result = self._scanner.scan_response(accumulated_text)
+                if result.has_threats:
+                    logger.info(
+                        f"Threat detected in Claude streaming response: {result.severity}"
+                    )
+            except Exception as e:
+                logger.error(f"Failed to scan streaming response: {e}")
+
+    def _extract_text_from_content(self, content: Any) -> str:
+        """Extract text from Anthropic content format.
+
+        Args:
+            content: Content in various formats (str, list, etc.)
+
+        Returns:
+            Extracted text string
+        """
+        # Handle string content
+        if isinstance(content, str):
+            return content
+
+        # Handle list of content blocks
+        if isinstance(content, list):
+            texts = []
+            for block in content:
+                if isinstance(block, dict) and "text" in block:
+                    texts.append(block["text"])
+                elif hasattr(block, "text"):
+                    texts.append(block.text)
+            return " ".join(texts)
+
+        # Handle single content block
+        if isinstance(content, dict) and "text" in content:
+            return content["text"]
+
+        if hasattr(content, "text"):
+            return content.text
+
+        return ""
+
+    def __getattr__(self, name):
+        """Proxy all other attributes to the Anthropic client.
+
+        This makes RaxeAnthropic a true drop-in replacement by forwarding
+        all attributes and methods we don't explicitly override.
+
+        Args:
+            name: Attribute name
+
+        Returns:
+            Attribute from underlying Anthropic client
+        """
+        return getattr(self._anthropic_client, name)
+
+    def __repr__(self) -> str:
+        """String representation of RaxeAnthropic client.
+
+        Returns:
+            Human-readable string
+        """
+        return (
+            f"RaxeAnthropic(block_on_threat={self.raxe_block_on_threat}, "
+            f"scan_responses={self.raxe_scan_responses})"
+        )