praisonai-code 0.0.1__py3-none-any.whl

praisonai_code/cli/features/checkpoints.py

@@ -0,0 +1,345 @@
+"""
+Checkpoints CLI Feature for PraisonAI.
+
+Provides CLI commands for managing file-level checkpoints.
+
+Commands:
+- praisonai checkpoint save <message>     # Save a checkpoint
+- praisonai checkpoint restore <id>       # Restore to a checkpoint
+- praisonai checkpoint list               # List all checkpoints
+- praisonai checkpoint diff [from] [to]   # Show diff between checkpoints
+- praisonai checkpoint delete             # Delete all checkpoints
+"""
+
+import os
+import asyncio
+from typing import Optional, List
+from dataclasses import dataclass
+
+
+class CheckpointsHandler:
+    """
+    Handler for checkpoints CLI commands.
+    
+    Provides functionality to:
+    - Save checkpoints before file modifications
+    - Restore workspace to previous checkpoints
+    - View diffs between checkpoints
+    - List and manage checkpoints
+    """
+    
+    def __init__(self, workspace_dir: Optional[str] = None, verbose: bool = False,
+                 storage_dir: Optional[str] = None):
+        self.workspace_dir = workspace_dir or os.getcwd()
+        self.verbose = verbose
+        self.storage_dir = storage_dir
+        self._service = None
+    
+    @property
+    def feature_name(self) -> str:
+        return "checkpoints"
+    
+    async def _get_service(self):
+        """Lazy load the checkpoint service."""
+        if self._service is None:
+            from praisonaiagents.checkpoints import CheckpointService
+            kwargs = {"workspace_dir": self.workspace_dir}
+            if self.storage_dir:
+                kwargs["storage_dir"] = self.storage_dir
+            self._service = CheckpointService(**kwargs)
+            await self._service.initialize()
+        return self._service
+    
+    async def save(self, message: str, allow_empty: bool = False, quiet: bool = False) -> bool:
+        """
+        Save a checkpoint.
+        
+        Args:
+            message: Checkpoint message
+            allow_empty: Allow checkpoint even if no changes
+            quiet: Suppress success/error console output (for auto-checkpointing
+                in machine-readable run modes)
+            
+        Returns:
+            True if successful
+        """
+        service = await self._get_service()
+        result = await service.save(message, allow_empty=allow_empty)
+        
+        if result.success:
+            if not quiet:
+                self._print_success(f"Checkpoint saved: {result.checkpoint.short_id}")
+                self._print_info(f"Message: {result.checkpoint.message}")
+            return True
+        else:
+            if not quiet:
+                self._print_error(f"Failed to save checkpoint: {result.error}")
+            return False
+    
+    async def restore(self, checkpoint_id: str) -> bool:
+        """
+        Restore to a checkpoint.
+        
+        Args:
+            checkpoint_id: Checkpoint ID to restore
+            
+        Returns:
+            True if successful
+        """
+        service = await self._get_service()
+        result = await service.restore(checkpoint_id)
+        
+        if result.success:
+            self._print_success(f"Restored to checkpoint: {result.checkpoint.short_id}")
+            return True
+        else:
+            self._print_error(f"Failed to restore: {result.error}")
+            return False
+    
+    async def list_checkpoints(self, limit: int = 20) -> List[dict]:
+        """
+        List all checkpoints.
+        
+        Args:
+            limit: Maximum number to show
+            
+        Returns:
+            List of checkpoint dictionaries
+        """
+        service = await self._get_service()
+        checkpoints = await service.list_checkpoints(limit=limit)
+        
+        if not checkpoints:
+            self._print_info("No checkpoints found")
+            return []
+        
+        self._print_checkpoints_table(checkpoints)
+        return [cp.to_dict() for cp in checkpoints]
+    
+    def _print_checkpoints_table(self, checkpoints):
+        """Print checkpoints in table format."""
+        try:
+            from rich.console import Console
+            from rich.table import Table
+            
+            console = Console()
+            table = Table(title="Checkpoints")
+            table.add_column("ID", style="cyan")
+            table.add_column("Message", style="green")
+            table.add_column("Timestamp", style="yellow")
+            
+            for cp in checkpoints:
+                table.add_row(
+                    cp.short_id,
+                    cp.message[:50] + "..." if len(cp.message) > 50 else cp.message,
+                    cp.timestamp.strftime("%Y-%m-%d %H:%M:%S")
+                )
+            
+            console.print(table)
+        except ImportError:
+            print("Checkpoints:")
+            for cp in checkpoints:
+                print(f"  {cp.short_id} - {cp.message} ({cp.timestamp})")
+    
+    async def diff(
+        self,
+        from_id: Optional[str] = None,
+        to_id: Optional[str] = None
+    ) -> dict:
+        """
+        Show diff between checkpoints.
+        
+        Args:
+            from_id: Starting checkpoint (default: previous)
+            to_id: Ending checkpoint (default: current)
+            
+        Returns:
+            Diff dictionary
+        """
+        service = await self._get_service()
+        diff = await service.diff(from_id, to_id)
+        
+        self._print_diff(diff)
+        return diff.to_dict()
+    
+    def _print_diff(self, diff):
+        """Print diff in a readable format."""
+        try:
+            from rich.console import Console
+            from rich.table import Table
+            
+            console = Console()
+            
+            if not diff.files:
+                console.print("[yellow]No changes[/yellow]")
+                return
+            
+            console.print(f"\n[bold]Changes from {diff.from_checkpoint} to {diff.to_checkpoint or 'working directory'}[/bold]\n")
+            
+            table = Table()
+            table.add_column("Status", style="cyan")
+            table.add_column("File", style="white")
+            table.add_column("+", style="green")
+            table.add_column("-", style="red")
+            
+            for f in diff.files:
+                status_color = {
+                    "added": "green",
+                    "modified": "yellow",
+                    "deleted": "red"
+                }.get(f.status, "white")
+                
+                table.add_row(
+                    f"[{status_color}]{f.status}[/{status_color}]",
+                    f.path,
+                    str(f.additions),
+                    str(f.deletions)
+                )
+            
+            console.print(table)
+            console.print(f"\n[green]+{diff.total_additions}[/green] [red]-{diff.total_deletions}[/red]")
+        except ImportError:
+            print(f"Changes from {diff.from_checkpoint} to {diff.to_checkpoint or 'working directory'}:")
+            for f in diff.files:
+                print(f"  {f.status}: {f.path} (+{f.additions}/-{f.deletions})")
+    
+    async def delete_all(self) -> bool:
+        """
+        Delete all checkpoints.
+        
+        Returns:
+            True if successful
+        """
+        service = await self._get_service()
+        result = await service.delete_all()
+        
+        if result:
+            self._print_success("All checkpoints deleted")
+            return True
+        else:
+            self._print_error("Failed to delete checkpoints")
+            return False
+    
+    def _print_success(self, message: str):
+        """Print success message."""
+        try:
+            from rich.console import Console
+            Console().print(f"[green]✓[/green] {message}")
+        except ImportError:
+            print(f"✓ {message}")
+    
+    def _print_error(self, message: str):
+        """Print error message."""
+        try:
+            from rich.console import Console
+            Console().print(f"[red]✗[/red] {message}")
+        except ImportError:
+            print(f"✗ {message}")
+    
+    def _print_info(self, message: str):
+        """Print info message."""
+        try:
+            from rich.console import Console
+            Console().print(f"[dim]{message}[/dim]")
+        except ImportError:
+            print(message)
+
+
+def handle_checkpoint_command(args: List[str], workspace_dir: Optional[str] = None, verbose: bool = False):
+    """
+    Handle checkpoint CLI commands.
+    
+    Usage:
+        praisonai checkpoint save <message>     # Save a checkpoint
+        praisonai checkpoint restore <id>       # Restore to a checkpoint
+        praisonai checkpoint list [--limit N]   # List checkpoints
+        praisonai checkpoint diff [from] [to]   # Show diff
+        praisonai checkpoint delete             # Delete all checkpoints
+    """
+    handler = CheckpointsHandler(workspace_dir=workspace_dir, verbose=verbose)
+    
+    if not args:
+        print("Usage: praisonai checkpoint <command> [options]")
+        print("\nCommands:")
+        print("  save <message>           Save a checkpoint with message")
+        print("  restore <id>             Restore to a checkpoint")
+        print("  list [--limit N]         List checkpoints (default: 20)")
+        print("  diff [from] [to]         Show diff between checkpoints")
+        print("  delete                   Delete all checkpoints")
+        return
+    
+    command = args[0]
+    
+    if command == "save":
+        if len(args) < 2:
+            print("Usage: praisonai checkpoint save <message>")
+            return
+        
+        message = " ".join(args[1:])
+        allow_empty = "--allow-empty" in args
+        if allow_empty:
+            message = message.replace("--allow-empty", "").strip()
+        
+        asyncio.run(handler.save(message, allow_empty=allow_empty))
+    
+    elif command == "restore":
+        if len(args) < 2:
+            print("Usage: praisonai checkpoint restore <checkpoint_id>")
+            return
+        
+        asyncio.run(handler.restore(args[1]))
+    
+    elif command == "list":
+        limit = 20
+        if "--limit" in args:
+            idx = args.index("--limit")
+            if idx + 1 < len(args):
+                try:
+                    limit = int(args[idx + 1])
+                except ValueError:
+                    pass
+        
+        asyncio.run(handler.list_checkpoints(limit=limit))
+    
+    elif command == "diff":
+        from_id = args[1] if len(args) > 1 else None
+        to_id = args[2] if len(args) > 2 else None
+        
+        asyncio.run(handler.diff(from_id, to_id))
+    
+    elif command == "delete":
+        # Confirm deletion
+        try:
+            from rich.console import Console
+            from rich.prompt import Confirm
+            
+            console = Console()
+            if Confirm.ask("[yellow]Delete all checkpoints?[/yellow]"):
+                asyncio.run(handler.delete_all())
+            else:
+                console.print("[dim]Cancelled[/dim]")
+        except ImportError:
+            response = input("Delete all checkpoints? [y/N] ")
+            if response.lower() == 'y':
+                asyncio.run(handler.delete_all())
+            else:
+                print("Cancelled")
+    
+    elif command == "help" or command == "--help":
+        print("Checkpoint CLI Commands:")
+        print("\n  praisonai checkpoint save <message>")
+        print("    Save a checkpoint with the given message")
+        print("    Options: --allow-empty  Allow checkpoint with no changes")
+        print("\n  praisonai checkpoint restore <id>")
+        print("    Restore workspace to a specific checkpoint")
+        print("\n  praisonai checkpoint list [--limit N]")
+        print("    List all checkpoints (default limit: 20)")
+        print("\n  praisonai checkpoint diff [from_id] [to_id]")
+        print("    Show diff between checkpoints")
+        print("    If no IDs given, shows diff from last checkpoint to current")
+        print("\n  praisonai checkpoint delete")
+        print("    Delete all checkpoints for this workspace")
+    
+    else:
+        print(f"Unknown command: {command}")
+        print("Use 'praisonai checkpoint help' for available commands")

praisonai_code/cli/features/cli_profiler.py

@@ -0,0 +1,335 @@
+"""
+CLI Profiler - Unified profiling for terminal-native CLI commands.
+
+Provides --profile and --profile-deep flags for eligible commands:
+- praisonai "prompt"
+- praisonai chat
+- praisonai chat "prompt"
+- praisonai code
+- praisonai code "prompt"
+- praisonai run agents.yaml
+
+NOT supported for:
+- praisonai ui ... (browser-based)
+- praisonai tui (Textual TUI)
+- praisonai call (long-running server)
+"""
+
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import typer
+
+
+# Commands that support profiling
+PROFILING_ELIGIBLE_COMMANDS = frozenset({
+    "chat",
+    "code",
+    "run",
+    # Direct prompt is handled separately in main.py
+})
+
+# Commands that explicitly reject profiling
+PROFILING_INELIGIBLE_COMMANDS = frozenset({
+    "ui",
+    "tui",
+    "call",
+    "realtime",
+    "serve",
+    "schedule",
+    "acp",
+    "mcp",
+    "lsp",
+})
+
+
+@dataclass
+class CLIProfileConfig:
+    """Configuration for CLI profiling."""
+    enabled: bool = False
+    deep: bool = False
+    output_format: str = "text"  # text or json
+    save_path: Optional[str] = None
+    
+    @classmethod
+    def from_flags(cls, profile: bool = False, profile_deep: bool = False) -> 'CLIProfileConfig':
+        """Create config from CLI flags."""
+        return cls(
+            enabled=profile or profile_deep,
+            deep=profile_deep,
+        )
+
+
+@dataclass
+class CLIProfileResult:
+    """Result from CLI profiling."""
+    import_time_ms: float = 0.0
+    agent_init_time_ms: float = 0.0
+    execution_time_ms: float = 0.0
+    total_time_ms: float = 0.0
+    response: str = ""
+    
+    # Deep profiling data (only if deep=True)
+    function_stats: Optional[list] = None
+    call_graph: Optional[dict] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        result = {
+            "import_time_ms": self.import_time_ms,
+            "agent_init_time_ms": self.agent_init_time_ms,
+            "execution_time_ms": self.execution_time_ms,
+            "total_time_ms": self.total_time_ms,
+        }
+        if self.function_stats:
+            result["function_stats"] = self.function_stats
+        if self.call_graph:
+            result["call_graph"] = self.call_graph
+        return result
+
+
+def print_profile_report(result: CLIProfileResult, config: CLIProfileConfig):
+    """Print profiling report to console."""
+    try:
+        from rich.console import Console
+        from rich.table import Table
+        from rich.panel import Panel
+        console = Console()
+    except ImportError:
+        # Fallback to simple output
+        print("\n─── Profiling ───")
+        print(f"Import:      {result.import_time_ms:.1f}ms")
+        print(f"Agent setup: {result.agent_init_time_ms:.1f}ms")
+        print(f"Execution:   {result.execution_time_ms:.1f}ms")
+        print(f"Total:       {result.total_time_ms:.1f}ms")
+        return
+    
+    # Create timing table
+    table = Table(show_header=False, box=None, padding=(0, 1))
+    table.add_column("Phase", style="cyan")
+    table.add_column("Time", style="green", justify="right")
+    
+    table.add_row("Import", f"{result.import_time_ms:.1f}ms")
+    table.add_row("Agent setup", f"{result.agent_init_time_ms:.1f}ms")
+    table.add_row("Execution", f"{result.execution_time_ms:.1f}ms")
+    table.add_row("─" * 12, "─" * 10)
+    table.add_row("Total", f"{result.total_time_ms:.1f}ms")
+    
+    console.print("\n")
+    console.print(Panel(table, title="[bold]Profiling[/bold]", border_style="dim"))
+    
+    # Deep profiling output
+    if config.deep and result.function_stats:
+        console.print("\n[bold yellow]⚠️  Deep profiling enabled - overhead included in measurements[/bold yellow]")
+        
+        func_table = Table(title="Top Functions by Cumulative Time")
+        func_table.add_column("Function", style="cyan")
+        func_table.add_column("Calls", justify="right")
+        func_table.add_column("Cumul (ms)", justify="right", style="green")
+        
+        for func in result.function_stats[:15]:
+            name = func.get("name", "")[:40]
+            calls = func.get("calls", 0)
+            cumtime = func.get("cumulative_time_ms", 0)
+            func_table.add_row(name, str(calls), f"{cumtime:.2f}")
+        
+        console.print(func_table)
+
+
+def reject_profiling_for_command(command_name: str):
+    """Raise error if profiling is attempted on ineligible command."""
+    if command_name in PROFILING_INELIGIBLE_COMMANDS:
+        raise typer.BadParameter(
+            f"Profiling is not supported for '{command_name}' command.\n"
+            f"Profiling is only available for terminal-native execution commands:\n"
+            f"  praisonai \"prompt\" --profile\n"
+            f"  praisonai chat \"prompt\" --profile\n"
+            f"  praisonai code \"prompt\" --profile\n"
+            f"  praisonai run agents.yaml --profile"
+        )
+
+
+class CLIProfiler:
+    """
+    CLI Profiler for wrapping command execution.
+    
+    Usage:
+        profiler = CLIProfiler(config)
+        with profiler.profile():
+            # execution code
+        profiler.print_report()
+    """
+    
+    def __init__(self, config: CLIProfileConfig):
+        self.config = config
+        self._result = CLIProfileResult()
+        self._start_time: float = 0.0
+        self._cprofile = None
+        
+        # Phase timestamps
+        self._import_start: float = 0.0
+        self._import_end: float = 0.0
+        self._init_start: float = 0.0
+        self._init_end: float = 0.0
+        self._exec_start: float = 0.0
+        self._exec_end: float = 0.0
+    
+    def start(self):
+        """Start profiling."""
+        if not self.config.enabled:
+            return
+        
+        self._start_time = time.perf_counter()
+        
+        if self.config.deep:
+            import cProfile
+            self._cprofile = cProfile.Profile()
+            self._cprofile.enable()
+    
+    def stop(self):
+        """Stop profiling and compute results."""
+        if not self.config.enabled:
+            return
+        
+        end_time = time.perf_counter()
+        
+        if self._cprofile:
+            self._cprofile.disable()
+            self._extract_cprofile_stats()
+        
+        self._result.total_time_ms = (end_time - self._start_time) * 1000
+    
+    def mark_import_start(self):
+        """Mark start of import phase."""
+        if self.config.enabled:
+            self._import_start = time.perf_counter()
+    
+    def mark_import_end(self):
+        """Mark end of import phase."""
+        if self.config.enabled:
+            self._import_end = time.perf_counter()
+            self._result.import_time_ms = (self._import_end - self._import_start) * 1000
+    
+    def mark_init_start(self):
+        """Mark start of agent initialization."""
+        if self.config.enabled:
+            self._init_start = time.perf_counter()
+    
+    def mark_init_end(self):
+        """Mark end of agent initialization."""
+        if self.config.enabled:
+            self._init_end = time.perf_counter()
+            self._result.agent_init_time_ms = (self._init_end - self._init_start) * 1000
+    
+    def mark_exec_start(self):
+        """Mark start of execution."""
+        if self.config.enabled:
+            self._exec_start = time.perf_counter()
+    
+    def mark_exec_end(self):
+        """Mark end of execution."""
+        if self.config.enabled:
+            self._exec_end = time.perf_counter()
+            self._result.execution_time_ms = (self._exec_end - self._exec_start) * 1000
+    
+    def set_response(self, response: str):
+        """Set the response for the profile result."""
+        self._result.response = response
+    
+    def _extract_cprofile_stats(self):
+        """Extract statistics from cProfile."""
+        if not self._cprofile:
+            return
+        
+        import pstats
+        
+        stats = pstats.Stats(self._cprofile)
+        stats.sort_stats("cumulative")
+        
+        function_stats = []
+        for (filename, line, name), (cc, nc, tt, ct, callers) in stats.stats.items():
+            if len(function_stats) >= 100:
+                break
+            function_stats.append({
+                "name": name,
+                "file": filename,
+                "line": line,
+                "calls": nc,
+                "total_time_ms": tt * 1000,
+                "cumulative_time_ms": ct * 1000,
+            })
+        
+        function_stats.sort(key=lambda x: x["cumulative_time_ms"], reverse=True)
+        self._result.function_stats = function_stats[:30]
+    
+    def get_result(self) -> CLIProfileResult:
+        """Get the profiling result."""
+        return self._result
+    
+    def print_report(self):
+        """Print the profiling report."""
+        if self.config.enabled:
+            print_profile_report(self._result, self.config)
+
+
+def run_with_profiling(
+    prompt: str,
+    config: CLIProfileConfig,
+    model: Optional[str] = None,
+    verbose: bool = False,
+    tools: Optional[str] = None,
+) -> tuple:
+    """
+    Run a prompt with profiling enabled.
+    
+    Returns:
+        Tuple of (response, CLIProfileResult)
+    """
+    profiler = CLIProfiler(config)
+    profiler.start()
+    
+    # Import phase
+    profiler.mark_import_start()
+    try:
+        from praisonaiagents import Agent
+    except ImportError:
+        raise ImportError("praisonaiagents not installed. Install with: pip install praisonaiagents")
+    profiler.mark_import_end()
+    
+    # Agent initialization phase
+    profiler.mark_init_start()
+    agent_config = {
+        "name": "ProfiledAgent",
+        "role": "Assistant",
+        "goal": "Complete the task",
+        "verbose": verbose,
+    }
+    if model:
+        agent_config["llm"] = model
+    
+    agent = Agent(**agent_config)
+    profiler.mark_init_end()
+    
+    # Execution phase
+    profiler.mark_exec_start()
+    response = agent.start(prompt)
+    profiler.mark_exec_end()
+    
+    profiler.stop()
+    profiler.set_response(str(response) if response else "")
+    
+    return response, profiler.get_result()
+
+
+# Export all
+__all__ = [
+    "CLIProfileConfig",
+    "CLIProfileResult",
+    "CLIProfiler",
+    "PROFILING_ELIGIBLE_COMMANDS",
+    "PROFILING_INELIGIBLE_COMMANDS",
+    "print_profile_report",
+    "reject_profiling_for_command",
+    "run_with_profiling",
+]