raxe 0.4.6__py3-none-any.whl

raxe/cli/profiler.py

@@ -0,0 +1,262 @@
+"""Performance profiling CLI commands.
+
+Commands for profiling scan performance and identifying bottlenecks.
+"""
+import click
+from rich.table import Table
+from rich.tree import Tree
+
+from raxe.cli.output import console, display_error
+from raxe.sdk.client import Raxe
+from raxe.utils.profiler import ScanProfiler
+
+
+@click.command(name="profile")
+@click.argument("text")
+@click.option(
+    "--l2/--no-l2",
+    default=True,
+    help="Include L2 profiling (default: True)",
+)
+@click.option(
+    "--format",
+    type=click.Choice(["tree", "table", "json"]),
+    default="tree",
+    help="Output format (default: tree)",
+)
+def profile_command(text: str, l2: bool, output_format: str) -> None:
+    """Profile scan performance.
+
+    Provides detailed performance breakdown including:
+    - Total time
+    - L1 time (per rule)
+    - L2 time (if enabled)
+    - Cache hit/miss
+    - Bottleneck identification
+    - Optimization recommendations
+
+    Examples:
+      raxe profile "test prompt"
+      raxe profile "Ignore previous instructions" --no-l2
+      raxe profile "test" --format table
+    """
+    from raxe.cli.branding import print_logo
+
+    # Show compact logo for text output
+    if format in ("tree", "table"):
+        print_logo(console, compact=True)
+        console.print()
+
+    try:
+        raxe = Raxe()
+    except Exception as e:
+        display_error("Failed to initialize RAXE", str(e))
+        raise click.Abort() from e
+
+    try:
+        # Get components using public API
+        components = raxe.get_profiling_components()
+        executor = components['executor']
+        l2_detector = components['l2_detector'] if l2 else None
+        rules = components['rules']
+
+        # Create profiler
+        profiler = ScanProfiler(
+            rule_executor=executor,
+            l2_detector=l2_detector,
+        )
+
+        # Profile scan
+        with console.status("[cyan]Profiling scan..."):
+            profile = profiler.profile_scan(text, rules, include_l2=l2)
+
+        # Display results
+        if output_format == "json":
+            _display_json(profile)
+        elif output_format == "table":
+            _display_table(profile)
+        else:  # tree
+            _display_tree(profile)
+
+    except Exception as e:
+        display_error("Profiling failed", str(e))
+        raise click.Abort() from e
+
+
+def _display_tree(profile) -> None:
+    """Display profile as a tree."""
+    console.print("\n[bold cyan]Scan Profile[/bold cyan]\n")
+
+    # Build tree
+    tree = Tree(f"[bold]Total Time:[/bold] {profile.total_time_ms:.2f}ms")
+
+    # L1 branch
+    l1_node = tree.add(
+        f"[cyan]L1 Detection[/cyan] ({profile.l1_profile.total_time_ms:.2f}ms, "
+        f"{profile.l1_percentage:.1f}%)"
+    )
+
+    # Show top slow rules
+    for rule_profile in profile.l1_profile.slowest_rules[:5]:
+        status = "[green]✓[/green]" if rule_profile.matched else "[dim]✗[/dim]"
+        cache = "[green]HIT[/green]" if rule_profile.cache_hit else "[red]MISS[/red]"
+        l1_node.add(
+            f"{status} Rule {rule_profile.rule_id}: {rule_profile.execution_time_ms:.3f}ms "
+            f"(cache: {cache})"
+        )
+
+    # Cache stats
+    l1_node.add(
+        f"[dim]Cache:[/dim] {profile.l1_profile.cache_hits} hits, "
+        f"{profile.l1_profile.cache_misses} misses "
+        f"({profile.l1_profile.cache_hit_rate * 100:.1f}% hit rate)"
+    )
+
+    # L2 branch (if present)
+    if profile.l2_profile:
+        l2_node = tree.add(
+            f"[magenta]L2 Detection[/magenta] ({profile.l2_profile.total_time_ms:.2f}ms, "
+            f"{profile.l2_percentage:.1f}%)"
+        )
+        l2_node.add("[dim]Model inference: Full scan[/dim]")
+
+    # Overhead
+    tree.add(
+        f"[yellow]Pipeline Overhead[/yellow] ({profile.overhead_ms:.2f}ms, "
+        f"{profile.overhead_percentage:.1f}%)"
+    )
+
+    console.print(tree)
+
+    # Bottlenecks
+    bottlenecks = profile.identify_bottlenecks()
+    if bottlenecks:
+        console.print("\n[bold red]Bottlenecks[/bold red]")
+        for bottleneck in bottlenecks:
+            console.print(f"  [red]⚠[/red]  {bottleneck}")
+
+    # Recommendations
+    recommendations = profile.get_recommendations()
+    if recommendations:
+        console.print("\n[bold green]Recommendations[/bold green]")
+        for rec in recommendations:
+            console.print(f"  [green]💡[/green] {rec}")
+
+    console.print()
+
+
+def _display_table(profile) -> None:
+    """Display profile as a table."""
+    console.print("\n[bold cyan]Scan Performance Profile[/bold cyan]\n")
+
+    # Summary table
+    summary = Table(show_header=True, header_style="bold cyan")
+    summary.add_column("Metric")
+    summary.add_column("Value", justify="right")
+    summary.add_column("Percentage", justify="right")
+
+    summary.add_row(
+        "Total Time",
+        f"{profile.total_time_ms:.2f}ms",
+        "100.0%",
+    )
+    summary.add_row(
+        "L1 Detection",
+        f"{profile.l1_profile.total_time_ms:.2f}ms",
+        f"{profile.l1_percentage:.1f}%",
+    )
+
+    if profile.l2_profile:
+        summary.add_row(
+            "L2 Detection",
+            f"{profile.l2_profile.total_time_ms:.2f}ms",
+            f"{profile.l2_percentage:.1f}%",
+        )
+
+    summary.add_row(
+        "Overhead",
+        f"{profile.overhead_ms:.2f}ms",
+        f"{profile.overhead_percentage:.1f}%",
+    )
+
+    console.print(summary)
+
+    # L1 rules table
+    if profile.l1_profile.rule_profiles:
+        console.print("\n[bold]L1 Rules (Top 10 Slowest)[/bold]\n")
+
+        rules_table = Table(show_header=True, header_style="bold cyan")
+        rules_table.add_column("Rule ID")
+        rules_table.add_column("Time (ms)", justify="right")
+        rules_table.add_column("Matched")
+        rules_table.add_column("Cache")
+
+        for rule_profile in profile.l1_profile.slowest_rules[:10]:
+            rules_table.add_row(
+                rule_profile.rule_id,
+                f"{rule_profile.execution_time_ms:.3f}",
+                "[green]Yes[/green]" if rule_profile.matched else "[dim]No[/dim]",
+                "[green]HIT[/green]" if rule_profile.cache_hit else "[red]MISS[/red]",
+            )
+
+        console.print(rules_table)
+
+    # Bottlenecks and recommendations
+    bottlenecks = profile.identify_bottlenecks()
+    recommendations = profile.get_recommendations()
+
+    if bottlenecks or recommendations:
+        console.print()
+
+    if bottlenecks:
+        console.print("[bold red]Bottlenecks:[/bold red]")
+        for b in bottlenecks:
+            console.print(f"  [red]⚠[/red]  {b}")
+
+    if recommendations:
+        console.print("\n[bold green]Recommendations:[/bold green]")
+        for r in recommendations:
+            console.print(f"  [green]💡[/green] {r}")
+
+    console.print()
+
+
+def _display_json(profile) -> None:
+    """Display profile as JSON."""
+    import json
+
+    data = {
+        "total_time_ms": profile.total_time_ms,
+        "text_length": profile.text_length,
+        "timestamp": profile.timestamp,
+        "layers": {
+            "l1": {
+                "total_time_ms": profile.l1_profile.total_time_ms,
+                "percentage": profile.l1_percentage,
+                "cache_hits": profile.l1_profile.cache_hits,
+                "cache_misses": profile.l1_profile.cache_misses,
+                "cache_hit_rate": profile.l1_profile.cache_hit_rate,
+                "slowest_rules": [
+                    {
+                        "rule_id": r.rule_id,
+                        "time_ms": r.execution_time_ms,
+                        "matched": r.matched,
+                        "cache_hit": r.cache_hit,
+                    }
+                    for r in profile.l1_profile.slowest_rules
+                ],
+            },
+        },
+        "overhead_ms": profile.overhead_ms,
+        "overhead_percentage": profile.overhead_percentage,
+        "bottlenecks": profile.identify_bottlenecks(),
+        "recommendations": profile.get_recommendations(),
+    }
+
+    if profile.l2_profile:
+        data["layers"]["l2"] = {
+            "total_time_ms": profile.l2_profile.total_time_ms,
+            "percentage": profile.l2_percentage,
+        }
+
+    console.print_json(json.dumps(data, indent=2))

raxe/cli/progress.py

@@ -0,0 +1,379 @@
+"""Progress indicators for RAXE initialization.
+
+Provides context-aware progress display for CLI:
+- Interactive terminals: Rich spinners with live updates
+- CI/CD environments: Plain timestamped text
+- Quiet mode: Silent (errors only)
+
+Usage:
+    from raxe.cli.progress import create_progress_indicator
+    from raxe.cli.progress_context import detect_progress_mode
+
+    mode = detect_progress_mode(quiet=False, no_color=False)
+    progress = create_progress_indicator(mode)
+
+    progress.start("Initializing RAXE...")
+    progress.update_component("rules", "complete", 633, {"count": 460})
+    progress.update_component("ml_model", "complete", 2150)
+    progress.complete(2933)
+"""
+
+import sys
+import time
+from abc import ABC, abstractmethod
+from typing import ClassVar, Literal
+
+from rich.console import Console
+from rich.live import Live
+from rich.panel import Panel
+from rich.text import Text
+
+ComponentStatus = Literal["loading", "complete", "error"]
+
+
+class ProgressIndicator(ABC):
+    """Abstract base for progress indicators.
+
+    All progress implementations must inherit from this class
+    and implement the required methods.
+    """
+
+    @abstractmethod
+    def start(self, message: str) -> None:
+        """Start showing progress.
+
+        Args:
+            message: Main progress message (e.g., "Initializing RAXE...")
+        """
+        pass
+
+    @abstractmethod
+    def update_component(
+        self,
+        name: str,
+        status: ComponentStatus,
+        duration_ms: float,
+        metadata: dict | None = None,
+    ) -> None:
+        """Update status of a component.
+
+        Args:
+            name: Component identifier (rules, ml_model, warmup)
+            status: Current status (loading, complete, error)
+            duration_ms: Time taken in milliseconds (0 if loading)
+            metadata: Optional metadata (e.g., {"count": 460} for rules)
+        """
+        pass
+
+    @abstractmethod
+    def complete(self, total_duration_ms: float) -> None:
+        """Mark initialization as complete.
+
+        Args:
+            total_duration_ms: Total initialization time
+        """
+        pass
+
+    @abstractmethod
+    def error(self, component: str, message: str) -> None:
+        """Report component failure.
+
+        Args:
+            component: Component that failed
+            message: Error message
+        """
+        pass
+
+
+class NullProgress(ProgressIndicator):
+    """No-op progress for when no callback provided.
+
+    Used internally when SDK is used without CLI progress.
+    """
+
+    def start(self, message: str) -> None:
+        pass
+
+    def update_component(
+        self,
+        name: str,
+        status: ComponentStatus,
+        duration_ms: float,
+        metadata: dict | None = None,
+    ) -> None:
+        pass
+
+    def complete(self, total_duration_ms: float) -> None:
+        pass
+
+    def error(self, component: str, message: str) -> None:
+        pass
+
+
+class InteractiveProgress(ProgressIndicator):
+    """Rich progress for interactive terminals.
+
+    Features:
+    - Live-updating spinners
+    - Color-coded status
+    - Transient display (clears after completion)
+    """
+
+    # Component display configuration
+    COMPONENT_LABELS: ClassVar[dict] = {
+        "rules": {
+            "loading": "Loading detection rules",
+            "complete": "Loaded {count} rules",
+        },
+        "ml_model": {
+            "loading": "Loading ML model",
+            "complete": "Loaded ML model",
+        },
+        "warmup": {
+            "loading": "Warming up components",
+            "complete": "Components ready",
+        },
+    }
+
+    def __init__(self):
+        self.console = Console()
+        self.components = {}
+        self.live = None
+        self.start_time = None
+        self.main_message = ""
+        self.rules_count = 0  # Track for display
+
+    def start(self, message: str) -> None:
+        """Start showing progress with live updates."""
+        self.start_time = time.time()
+        self.main_message = message
+
+        # Initialize components in order
+        self.components = {
+            "rules": {"status": "loading", "duration_ms": 0},
+            "ml_model": {"status": "loading", "duration_ms": 0},
+            "warmup": {"status": "loading", "duration_ms": 0},
+        }
+
+        # Start live display (transient=True makes it disappear after stop)
+        self.live = Live(
+            self._render(),
+            console=self.console,
+            refresh_per_second=10,  # 10 FPS for smooth spinner
+            transient=True,
+        )
+        self.live.start()
+
+    def update_component(
+        self,
+        name: str,
+        status: ComponentStatus,
+        duration_ms: float,
+        metadata: dict | None = None,
+    ) -> None:
+        """Update component status with live refresh."""
+        if name not in self.components:
+            return
+
+        self.components[name]["status"] = status
+        self.components[name]["duration_ms"] = duration_ms
+
+        # Store metadata for display (e.g., rules count)
+        if metadata:
+            if name == "rules" and "count" in metadata:
+                self.rules_count = metadata["count"]
+
+        if self.live:
+            self.live.update(self._render())
+
+    def complete(self, total_duration_ms: float) -> None:
+        """Show completion and clear after delay."""
+        if not self.live:
+            return
+
+        # Update to completion state
+        final_render = self._render_complete(total_duration_ms)
+        self.live.update(final_render)
+
+        # Wait for user to read (500ms)
+        time.sleep(0.5)
+
+        # Stop live display (transient will auto-clear)
+        self.live.stop()
+
+    def error(self, component: str, message: str) -> None:
+        """Show error panel."""
+        if self.live:
+            self.live.stop()
+
+        error_text = Text()
+        error_text.append("✗ ", style="red bold")
+        error_text.append(f"Initialization failed: {component}\n", style="red")
+        error_text.append(message, style="dim")
+
+        self.console.print(Panel(error_text, border_style="red"))
+
+    def _render(self) -> Panel:
+        """Render current progress state."""
+        content = Text()
+
+        # Header
+        content.append("🔧 ", style="cyan")
+        content.append(self.main_message, style="cyan bold")
+        content.append("\n")
+
+        # Component status lines
+        for name, data in self.components.items():
+            content.append("  ")  # Indent
+
+            if data["status"] == "loading":
+                content.append("⏳ ", style="cyan")
+                label = self._get_label(name, "loading")
+                content.append(label, style="white")
+                content.append("...\n")
+
+            elif data["status"] == "complete":
+                content.append("✓ ", style="green")
+                label = self._get_label(name, "complete")
+                content.append(label, style="green")
+                content.append(
+                    f" ({data['duration_ms']:.0f}ms)", style="dim white"
+                )
+                content.append("\n")
+
+            elif data["status"] == "error":
+                content.append("✗ ", style="red")
+                label = self._get_label(name, "error")
+                content.append(label, style="red")
+                content.append("\n")
+
+        return Panel(content, border_style="cyan", padding=(0, 1))
+
+    def _render_complete(self, total_ms: float) -> Panel:
+        """Render completion state."""
+        content = Text()
+        content.append("✓ ", style="green bold")
+        content.append("Ready to scan", style="green bold")
+        content.append(
+            f" (Total: {total_ms:.0f}ms, one-time)", style="dim white"
+        )
+
+        return Panel(content, border_style="green", padding=(0, 1))
+
+    def _get_label(self, name: str, status: str) -> str:
+        """Get human-readable label for component."""
+        labels = self.COMPONENT_LABELS.get(name, {})
+        label_template = labels.get(status, name)
+
+        # Format with variables (e.g., rules count)
+        if name == "rules" and status == "complete":
+            return label_template.format(count=self.rules_count)
+
+        return label_template
+
+
+class SimpleProgress(ProgressIndicator):
+    """Plain text progress for CI/CD environments.
+
+    Features:
+    - No ANSI codes
+    - Timestamped messages
+    - Minimal output
+    - Log-parser friendly
+    """
+
+    def __init__(self):
+        self.start_time = None
+
+    def start(self, message: str) -> None:
+        """Print start message with timestamp."""
+        self._log(message)
+
+    def update_component(
+        self,
+        name: str,
+        status: ComponentStatus,
+        duration_ms: float,
+        metadata: dict | None = None,
+    ) -> None:
+        """Print component completion (skip loading state)."""
+        if status == "complete":
+            label = self._get_label(name, metadata)
+            self._log(f"{label} ({duration_ms:.0f}ms)")
+
+    def complete(self, total_duration_ms: float) -> None:
+        """Print completion message."""
+        self._log(
+            f"Initialization complete ({total_duration_ms:.0f}ms, one-time)"
+        )
+
+    def error(self, component: str, message: str) -> None:
+        """Print error message."""
+        self._log(f"ERROR: {component} - {message}")
+
+    def _log(self, message: str) -> None:
+        """Print timestamped message to stderr."""
+        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
+        print(f"[{timestamp}] {message}", file=sys.stderr)
+
+    def _get_label(self, name: str, metadata: dict | None = None) -> str:
+        """Get simple label for component."""
+        labels = {
+            "rules": "Loaded detection rules",
+            "ml_model": "Loaded ML model",
+            "warmup": "Components ready",
+        }
+
+        label = labels.get(name, name)
+
+        # Add count for rules
+        if name == "rules" and metadata and "count" in metadata:
+            label = f"Loaded {metadata['count']} rules"
+
+        return label
+
+
+class QuietProgress(ProgressIndicator):
+    """Silent progress for --quiet mode.
+
+    Only errors are shown (can't suppress errors).
+    """
+
+    def start(self, message: str) -> None:
+        pass
+
+    def update_component(
+        self,
+        name: str,
+        status: ComponentStatus,
+        duration_ms: float,
+        metadata: dict | None = None,
+    ) -> None:
+        pass
+
+    def complete(self, total_duration_ms: float) -> None:
+        pass
+
+    def error(self, component: str, message: str) -> None:
+        # Errors must always be shown
+        pass
+
+
+def create_progress_indicator(mode: str) -> ProgressIndicator:
+    """Factory function to create appropriate progress indicator.
+
+    Args:
+        mode: Progress mode (interactive, simple, quiet)
+
+    Returns:
+        Appropriate ProgressIndicator instance
+    """
+    if mode == "interactive":
+        return InteractiveProgress()
+    elif mode == "simple":
+        return SimpleProgress()
+    elif mode == "quiet":
+        return QuietProgress()
+    else:
+        # Safe default for unknown modes
+        return SimpleProgress()