raxe 0.4.6__py3-none-any.whl

raxe/async_sdk/wrappers/openai.py

@@ -0,0 +1,238 @@
+"""Async OpenAI client wrapper with automatic RAXE scanning.
+
+Drop-in async replacement for openai.AsyncOpenAI that scans all prompts and responses.
+
+This wrapper allows users to replace:
+    from openai import AsyncOpenAI
+    client = AsyncOpenAI()
+
+With:
+    from raxe.async_sdk import AsyncRaxeOpenAI
+    client = AsyncRaxeOpenAI()  # All calls automatically scanned
+
+The wrapper intercepts chat.completions.create calls, scans user messages
+before sending to OpenAI, and optionally scans responses.
+"""
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from raxe.async_sdk.client import AsyncRaxe
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncRaxeOpenAI:
+    """Drop-in async replacement for openai.AsyncOpenAI with automatic scanning.
+
+    This wrapper wraps the AsyncOpenAI client and intercepts all
+    chat.completions.create calls to scan prompts and responses.
+
+    Usage:
+        # Instead of:
+        from openai import AsyncOpenAI
+        client = AsyncOpenAI(api_key="sk-...")
+
+        # Use:
+        from raxe.async_sdk import AsyncRaxeOpenAI
+        async with AsyncRaxeOpenAI(api_key="sk-...") as client:
+            response = await client.chat.completions.create(
+                model="gpt-4",
+                messages=[{"role": "user", "content": "Hello"}]
+            )
+
+    Attributes:
+        raxe: AsyncRaxe client instance for scanning
+        raxe_block_on_threat: Whether to block requests on threat detection
+        raxe_scan_responses: Whether to scan OpenAI responses
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        raxe: Optional["AsyncRaxe"] = None,
+        raxe_block_on_threat: bool = True,
+        raxe_scan_responses: bool = True,
+        raxe_cache_size: int = 1000,
+        **kwargs: Any
+    ):
+        """Initialize AsyncRaxeOpenAI client.
+
+        Args:
+            *args: Passed to AsyncOpenAI.__init__
+            raxe: Optional AsyncRaxe client (creates default if not provided)
+            raxe_block_on_threat: Block requests on threat detection
+            raxe_scan_responses: Also scan OpenAI responses
+            raxe_cache_size: Cache size for AsyncRaxe (default: 1000)
+            **kwargs: Passed to AsyncOpenAI.__init__
+
+        Example:
+            # With default AsyncRaxe client
+            async with AsyncRaxeOpenAI(api_key="sk-...") as client:
+                response = await client.chat.completions.create(...)
+
+            # With custom AsyncRaxe client
+            from raxe.async_sdk import AsyncRaxe
+            raxe = AsyncRaxe(telemetry=False)
+            client = AsyncRaxeOpenAI(api_key="sk-...", raxe=raxe)
+
+            # Disable blocking (just monitor)
+            client = AsyncRaxeOpenAI(
+                api_key="sk-...",
+                raxe_block_on_threat=False
+            )
+        """
+        # Try to import AsyncOpenAI
+        try:
+            from openai import AsyncOpenAI
+        except ImportError:
+            raise ImportError(
+                "openai package is required for AsyncRaxeOpenAI. "
+                "Install with: pip install openai"
+            ) from None
+
+        # Initialize AsyncOpenAI client
+        self._openai_client = AsyncOpenAI(*args, **kwargs)
+
+        # Create or use provided AsyncRaxe client
+        if raxe is None:
+            from raxe.async_sdk.client import AsyncRaxe
+            raxe = AsyncRaxe(cache_size=raxe_cache_size)
+
+        self.raxe = raxe
+        self.raxe_block_on_threat = raxe_block_on_threat
+        self.raxe_scan_responses = raxe_scan_responses
+
+        # Wrap the chat completions resource
+        self._wrap_chat_completions()
+
+        logger.debug(
+            f"AsyncRaxeOpenAI initialized: block={raxe_block_on_threat}, "
+            f"scan_responses={raxe_scan_responses}"
+        )
+
+    def _wrap_chat_completions(self) -> None:
+        """Wrap chat.completions.create method with scanning."""
+        # Get original create method
+        original_create = self._openai_client.chat.completions.create
+
+        async def wrapped_create(*args: Any, **kwargs: Any) -> Any:
+            """Wrapped chat.completions.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:
+                            await self._scan_message(content)
+                elif hasattr(message, "role") and hasattr(message, "content"):
+                    # Handle object-style messages
+                    if message.role == "user":
+                        await self._scan_message(message.content)
+
+            # Call original AsyncOpenAI
+            response = await original_create(*args, **kwargs)
+
+            # Optionally scan response
+            if self.raxe_scan_responses and hasattr(response, "choices"):
+                await self._scan_response(response)
+
+            return response
+
+        # Replace the method
+        self._openai_client.chat.completions.create = wrapped_create  # type: ignore
+
+    async def _scan_message(self, content: str) -> None:
+        """Scan a user message for threats.
+
+        Args:
+            content: Message content to scan
+
+        Raises:
+            RaxeBlockedError: If threat detected and blocking enabled
+        """
+        # Use AsyncRaxe.scan() - single entry point
+        result = await self.raxe.scan(
+            content,
+            block_on_threat=self.raxe_block_on_threat
+        )
+
+        # If we get here and blocking is enabled but not raised,
+        # it means no threat was detected or scan() didn't raise
+        # Just log for monitoring
+        if result.has_threats:
+            logger.warning(
+                f"Threat detected in user message: {result.severity} "
+                f"(block={self.raxe_block_on_threat})"
+            )
+
+    async def _scan_response(self, response: Any) -> None:
+        """Scan OpenAI response for threats.
+
+        Args:
+            response: OpenAI response object
+        """
+        try:
+            for choice in response.choices:
+                if hasattr(choice, "message") and hasattr(choice.message, "content"):
+                    content = choice.message.content
+                    if content:
+                        # Scan but don't block on responses
+                        # (just monitor for policy violations)
+                        result = await self.raxe.scan(content, block_on_threat=False)
+                        if result.has_threats:
+                            logger.info(
+                                f"Threat detected in OpenAI response: {result.severity}"
+                            )
+        except Exception as e:
+            # Don't fail on response scanning
+            logger.error(f"Failed to scan response: {e}")
+
+    def __getattr__(self, name: str) -> Any:
+        """Proxy all other attributes to the AsyncOpenAI client.
+
+        This makes AsyncRaxeOpenAI a true drop-in replacement by forwarding
+        all attributes and methods we don't explicitly override.
+
+        Args:
+            name: Attribute name
+
+        Returns:
+            Attribute from underlying AsyncOpenAI client
+        """
+        return getattr(self._openai_client, name)
+
+    async def __aenter__(self) -> "AsyncRaxeOpenAI":
+        """Enter async context manager.
+
+        Returns:
+            Self for use in async with statement
+        """
+        await self._openai_client.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit async context manager and cleanup.
+
+        Args:
+            exc_type: Exception type (if any)
+            exc_val: Exception value (if any)
+            exc_tb: Exception traceback (if any)
+        """
+        await self._openai_client.__aexit__(exc_type, exc_val, exc_tb)
+        await self.raxe.close()
+
+    def __repr__(self) -> str:
+        """String representation of AsyncRaxeOpenAI client.
+
+        Returns:
+            Human-readable string
+        """
+        return (
+            f"AsyncRaxeOpenAI(block_on_threat={self.raxe_block_on_threat}, "
+            f"scan_responses={self.raxe_scan_responses})"
+        )

raxe/cli/__init__.py

@@ -0,0 +1,21 @@
+"""
+CLI Layer - Command-Line Interface
+
+Built with Click framework for excellent UX.
+
+Commands (to be implemented):
+    - raxe init: Initialize RAXE configuration
+    - raxe scan: Scan a prompt or file
+    - raxe config: Manage configuration
+    - raxe upgrade: Upgrade to paid tier
+    - raxe completion: Generate shell completions
+
+Design Goals:
+    - <60 seconds from install to first detection
+    - Helpful error messages
+    - Shell completions for all major shells
+    - Progress indicators for long operations
+    - Colored output for better UX
+"""
+
+__all__ = []