raxe 0.4.6__py3-none-any.whl

raxe/application/scan_pipeline_async.py

@@ -0,0 +1,403 @@
+"""Async parallel scan pipeline with concurrent L1/L2 execution.
+
+This module provides async scanning that runs L1 (regex) and L2 (ML) detection
+in parallel, reducing total latency from ~53ms (sequential) to ~50ms (parallel).
+
+Key benefits:
+- L1 and L2 run concurrently (not sequentially)
+- Total latency = max(L1, L2) instead of L1 + L2
+- Smart cancellation: L2 cancelled if CRITICAL detected in L1
+- Graceful degradation: timeouts don't block the scan
+- Better resource utilization: CPU and GPU work in parallel
+
+Performance:
+- Sequential (current): L1 (3ms) + L2 (50ms) = 53ms
+- Parallel (async): max(L1: 3ms, L2: 50ms) = 50ms (5.6% faster)
+- CRITICAL fast path: 3ms (L2 cancelled)
+
+Example:
+    async with AsyncRaxe() as raxe:
+        result = await raxe.scan("Ignore all instructions")
+        print(f"Latency: {result.duration_ms}ms")  # ~50ms
+"""
+
+import asyncio
+import hashlib
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any
+
+from raxe.application.scan_merger import CombinedScanResult, ScanMerger
+from raxe.domain.engine.executor import RuleExecutor, ScanResult
+from raxe.domain.ml.protocol import L2Detector, L2Result
+from raxe.domain.rules.models import Severity
+from raxe.infrastructure.packs.registry import PackRegistry
+from raxe.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass(frozen=True)
+class AsyncScanMetrics:
+    """Performance metrics for async parallel scan."""
+    l1_start_time: float
+    l1_end_time: float
+    l1_duration_ms: float
+    l2_start_time: float | None
+    l2_end_time: float | None
+    l2_duration_ms: float
+    l2_cancelled: bool
+    l2_timeout: bool
+    parallel_speedup: float  # (L1 + L2) / max(L1, L2)
+    total_duration_ms: float
+
+
+@dataclass(frozen=True)
+class AsyncScanPipelineResult:
+    """Result from async parallel scan pipeline."""
+    scan_result: CombinedScanResult
+    duration_ms: float
+    text_hash: str
+    metadata: dict[str, Any]
+    metrics: AsyncScanMetrics | None = None
+
+    @property
+    def has_threats(self) -> bool:
+        """True if any threats detected."""
+        return self.scan_result.has_threats
+
+    @property
+    def l1_detections(self) -> int:
+        """Count of L1 detections."""
+        return len(self.scan_result.l1_detections or [])
+
+    @property
+    def l2_detections(self) -> int:
+        """Count of L2 predictions."""
+        return len(self.scan_result.l2_predictions or [])
+
+
+class AsyncScanPipeline:
+    """Async parallel scan pipeline with concurrent L1/L2 execution.
+
+    This pipeline runs L1 (regex) and L2 (ML) detection concurrently to minimize
+    total latency. Results are merged after both complete.
+
+    Architecture:
+        1. Start L1 and L2 tasks in parallel
+        2. Wait for L1 first (fast path)
+        3. Check if L2 should be cancelled (CRITICAL optimization)
+        4. Wait for L2 or timeout
+        5. Merge results and return
+
+    Performance targets:
+        - fast mode: <5ms (L1 only)
+        - balanced mode: <55ms (parallel L1+L2)
+        - thorough mode: <160ms (parallel L1+L2 with timeout)
+    """
+
+    def __init__(
+        self,
+        pack_registry: PackRegistry,
+        rule_executor: RuleExecutor,
+        l2_detector: L2Detector,
+        scan_merger: ScanMerger,
+        *,
+        enable_l2: bool = True,
+        fail_fast_on_critical: bool = True,
+        min_confidence_for_skip: float = 0.7,
+        l1_timeout_ms: float = 10.0,
+        l2_timeout_ms: float = 150.0,
+    ):
+        """Initialize async scan pipeline.
+
+        Args:
+            pack_registry: Pack registry for loading rules
+            rule_executor: L1 rule execution engine
+            l2_detector: L2 ML detector
+            scan_merger: Result merger
+            enable_l2: Enable L2 analysis (default: True)
+            fail_fast_on_critical: Cancel L2 if CRITICAL detected (optimization)
+            min_confidence_for_skip: Minimum L1 confidence to skip L2 on CRITICAL
+            l1_timeout_ms: L1 timeout in milliseconds (default: 10ms)
+            l2_timeout_ms: L2 timeout in milliseconds (default: 150ms)
+        """
+        self.pack_registry = pack_registry
+        self.rule_executor = rule_executor
+        self.l2_detector = l2_detector
+        self.scan_merger = scan_merger
+        self.enable_l2 = enable_l2
+        self.fail_fast_on_critical = fail_fast_on_critical
+        self.min_confidence_for_skip = min_confidence_for_skip
+        self.l1_timeout_ms = l1_timeout_ms
+        self.l2_timeout_ms = l2_timeout_ms
+
+        logger.info(
+            "AsyncScanPipeline initialized",
+            enable_l2=enable_l2,
+            fail_fast_on_critical=fail_fast_on_critical,
+            l1_timeout_ms=l1_timeout_ms,
+            l2_timeout_ms=l2_timeout_ms,
+        )
+
+    async def scan(
+        self,
+        text: str,
+        *,
+        customer_id: str | None = None,
+        context: dict[str, Any] | None = None,
+        l1_enabled: bool = True,
+        l2_enabled: bool = True,
+        mode: str = "balanced",
+    ) -> AsyncScanPipelineResult:
+        """Execute async parallel scan with L1 and L2 running concurrently.
+
+        Args:
+            text: Text to scan
+            customer_id: Optional customer ID
+            context: Optional context metadata
+            l1_enabled: Enable L1 detection (default: True)
+            l2_enabled: Enable L2 detection (default: True)
+            mode: Performance mode (fast/balanced/thorough)
+
+        Returns:
+            AsyncScanPipelineResult with scan results and metrics
+
+        Raises:
+            ValueError: If text is empty or mode is invalid
+        """
+        if not text:
+            raise ValueError("Text cannot be empty")
+
+        if mode not in ("fast", "balanced", "thorough"):
+            raise ValueError(f"Invalid mode: {mode}")
+
+        # Apply mode-specific settings
+        if mode == "fast":
+            l1_enabled = True
+            l2_enabled = False
+        elif mode == "thorough":
+            l1_enabled = True
+            l2_enabled = True
+
+        start_time = time.perf_counter()
+        scan_timestamp = datetime.now(timezone.utc).isoformat()
+
+        # Load rules
+        rules = self.pack_registry.get_all_rules()
+
+        # Track metrics
+        l1_start = time.perf_counter()
+        l1_end = 0.0
+        l2_start = 0.0
+        l2_end = 0.0
+        l2_cancelled = False
+        l2_timeout = False
+
+        # Create tasks for parallel execution
+        tasks = {}
+
+        if l1_enabled:
+            tasks["l1"] = asyncio.create_task(
+                self._run_l1_async(text, rules),
+                name="L1-detection"
+            )
+
+        if l2_enabled and self.enable_l2 and mode != "fast":
+            l2_start = time.perf_counter()
+            tasks["l2"] = asyncio.create_task(
+                self._run_l2_async(text, context),
+                name="L2-detection"
+            )
+
+        # Wait for L1 first (fast path)
+        l1_result = None
+        if "l1" in tasks:
+            try:
+                l1_result = await asyncio.wait_for(
+                    tasks["l1"],
+                    timeout=self.l1_timeout_ms / 1000
+                )
+                l1_end = time.perf_counter()
+            except asyncio.TimeoutError:
+                logger.warning(f"L1 timeout after {self.l1_timeout_ms}ms")
+                l1_result = self._create_empty_l1_result(text, scan_timestamp)
+                l1_end = time.perf_counter()
+        else:
+            # L1 disabled - create empty result
+            l1_result = self._create_empty_l1_result(text, scan_timestamp)
+
+        # Check if we should cancel L2 (CRITICAL optimization)
+        l2_result = None
+        if "l2" in tasks:
+            should_cancel = self._should_cancel_l2(l1_result)
+
+            if should_cancel:
+                # Cancel L2 task
+                tasks["l2"].cancel()
+                l2_cancelled = True
+                l2_end = time.perf_counter()
+                logger.info(
+                    "l2_cancelled_critical",
+                    reason="high_confidence_critical_detected",
+                    l1_severity="CRITICAL",
+                    text_hash=self._hash_text(text),
+                )
+            else:
+                # Wait for L2 with timeout
+                try:
+                    l2_result = await asyncio.wait_for(
+                        tasks["l2"],
+                        timeout=self.l2_timeout_ms / 1000
+                    )
+                    l2_end = time.perf_counter()
+                except asyncio.TimeoutError:
+                    logger.warning(f"L2 timeout after {self.l2_timeout_ms}ms")
+                    l2_timeout = True
+                    l2_end = time.perf_counter()
+                except asyncio.CancelledError:
+                    l2_cancelled = True
+                    l2_end = time.perf_counter()
+
+        # Calculate timings
+        l1_duration_ms = (l1_end - l1_start) * 1000 if l1_end > 0 else 0.0
+        l2_duration_ms = (l2_end - l2_start) * 1000 if l2_start > 0 and l2_end > 0 else 0.0
+
+        # Calculate parallel speedup
+        sequential_time = l1_duration_ms + l2_duration_ms
+        parallel_time = max(l1_duration_ms, l2_duration_ms)
+        parallel_speedup = sequential_time / parallel_time if parallel_time > 0 else 1.0
+
+        # Merge results
+        metadata: dict[str, Any] = {
+            "customer_id": customer_id,
+            "scan_timestamp": scan_timestamp,
+            "rules_loaded": len(rules),
+            "l2_skipped": self.enable_l2 and l2_result is None,
+            "l1_duration_ms": l1_duration_ms,
+            "l2_duration_ms": l2_duration_ms,
+            "mode": mode,
+            "l1_enabled": l1_enabled,
+            "l2_enabled": l2_enabled,
+            "l2_cancelled": l2_cancelled,
+            "l2_timeout": l2_timeout,
+            "execution_mode": "async_parallel",
+        }
+
+        if context:
+            metadata["context"] = context
+
+        combined_result = self.scan_merger.merge(
+            l1_result=l1_result,
+            l2_result=l2_result,
+            metadata=metadata,
+        )
+
+        # Calculate total duration
+        total_duration_ms = (time.perf_counter() - start_time) * 1000
+
+        # Create metrics
+        metrics = AsyncScanMetrics(
+            l1_start_time=l1_start,
+            l1_end_time=l1_end,
+            l1_duration_ms=l1_duration_ms,
+            l2_start_time=l2_start if l2_start > 0 else None,
+            l2_end_time=l2_end if l2_end > 0 else None,
+            l2_duration_ms=l2_duration_ms,
+            l2_cancelled=l2_cancelled,
+            l2_timeout=l2_timeout,
+            parallel_speedup=parallel_speedup,
+            total_duration_ms=total_duration_ms,
+        )
+
+        # Log async performance
+        logger.info(
+            "async_scan_complete",
+            total_duration_ms=total_duration_ms,
+            l1_duration_ms=l1_duration_ms,
+            l2_duration_ms=l2_duration_ms,
+            parallel_speedup=parallel_speedup,
+            l2_cancelled=l2_cancelled,
+            l2_timeout=l2_timeout,
+            has_threats=combined_result.has_threats,
+        )
+
+        return AsyncScanPipelineResult(
+            scan_result=combined_result,
+            duration_ms=total_duration_ms,
+            text_hash=self._hash_text(text),
+            metadata=metadata,
+            metrics=metrics,
+        )
+
+    async def _run_l1_async(self, text: str, rules: list) -> ScanResult:
+        """Run L1 detection asynchronously in thread pool.
+
+        L1 is CPU-bound (regex), so we run it in a thread pool executor
+        to avoid blocking the event loop.
+        """
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,  # Use default thread pool executor
+            self.rule_executor.execute_rules,
+            text,
+            rules
+        )
+
+    async def _run_l2_async(self, text: str, context: dict[str, Any] | None) -> L2Result:
+        """Run L2 ML detection asynchronously in thread pool.
+
+        L2 embedding generation is CPU-bound, so we run it in a thread pool
+        executor to avoid blocking the event loop.
+
+        Note: L1 results are not needed by the current bundle detector,
+        so we pass None to avoid waiting for L1 to complete.
+        """
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,  # Use default thread pool executor
+            self.l2_detector.analyze,
+            text,
+            None,  # L1 results not needed for bundle detector
+            context
+        )
+
+    def _should_cancel_l2(self, l1_result: ScanResult | None) -> bool:
+        """Check if L2 should be cancelled based on L1 results.
+
+        Cancels L2 if:
+        - L1 detected CRITICAL severity
+        - L1 has high confidence (>= min_confidence_for_skip)
+        - fail_fast_on_critical is enabled
+
+        This optimization saves ~50ms when high-confidence CRITICAL is detected.
+        """
+        if not l1_result or not self.fail_fast_on_critical:
+            return False
+
+        if l1_result.highest_severity == Severity.CRITICAL:
+            # Check confidence of CRITICAL detections
+            max_confidence = max(
+                (d.confidence for d in l1_result.detections
+                 if d.severity == Severity.CRITICAL),
+                default=0.0
+            )
+
+            return max_confidence >= self.min_confidence_for_skip
+
+        return False
+
+    def _create_empty_l1_result(self, text: str, timestamp: str) -> ScanResult:
+        """Create empty L1 result when L1 is disabled or times out."""
+        return ScanResult(
+            detections=[],
+            scanned_at=timestamp,
+            text_length=len(text),
+            rules_checked=0,
+            scan_duration_ms=0.0,
+        )
+
+    def _hash_text(self, text: str) -> str:
+        """Create privacy-preserving hash of text."""
+        return hashlib.sha256(text.encode("utf-8")).hexdigest()