From 397874ccfc31899daaa6524c7b02c2c225efd417 Mon Sep 17 00:00:00 2001 From: Mervin Praison Date: Fri, 19 Dec 2025 12:16:12 +0000 Subject: [PATCH] Release v2.3.45 --- README.md | 8 + docker/Dockerfile.chat | 2 +- docker/Dockerfile.dev | 2 +- docker/Dockerfile.ui | 2 +- examples/python/cli/autonomy_modes_example.py | 38 + examples/python/cli/cost_tracking_example.py | 34 + .../python/cli/git_integration_example.py | 45 + .../python/cli/interactive_tui_example.py | 49 + examples/python/cli/repo_map_example.py | 24 + .../python/cli/sandbox_execution_example.py | 42 + examples/python/cli/slash_commands_example.py | 51 ++ src/praisonai/README.md | 76 +- src/praisonai/praisonai.rb | 4 +- src/praisonai/praisonai/__init__.py | 4 + .../praisonai/cli/features/__init__.py | 33 + .../praisonai/cli/features/autonomy_mode.py | 490 ++++++++++ .../praisonai/cli/features/cost_tracker.py | 514 +++++++++++ .../praisonai/cli/features/git_integration.py | 651 +++++++++++++ .../praisonai/cli/features/interactive_tui.py | 603 ++++++++++++ .../praisonai/cli/features/repo_map.py | 860 ++++++++++++++++++ .../cli/features/sandbox_executor.py | 515 +++++++++++ .../praisonai/cli/features/slash_commands.py | 650 +++++++++++++ src/praisonai/praisonai/deploy.py | 2 +- src/praisonai/praisonai/version.py | 2 +- src/praisonai/test_routing_live.py | 231 ----- .../tests/e2e/test_cli_features_e2e.py | 309 +++++++ .../tests/unit/cli/test_autonomy_mode.py | 590 ++++++++++++ .../tests/unit/cli/test_cost_tracker.py | 620 +++++++++++++ .../tests/unit/cli/test_git_integration.py | 368 ++++++++ .../tests/unit/cli/test_interactive_tui.py | 460 ++++++++++ src/praisonai/tests/unit/cli/test_repo_map.py | 611 +++++++++++++ .../tests/unit/cli/test_sandbox_executor.py | 373 ++++++++ .../tests/unit/cli/test_slash_commands.py | 608 +++++++++++++ 33 files changed, 8629 insertions(+), 242 deletions(-) create mode 100644 examples/python/cli/autonomy_modes_example.py create mode 100644 examples/python/cli/cost_tracking_example.py create mode 100644 examples/python/cli/git_integration_example.py create mode 100644 examples/python/cli/interactive_tui_example.py create mode 100644 examples/python/cli/repo_map_example.py create mode 100644 examples/python/cli/sandbox_execution_example.py create mode 100644 examples/python/cli/slash_commands_example.py create mode 100644 src/praisonai/praisonai/cli/features/autonomy_mode.py create mode 100644 src/praisonai/praisonai/cli/features/cost_tracker.py create mode 100644 src/praisonai/praisonai/cli/features/git_integration.py create mode 100644 src/praisonai/praisonai/cli/features/interactive_tui.py create mode 100644 src/praisonai/praisonai/cli/features/repo_map.py create mode 100644 src/praisonai/praisonai/cli/features/sandbox_executor.py create mode 100644 src/praisonai/praisonai/cli/features/slash_commands.py delete mode 100644 src/praisonai/test_routing_live.py create mode 100644 src/praisonai/tests/e2e/test_cli_features_e2e.py create mode 100644 src/praisonai/tests/unit/cli/test_autonomy_mode.py create mode 100644 src/praisonai/tests/unit/cli/test_cost_tracker.py create mode 100644 src/praisonai/tests/unit/cli/test_git_integration.py create mode 100644 src/praisonai/tests/unit/cli/test_interactive_tui.py create mode 100644 src/praisonai/tests/unit/cli/test_repo_map.py create mode 100644 src/praisonai/tests/unit/cli/test_sandbox_executor.py create mode 100644 src/praisonai/tests/unit/cli/test_slash_commands.py diff --git a/README.md b/README.md index 5d2273863..aaa3f56de 100644 --- a/README.md +++ b/README.md @@ -258,6 +258,14 @@ npm install praisonai | @ @Mentions in Prompts | [Example](#mentions-in-prompts) | [📖](https://docs.praison.ai/cli/mentions) | | 💾 Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/cli/session) | | 📜 History in Context | [Example](#session-management-python) | [📖](https://docs.praison.ai/cli/session) | +| **🖥️ CLI Features** | | | +| ↳ Slash Commands | [Example](examples/python/cli/slash_commands_example.py) | [📖](https://docs.praison.ai/cli/slash-commands) | +| ↳ Autonomy Modes | [Example](examples/python/cli/autonomy_modes_example.py) | [📖](https://docs.praison.ai/cli/autonomy-modes) | +| ↳ Cost Tracking | [Example](examples/python/cli/cost_tracking_example.py) | [📖](https://docs.praison.ai/cli/cost-tracking) | +| ↳ Repository Map | [Example](examples/python/cli/repo_map_example.py) | [📖](https://docs.praison.ai/cli/repo-map) | +| ↳ Interactive TUI | [Example](examples/python/cli/interactive_tui_example.py) | [📖](https://docs.praison.ai/cli/interactive-tui) | +| ↳ Git Integration | [Example](examples/python/cli/git_integration_example.py) | [📖](https://docs.praison.ai/cli/git-integration) | +| ↳ Sandbox Execution | [Example](examples/python/cli/sandbox_execution_example.py) | [📖](https://docs.praison.ai/cli/sandbox-execution) | --- diff --git a/docker/Dockerfile.chat b/docker/Dockerfile.chat index 6a7a206e6..34d39aad3 100644 --- a/docker/Dockerfile.chat +++ b/docker/Dockerfile.chat @@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison # Install Python packages (using latest versions) RUN pip install --no-cache-dir \ praisonai_tools \ - "praisonai>=2.3.44" \ + "praisonai>=2.3.45" \ "praisonai[chat]" \ "embedchain[github,youtube]" diff --git a/docker/Dockerfile.dev b/docker/Dockerfile.dev index 11bffea5e..63a99e056 100644 --- a/docker/Dockerfile.dev +++ b/docker/Dockerfile.dev @@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison # Install Python packages (using latest versions) RUN pip install --no-cache-dir \ praisonai_tools \ - "praisonai>=2.3.44" \ + "praisonai>=2.3.45" \ "praisonai[ui]" \ "praisonai[chat]" \ "praisonai[realtime]" \ diff --git a/docker/Dockerfile.ui b/docker/Dockerfile.ui index 44d334c06..9240ac44e 100644 --- a/docker/Dockerfile.ui +++ b/docker/Dockerfile.ui @@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison # Install Python packages (using latest versions) RUN pip install --no-cache-dir \ praisonai_tools \ - "praisonai>=2.3.44" \ + "praisonai>=2.3.45" \ "praisonai[ui]" \ "praisonai[crewai]" diff --git a/examples/python/cli/autonomy_modes_example.py b/examples/python/cli/autonomy_modes_example.py new file mode 100644 index 000000000..66d3ee21c --- /dev/null +++ b/examples/python/cli/autonomy_modes_example.py @@ -0,0 +1,38 @@ +""" +Autonomy Modes Example for PraisonAI CLI. + +Control AI autonomy: suggest, auto_edit, full_auto. +Docs: https://docs.praison.ai/cli/autonomy-modes +""" + +from praisonai.cli.features import AutonomyModeHandler +from praisonai.cli.features.autonomy_mode import ActionRequest, ActionType + +# Initialize with suggest mode (safest) +handler = AutonomyModeHandler() +manager = handler.initialize(mode="suggest") + +print(f"Current mode: {handler.get_mode()}") +print("Available modes: suggest, auto_edit, full_auto") + +# Create an action request +action = ActionRequest( + action_type=ActionType.FILE_WRITE, + description="Edit src/main.py to add logging", + details={"file": "src/main.py"} +) + +# In suggest mode, this would require approval +# For demo, we'll just show the action +print(f"\nAction: {action.description}") +print(f"Type: {action.action_type.value}") +print("Would require approval in 'suggest' mode") + +# Change to auto_edit mode +handler.set_mode("auto_edit") +print(f"\nChanged to: {handler.get_mode()}") +print("File edits now auto-approved, shell commands still require approval") + +# Get statistics +stats = handler.get_stats() +print(f"\nStats: {stats}") diff --git a/examples/python/cli/cost_tracking_example.py b/examples/python/cli/cost_tracking_example.py new file mode 100644 index 000000000..b0c95c5e3 --- /dev/null +++ b/examples/python/cli/cost_tracking_example.py @@ -0,0 +1,34 @@ +""" +Cost Tracking Example for PraisonAI CLI. + +Real-time token usage and cost monitoring. +Docs: https://docs.praison.ai/cli/cost-tracking +""" + +from praisonai.cli.features import CostTrackerHandler + +# Initialize tracker +handler = CostTrackerHandler() +tracker = handler.initialize(session_id="demo-session") + +# Track some requests +handler.track_request("gpt-4o", input_tokens=1000, output_tokens=500) +handler.track_request("gpt-4o-mini", input_tokens=2000, output_tokens=800) +handler.track_request("claude-3-5-sonnet", input_tokens=1500, output_tokens=600) + +# Get totals +print(f"Total tokens: {handler.get_tokens():,}") +print(f"Total cost: ${handler.get_cost():.4f}") + +# Get detailed summary +summary = handler.get_summary() +print("\n=== Session Summary ===") +print(f"Session ID: {summary['session_id']}") +print(f"Total requests: {summary['total_requests']}") +print(f"Input tokens: {summary['total_input_tokens']:,}") +print(f"Output tokens: {summary['total_output_tokens']:,}") +print(f"Total cost: ${summary['total_cost']:.4f}") +print(f"Avg cost/request: ${summary['avg_cost_per_request']:.4f}") + +# Display formatted (uses Rich if available) +handler.display_summary() diff --git a/examples/python/cli/git_integration_example.py b/examples/python/cli/git_integration_example.py new file mode 100644 index 000000000..304202ce5 --- /dev/null +++ b/examples/python/cli/git_integration_example.py @@ -0,0 +1,45 @@ +""" +Git Integration Example for PraisonAI CLI. + +Auto-commit with AI messages, diff viewing, undo. +Docs: https://docs.praison.ai/cli/git-integration +""" + +from praisonai.cli.features import GitIntegrationHandler + +# Initialize +handler = GitIntegrationHandler(verbose=True) +git = handler.initialize(repo_path=".") + +# Check if it's a git repo +if not git.is_repo: + print("Not a git repository!") + exit(1) + +# Show status +print("=== Git Status ===") +status = handler.show_status() +print(f"Branch: {status.branch}") +print(f"Staged: {len(status.staged_files)}") +print(f"Modified: {len(status.modified_files)}") +print(f"Untracked: {len(status.untracked_files)}") + +# Show diff (if any changes) +if status.has_changes: + print("\n=== Diff ===") + diff = handler.show_diff() + if diff: + print(diff[:300] + "..." if len(diff) > 300 else diff) + +# Show recent commits +print("\n=== Recent Commits ===") +commits = handler.show_log(count=5) +for c in commits: + print(f" {c.short_hash} {c.message[:50]}") + +# To commit with AI message: +# commit = handler.commit() # Auto-generates message +# commit = handler.commit(message="Custom message") + +# To undo last commit: +# handler.undo(soft=True) # Keeps changes staged diff --git a/examples/python/cli/interactive_tui_example.py b/examples/python/cli/interactive_tui_example.py new file mode 100644 index 000000000..9d1fc2c8a --- /dev/null +++ b/examples/python/cli/interactive_tui_example.py @@ -0,0 +1,49 @@ +""" +Interactive TUI Example for PraisonAI CLI. + +Rich terminal interface with completions and history. +Docs: https://docs.praison.ai/cli/interactive-tui +""" + +from praisonai.cli.features import InteractiveTUIHandler +from praisonai.cli.features.interactive_tui import InteractiveConfig + +# Custom configuration +config = InteractiveConfig( + prompt="ai> ", + multiline=True, + enable_completions=True, + show_status_bar=True +) + +# Define callbacks +def on_input(text): + """Handle regular input.""" + return f"You said: {text}" + +def on_command(cmd): + """Handle slash commands.""" + if cmd == "/exit": + return {"type": "exit"} + return {"type": "command", "message": f"Executed: {cmd}"} + +# Initialize +handler = InteractiveTUIHandler() +session = handler.initialize( + config=config, + on_input=on_input, + on_command=on_command +) + +# Add commands for completion +session.add_commands(["help", "exit", "cost", "model", "plan"]) + +# Add symbols from codebase +session.add_symbols(["MyClass", "my_function"]) + +print("Interactive TUI initialized!") +print("Commands available: /help, /exit, /cost, /model, /plan") +print("Run handler.run() to start interactive session") + +# Uncomment to run interactive session: +# handler.run() diff --git a/examples/python/cli/repo_map_example.py b/examples/python/cli/repo_map_example.py new file mode 100644 index 000000000..81c578cac --- /dev/null +++ b/examples/python/cli/repo_map_example.py @@ -0,0 +1,24 @@ +""" +Repository Map Example for PraisonAI CLI. + +Intelligent codebase mapping with symbol extraction. +Docs: https://docs.praison.ai/cli/repo-map +""" + +from praisonai.cli.features import RepoMapHandler + +# Initialize with current directory +handler = RepoMapHandler(verbose=True) +repo_map = handler.initialize(root=".") + +# Get the repository map +print("=== Repository Map ===") +map_str = handler.get_map() +print(map_str[:500] if len(map_str) > 500 else map_str) +print("..." if len(map_str) > 500 else "") + +# Refresh after file changes +handler.refresh() + +# Get symbols for a specific file (if exists) +# symbols = handler.get_context("MyClass") diff --git a/examples/python/cli/sandbox_execution_example.py b/examples/python/cli/sandbox_execution_example.py new file mode 100644 index 000000000..2b278aa9d --- /dev/null +++ b/examples/python/cli/sandbox_execution_example.py @@ -0,0 +1,42 @@ +""" +Sandbox Execution Example for PraisonAI CLI. + +Secure isolated command execution (--sandbox flag). +Docs: https://docs.praison.ai/cli/sandbox-execution +""" + +from praisonai.cli.features import SandboxExecutorHandler +from praisonai.cli.features.sandbox_executor import SandboxMode + +# Initialize with basic sandbox +handler = SandboxExecutorHandler(verbose=True) +sandbox = handler.initialize(mode="basic") + +print(f"Sandbox mode: {handler.get_mode()}") +print(f"Sandbox enabled: {handler.is_enabled}") + +# Execute a safe command +print("\n=== Safe Command ===") +result = handler.execute("echo 'Hello from sandbox!'") +print(f"Success: {result.success}") +print(f"Output: {result.stdout.strip()}") +print(f"Sandboxed: {result.was_sandboxed}") + +# Try a blocked command +print("\n=== Blocked Command ===") +result = handler.execute("rm -rf /") +print(f"Success: {result.success}") +print(f"Violations: {result.policy_violations}") + +# Validate command before execution +print("\n=== Command Validation ===") +violations = handler.validate_command("sudo apt install something") +if violations: + print(f"Command blocked: {violations}") +else: + print("Command allowed") + +# Available modes +print("\n=== Available Modes ===") +for mode in SandboxMode: + print(f" {mode.value}") diff --git a/examples/python/cli/slash_commands_example.py b/examples/python/cli/slash_commands_example.py new file mode 100644 index 000000000..9802d15b6 --- /dev/null +++ b/examples/python/cli/slash_commands_example.py @@ -0,0 +1,51 @@ +""" +Slash Commands Example for PraisonAI CLI. + +Demonstrates interactive slash commands like /help, /cost, /model. +Docs: https://docs.praison.ai/cli/slash-commands +""" + +from praisonai.cli.features import SlashCommandHandler +from praisonai.cli.features.slash_commands import SlashCommand, CommandKind + +# Create handler +handler = SlashCommandHandler() + +# Execute commands +print("=== /help ===") +result = handler.execute("/help") +print(result.get("message", "")) + +print("\n=== /cost ===") +result = handler.execute("/cost") +print(result.get("message", "")) + +print("\n=== /model ===") +result = handler.execute("/model") +print(result.get("message", "")) + +# Check if input is a command +print("\n=== Command Detection ===") +print(f"/help is command: {handler.is_command('/help')}") +print(f"hello is command: {handler.is_command('hello')}") + +# Get completions for auto-complete +print("\n=== Completions for '/he' ===") +completions = handler.get_completions("/he") +print(completions) + +# Register custom command +def my_custom_handler(args, context): + return {"type": "custom", "message": f"Custom command with args: {args}"} + +custom_cmd = SlashCommand( + name="mycommand", + description="My custom command", + handler=my_custom_handler, + kind=CommandKind.ACTION, + aliases=["mc"] +) + +handler.register(custom_cmd) +result = handler.execute("/mycommand arg1 arg2") +print(f"\n=== Custom Command ===\n{result.get('message', '')}") diff --git a/src/praisonai/README.md b/src/praisonai/README.md index 43d9adc06..aaa3f56de 100644 --- a/src/praisonai/README.md +++ b/src/praisonai/README.md @@ -191,8 +191,14 @@ npm install praisonai | 📚 Add Custom Knowledge | [Example](examples/python/concepts/knowledge-agents.py) | [📖](https://docs.praison.ai/features/knowledge) | | 🧠 Memory (Short & Long Term) | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) | | 📄 Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [📖](https://docs.praison.ai/features/chat-with-pdf) | -| 💻 Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [📖](https://docs.praison.ai/features/codeagent) | -| ✏️ AI Code Editing Tools | [Example](#ai-code-editing) | [📖](https://docs.praison.ai/features/code) | +| **💻 Code** | | | +| ↳ Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [📖](https://docs.praison.ai/features/codeagent) | +| ↳ AI Code Editing Tools | [Example](examples/python/code/code_editing_example.py) | [📖](https://docs.praison.ai/code/editing) | +| ↳ External Agents (All) | [Example](examples/python/code/external_agents_example.py) | [📖](https://docs.praison.ai/code/external-agents) | +| ↳ Claude Code CLI | [Example](examples/python/code/claude_code_example.py) | [📖](https://docs.praison.ai/code/claude-code) | +| ↳ Gemini CLI | [Example](examples/python/code/gemini_cli_example.py) | [📖](https://docs.praison.ai/code/gemini-cli) | +| ↳ Codex CLI | [Example](examples/python/code/codex_cli_example.py) | [📖](https://docs.praison.ai/code/codex-cli) | +| ↳ Cursor CLI | [Example](examples/python/code/cursor_cli_example.py) | [📖](https://docs.praison.ai/code/cursor-cli) | | 📚 RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [📖](https://docs.praison.ai/features/rag) | | 🤔 Async & Parallel Processing | [Example](examples/python/general/async_example.py) | [📖](https://docs.praison.ai/features/async) | | 🔢 Math Agents | [Example](examples/python/agents/math-agent.py) | [📖](https://docs.praison.ai/features/mathagent) | @@ -230,7 +236,7 @@ npm install praisonai | 💻 Programming Agent | [Example](examples/python/agents/programming-agent.py) | [📖](https://docs.praison.ai/agents/programming) | | 📝 Markdown Agent | [Example](examples/python/agents/markdown-agent.py) | [📖](https://docs.praison.ai/agents/markdown) | | 📝 Prompt Expander Agent | [Example](#prompt-expansion) | [📖](https://docs.praison.ai/agents/prompt-expander) | -| 🔀 Router Agent | [Example](examples/python/agents/router-agent-cost-optimization.py) | [📖](https://docs.praison.ai/features/routing) | +| 🔀 Model Router | [Example](examples/python/agents/router-agent-cost-optimization.py) | [📖](https://docs.praison.ai/features/model-router) | | ⛓️ Prompt Chaining | [Example](examples/python/general/prompt_chaining.py) | [📖](https://docs.praison.ai/features/promptchaining) | | 🔍 Evaluator Optimiser | [Example](examples/python/general/evaluator-optimiser.py) | [📖](https://docs.praison.ai/features/evaluator-optimiser) | | 👷 Orchestrator Workers | [Example](examples/python/general/orchestrator-workers.py) | [📖](https://docs.praison.ai/features/orchestrator-worker) | @@ -252,6 +258,14 @@ npm install praisonai | @ @Mentions in Prompts | [Example](#mentions-in-prompts) | [📖](https://docs.praison.ai/cli/mentions) | | 💾 Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/cli/session) | | 📜 History in Context | [Example](#session-management-python) | [📖](https://docs.praison.ai/cli/session) | +| **🖥️ CLI Features** | | | +| ↳ Slash Commands | [Example](examples/python/cli/slash_commands_example.py) | [📖](https://docs.praison.ai/cli/slash-commands) | +| ↳ Autonomy Modes | [Example](examples/python/cli/autonomy_modes_example.py) | [📖](https://docs.praison.ai/cli/autonomy-modes) | +| ↳ Cost Tracking | [Example](examples/python/cli/cost_tracking_example.py) | [📖](https://docs.praison.ai/cli/cost-tracking) | +| ↳ Repository Map | [Example](examples/python/cli/repo_map_example.py) | [📖](https://docs.praison.ai/cli/repo-map) | +| ↳ Interactive TUI | [Example](examples/python/cli/interactive_tui_example.py) | [📖](https://docs.praison.ai/cli/interactive-tui) | +| ↳ Git Integration | [Example](examples/python/cli/git_integration_example.py) | [📖](https://docs.praison.ai/cli/git-integration) | +| ↳ Sandbox Execution | [Example](examples/python/cli/sandbox_execution_example.py) | [📖](https://docs.praison.ai/cli/sandbox-execution) | --- @@ -1374,8 +1388,13 @@ praisonai "Complex analysis" --router --router-provider anthropic # Router automatically selects: # - Simple tasks → gpt-4o-mini, claude-3-haiku # - Complex tasks → gpt-4-turbo, claude-3-opus + +# Create workflow with model routing template +praisonai workflow create --template model-routing --output my_workflow.yaml ``` +Custom models can be configured in `agents.yaml`. See [Model Router Docs](https://docs.praison.ai/features/model-router) for details. + ### Flow Display CLI: ```bash # Enable visual workflow tracking @@ -1464,6 +1483,55 @@ export N8N_API_KEY="your-api-key" praisonai agents.yaml --n8n ``` +### External Agents CLI: + +Use external AI coding CLI tools (Claude Code, Gemini CLI, Codex CLI, Cursor CLI) as agent tools: + +```bash +# Use Claude Code for coding tasks +praisonai "Refactor the auth module" --external-agent claude + +# Use Gemini CLI for code analysis +praisonai "Analyze codebase architecture" --external-agent gemini + +# Use OpenAI Codex CLI +praisonai "Fix all bugs in src/" --external-agent codex + +# Use Cursor CLI +praisonai "Add comprehensive tests" --external-agent cursor +``` + +**Python API:** +```python +from praisonai.integrations import ( + ClaudeCodeIntegration, + GeminiCLIIntegration, + CodexCLIIntegration, + CursorCLIIntegration +) + +# Create integration +claude = ClaudeCodeIntegration(workspace="/project") + +# Execute a coding task +result = await claude.execute("Refactor the auth module") + +# Use as agent tool +from praisonai import Agent +tool = claude.as_tool() +agent = Agent(tools=[tool]) +``` + +**Environment Variables:** +```bash +export ANTHROPIC_API_KEY=your-key # Claude Code +export GEMINI_API_KEY=your-key # Gemini CLI +export OPENAI_API_KEY=your-key # Codex CLI +export CURSOR_API_KEY=your-key # Cursor CLI +``` + +See [External Agents Documentation](https://docs.praison.ai/code/external-agents) for more details. + ### @Mentions in Prompts: ```bash # Include file content in prompt @@ -2266,7 +2334,7 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex - 📖 **Wikipedia Agent** - Search and extract Wikipedia information - 💻 **Programming Agent** - Code development and analysis - 📝 **Markdown Agent** - Generate and format Markdown content -- 🔀 **Router Agent** - Dynamic task routing with cost optimisation +- 🔀 **Model Router** - Smart model selection based on task complexity ### MCP Protocol diff --git a/src/praisonai/praisonai.rb b/src/praisonai/praisonai.rb index ddf76687d..15b75cdee 100644 --- a/src/praisonai/praisonai.rb +++ b/src/praisonai/praisonai.rb @@ -3,8 +3,8 @@ class Praisonai < Formula desc "AI tools for various AI applications" homepage "https://github.com/MervinPraison/PraisonAI" - url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.44.tar.gz" - sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.44.tar.gz | shasum -a 256`.split.first + url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.45.tar.gz" + sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.45.tar.gz | shasum -a 256`.split.first license "MIT" depends_on "python@3.11" diff --git a/src/praisonai/praisonai/__init__.py b/src/praisonai/praisonai/__init__.py index 3710a6b4d..8333aade5 100644 --- a/src/praisonai/praisonai/__init__.py +++ b/src/praisonai/praisonai/__init__.py @@ -1,3 +1,7 @@ +# Suppress crewai.cli.config logger BEFORE any imports to prevent INFO log +import logging +logging.getLogger('crewai.cli.config').setLevel(logging.ERROR) + # Disable OpenTelemetry SDK import os os.environ["OTEL_SDK_DISABLED"] = "true" diff --git a/src/praisonai/praisonai/cli/features/__init__.py b/src/praisonai/praisonai/cli/features/__init__.py index c19c14a1e..7e8aba004 100644 --- a/src/praisonai/praisonai/cli/features/__init__.py +++ b/src/praisonai/praisonai/cli/features/__init__.py @@ -22,6 +22,9 @@ Features: - workflow: YAML workflow management - n8n: Export workflows to n8n visual editor - external_agents: External AI CLI tool integrations (Claude, Gemini, Codex, Cursor) +- slash_commands: Interactive slash commands (/help, /cost, /plan, etc.) +- autonomy_mode: Autonomy levels (suggest, auto-edit, full-auto) +- cost_tracker: Real-time cost and token tracking """ # Type hints available for IDE support @@ -45,6 +48,14 @@ __all__ = [ 'WorkflowHandler', 'N8nHandler', 'ExternalAgentsHandler', + # New CLI enhancement features + 'SlashCommandHandler', + 'AutonomyModeHandler', + 'CostTrackerHandler', + 'RepoMapHandler', + 'InteractiveTUIHandler', + 'GitIntegrationHandler', + 'SandboxExecutorHandler', ] def __getattr__(name): @@ -100,4 +111,26 @@ def __getattr__(name): elif name == 'ExternalAgentsHandler': from .external_agents import ExternalAgentsHandler return ExternalAgentsHandler + # New CLI enhancement features + elif name == 'SlashCommandHandler': + from .slash_commands import SlashCommandHandler + return SlashCommandHandler + elif name == 'AutonomyModeHandler': + from .autonomy_mode import AutonomyModeHandler + return AutonomyModeHandler + elif name == 'CostTrackerHandler': + from .cost_tracker import CostTrackerHandler + return CostTrackerHandler + elif name == 'RepoMapHandler': + from .repo_map import RepoMapHandler + return RepoMapHandler + elif name == 'InteractiveTUIHandler': + from .interactive_tui import InteractiveTUIHandler + return InteractiveTUIHandler + elif name == 'GitIntegrationHandler': + from .git_integration import GitIntegrationHandler + return GitIntegrationHandler + elif name == 'SandboxExecutorHandler': + from .sandbox_executor import SandboxExecutorHandler + return SandboxExecutorHandler raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/praisonai/praisonai/cli/features/autonomy_mode.py b/src/praisonai/praisonai/cli/features/autonomy_mode.py new file mode 100644 index 000000000..06a4e8043 --- /dev/null +++ b/src/praisonai/praisonai/cli/features/autonomy_mode.py @@ -0,0 +1,490 @@ +""" +Autonomy Modes for PraisonAI CLI. + +Inspired by Codex CLI's approval modes and Gemini CLI's approval settings. +Provides configurable autonomy levels: suggest, auto-edit, full-auto. + +Architecture: +- AutonomyMode: Enum for different autonomy levels +- AutonomyPolicy: Policy configuration for each mode +- AutonomyManager: Manages mode transitions and approvals +""" + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Optional, Set +from enum import Enum +import logging + +logger = logging.getLogger(__name__) + + +class AutonomyMode(Enum): + """ + Autonomy levels for agent execution. + + Inspired by Codex CLI's approval modes: + - SUGGEST: Read-only, requires approval for all changes + - AUTO_EDIT: Auto-approve file edits, require approval for commands + - FULL_AUTO: Auto-approve everything (YOLO mode) + """ + SUGGEST = "suggest" # Default: ask for approval on everything + AUTO_EDIT = "auto_edit" # Auto-approve file edits only + FULL_AUTO = "full_auto" # Auto-approve everything (dangerous) + + @classmethod + def from_string(cls, value: str) -> "AutonomyMode": + """Parse mode from string.""" + value = value.lower().strip().replace("-", "_") + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"Unknown autonomy mode: {value}") + + +class ActionType(Enum): + """Types of actions that require approval.""" + FILE_READ = "file_read" + FILE_WRITE = "file_write" + FILE_DELETE = "file_delete" + SHELL_COMMAND = "shell_command" + NETWORK_REQUEST = "network_request" + CODE_EXECUTION = "code_execution" + GIT_OPERATION = "git_operation" + TOOL_CALL = "tool_call" + + +@dataclass +class ActionRequest: + """ + Request for an action that may require approval. + + Attributes: + action_type: Type of action + description: Human-readable description + details: Additional details (file path, command, etc.) + risk_level: Estimated risk (low, medium, high) + reversible: Whether the action can be undone + """ + action_type: ActionType + description: str + details: Dict[str, Any] = field(default_factory=dict) + risk_level: str = "medium" + reversible: bool = True + + def __str__(self) -> str: + return f"{self.action_type.value}: {self.description}" + + +@dataclass +class ApprovalResult: + """Result of an approval request.""" + approved: bool + reason: Optional[str] = None + modified_action: Optional[ActionRequest] = None + remember_choice: bool = False + + +@dataclass +class AutonomyPolicy: + """ + Policy configuration for autonomy mode. + + Defines which actions are auto-approved, require approval, + or are blocked entirely. + """ + mode: AutonomyMode + + # Actions that are auto-approved + auto_approve: Set[ActionType] = field(default_factory=set) + + # Actions that always require approval + require_approval: Set[ActionType] = field(default_factory=set) + + # Actions that are blocked entirely + blocked: Set[ActionType] = field(default_factory=set) + + # Trusted paths for file operations + trusted_paths: Set[str] = field(default_factory=set) + + # Trusted commands (patterns) + trusted_commands: Set[str] = field(default_factory=set) + + # Maximum auto-approved actions before requiring confirmation + max_auto_actions: int = 10 + + # Whether to show what would be done in suggest mode + show_preview: bool = True + + @classmethod + def for_mode(cls, mode: AutonomyMode) -> "AutonomyPolicy": + """Create default policy for a given mode.""" + if mode == AutonomyMode.SUGGEST: + return cls( + mode=mode, + auto_approve={ActionType.FILE_READ}, + require_approval={ + ActionType.FILE_WRITE, + ActionType.FILE_DELETE, + ActionType.SHELL_COMMAND, + ActionType.CODE_EXECUTION, + ActionType.GIT_OPERATION, + ActionType.NETWORK_REQUEST, + ActionType.TOOL_CALL, + }, + blocked=set(), + show_preview=True + ) + elif mode == AutonomyMode.AUTO_EDIT: + return cls( + mode=mode, + auto_approve={ + ActionType.FILE_READ, + ActionType.FILE_WRITE, + ActionType.TOOL_CALL, + }, + require_approval={ + ActionType.FILE_DELETE, + ActionType.SHELL_COMMAND, + ActionType.CODE_EXECUTION, + ActionType.GIT_OPERATION, + ActionType.NETWORK_REQUEST, + }, + blocked=set(), + show_preview=True + ) + elif mode == AutonomyMode.FULL_AUTO: + return cls( + mode=mode, + auto_approve={ + ActionType.FILE_READ, + ActionType.FILE_WRITE, + ActionType.FILE_DELETE, + ActionType.SHELL_COMMAND, + ActionType.CODE_EXECUTION, + ActionType.GIT_OPERATION, + ActionType.NETWORK_REQUEST, + ActionType.TOOL_CALL, + }, + require_approval=set(), + blocked=set(), + show_preview=False, + max_auto_actions=100 + ) + else: + raise ValueError(f"Unknown mode: {mode}") + + +class AutonomyManager: + """ + Manages autonomy mode and approval flow. + + Handles: + - Mode transitions + - Approval requests + - Action tracking + - Policy enforcement + """ + + def __init__( + self, + mode: AutonomyMode = AutonomyMode.SUGGEST, + policy: Optional[AutonomyPolicy] = None, + approval_callback: Optional[Callable[[ActionRequest], ApprovalResult]] = None, + verbose: bool = False + ): + self.mode = mode + self.policy = policy or AutonomyPolicy.for_mode(mode) + self.approval_callback = approval_callback or self._default_approval + self.verbose = verbose + + # Tracking + self._action_count = 0 + self._auto_approved_count = 0 + self._denied_count = 0 + self._remembered_approvals: Dict[str, bool] = {} + + def set_mode(self, mode: AutonomyMode) -> None: + """Change the autonomy mode.""" + old_mode = self.mode + self.mode = mode + self.policy = AutonomyPolicy.for_mode(mode) + + if self.verbose: + logger.info(f"Autonomy mode changed: {old_mode.value} -> {mode.value}") + + def request_approval(self, action: ActionRequest) -> ApprovalResult: + """ + Request approval for an action. + + Returns: + ApprovalResult with approval decision + """ + self._action_count += 1 + + # Check if blocked + if action.action_type in self.policy.blocked: + self._denied_count += 1 + return ApprovalResult( + approved=False, + reason=f"Action type '{action.action_type.value}' is blocked by policy" + ) + + # Check remembered approvals + action_key = self._get_action_key(action) + if action_key in self._remembered_approvals: + approved = self._remembered_approvals[action_key] + if approved: + self._auto_approved_count += 1 + else: + self._denied_count += 1 + return ApprovalResult(approved=approved, reason="Remembered choice") + + # Check if auto-approved + if action.action_type in self.policy.auto_approve: + # Check max auto actions + if self._auto_approved_count >= self.policy.max_auto_actions: + # Force approval after max auto actions + return self._request_user_approval(action) + + self._auto_approved_count += 1 + return ApprovalResult(approved=True, reason="Auto-approved by policy") + + # Check if requires approval + if action.action_type in self.policy.require_approval: + return self._request_user_approval(action) + + # Default: require approval for unknown action types + return self._request_user_approval(action) + + def _request_user_approval(self, action: ActionRequest) -> ApprovalResult: + """Request approval from user via callback.""" + result = self.approval_callback(action) + + if result.remember_choice: + action_key = self._get_action_key(action) + self._remembered_approvals[action_key] = result.approved + + if result.approved: + self._auto_approved_count += 1 + else: + self._denied_count += 1 + + return result + + def _get_action_key(self, action: ActionRequest) -> str: + """Generate a key for remembering approval decisions.""" + # Use action type and relevant details + key_parts = [action.action_type.value] + + if "path" in action.details: + key_parts.append(action.details["path"]) + if "command" in action.details: + key_parts.append(action.details["command"]) + + return ":".join(key_parts) + + def _default_approval(self, action: ActionRequest) -> ApprovalResult: + """Default approval callback using console input.""" + from rich.console import Console + from rich.panel import Panel + from rich.prompt import Confirm + + console = Console() + + # Build action description + risk_color = { + "low": "green", + "medium": "yellow", + "high": "red" + }.get(action.risk_level, "white") + + details_str = "\n".join( + f" {k}: {v}" for k, v in action.details.items() + ) if action.details else " (no details)" + + panel_content = f""" +[bold]{action.description}[/bold] + +Type: {action.action_type.value} +Risk: [{risk_color}]{action.risk_level}[/{risk_color}] +Reversible: {"Yes" if action.reversible else "No"} + +Details: +{details_str} +""" + + console.print(Panel( + panel_content, + title="🔒 Approval Required", + border_style="yellow" + )) + + approved = Confirm.ask("Approve this action?", default=False) + remember = False + + if approved: + remember = Confirm.ask("Remember this choice for similar actions?", default=False) + + return ApprovalResult( + approved=approved, + remember_choice=remember + ) + + def get_stats(self) -> Dict[str, int]: + """Get approval statistics.""" + return { + "total_actions": self._action_count, + "auto_approved": self._auto_approved_count, + "denied": self._denied_count, + "remembered": len(self._remembered_approvals) + } + + def reset_stats(self) -> None: + """Reset approval statistics.""" + self._action_count = 0 + self._auto_approved_count = 0 + self._denied_count = 0 + + def clear_remembered(self) -> None: + """Clear remembered approval decisions.""" + self._remembered_approvals.clear() + + +# ============================================================================ +# CLI Integration Handler +# ============================================================================ + +class AutonomyModeHandler: + """ + Handler for integrating autonomy modes with PraisonAI CLI. + """ + + def __init__(self, verbose: bool = False): + self.verbose = verbose + self._manager: Optional[AutonomyManager] = None + + @property + def feature_name(self) -> str: + return "autonomy_mode" + + def initialize( + self, + mode: str = "suggest", + approval_callback: Optional[Callable[[ActionRequest], ApprovalResult]] = None + ) -> AutonomyManager: + """ + Initialize the autonomy manager. + + Args: + mode: Autonomy mode string (suggest, auto_edit, full_auto) + approval_callback: Optional custom approval callback + + Returns: + Configured AutonomyManager + """ + try: + autonomy_mode = AutonomyMode.from_string(mode) + except ValueError: + logger.warning(f"Unknown mode '{mode}', defaulting to 'suggest'") + autonomy_mode = AutonomyMode.SUGGEST + + self._manager = AutonomyManager( + mode=autonomy_mode, + approval_callback=approval_callback, + verbose=self.verbose + ) + + if self.verbose: + from rich import print as rprint + rprint(f"[cyan]Autonomy mode: {autonomy_mode.value}[/cyan]") + + return self._manager + + def get_manager(self) -> Optional[AutonomyManager]: + """Get the current autonomy manager.""" + return self._manager + + def request_approval(self, action: ActionRequest) -> ApprovalResult: + """Request approval for an action.""" + if not self._manager: + self._manager = self.initialize() + + return self._manager.request_approval(action) + + def set_mode(self, mode: str) -> None: + """Change the autonomy mode.""" + if not self._manager: + self.initialize(mode) + else: + try: + autonomy_mode = AutonomyMode.from_string(mode) + self._manager.set_mode(autonomy_mode) + except ValueError as e: + logger.error(f"Failed to set mode: {e}") + + def get_mode(self) -> str: + """Get the current autonomy mode.""" + if self._manager: + return self._manager.mode.value + return AutonomyMode.SUGGEST.value + + +# ============================================================================ +# Convenience Functions +# ============================================================================ + +def create_file_write_action(path: str, content_preview: str = "") -> ActionRequest: + """Create an action request for file write.""" + return ActionRequest( + action_type=ActionType.FILE_WRITE, + description=f"Write to file: {path}", + details={ + "path": path, + "preview": content_preview[:200] if content_preview else "(no preview)" + }, + risk_level="medium", + reversible=True + ) + + +def create_shell_command_action(command: str, cwd: str = ".") -> ActionRequest: + """Create an action request for shell command.""" + # Assess risk based on command + dangerous_patterns = ["rm ", "sudo ", "chmod ", "chown ", "> /", "| sh", "curl |"] + risk = "high" if any(p in command for p in dangerous_patterns) else "medium" + + return ActionRequest( + action_type=ActionType.SHELL_COMMAND, + description=f"Execute command: {command}", + details={ + "command": command, + "cwd": cwd + }, + risk_level=risk, + reversible=False + ) + + +def create_git_action(operation: str, details: Dict[str, Any] = None) -> ActionRequest: + """Create an action request for git operation.""" + return ActionRequest( + action_type=ActionType.GIT_OPERATION, + description=f"Git {operation}", + details=details or {}, + risk_level="low" if operation in ["status", "log", "diff"] else "medium", + reversible=operation not in ["push", "force-push"] + ) + + +def create_tool_call_action(tool_name: str, args: Dict[str, Any] = None) -> ActionRequest: + """Create an action request for tool call.""" + return ActionRequest( + action_type=ActionType.TOOL_CALL, + description=f"Call tool: {tool_name}", + details={ + "tool": tool_name, + "args": args or {} + }, + risk_level="low", + reversible=True + ) diff --git a/src/praisonai/praisonai/cli/features/cost_tracker.py b/src/praisonai/praisonai/cli/features/cost_tracker.py new file mode 100644 index 000000000..3cd020cd6 --- /dev/null +++ b/src/praisonai/praisonai/cli/features/cost_tracker.py @@ -0,0 +1,514 @@ +""" +Cost Tracking System for PraisonAI CLI. + +Inspired by Aider's cost tracking and Gemini CLI's stats command. +Provides real-time cost and token usage tracking. + +Architecture: +- CostTracker: Tracks tokens and costs per session +- ModelPricing: Pricing data for different models +- UsageStats: Aggregated usage statistics +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional +from datetime import datetime +import logging +import json + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Model Pricing Data +# ============================================================================ + +@dataclass +class ModelPricing: + """ + Pricing information for a model. + + Prices are per 1M tokens (as commonly quoted). + """ + model_name: str + input_price_per_1m: float # Price per 1M input tokens + output_price_per_1m: float # Price per 1M output tokens + context_window: int = 128000 + + def calculate_cost(self, input_tokens: int, output_tokens: int) -> float: + """Calculate cost for given token counts.""" + input_cost = (input_tokens / 1_000_000) * self.input_price_per_1m + output_cost = (output_tokens / 1_000_000) * self.output_price_per_1m + return input_cost + output_cost + + +# Default pricing for common models (as of late 2024) +DEFAULT_PRICING: Dict[str, ModelPricing] = { + # OpenAI models + "gpt-4o": ModelPricing("gpt-4o", 2.50, 10.00, 128000), + "gpt-4o-mini": ModelPricing("gpt-4o-mini", 0.15, 0.60, 128000), + "gpt-4-turbo": ModelPricing("gpt-4-turbo", 10.00, 30.00, 128000), + "gpt-4": ModelPricing("gpt-4", 30.00, 60.00, 8192), + "gpt-3.5-turbo": ModelPricing("gpt-3.5-turbo", 0.50, 1.50, 16385), + "o1": ModelPricing("o1", 15.00, 60.00, 200000), + "o1-mini": ModelPricing("o1-mini", 3.00, 12.00, 128000), + "o1-preview": ModelPricing("o1-preview", 15.00, 60.00, 128000), + "o3-mini": ModelPricing("o3-mini", 1.10, 4.40, 200000), + + # Anthropic models + "claude-3-5-sonnet-20241022": ModelPricing("claude-3-5-sonnet", 3.00, 15.00, 200000), + "claude-3-5-sonnet": ModelPricing("claude-3-5-sonnet", 3.00, 15.00, 200000), + "claude-3-opus": ModelPricing("claude-3-opus", 15.00, 75.00, 200000), + "claude-3-sonnet": ModelPricing("claude-3-sonnet", 3.00, 15.00, 200000), + "claude-3-haiku": ModelPricing("claude-3-haiku", 0.25, 1.25, 200000), + + # Google models + "gemini-2.0-flash": ModelPricing("gemini-2.0-flash", 0.10, 0.40, 1000000), + "gemini-1.5-pro": ModelPricing("gemini-1.5-pro", 1.25, 5.00, 2000000), + "gemini-1.5-flash": ModelPricing("gemini-1.5-flash", 0.075, 0.30, 1000000), + + # Default fallback + "default": ModelPricing("default", 1.00, 3.00, 128000), +} + + +def get_pricing(model_name: str) -> ModelPricing: + """Get pricing for a model, with fallback to default.""" + # Try exact match + if model_name in DEFAULT_PRICING: + return DEFAULT_PRICING[model_name] + + # Try partial match + model_lower = model_name.lower() + for key, pricing in DEFAULT_PRICING.items(): + if key in model_lower or model_lower in key: + return pricing + + # Return default + logger.debug(f"No pricing found for model '{model_name}', using default") + return DEFAULT_PRICING["default"] + + +# ============================================================================ +# Usage Statistics +# ============================================================================ + +@dataclass +class TokenUsage: + """Token usage for a single request.""" + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + cached_tokens: int = 0 + + def __post_init__(self): + if self.total_tokens == 0: + self.total_tokens = self.input_tokens + self.output_tokens + + +@dataclass +class RequestStats: + """Statistics for a single request.""" + timestamp: datetime + model: str + usage: TokenUsage + cost: float + duration_ms: float = 0.0 + prompt_preview: str = "" + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for serialization.""" + return { + "timestamp": self.timestamp.isoformat(), + "model": self.model, + "input_tokens": self.usage.input_tokens, + "output_tokens": self.usage.output_tokens, + "total_tokens": self.usage.total_tokens, + "cached_tokens": self.usage.cached_tokens, + "cost": self.cost, + "duration_ms": self.duration_ms, + } + + +@dataclass +class SessionStats: + """Aggregated statistics for a session.""" + session_id: str + start_time: datetime + end_time: Optional[datetime] = None + + total_requests: int = 0 + total_input_tokens: int = 0 + total_output_tokens: int = 0 + total_tokens: int = 0 + total_cached_tokens: int = 0 + total_cost: float = 0.0 + total_duration_ms: float = 0.0 + + models_used: Dict[str, int] = field(default_factory=dict) + + def add_request(self, stats: RequestStats) -> None: + """Add request statistics to session totals.""" + self.total_requests += 1 + self.total_input_tokens += stats.usage.input_tokens + self.total_output_tokens += stats.usage.output_tokens + self.total_tokens += stats.usage.total_tokens + self.total_cached_tokens += stats.usage.cached_tokens + self.total_cost += stats.cost + self.total_duration_ms += stats.duration_ms + + # Track model usage + if stats.model not in self.models_used: + self.models_used[stats.model] = 0 + self.models_used[stats.model] += 1 + + @property + def duration_seconds(self) -> float: + """Get session duration in seconds.""" + end = self.end_time or datetime.now() + return (end - self.start_time).total_seconds() + + @property + def avg_tokens_per_request(self) -> float: + """Get average tokens per request.""" + if self.total_requests == 0: + return 0.0 + return self.total_tokens / self.total_requests + + @property + def avg_cost_per_request(self) -> float: + """Get average cost per request.""" + if self.total_requests == 0: + return 0.0 + return self.total_cost / self.total_requests + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for serialization.""" + return { + "session_id": self.session_id, + "start_time": self.start_time.isoformat(), + "end_time": self.end_time.isoformat() if self.end_time else None, + "duration_seconds": self.duration_seconds, + "total_requests": self.total_requests, + "total_input_tokens": self.total_input_tokens, + "total_output_tokens": self.total_output_tokens, + "total_tokens": self.total_tokens, + "total_cached_tokens": self.total_cached_tokens, + "total_cost": self.total_cost, + "avg_tokens_per_request": self.avg_tokens_per_request, + "avg_cost_per_request": self.avg_cost_per_request, + "models_used": self.models_used, + } + + +# ============================================================================ +# Cost Tracker +# ============================================================================ + +class CostTracker: + """ + Tracks costs and token usage for a session. + + Features: + - Real-time cost tracking + - Per-model statistics + - Session history + - Export to JSON + """ + + def __init__( + self, + session_id: Optional[str] = None, + pricing: Optional[Dict[str, ModelPricing]] = None, + verbose: bool = False + ): + self.session_id = session_id or self._generate_session_id() + self.pricing = pricing or DEFAULT_PRICING + self.verbose = verbose + + self.session_stats = SessionStats( + session_id=self.session_id, + start_time=datetime.now() + ) + self.request_history: List[RequestStats] = [] + + def _generate_session_id(self) -> str: + """Generate a unique session ID.""" + import uuid + return str(uuid.uuid4())[:8] + + def track_request( + self, + model: str, + input_tokens: int, + output_tokens: int, + cached_tokens: int = 0, + duration_ms: float = 0.0, + prompt_preview: str = "" + ) -> RequestStats: + """ + Track a single request. + + Returns: + RequestStats for the tracked request + """ + usage = TokenUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + cached_tokens=cached_tokens + ) + + # Calculate cost + pricing = get_pricing(model) + cost = pricing.calculate_cost(input_tokens, output_tokens) + + # Create request stats + stats = RequestStats( + timestamp=datetime.now(), + model=model, + usage=usage, + cost=cost, + duration_ms=duration_ms, + prompt_preview=prompt_preview[:100] if prompt_preview else "" + ) + + # Update session stats + self.session_stats.add_request(stats) + self.request_history.append(stats) + + if self.verbose: + logger.info( + f"Request tracked: {model} - " + f"{usage.total_tokens} tokens, ${cost:.4f}" + ) + + return stats + + def track_from_response( + self, + model: str, + response: Any, + duration_ms: float = 0.0 + ) -> Optional[RequestStats]: + """ + Track request from LLM response object. + + Handles various response formats from different providers. + """ + try: + # Try to extract usage from response + usage = None + + # OpenAI format + if hasattr(response, 'usage'): + usage = response.usage + input_tokens = getattr(usage, 'prompt_tokens', 0) + output_tokens = getattr(usage, 'completion_tokens', 0) + cached_tokens = getattr(usage, 'cached_tokens', 0) + # Dict format + elif isinstance(response, dict) and 'usage' in response: + usage = response['usage'] + input_tokens = usage.get('prompt_tokens', 0) + output_tokens = usage.get('completion_tokens', 0) + cached_tokens = usage.get('cached_tokens', 0) + else: + logger.debug("Could not extract usage from response") + return None + + return self.track_request( + model=model, + input_tokens=input_tokens, + output_tokens=output_tokens, + cached_tokens=cached_tokens, + duration_ms=duration_ms + ) + except Exception as e: + logger.error(f"Error tracking response: {e}") + return None + + def get_current_stats(self) -> SessionStats: + """Get current session statistics.""" + return self.session_stats + + def get_total_cost(self) -> float: + """Get total cost for the session.""" + return self.session_stats.total_cost + + def get_total_tokens(self) -> int: + """Get total tokens for the session.""" + return self.session_stats.total_tokens + + def get_request_count(self) -> int: + """Get total request count.""" + return self.session_stats.total_requests + + def format_summary(self) -> str: + """Format a summary of the session statistics.""" + stats = self.session_stats + + lines = [ + f"Session: {stats.session_id}", + f"Duration: {stats.duration_seconds:.1f}s", + f"Requests: {stats.total_requests}", + "", + "Tokens:", + f" Input: {stats.total_input_tokens:,}", + f" Output: {stats.total_output_tokens:,}", + f" Total: {stats.total_tokens:,}", + f" Cached: {stats.total_cached_tokens:,}", + "", + f"Cost: ${stats.total_cost:.4f}", + f"Avg per request: ${stats.avg_cost_per_request:.4f}", + ] + + if stats.models_used: + lines.append("") + lines.append("Models used:") + for model, count in stats.models_used.items(): + lines.append(f" {model}: {count} requests") + + return "\n".join(lines) + + def print_summary(self) -> None: + """Print a formatted summary to console.""" + from rich.console import Console + from rich.panel import Panel + from rich.table import Table + + console = Console() + stats = self.session_stats + + # Create summary table + table = Table(show_header=True, header_style="bold cyan") + table.add_column("Metric") + table.add_column("Value", justify="right") + + table.add_row("Session ID", stats.session_id) + table.add_row("Duration", f"{stats.duration_seconds:.1f}s") + table.add_row("Requests", str(stats.total_requests)) + table.add_row("", "") + table.add_row("Input Tokens", f"{stats.total_input_tokens:,}") + table.add_row("Output Tokens", f"{stats.total_output_tokens:,}") + table.add_row("Total Tokens", f"{stats.total_tokens:,}") + table.add_row("Cached Tokens", f"{stats.total_cached_tokens:,}") + table.add_row("", "") + table.add_row("Total Cost", f"${stats.total_cost:.4f}") + table.add_row("Avg Cost/Request", f"${stats.avg_cost_per_request:.4f}") + + console.print(Panel(table, title="💰 Session Statistics", border_style="green")) + + # Print model breakdown if multiple models used + if len(stats.models_used) > 1: + model_table = Table(show_header=True, header_style="bold cyan") + model_table.add_column("Model") + model_table.add_column("Requests", justify="right") + + for model, count in stats.models_used.items(): + model_table.add_row(model, str(count)) + + console.print(Panel(model_table, title="Models Used", border_style="blue")) + + def export_json(self, filepath: Optional[str] = None) -> str: + """ + Export session data to JSON. + + Args: + filepath: Optional file path to write to + + Returns: + JSON string + """ + data = { + "session": self.session_stats.to_dict(), + "requests": [r.to_dict() for r in self.request_history] + } + + json_str = json.dumps(data, indent=2) + + if filepath: + with open(filepath, 'w') as f: + f.write(json_str) + logger.info(f"Exported session data to {filepath}") + + return json_str + + def end_session(self) -> SessionStats: + """Mark the session as ended.""" + self.session_stats.end_time = datetime.now() + return self.session_stats + + +# ============================================================================ +# CLI Integration Handler +# ============================================================================ + +class CostTrackerHandler: + """ + Handler for integrating cost tracking with PraisonAI CLI. + """ + + def __init__(self, verbose: bool = False): + self.verbose = verbose + self._tracker: Optional[CostTracker] = None + + @property + def feature_name(self) -> str: + return "cost_tracker" + + def initialize(self, session_id: Optional[str] = None) -> CostTracker: + """Initialize the cost tracker.""" + self._tracker = CostTracker( + session_id=session_id, + verbose=self.verbose + ) + + if self.verbose: + from rich import print as rprint + rprint(f"[cyan]Cost tracking enabled for session: {self._tracker.session_id}[/cyan]") + + return self._tracker + + def get_tracker(self) -> Optional[CostTracker]: + """Get the current cost tracker.""" + return self._tracker + + def track_request( + self, + model: str, + input_tokens: int, + output_tokens: int, + **kwargs + ) -> Optional[RequestStats]: + """Track a request.""" + if not self._tracker: + self._tracker = self.initialize() + + return self._tracker.track_request( + model=model, + input_tokens=input_tokens, + output_tokens=output_tokens, + **kwargs + ) + + def get_summary(self) -> Dict[str, Any]: + """Get session summary as dict.""" + if not self._tracker: + return {} + return self._tracker.session_stats.to_dict() + + def print_summary(self) -> None: + """Print session summary.""" + if self._tracker: + self._tracker.print_summary() + + def get_cost(self) -> float: + """Get total cost.""" + if self._tracker: + return self._tracker.get_total_cost() + return 0.0 + + def get_tokens(self) -> int: + """Get total tokens.""" + if self._tracker: + return self._tracker.get_total_tokens() + return 0 diff --git a/src/praisonai/praisonai/cli/features/git_integration.py b/src/praisonai/praisonai/cli/features/git_integration.py new file mode 100644 index 000000000..8fc02f8c5 --- /dev/null +++ b/src/praisonai/praisonai/cli/features/git_integration.py @@ -0,0 +1,651 @@ +""" +Git Integration System for PraisonAI CLI. + +Inspired by Aider's Git integration for auto-commits and change tracking. +Provides seamless Git operations with AI-generated commit messages. + +Architecture: +- GitManager: Core Git operations using subprocess +- CommitGenerator: AI-powered commit message generation +- DiffViewer: Rich diff display +""" + +from dataclasses import dataclass, field +from typing import List, Optional +from pathlib import Path +import subprocess +import logging +import re + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Data Classes +# ============================================================================ + +@dataclass +class GitStatus: + """Git repository status.""" + branch: str = "" + is_clean: bool = True + staged_files: List[str] = field(default_factory=list) + modified_files: List[str] = field(default_factory=list) + untracked_files: List[str] = field(default_factory=list) + ahead: int = 0 + behind: int = 0 + + @property + def has_changes(self) -> bool: + """Check if there are any changes.""" + return bool(self.staged_files or self.modified_files or self.untracked_files) + + +@dataclass +class GitCommit: + """Git commit information.""" + hash: str + short_hash: str + message: str + author: str + date: str + files_changed: int = 0 + insertions: int = 0 + deletions: int = 0 + + +@dataclass +class GitDiff: + """Git diff information.""" + file_path: str + additions: int = 0 + deletions: int = 0 + content: str = "" + is_new: bool = False + is_deleted: bool = False + + +# ============================================================================ +# Git Manager +# ============================================================================ + +class GitManager: + """ + Manages Git operations using subprocess. + + Uses subprocess instead of GitPython for lighter dependencies. + """ + + def __init__(self, repo_path: Optional[str] = None, verbose: bool = False): + self.repo_path = Path(repo_path) if repo_path else Path.cwd() + self.verbose = verbose + self._git_available = self._check_git() + + def _check_git(self) -> bool: + """Check if git is available.""" + try: + result = subprocess.run( + ["git", "--version"], + capture_output=True, + text=True, + cwd=self.repo_path + ) + return result.returncode == 0 + except FileNotFoundError: + return False + + def _run_git(self, *args: str, check: bool = True) -> subprocess.CompletedProcess: + """Run a git command.""" + cmd = ["git"] + list(args) + if self.verbose: + logger.debug(f"Running: {' '.join(cmd)}") + + result = subprocess.run( + cmd, + capture_output=True, + text=True, + cwd=self.repo_path + ) + + if check and result.returncode != 0: + logger.debug(f"Git error: {result.stderr}") + + return result + + @property + def is_repo(self) -> bool: + """Check if current directory is a git repository.""" + if not self._git_available: + return False + result = self._run_git("rev-parse", "--git-dir", check=False) + return result.returncode == 0 + + def get_status(self) -> GitStatus: + """Get repository status.""" + status = GitStatus() + + if not self.is_repo: + return status + + # Get branch + result = self._run_git("branch", "--show-current", check=False) + if result.returncode == 0: + status.branch = result.stdout.strip() + + # Get status + result = self._run_git("status", "--porcelain", check=False) + if result.returncode == 0: + for line in result.stdout.strip().split("\n"): + if not line: + continue + + status_code = line[:2] + file_path = line[3:] + + if status_code[0] in "MADRCU": + status.staged_files.append(file_path) + if status_code[1] in "MD": + status.modified_files.append(file_path) + if status_code == "??": + status.untracked_files.append(file_path) + + status.is_clean = not status.has_changes + + # Get ahead/behind + result = self._run_git("rev-list", "--left-right", "--count", "HEAD...@{u}", check=False) + if result.returncode == 0: + parts = result.stdout.strip().split() + if len(parts) == 2: + status.ahead = int(parts[0]) + status.behind = int(parts[1]) + + return status + + def get_diff(self, staged: bool = False, file_path: Optional[str] = None) -> List[GitDiff]: + """Get diff information.""" + diffs = [] + + if not self.is_repo: + return diffs + + args = ["diff"] + if staged: + args.append("--staged") + args.append("--numstat") + + if file_path: + args.append("--") + args.append(file_path) + + result = self._run_git(*args, check=False) + if result.returncode != 0: + return diffs + + for line in result.stdout.strip().split("\n"): + if not line: + continue + + parts = line.split("\t") + if len(parts) >= 3: + additions = int(parts[0]) if parts[0] != "-" else 0 + deletions = int(parts[1]) if parts[1] != "-" else 0 + path = parts[2] + + diffs.append(GitDiff( + file_path=path, + additions=additions, + deletions=deletions + )) + + return diffs + + def get_diff_content(self, staged: bool = False, file_path: Optional[str] = None) -> str: + """Get full diff content.""" + if not self.is_repo: + return "" + + args = ["diff"] + if staged: + args.append("--staged") + + if file_path: + args.append("--") + args.append(file_path) + + result = self._run_git(*args, check=False) + return result.stdout if result.returncode == 0 else "" + + def stage_files(self, files: Optional[List[str]] = None) -> bool: + """Stage files for commit.""" + if not self.is_repo: + return False + + if files: + result = self._run_git("add", *files, check=False) + else: + result = self._run_git("add", "-A", check=False) + + return result.returncode == 0 + + def commit(self, message: str, allow_empty: bool = False) -> Optional[GitCommit]: + """Create a commit.""" + if not self.is_repo: + return None + + args = ["commit", "-m", message] + if allow_empty: + args.append("--allow-empty") + + result = self._run_git(*args, check=False) + + if result.returncode != 0: + logger.debug(f"Commit failed: {result.stderr}") + return None + + # Get commit info + return self.get_last_commit() + + def get_last_commit(self) -> Optional[GitCommit]: + """Get the last commit.""" + if not self.is_repo: + return None + + result = self._run_git( + "log", "-1", + "--format=%H|%h|%s|%an|%ai", + check=False + ) + + if result.returncode != 0: + return None + + parts = result.stdout.strip().split("|") + if len(parts) >= 5: + return GitCommit( + hash=parts[0], + short_hash=parts[1], + message=parts[2], + author=parts[3], + date=parts[4] + ) + + return None + + def get_log(self, count: int = 10) -> List[GitCommit]: + """Get commit log.""" + commits = [] + + if not self.is_repo: + return commits + + result = self._run_git( + "log", f"-{count}", + "--format=%H|%h|%s|%an|%ai", + check=False + ) + + if result.returncode != 0: + return commits + + for line in result.stdout.strip().split("\n"): + if not line: + continue + + parts = line.split("|") + if len(parts) >= 5: + commits.append(GitCommit( + hash=parts[0], + short_hash=parts[1], + message=parts[2], + author=parts[3], + date=parts[4] + )) + + return commits + + def undo_last_commit(self, soft: bool = True) -> bool: + """Undo the last commit.""" + if not self.is_repo: + return False + + reset_type = "--soft" if soft else "--hard" + result = self._run_git("reset", reset_type, "HEAD~1", check=False) + + return result.returncode == 0 + + def create_branch(self, name: str, checkout: bool = True) -> bool: + """Create a new branch.""" + if not self.is_repo: + return False + + if checkout: + result = self._run_git("checkout", "-b", name, check=False) + else: + result = self._run_git("branch", name, check=False) + + return result.returncode == 0 + + def checkout_branch(self, name: str) -> bool: + """Checkout a branch.""" + if not self.is_repo: + return False + + result = self._run_git("checkout", name, check=False) + return result.returncode == 0 + + def get_branches(self) -> List[str]: + """Get list of branches.""" + if not self.is_repo: + return [] + + result = self._run_git("branch", "--list", check=False) + if result.returncode != 0: + return [] + + branches = [] + for line in result.stdout.strip().split("\n"): + if line: + # Remove * and whitespace + branch = line.strip().lstrip("* ") + branches.append(branch) + + return branches + + def stash(self, message: Optional[str] = None) -> bool: + """Stash changes.""" + if not self.is_repo: + return False + + args = ["stash", "push"] + if message: + args.extend(["-m", message]) + + result = self._run_git(*args, check=False) + return result.returncode == 0 + + def stash_pop(self) -> bool: + """Pop stashed changes.""" + if not self.is_repo: + return False + + result = self._run_git("stash", "pop", check=False) + return result.returncode == 0 + + +# ============================================================================ +# Commit Message Generator +# ============================================================================ + +class CommitMessageGenerator: + """ + Generates commit messages using AI or templates. + """ + + def __init__(self, use_ai: bool = True): + self.use_ai = use_ai + + def generate( + self, + diff_content: str, + context: Optional[str] = None, + style: str = "conventional" + ) -> str: + """ + Generate a commit message. + + Args: + diff_content: The diff to describe + context: Optional context about the changes + style: Message style (conventional, simple, detailed) + + Returns: + Generated commit message + """ + if not diff_content: + return "Empty commit" + + if self.use_ai: + return self._generate_with_ai(diff_content, context, style) + else: + return self._generate_from_diff(diff_content, style) + + def _generate_with_ai( + self, + diff_content: str, + context: Optional[str], + style: str + ) -> str: + """Generate commit message using AI.""" + # This would integrate with the agent system + # For now, return a template-based message + return self._generate_from_diff(diff_content, style) + + def _generate_from_diff(self, diff_content: str, style: str) -> str: + """Generate commit message from diff analysis.""" + # Analyze diff + files_changed = set() + additions = 0 + deletions = 0 + + for line in diff_content.split("\n"): + if line.startswith("diff --git"): + match = re.search(r"b/(.+)$", line) + if match: + files_changed.add(match.group(1)) + elif line.startswith("+") and not line.startswith("+++"): + additions += 1 + elif line.startswith("-") and not line.startswith("---"): + deletions += 1 + + # Determine change type + if additions > deletions * 2: + change_type = "feat" + action = "Add" + elif deletions > additions * 2: + change_type = "refactor" + action = "Remove" + else: + change_type = "fix" + action = "Update" + + # Build message + if len(files_changed) == 1: + file_name = list(files_changed)[0] + scope = Path(file_name).stem + message = f"{change_type}({scope}): {action} {file_name}" + elif len(files_changed) <= 3: + files_str = ", ".join(sorted(files_changed)) + message = f"{change_type}: {action} {files_str}" + else: + message = f"{change_type}: {action} {len(files_changed)} files" + + if style == "detailed": + message += f"\n\n+{additions} -{deletions} lines" + + return message + + +# ============================================================================ +# Diff Viewer +# ============================================================================ + +class DiffViewer: + """ + Rich-based diff viewer. + """ + + def __init__(self): + self._console = None + + @property + def console(self): + """Lazy load Rich console.""" + if self._console is None: + try: + from rich.console import Console + self._console = Console() + except ImportError: + self._console = None + return self._console + + def display_diff(self, diff_content: str, title: str = "Changes") -> None: + """Display diff with syntax highlighting.""" + if not self.console: + print(diff_content) + return + + from rich.panel import Panel + from rich.syntax import Syntax + + syntax = Syntax(diff_content, "diff", theme="monokai", line_numbers=True) + self.console.print(Panel(syntax, title=title, border_style="blue")) + + def display_status(self, status: GitStatus) -> None: + """Display git status.""" + if not self.console: + print(f"Branch: {status.branch}") + print(f"Staged: {len(status.staged_files)}") + print(f"Modified: {len(status.modified_files)}") + print(f"Untracked: {len(status.untracked_files)}") + return + + from rich.panel import Panel + from rich.table import Table + + table = Table(show_header=True, header_style="bold cyan") + table.add_column("Category") + table.add_column("Files", justify="right") + + table.add_row("Branch", status.branch or "detached") + table.add_row("Staged", str(len(status.staged_files))) + table.add_row("Modified", str(len(status.modified_files))) + table.add_row("Untracked", str(len(status.untracked_files))) + + if status.ahead or status.behind: + table.add_row("Ahead/Behind", f"+{status.ahead} / -{status.behind}") + + self.console.print(Panel(table, title="📊 Git Status", border_style="green")) + + def display_log(self, commits: List[GitCommit]) -> None: + """Display commit log.""" + if not self.console: + for commit in commits: + print(f"{commit.short_hash} {commit.message}") + return + + from rich.table import Table + + table = Table(show_header=True, header_style="bold cyan") + table.add_column("Hash") + table.add_column("Message") + table.add_column("Author") + table.add_column("Date") + + for commit in commits: + table.add_row( + commit.short_hash, + commit.message[:50] + "..." if len(commit.message) > 50 else commit.message, + commit.author, + commit.date[:10] + ) + + self.console.print(table) + + +# ============================================================================ +# CLI Integration Handler +# ============================================================================ + +class GitIntegrationHandler: + """ + Handler for integrating Git with PraisonAI CLI. + """ + + def __init__(self, verbose: bool = False): + self.verbose = verbose + self._git: Optional[GitManager] = None + self._viewer: Optional[DiffViewer] = None + self._generator: Optional[CommitMessageGenerator] = None + + @property + def feature_name(self) -> str: + return "git_integration" + + def initialize(self, repo_path: Optional[str] = None) -> GitManager: + """Initialize Git manager.""" + self._git = GitManager(repo_path=repo_path, verbose=self.verbose) + self._viewer = DiffViewer() + self._generator = CommitMessageGenerator() + + if self.verbose and self._git.is_repo: + from rich import print as rprint + status = self._git.get_status() + rprint(f"[cyan]Git initialized on branch: {status.branch}[/cyan]") + + return self._git + + def get_git(self) -> Optional[GitManager]: + """Get the Git manager.""" + return self._git + + def show_status(self) -> GitStatus: + """Show git status.""" + if not self._git: + self._git = self.initialize() + + status = self._git.get_status() + if self._viewer: + self._viewer.display_status(status) + + return status + + def show_diff(self, staged: bool = False) -> str: + """Show diff.""" + if not self._git: + self._git = self.initialize() + + diff_content = self._git.get_diff_content(staged=staged) + if self._viewer and diff_content: + self._viewer.display_diff(diff_content) + + return diff_content + + def commit( + self, + message: Optional[str] = None, + auto_stage: bool = True + ) -> Optional[GitCommit]: + """Create a commit.""" + if not self._git: + self._git = self.initialize() + + if auto_stage: + self._git.stage_files() + + if not message: + # Generate message from diff + diff_content = self._git.get_diff_content(staged=True) + if self._generator: + message = self._generator.generate(diff_content) + else: + message = "Update files" + + return self._git.commit(message) + + def undo(self, soft: bool = True) -> bool: + """Undo last commit.""" + if not self._git: + self._git = self.initialize() + + return self._git.undo_last_commit(soft=soft) + + def show_log(self, count: int = 10) -> List[GitCommit]: + """Show commit log.""" + if not self._git: + self._git = self.initialize() + + commits = self._git.get_log(count=count) + if self._viewer: + self._viewer.display_log(commits) + + return commits diff --git a/src/praisonai/praisonai/cli/features/interactive_tui.py b/src/praisonai/praisonai/cli/features/interactive_tui.py new file mode 100644 index 000000000..0efa5fdf5 --- /dev/null +++ b/src/praisonai/praisonai/cli/features/interactive_tui.py @@ -0,0 +1,603 @@ +""" +Interactive TUI System for PraisonAI CLI. + +Inspired by Aider's interactive mode and prompt_toolkit usage. +Provides rich interactive terminal experience with completions and history. + +Architecture: +- InteractiveSession: Main interactive session manager +- CommandCompleter: Auto-completion for commands and files +- HistoryManager: Persistent command history +- StatusDisplay: Rich-based status bar and output +""" + +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional +from pathlib import Path +import logging + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Configuration +# ============================================================================ + +@dataclass +class InteractiveConfig: + """Configuration for interactive session.""" + prompt: str = ">>> " + multiline: bool = True + history_file: Optional[str] = None + max_history: int = 1000 + enable_completions: bool = True + enable_syntax_highlighting: bool = True + vi_mode: bool = False + auto_suggest: bool = True + show_status_bar: bool = True + color_scheme: str = "monokai" + + +# ============================================================================ +# Command Completer +# ============================================================================ + +class CommandCompleter: + """ + Provides auto-completion for commands, files, and symbols. + """ + + def __init__( + self, + commands: Optional[List[str]] = None, + file_patterns: Optional[List[str]] = None + ): + self.commands = commands or [] + self.file_patterns = file_patterns or ["*.py", "*.js", "*.ts"] + self._file_cache: List[str] = [] + self._symbol_cache: List[str] = [] + + def add_commands(self, commands: List[str]) -> None: + """Add commands to completion list.""" + self.commands.extend(commands) + + def add_symbols(self, symbols: List[str]) -> None: + """Add symbols to completion list.""" + self._symbol_cache.extend(symbols) + + def refresh_files(self, root: Optional[Path] = None) -> None: + """Refresh file cache.""" + root = root or Path.cwd() + self._file_cache = [] + + for pattern in self.file_patterns: + for file_path in root.rglob(pattern): + try: + rel_path = str(file_path.relative_to(root)) + if not any(excl in rel_path for excl in ["__pycache__", "node_modules", ".git"]): + self._file_cache.append(rel_path) + except ValueError: + pass + + def get_completions(self, text: str, cursor_pos: int) -> List[str]: + """ + Get completions for the current input. + + Args: + text: Current input text + cursor_pos: Cursor position + + Returns: + List of completion suggestions + """ + # Get word before cursor + before_cursor = text[:cursor_pos] + words = before_cursor.split() + current_word = words[-1] if words else "" + + completions = [] + + # Slash commands + if current_word.startswith("/"): + cmd_part = current_word[1:].lower() + for cmd in self.commands: + if cmd.lower().startswith(cmd_part): + completions.append(f"/{cmd}") + + # @ mentions for files + elif current_word.startswith("@"): + file_part = current_word[1:].lower() + for file_path in self._file_cache: + if file_part in file_path.lower(): + completions.append(f"@{file_path}") + + # General completions + else: + word_lower = current_word.lower() + + # Check symbols + for symbol in self._symbol_cache: + if symbol.lower().startswith(word_lower): + completions.append(symbol) + + # Check files + for file_path in self._file_cache: + if word_lower in file_path.lower(): + completions.append(file_path) + + return completions[:20] # Limit results + + def create_prompt_toolkit_completer(self) -> Any: + """Create a prompt_toolkit completer.""" + try: + from prompt_toolkit.completion import Completer, Completion + + outer_self = self + + class PTCompleter(Completer): + def get_completions(self, document, complete_event): + text = document.text + cursor_pos = document.cursor_position + + for completion in outer_self.get_completions(text, cursor_pos): + # Calculate start position + words = text[:cursor_pos].split() + current_word = words[-1] if words else "" + start_pos = -len(current_word) + + yield Completion(completion, start_position=start_pos) + + return PTCompleter() + except ImportError: + return None + + +# ============================================================================ +# History Manager +# ============================================================================ + +class HistoryManager: + """ + Manages command history with persistence. + """ + + def __init__( + self, + history_file: Optional[str] = None, + max_entries: int = 1000 + ): + self.history_file = Path(history_file) if history_file else None + self.max_entries = max_entries + self._history: List[str] = [] + self._position: int = 0 + + if self.history_file: + self._load_history() + + def _load_history(self) -> None: + """Load history from file.""" + if self.history_file and self.history_file.exists(): + try: + with open(self.history_file, 'r', encoding='utf-8') as f: + self._history = [ + line.strip() for line in f.readlines() + if line.strip() + ][-self.max_entries:] + except Exception as e: + logger.debug(f"Could not load history: {e}") + + def _save_history(self) -> None: + """Save history to file.""" + if self.history_file: + try: + self.history_file.parent.mkdir(parents=True, exist_ok=True) + with open(self.history_file, 'w', encoding='utf-8') as f: + for entry in self._history[-self.max_entries:]: + f.write(entry + "\n") + except Exception as e: + logger.debug(f"Could not save history: {e}") + + def add(self, entry: str) -> None: + """Add an entry to history.""" + entry = entry.strip() + if entry and (not self._history or self._history[-1] != entry): + self._history.append(entry) + self._position = len(self._history) + self._save_history() + + def get_previous(self) -> Optional[str]: + """Get previous history entry.""" + if self._position > 0: + self._position -= 1 + return self._history[self._position] + return None + + def get_next(self) -> Optional[str]: + """Get next history entry.""" + if self._position < len(self._history) - 1: + self._position += 1 + return self._history[self._position] + elif self._position == len(self._history) - 1: + self._position = len(self._history) + return "" + return None + + def search(self, prefix: str) -> List[str]: + """Search history for entries starting with prefix.""" + return [ + entry for entry in self._history + if entry.lower().startswith(prefix.lower()) + ] + + def get_all(self) -> List[str]: + """Get all history entries.""" + return self._history.copy() + + def clear(self) -> None: + """Clear history.""" + self._history.clear() + self._position = 0 + if self.history_file and self.history_file.exists(): + self.history_file.unlink() + + def create_prompt_toolkit_history(self) -> Any: + """Create a prompt_toolkit history object.""" + try: + from prompt_toolkit.history import FileHistory, InMemoryHistory + + if self.history_file: + return FileHistory(str(self.history_file)) + return InMemoryHistory() + except ImportError: + return None + + +# ============================================================================ +# Status Display +# ============================================================================ + +class StatusDisplay: + """ + Rich-based status display for interactive session. + """ + + def __init__(self, show_status_bar: bool = True): + self.show_status_bar = show_status_bar + self._console = None + self._status_items: Dict[str, str] = {} + + @property + def console(self): + """Lazy load Rich console.""" + if self._console is None: + try: + from rich.console import Console + self._console = Console() + except ImportError: + self._console = None + return self._console + + def set_status(self, key: str, value: str) -> None: + """Set a status item.""" + self._status_items[key] = value + + def clear_status(self, key: str) -> None: + """Clear a status item.""" + self._status_items.pop(key, None) + + def print_status_bar(self) -> None: + """Print the status bar.""" + if not self.show_status_bar or not self.console: + return + + from rich.columns import Columns + from rich.text import Text + + items = [] + for key, value in self._status_items.items(): + items.append(Text(f"{key}: {value}", style="dim")) + + if items: + self.console.print(Columns(items, equal=True, expand=True)) + + def print_welcome(self, version: str = "1.0.0") -> None: + """Print welcome message.""" + if not self.console: + print(f"PraisonAI CLI v{version}") + return + + from rich.panel import Panel + from rich.text import Text + + welcome = Text() + welcome.append("PraisonAI CLI", style="bold cyan") + welcome.append(f" v{version}\n", style="dim") + welcome.append("Type ", style="dim") + welcome.append("/help", style="bold green") + welcome.append(" for commands, ", style="dim") + welcome.append("/exit", style="bold yellow") + welcome.append(" to quit", style="dim") + + self.console.print(Panel(welcome, border_style="blue")) + + def print_response(self, response: str, title: str = "Response") -> None: + """Print a response with formatting.""" + if not self.console: + print(response) + return + + from rich.panel import Panel + from rich.markdown import Markdown + + try: + md = Markdown(response) + self.console.print(Panel(md, title=title, border_style="green")) + except Exception: + self.console.print(Panel(response, title=title, border_style="green")) + + def print_error(self, error: str) -> None: + """Print an error message.""" + if not self.console: + print(f"Error: {error}") + return + + from rich.panel import Panel + self.console.print(Panel(error, title="Error", border_style="red")) + + def print_info(self, message: str) -> None: + """Print an info message.""" + if not self.console: + print(message) + return + + self.console.print(f"[cyan]ℹ[/cyan] {message}") + + def print_success(self, message: str) -> None: + """Print a success message.""" + if not self.console: + print(message) + return + + self.console.print(f"[green]✓[/green] {message}") + + +# ============================================================================ +# Interactive Session +# ============================================================================ + +class InteractiveSession: + """ + Main interactive session manager. + + Provides a rich interactive terminal experience with: + - Command completion + - History + - Syntax highlighting + - Status bar + - Keyboard shortcuts + """ + + def __init__( + self, + config: Optional[InteractiveConfig] = None, + on_input: Optional[Callable[[str], Optional[str]]] = None, + on_command: Optional[Callable[[str], Optional[Dict[str, Any]]]] = None + ): + self.config = config or InteractiveConfig() + self.on_input = on_input + self.on_command = on_command + + self.completer = CommandCompleter() + self.history = HistoryManager( + history_file=self.config.history_file, + max_entries=self.config.max_history + ) + self.display = StatusDisplay(show_status_bar=self.config.show_status_bar) + + self._running = False + self._prompt_session = None + + def _create_prompt_session(self) -> Any: + """Create a prompt_toolkit session.""" + try: + from prompt_toolkit import PromptSession + from prompt_toolkit.lexers import PygmentsLexer + from prompt_toolkit.auto_suggest import AutoSuggestFromHistory + from pygments.lexers import MarkdownLexer + + session_kwargs = { + "message": self.config.prompt, + "multiline": self.config.multiline, + } + + # Add completer + if self.config.enable_completions: + completer = self.completer.create_prompt_toolkit_completer() + if completer: + session_kwargs["completer"] = completer + + # Add history + history = self.history.create_prompt_toolkit_history() + if history: + session_kwargs["history"] = history + + # Add syntax highlighting + if self.config.enable_syntax_highlighting: + session_kwargs["lexer"] = PygmentsLexer(MarkdownLexer) + + # Add auto-suggest + if self.config.auto_suggest: + session_kwargs["auto_suggest"] = AutoSuggestFromHistory() + + # VI mode + if self.config.vi_mode: + from prompt_toolkit.enums import EditingMode + session_kwargs["editing_mode"] = EditingMode.VI + + return PromptSession(**session_kwargs) + + except ImportError as e: + logger.debug(f"prompt_toolkit not available: {e}") + return None + + def _get_input_fallback(self) -> str: + """Fallback input method without prompt_toolkit.""" + try: + if self.config.multiline: + print(f"{self.config.prompt}(Enter empty line to submit)") + lines = [] + while True: + line = input() + if not line: + break + lines.append(line) + return "\n".join(lines) + else: + return input(self.config.prompt) + except EOFError: + return "/exit" + except KeyboardInterrupt: + print() + return "" + + def get_input(self) -> str: + """Get input from user.""" + if self._prompt_session: + try: + return self._prompt_session.prompt() + except KeyboardInterrupt: + return "" + except EOFError: + return "/exit" + else: + return self._get_input_fallback() + + def process_input(self, user_input: str) -> Optional[str]: + """ + Process user input. + + Args: + user_input: Raw user input + + Returns: + Response string or None + """ + user_input = user_input.strip() + + if not user_input: + return None + + # Add to history + self.history.add(user_input) + + # Check for slash commands + if user_input.startswith("/"): + if self.on_command: + result = self.on_command(user_input) + if result: + if result.get("type") == "exit": + self._running = False + return None + return result.get("message") + return None + + # Regular input + if self.on_input: + return self.on_input(user_input) + + return None + + def run(self) -> None: + """Run the interactive session.""" + self._running = True + self._prompt_session = self._create_prompt_session() + + # Print welcome + self.display.print_welcome() + + while self._running: + try: + # Get input + user_input = self.get_input() + + # Process + response = self.process_input(user_input) + + # Display response + if response: + self.display.print_response(response) + + except KeyboardInterrupt: + self.display.print_info("Use /exit to quit") + except Exception as e: + self.display.print_error(str(e)) + logger.exception("Error in interactive session") + + def stop(self) -> None: + """Stop the interactive session.""" + self._running = False + + def add_commands(self, commands: List[str]) -> None: + """Add commands for completion.""" + self.completer.add_commands(commands) + + def add_symbols(self, symbols: List[str]) -> None: + """Add symbols for completion.""" + self.completer.add_symbols(symbols) + + def refresh_files(self, root: Optional[Path] = None) -> None: + """Refresh file completions.""" + self.completer.refresh_files(root) + + +# ============================================================================ +# CLI Integration Handler +# ============================================================================ + +class InteractiveTUIHandler: + """ + Handler for integrating Interactive TUI with PraisonAI CLI. + """ + + def __init__(self, verbose: bool = False): + self.verbose = verbose + self._session: Optional[InteractiveSession] = None + + @property + def feature_name(self) -> str: + return "interactive_tui" + + def initialize( + self, + config: Optional[InteractiveConfig] = None, + on_input: Optional[Callable[[str], Optional[str]]] = None, + on_command: Optional[Callable[[str], Optional[Dict[str, Any]]]] = None + ) -> InteractiveSession: + """Initialize the interactive session.""" + self._session = InteractiveSession( + config=config, + on_input=on_input, + on_command=on_command + ) + + if self.verbose: + from rich import print as rprint + rprint("[cyan]Interactive TUI initialized[/cyan]") + + return self._session + + def get_session(self) -> Optional[InteractiveSession]: + """Get the current session.""" + return self._session + + def run(self) -> None: + """Run the interactive session.""" + if not self._session: + self._session = self.initialize() + + self._session.run() + + def stop(self) -> None: + """Stop the interactive session.""" + if self._session: + self._session.stop() diff --git a/src/praisonai/praisonai/cli/features/repo_map.py b/src/praisonai/praisonai/cli/features/repo_map.py new file mode 100644 index 000000000..3351a31d6 --- /dev/null +++ b/src/praisonai/praisonai/cli/features/repo_map.py @@ -0,0 +1,860 @@ +""" +Repository Map System for PraisonAI CLI. + +Inspired by Aider's RepoMap using tree-sitter for code parsing. +Provides intelligent codebase context for LLM interactions. + +Architecture: +- RepoMap: Main class for building repository maps +- SymbolExtractor: Extracts symbols using tree-sitter or regex fallback +- SymbolRanker: Ranks symbols by importance using graph analysis +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Set +from pathlib import Path +from collections import defaultdict +import logging +import re + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Data Classes +# ============================================================================ + +@dataclass +class Symbol: + """ + Represents a code symbol (class, function, method, variable). + """ + name: str + kind: str # class, function, method, variable, import + file_path: str + line_number: int + signature: str = "" # Full signature for functions/methods + parent: Optional[str] = None # Parent class for methods + references: int = 0 # Number of references in codebase + + @property + def qualified_name(self) -> str: + """Get fully qualified name.""" + if self.parent: + return f"{self.parent}.{self.name}" + return self.name + + def __hash__(self): + return hash((self.name, self.kind, self.file_path, self.line_number)) + + +@dataclass +class FileMap: + """ + Map of symbols in a single file. + """ + file_path: str + symbols: List[Symbol] = field(default_factory=list) + imports: List[str] = field(default_factory=list) + + def add_symbol(self, symbol: Symbol) -> None: + """Add a symbol to the file map.""" + self.symbols.append(symbol) + + def get_summary(self, max_symbols: int = 10) -> str: + """Get a summary of the file's symbols.""" + lines = [f"{self.file_path}:"] + + # Sort by importance (references) and take top symbols + sorted_symbols = sorted( + self.symbols, + key=lambda s: (s.references, s.kind == "class"), + reverse=True + )[:max_symbols] + + for symbol in sorted_symbols: + if symbol.signature: + lines.append(f" {symbol.signature}") + else: + lines.append(f" {symbol.kind} {symbol.name}") + + return "\n".join(lines) + + +@dataclass +class RepoMapConfig: + """Configuration for repository mapping.""" + max_tokens: int = 1024 + max_files: int = 50 + max_symbols_per_file: int = 20 + include_imports: bool = True + include_docstrings: bool = False + file_extensions: Set[str] = field(default_factory=lambda: { + ".py", ".js", ".ts", ".jsx", ".tsx", ".java", ".go", ".rs", + ".cpp", ".c", ".h", ".hpp", ".rb", ".php", ".swift", ".kt" + }) + exclude_patterns: Set[str] = field(default_factory=lambda: { + "__pycache__", "node_modules", ".git", ".venv", "venv", + "build", "dist", ".egg-info", "__init__.py" + }) + + +# ============================================================================ +# Symbol Extraction +# ============================================================================ + +class SymbolExtractor: + """ + Extracts symbols from source code. + + Uses tree-sitter when available, falls back to regex patterns. + """ + + def __init__(self, use_tree_sitter: bool = True): + self.use_tree_sitter = use_tree_sitter + self._tree_sitter_available = False + self._parsers: Dict[str, Any] = {} + + if use_tree_sitter: + self._init_tree_sitter() + + def _init_tree_sitter(self) -> None: + """Initialize tree-sitter if available.""" + try: + # Try to import tree-sitter-languages (pip installable) + import tree_sitter_languages + self._tree_sitter_available = True + self._ts_languages = tree_sitter_languages + logger.debug("tree-sitter-languages available") + except ImportError: + try: + # Fallback to tree-sitter with manual language setup + import tree_sitter + self._tree_sitter_available = True + self._tree_sitter = tree_sitter + logger.debug("tree-sitter available (manual setup required)") + except ImportError: + logger.debug("tree-sitter not available, using regex fallback") + self._tree_sitter_available = False + + def extract_symbols(self, file_path: str, content: str) -> List[Symbol]: + """ + Extract symbols from file content. + + Args: + file_path: Path to the file + content: File content + + Returns: + List of extracted symbols + """ + ext = Path(file_path).suffix.lower() + + if self._tree_sitter_available: + try: + return self._extract_with_tree_sitter(file_path, content, ext) + except Exception as e: + logger.debug(f"tree-sitter extraction failed: {e}") + + # Fallback to regex + return self._extract_with_regex(file_path, content, ext) + + def _extract_with_tree_sitter( + self, file_path: str, content: str, ext: str + ) -> List[Symbol]: + """Extract symbols using tree-sitter.""" + symbols = [] + + # Map extension to language + lang_map = { + ".py": "python", + ".js": "javascript", + ".ts": "typescript", + ".jsx": "javascript", + ".tsx": "typescript", + ".java": "java", + ".go": "go", + ".rs": "rust", + ".rb": "ruby", + ".cpp": "cpp", + ".c": "c", + ".h": "c", + } + + lang = lang_map.get(ext) + if not lang: + return self._extract_with_regex(file_path, content, ext) + + try: + parser = self._ts_languages.get_parser(lang) + tree = parser.parse(content.encode()) + + # Extract based on language + if lang == "python": + symbols = self._extract_python_symbols(file_path, tree, content) + else: + # Generic extraction for other languages + symbols = self._extract_generic_symbols(file_path, tree, content, lang) + except Exception as e: + logger.debug(f"tree-sitter parsing failed for {lang}: {e}") + return self._extract_with_regex(file_path, content, ext) + + return symbols + + def _extract_python_symbols( + self, file_path: str, tree: Any, content: str + ) -> List[Symbol]: + """Extract Python symbols from tree-sitter AST.""" + symbols = [] + lines = content.split("\n") + + def visit_node(node, parent_class=None): + if node.type == "class_definition": + name_node = node.child_by_field_name("name") + if name_node: + class_name = content[name_node.start_byte:name_node.end_byte] + line_num = node.start_point[0] + 1 + + # Get class signature + sig_end = node.children[0].end_byte if node.children else node.start_byte + 50 + signature = content[node.start_byte:sig_end].split("\n")[0] + + symbols.append(Symbol( + name=class_name, + kind="class", + file_path=file_path, + line_number=line_num, + signature=signature.strip() + )) + + # Visit children with class context + for child in node.children: + visit_node(child, class_name) + return + + elif node.type == "function_definition": + name_node = node.child_by_field_name("name") + if name_node: + func_name = content[name_node.start_byte:name_node.end_byte] + line_num = node.start_point[0] + 1 + + # Get function signature + sig_line = lines[line_num - 1] if line_num <= len(lines) else "" + + kind = "method" if parent_class else "function" + symbols.append(Symbol( + name=func_name, + kind=kind, + file_path=file_path, + line_number=line_num, + signature=sig_line.strip(), + parent=parent_class + )) + + # Visit children + for child in node.children: + visit_node(child, parent_class) + + visit_node(tree.root_node) + return symbols + + def _extract_generic_symbols( + self, file_path: str, tree: Any, content: str, lang: str + ) -> List[Symbol]: + """Generic symbol extraction for other languages.""" + symbols = [] + + # Node types for different languages + class_types = {"class_definition", "class_declaration", "struct_definition"} + func_types = {"function_definition", "function_declaration", "method_definition"} + + def visit_node(node, parent=None): + if node.type in class_types: + name = self._get_node_name(node, content) + if name: + symbols.append(Symbol( + name=name, + kind="class", + file_path=file_path, + line_number=node.start_point[0] + 1 + )) + for child in node.children: + visit_node(child, name) + return + + elif node.type in func_types: + name = self._get_node_name(node, content) + if name: + symbols.append(Symbol( + name=name, + kind="method" if parent else "function", + file_path=file_path, + line_number=node.start_point[0] + 1, + parent=parent + )) + + for child in node.children: + visit_node(child, parent) + + visit_node(tree.root_node) + return symbols + + def _get_node_name(self, node: Any, content: str) -> Optional[str]: + """Get name from a tree-sitter node.""" + name_node = node.child_by_field_name("name") + if name_node: + return content[name_node.start_byte:name_node.end_byte] + return None + + def _extract_with_regex( + self, file_path: str, content: str, ext: str + ) -> List[Symbol]: + """Extract symbols using regex patterns (fallback).""" + symbols = [] + lines = content.split("\n") + + if ext == ".py": + symbols = self._extract_python_regex(file_path, lines) + elif ext in {".js", ".ts", ".jsx", ".tsx"}: + symbols = self._extract_js_regex(file_path, lines) + elif ext in {".java", ".kt"}: + symbols = self._extract_java_regex(file_path, lines) + elif ext == ".go": + symbols = self._extract_go_regex(file_path, lines) + elif ext == ".rs": + symbols = self._extract_rust_regex(file_path, lines) + else: + # Generic extraction + symbols = self._extract_generic_regex(file_path, lines) + + return symbols + + def _extract_python_regex(self, file_path: str, lines: List[str]) -> List[Symbol]: + """Extract Python symbols using regex.""" + symbols = [] + current_class = None + + class_pattern = re.compile(r'^class\s+(\w+)') + func_pattern = re.compile(r'^(\s*)def\s+(\w+)\s*\(([^)]*)\)') + + for i, line in enumerate(lines): + line_num = i + 1 + + # Check for class + class_match = class_pattern.match(line) + if class_match: + current_class = class_match.group(1) + symbols.append(Symbol( + name=current_class, + kind="class", + file_path=file_path, + line_number=line_num, + signature=line.strip() + )) + continue + + # Check for function/method + func_match = func_pattern.match(line) + if func_match: + indent = len(func_match.group(1)) + func_name = func_match.group(2) + + # If indented, it's a method + if indent > 0 and current_class: + symbols.append(Symbol( + name=func_name, + kind="method", + file_path=file_path, + line_number=line_num, + signature=line.strip(), + parent=current_class + )) + else: + current_class = None # Reset class context + symbols.append(Symbol( + name=func_name, + kind="function", + file_path=file_path, + line_number=line_num, + signature=line.strip() + )) + + return symbols + + def _extract_js_regex(self, file_path: str, lines: List[str]) -> List[Symbol]: + """Extract JavaScript/TypeScript symbols using regex.""" + symbols = [] + + class_pattern = re.compile(r'(?:export\s+)?class\s+(\w+)') + func_pattern = re.compile( + r'(?:export\s+)?(?:async\s+)?function\s+(\w+)|' + r'(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\([^)]*\)\s*=>' + ) + + for i, line in enumerate(lines): + line_num = i + 1 + + class_match = class_pattern.search(line) + if class_match: + symbols.append(Symbol( + name=class_match.group(1), + kind="class", + file_path=file_path, + line_number=line_num, + signature=line.strip()[:80] + )) + + func_match = func_pattern.search(line) + if func_match: + name = func_match.group(1) or func_match.group(2) + if name: + symbols.append(Symbol( + name=name, + kind="function", + file_path=file_path, + line_number=line_num, + signature=line.strip()[:80] + )) + + return symbols + + def _extract_java_regex(self, file_path: str, lines: List[str]) -> List[Symbol]: + """Extract Java/Kotlin symbols using regex.""" + symbols = [] + + class_pattern = re.compile( + r'(?:public|private|protected)?\s*(?:static\s+)?' + r'(?:class|interface|enum)\s+(\w+)' + ) + method_pattern = re.compile( + r'(?:public|private|protected)?\s*(?:static\s+)?' + r'(?:\w+(?:<[^>]+>)?)\s+(\w+)\s*\(' + ) + + for i, line in enumerate(lines): + line_num = i + 1 + + class_match = class_pattern.search(line) + if class_match: + symbols.append(Symbol( + name=class_match.group(1), + kind="class", + file_path=file_path, + line_number=line_num + )) + + method_match = method_pattern.search(line) + if method_match and not class_match: + symbols.append(Symbol( + name=method_match.group(1), + kind="method", + file_path=file_path, + line_number=line_num + )) + + return symbols + + def _extract_go_regex(self, file_path: str, lines: List[str]) -> List[Symbol]: + """Extract Go symbols using regex.""" + symbols = [] + + func_pattern = re.compile(r'^func\s+(?:\([^)]+\)\s+)?(\w+)') + type_pattern = re.compile(r'^type\s+(\w+)\s+(?:struct|interface)') + + for i, line in enumerate(lines): + line_num = i + 1 + + type_match = type_pattern.match(line) + if type_match: + symbols.append(Symbol( + name=type_match.group(1), + kind="class", + file_path=file_path, + line_number=line_num + )) + + func_match = func_pattern.match(line) + if func_match: + symbols.append(Symbol( + name=func_match.group(1), + kind="function", + file_path=file_path, + line_number=line_num + )) + + return symbols + + def _extract_rust_regex(self, file_path: str, lines: List[str]) -> List[Symbol]: + """Extract Rust symbols using regex.""" + symbols = [] + + fn_pattern = re.compile(r'(?:pub\s+)?(?:async\s+)?fn\s+(\w+)') + struct_pattern = re.compile(r'(?:pub\s+)?struct\s+(\w+)') + + for i, line in enumerate(lines): + line_num = i + 1 + + struct_match = struct_pattern.match(line) + if struct_match: + symbols.append(Symbol( + name=struct_match.group(1), + kind="class", + file_path=file_path, + line_number=line_num + )) + + fn_match = fn_pattern.search(line) + if fn_match: + symbols.append(Symbol( + name=fn_match.group(1), + kind="function", + file_path=file_path, + line_number=line_num + )) + + return symbols + + def _extract_generic_regex(self, file_path: str, lines: List[str]) -> List[Symbol]: + """Generic symbol extraction for unknown languages.""" + symbols = [] + + # Very basic patterns + class_pattern = re.compile(r'(?:class|struct|interface)\s+(\w+)') + func_pattern = re.compile(r'(?:function|def|fn|func)\s+(\w+)') + + for i, line in enumerate(lines): + line_num = i + 1 + + class_match = class_pattern.search(line) + if class_match: + symbols.append(Symbol( + name=class_match.group(1), + kind="class", + file_path=file_path, + line_number=line_num + )) + + func_match = func_pattern.search(line) + if func_match: + symbols.append(Symbol( + name=func_match.group(1), + kind="function", + file_path=file_path, + line_number=line_num + )) + + return symbols + + +# ============================================================================ +# Symbol Ranking +# ============================================================================ + +class SymbolRanker: + """ + Ranks symbols by importance using reference counting and graph analysis. + """ + + def __init__(self): + self.symbol_refs: Dict[str, int] = defaultdict(int) + self.file_refs: Dict[str, Set[str]] = defaultdict(set) + + def analyze_references( + self, file_maps: Dict[str, FileMap], all_content: Dict[str, str] + ) -> None: + """ + Analyze references between symbols across files. + + Args: + file_maps: Map of file path to FileMap + all_content: Map of file path to content + """ + # Build symbol name set + all_symbols = set() + for file_map in file_maps.values(): + for symbol in file_map.symbols: + all_symbols.add(symbol.name) + + # Count references + for file_path, content in all_content.items(): + for symbol_name in all_symbols: + # Simple word boundary match + pattern = rf'\b{re.escape(symbol_name)}\b' + matches = len(re.findall(pattern, content)) + if matches > 0: + self.symbol_refs[symbol_name] += matches + self.file_refs[file_path].add(symbol_name) + + # Update symbol reference counts + for file_map in file_maps.values(): + for symbol in file_map.symbols: + symbol.references = self.symbol_refs.get(symbol.name, 0) + + def get_top_symbols( + self, file_maps: Dict[str, FileMap], max_symbols: int = 50 + ) -> List[Symbol]: + """Get the most important symbols across the codebase.""" + all_symbols = [] + for file_map in file_maps.values(): + all_symbols.extend(file_map.symbols) + + # Sort by references and kind priority + kind_priority = {"class": 3, "function": 2, "method": 1, "variable": 0} + + sorted_symbols = sorted( + all_symbols, + key=lambda s: (s.references, kind_priority.get(s.kind, 0)), + reverse=True + ) + + return sorted_symbols[:max_symbols] + + +# ============================================================================ +# Repository Map +# ============================================================================ + +class RepoMap: + """ + Main class for building repository maps. + + Provides intelligent codebase context for LLM interactions. + """ + + def __init__( + self, + root: Optional[str] = None, + config: Optional[RepoMapConfig] = None, + verbose: bool = False + ): + self.root = Path(root) if root else Path.cwd() + self.config = config or RepoMapConfig() + self.verbose = verbose + + self.extractor = SymbolExtractor() + self.ranker = SymbolRanker() + + self._file_maps: Dict[str, FileMap] = {} + self._all_content: Dict[str, str] = {} + self._last_map: Optional[str] = None + + def scan(self, paths: Optional[List[str]] = None) -> None: + """ + Scan the repository for symbols. + + Args: + paths: Optional list of specific paths to scan + """ + if paths: + files_to_scan = [Path(p) for p in paths] + else: + files_to_scan = self._find_files() + + for file_path in files_to_scan: + try: + self._scan_file(file_path) + except Exception as e: + logger.debug(f"Error scanning {file_path}: {e}") + + # Analyze references + self.ranker.analyze_references(self._file_maps, self._all_content) + + if self.verbose: + logger.info(f"Scanned {len(self._file_maps)} files") + + def _find_files(self) -> List[Path]: + """Find all relevant source files in the repository.""" + files = [] + + for ext in self.config.file_extensions: + for file_path in self.root.rglob(f"*{ext}"): + # Check exclusions + if any(excl in str(file_path) for excl in self.config.exclude_patterns): + continue + files.append(file_path) + + return files[:self.config.max_files] + + def _scan_file(self, file_path: Path) -> None: + """Scan a single file for symbols.""" + try: + content = file_path.read_text(encoding="utf-8", errors="ignore") + except Exception as e: + logger.debug(f"Could not read {file_path}: {e}") + return + + rel_path = str(file_path.relative_to(self.root)) + self._all_content[rel_path] = content + + symbols = self.extractor.extract_symbols(rel_path, content) + + file_map = FileMap(file_path=rel_path) + for symbol in symbols[:self.config.max_symbols_per_file]: + file_map.add_symbol(symbol) + + self._file_maps[rel_path] = file_map + + def get_map( + self, + focus_files: Optional[List[str]] = None, + max_tokens: Optional[int] = None + ) -> str: + """ + Get the repository map as a string. + + Args: + focus_files: Files to prioritize in the map + max_tokens: Maximum tokens for the map + + Returns: + Formatted repository map string + """ + if not self._file_maps: + self.scan() + + max_tokens = max_tokens or self.config.max_tokens + + # Get top symbols + top_symbols = self.ranker.get_top_symbols(self._file_maps) + + # Group by file + file_symbols: Dict[str, List[Symbol]] = defaultdict(list) + for symbol in top_symbols: + file_symbols[symbol.file_path].append(symbol) + + # Prioritize focus files + if focus_files: + focus_set = set(focus_files) + sorted_files = sorted( + file_symbols.keys(), + key=lambda f: (f not in focus_set, f) + ) + else: + sorted_files = sorted(file_symbols.keys()) + + # Build map + lines = [] + estimated_tokens = 0 + + for file_path in sorted_files: + symbols = file_symbols[file_path] + + file_lines = [f"{file_path}:"] + for symbol in symbols: + if symbol.signature: + file_lines.append(f" │{symbol.signature}") + else: + file_lines.append(f" │{symbol.kind} {symbol.name}") + file_lines.append(" ⋮...") + + file_text = "\n".join(file_lines) + file_tokens = len(file_text) // 4 # Rough token estimate + + if estimated_tokens + file_tokens > max_tokens: + break + + lines.extend(file_lines) + estimated_tokens += file_tokens + + self._last_map = "\n".join(lines) + return self._last_map + + def get_file_symbols(self, file_path: str) -> List[Symbol]: + """Get symbols for a specific file.""" + if file_path in self._file_maps: + return self._file_maps[file_path].symbols + return [] + + def get_symbol_context(self, symbol_name: str) -> Optional[str]: + """Get context around a specific symbol.""" + for file_path, file_map in self._file_maps.items(): + for symbol in file_map.symbols: + if symbol.name == symbol_name: + content = self._all_content.get(file_path, "") + lines = content.split("\n") + + # Get lines around the symbol + start = max(0, symbol.line_number - 3) + end = min(len(lines), symbol.line_number + 10) + + context_lines = lines[start:end] + return f"{file_path}:{symbol.line_number}\n" + "\n".join(context_lines) + + return None + + def refresh(self) -> None: + """Refresh the repository map.""" + self._file_maps.clear() + self._all_content.clear() + self.ranker = SymbolRanker() + self.scan() + + +# ============================================================================ +# CLI Integration Handler +# ============================================================================ + +class RepoMapHandler: + """ + Handler for integrating RepoMap with PraisonAI CLI. + """ + + def __init__(self, verbose: bool = False): + self.verbose = verbose + self._repo_map: Optional[RepoMap] = None + + @property + def feature_name(self) -> str: + return "repo_map" + + def initialize( + self, + root: Optional[str] = None, + config: Optional[RepoMapConfig] = None + ) -> RepoMap: + """Initialize the repository map.""" + self._repo_map = RepoMap( + root=root, + config=config, + verbose=self.verbose + ) + + if self.verbose: + from rich import print as rprint + rprint(f"[cyan]RepoMap initialized for: {self._repo_map.root}[/cyan]") + + return self._repo_map + + def get_map(self, focus_files: Optional[List[str]] = None) -> str: + """Get the repository map.""" + if not self._repo_map: + self._repo_map = self.initialize() + + return self._repo_map.get_map(focus_files=focus_files) + + def get_context(self, symbol_name: str) -> Optional[str]: + """Get context for a symbol.""" + if not self._repo_map: + self._repo_map = self.initialize() + + return self._repo_map.get_symbol_context(symbol_name) + + def refresh(self) -> None: + """Refresh the map.""" + if self._repo_map: + self._repo_map.refresh() + + def print_map(self) -> None: + """Print the repository map.""" + from rich.console import Console + from rich.panel import Panel + from rich.syntax import Syntax + + console = Console() + map_str = self.get_map() + + console.print(Panel( + Syntax(map_str, "text", theme="monokai"), + title="📁 Repository Map", + border_style="blue" + )) diff --git a/src/praisonai/praisonai/cli/features/sandbox_executor.py b/src/praisonai/praisonai/cli/features/sandbox_executor.py new file mode 100644 index 000000000..89a2bb2e4 --- /dev/null +++ b/src/praisonai/praisonai/cli/features/sandbox_executor.py @@ -0,0 +1,515 @@ +""" +Sandboxed Execution System for PraisonAI CLI. + +Inspired by Codex CLI's sandbox modes for secure command execution. +Provides isolated execution environment when enabled via --sandbox flag. + +Architecture: +- SandboxExecutor: Base class for sandboxed execution +- SubprocessSandbox: Subprocess-based isolation with resource limits +- SandboxPolicy: Configurable security policies + +Note: Sandbox is ONLY activated when explicitly requested via CLI flag. +""" + +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Set +from pathlib import Path +from enum import Enum +import subprocess +import logging +import tempfile +import shutil +import time + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Sandbox Modes +# ============================================================================ + +class SandboxMode(Enum): + """ + Sandbox execution modes. + + - DISABLED: No sandboxing (default) + - BASIC: Basic isolation with resource limits + - STRICT: Strict isolation with filesystem restrictions + - NETWORK_ISOLATED: No network access + """ + DISABLED = "disabled" + BASIC = "basic" + STRICT = "strict" + NETWORK_ISOLATED = "network_isolated" + + @classmethod + def from_string(cls, value: str) -> "SandboxMode": + """Parse mode from string.""" + value = value.lower().strip().replace("-", "_") + for mode in cls: + if mode.value == value: + return mode + raise ValueError(f"Unknown sandbox mode: {value}") + + +# ============================================================================ +# Data Classes +# ============================================================================ + +@dataclass +class SandboxPolicy: + """ + Security policy for sandbox execution. + """ + mode: SandboxMode = SandboxMode.DISABLED + + # Resource limits + max_memory_mb: int = 512 + max_cpu_seconds: int = 30 + max_file_size_mb: int = 10 + max_processes: int = 10 + + # Filesystem restrictions + allowed_paths: Set[str] = field(default_factory=set) + blocked_paths: Set[str] = field(default_factory=lambda: { + "/etc", "/var", "/usr", "/bin", "/sbin", + "/root", "/home", "/sys", "/proc" + }) + temp_dir: Optional[str] = None + + # Network restrictions + allow_network: bool = True + allowed_hosts: Set[str] = field(default_factory=set) + + # Command restrictions + blocked_commands: Set[str] = field(default_factory=lambda: { + "rm", "rmdir", "mv", "dd", "mkfs", "fdisk", + "sudo", "su", "chmod", "chown", "kill", "pkill" + }) + + @classmethod + def for_mode(cls, mode: SandboxMode) -> "SandboxPolicy": + """Create policy for a given mode.""" + if mode == SandboxMode.DISABLED: + return cls(mode=mode) + + elif mode == SandboxMode.BASIC: + return cls( + mode=mode, + max_memory_mb=512, + max_cpu_seconds=60, + allow_network=True + ) + + elif mode == SandboxMode.STRICT: + return cls( + mode=mode, + max_memory_mb=256, + max_cpu_seconds=30, + max_processes=5, + allow_network=True, + blocked_commands={ + "rm", "rmdir", "mv", "dd", "mkfs", "fdisk", + "sudo", "su", "chmod", "chown", "kill", "pkill", + "curl", "wget", "nc", "netcat", "ssh", "scp" + } + ) + + elif mode == SandboxMode.NETWORK_ISOLATED: + return cls( + mode=mode, + max_memory_mb=512, + max_cpu_seconds=60, + allow_network=False + ) + + return cls(mode=mode) + + +@dataclass +class ExecutionResult: + """Result of sandboxed execution.""" + success: bool + exit_code: int + stdout: str + stderr: str + duration_ms: float + was_sandboxed: bool + policy_violations: List[str] = field(default_factory=list) + + @property + def output(self) -> str: + """Get combined output.""" + return self.stdout + self.stderr + + +# ============================================================================ +# Command Validator +# ============================================================================ + +class CommandValidator: + """ + Validates commands against sandbox policy. + """ + + def __init__(self, policy: SandboxPolicy): + self.policy = policy + + def validate(self, command: str) -> List[str]: + """ + Validate a command against the policy. + + Returns: + List of policy violations (empty if valid) + """ + violations = [] + + if self.policy.mode == SandboxMode.DISABLED: + return violations + + # Parse command + parts = command.split() + if not parts: + return violations + + cmd_name = Path(parts[0]).name + + # Check blocked commands + if cmd_name in self.policy.blocked_commands: + violations.append(f"Command '{cmd_name}' is blocked by sandbox policy") + + # Check for dangerous patterns + dangerous_patterns = [ + ("rm -rf", "Recursive force delete is blocked"), + ("> /dev/", "Writing to device files is blocked"), + ("| sh", "Piping to shell is blocked"), + ("| bash", "Piping to bash is blocked"), + ("$(", "Command substitution is blocked"), + ("`", "Backtick command substitution is blocked"), + ] + + for pattern, message in dangerous_patterns: + if pattern in command: + violations.append(message) + + # Check path access + for part in parts: + if part.startswith("/"): + for blocked in self.policy.blocked_paths: + if part.startswith(blocked): + violations.append(f"Access to '{blocked}' is blocked") + break + + return violations + + def is_allowed(self, command: str) -> bool: + """Check if command is allowed.""" + return len(self.validate(command)) == 0 + + +# ============================================================================ +# Subprocess Sandbox +# ============================================================================ + +class SubprocessSandbox: + """ + Subprocess-based sandbox execution. + + Provides basic isolation through: + - Resource limits (timeout, memory via ulimit) + - Working directory isolation + - Environment variable control + - Command validation + """ + + def __init__( + self, + policy: Optional[SandboxPolicy] = None, + working_dir: Optional[str] = None, + verbose: bool = False + ): + self.policy = policy or SandboxPolicy() + self.working_dir = Path(working_dir) if working_dir else Path.cwd() + self.verbose = verbose + self.validator = CommandValidator(self.policy) + + self._temp_dir: Optional[Path] = None + + def _setup_temp_dir(self) -> Path: + """Set up temporary directory for isolated execution.""" + if self._temp_dir is None: + self._temp_dir = Path(tempfile.mkdtemp(prefix="praisonai_sandbox_")) + return self._temp_dir + + def _cleanup_temp_dir(self) -> None: + """Clean up temporary directory.""" + if self._temp_dir and self._temp_dir.exists(): + try: + shutil.rmtree(self._temp_dir) + except Exception as e: + logger.debug(f"Failed to cleanup temp dir: {e}") + self._temp_dir = None + + def _build_environment(self) -> Dict[str, str]: + """Build restricted environment variables.""" + import os + + # Start with minimal environment + env = { + "PATH": "/usr/local/bin:/usr/bin:/bin", + "HOME": str(self._temp_dir or self.working_dir), + "TERM": "xterm", + "LANG": "en_US.UTF-8", + } + + # Add some safe variables from current environment + safe_vars = ["USER", "SHELL", "PYTHONPATH"] + for var in safe_vars: + if var in os.environ: + env[var] = os.environ[var] + + return env + + def execute( + self, + command: str, + timeout: Optional[float] = None, + capture_output: bool = True + ) -> ExecutionResult: + """ + Execute a command in the sandbox. + + Args: + command: Command to execute + timeout: Timeout in seconds (overrides policy) + capture_output: Whether to capture stdout/stderr + + Returns: + ExecutionResult with execution details + """ + start_time = time.time() + + # Check if sandbox is enabled + if self.policy.mode == SandboxMode.DISABLED: + return self._execute_unsandboxed(command, timeout, capture_output) + + # Validate command + violations = self.validator.validate(command) + if violations: + return ExecutionResult( + success=False, + exit_code=-1, + stdout="", + stderr="Command blocked by sandbox policy:\n" + "\n".join(violations), + duration_ms=(time.time() - start_time) * 1000, + was_sandboxed=True, + policy_violations=violations + ) + + # Set up sandbox + if self.policy.mode == SandboxMode.STRICT: + self._setup_temp_dir() + cwd = self._temp_dir + else: + cwd = self.working_dir + + # Build environment + env = self._build_environment() + + # Set timeout + if timeout is None: + timeout = float(self.policy.max_cpu_seconds) + + # Execute + try: + result = subprocess.run( + command, + shell=True, + cwd=cwd, + env=env, + capture_output=capture_output, + text=True, + timeout=timeout + ) + + return ExecutionResult( + success=result.returncode == 0, + exit_code=result.returncode, + stdout=result.stdout if capture_output else "", + stderr=result.stderr if capture_output else "", + duration_ms=(time.time() - start_time) * 1000, + was_sandboxed=True + ) + + except subprocess.TimeoutExpired: + return ExecutionResult( + success=False, + exit_code=-1, + stdout="", + stderr=f"Command timed out after {timeout}s", + duration_ms=(time.time() - start_time) * 1000, + was_sandboxed=True, + policy_violations=["Timeout exceeded"] + ) + + except Exception as e: + return ExecutionResult( + success=False, + exit_code=-1, + stdout="", + stderr=str(e), + duration_ms=(time.time() - start_time) * 1000, + was_sandboxed=True, + policy_violations=[str(e)] + ) + + finally: + if self.policy.mode == SandboxMode.STRICT: + self._cleanup_temp_dir() + + def _execute_unsandboxed( + self, + command: str, + timeout: Optional[float], + capture_output: bool + ) -> ExecutionResult: + """Execute without sandbox (when disabled).""" + start_time = time.time() + + try: + result = subprocess.run( + command, + shell=True, + cwd=self.working_dir, + capture_output=capture_output, + text=True, + timeout=timeout + ) + + return ExecutionResult( + success=result.returncode == 0, + exit_code=result.returncode, + stdout=result.stdout if capture_output else "", + stderr=result.stderr if capture_output else "", + duration_ms=(time.time() - start_time) * 1000, + was_sandboxed=False + ) + + except subprocess.TimeoutExpired: + return ExecutionResult( + success=False, + exit_code=-1, + stdout="", + stderr=f"Command timed out after {timeout}s", + duration_ms=(time.time() - start_time) * 1000, + was_sandboxed=False + ) + + except Exception as e: + return ExecutionResult( + success=False, + exit_code=-1, + stdout="", + stderr=str(e), + duration_ms=(time.time() - start_time) * 1000, + was_sandboxed=False + ) + + +# ============================================================================ +# CLI Integration Handler +# ============================================================================ + +class SandboxExecutorHandler: + """ + Handler for integrating Sandbox Execution with PraisonAI CLI. + + Note: Sandbox is ONLY activated when --sandbox flag is passed. + """ + + def __init__(self, verbose: bool = False): + self.verbose = verbose + self._sandbox: Optional[SubprocessSandbox] = None + self._enabled: bool = False + + @property + def feature_name(self) -> str: + return "sandbox_executor" + + @property + def is_enabled(self) -> bool: + """Check if sandbox is enabled.""" + return self._enabled + + def initialize( + self, + mode: str = "disabled", + working_dir: Optional[str] = None, + policy: Optional[SandboxPolicy] = None + ) -> SubprocessSandbox: + """ + Initialize the sandbox. + + Args: + mode: Sandbox mode (disabled, basic, strict, network_isolated) + working_dir: Working directory for execution + policy: Optional custom policy + + Returns: + Configured SubprocessSandbox + """ + try: + sandbox_mode = SandboxMode.from_string(mode) + except ValueError: + logger.warning(f"Unknown sandbox mode '{mode}', defaulting to 'disabled'") + sandbox_mode = SandboxMode.DISABLED + + if policy is None: + policy = SandboxPolicy.for_mode(sandbox_mode) + + self._sandbox = SubprocessSandbox( + policy=policy, + working_dir=working_dir, + verbose=self.verbose + ) + + self._enabled = sandbox_mode != SandboxMode.DISABLED + + if self.verbose and self._enabled: + from rich import print as rprint + rprint(f"[yellow]⚠ Sandbox enabled: {sandbox_mode.value}[/yellow]") + + return self._sandbox + + def get_sandbox(self) -> Optional[SubprocessSandbox]: + """Get the sandbox executor.""" + return self._sandbox + + def execute( + self, + command: str, + timeout: Optional[float] = None + ) -> ExecutionResult: + """ + Execute a command. + + If sandbox is not initialized, runs without sandboxing. + """ + if not self._sandbox: + self._sandbox = self.initialize() + + return self._sandbox.execute(command, timeout=timeout) + + def validate_command(self, command: str) -> List[str]: + """Validate a command against sandbox policy.""" + if not self._sandbox: + return [] + + return self._sandbox.validator.validate(command) + + def get_mode(self) -> str: + """Get current sandbox mode.""" + if self._sandbox: + return self._sandbox.policy.mode.value + return SandboxMode.DISABLED.value diff --git a/src/praisonai/praisonai/cli/features/slash_commands.py b/src/praisonai/praisonai/cli/features/slash_commands.py new file mode 100644 index 000000000..4f9f9ada4 --- /dev/null +++ b/src/praisonai/praisonai/cli/features/slash_commands.py @@ -0,0 +1,650 @@ +""" +Slash Commands System for PraisonAI CLI. + +Inspired by Aider's command system and Gemini CLI's slash commands. +Provides interactive commands like /help, /cost, /plan, /model, /clear, etc. + +Architecture: +- SlashCommand: Base class for all commands +- SlashCommandRegistry: Manages command registration and lookup +- SlashCommandParser: Parses user input for slash commands +""" + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional +from enum import Enum +import logging + +logger = logging.getLogger(__name__) + + +class CommandKind(Enum): + """Type of slash command.""" + BUILT_IN = "built_in" + CUSTOM = "custom" + MCP = "mcp" + + +@dataclass +class SlashCommand: + """ + Represents a slash command. + + Attributes: + name: Primary command name (e.g., "help", "cost") + description: Short description for help text + action: Callable that executes the command + alt_names: Alternative names/aliases + sub_commands: Nested sub-commands + kind: Type of command (built-in, custom, MCP) + auto_execute: Whether to execute without confirmation + """ + name: str + description: str + action: Optional[Callable[["CommandContext", str], Any]] = None + alt_names: List[str] = field(default_factory=list) + sub_commands: List["SlashCommand"] = field(default_factory=list) + kind: CommandKind = CommandKind.BUILT_IN + auto_execute: bool = True + + def __post_init__(self): + """Validate command after initialization.""" + if not self.name: + raise ValueError("Command name cannot be empty") + if not self.name.isalnum() and "_" not in self.name and "-" not in self.name: + raise ValueError(f"Invalid command name: {self.name}") + + +@dataclass +class CommandContext: + """ + Context passed to command actions. + + Contains all services and state needed by commands. + """ + # Core services + agent: Any = None # PraisonAI Agent instance + config: Dict[str, Any] = field(default_factory=dict) + + # Session state + session_id: Optional[str] = None + session_start_time: Optional[float] = None + + # Metrics + total_tokens: int = 0 + total_cost: float = 0.0 + prompt_count: int = 0 + + # UI callbacks + print_fn: Callable[[str], None] = print + input_fn: Callable[[str], str] = input + + # Command invocation info + raw_input: str = "" + command_name: str = "" + args: str = "" + + +@dataclass +class ParsedSlashCommand: + """Result of parsing a slash command.""" + command: Optional[SlashCommand] = None + args: str = "" + canonical_path: List[str] = field(default_factory=list) + error: Optional[str] = None + + +class SlashCommandRegistry: + """ + Registry for slash commands. + + Manages command registration, lookup, and help generation. + """ + + def __init__(self): + self._commands: Dict[str, SlashCommand] = {} + self._aliases: Dict[str, str] = {} # alias -> primary name + + def register(self, command: SlashCommand) -> None: + """Register a slash command.""" + if command.name in self._commands: + logger.warning(f"Overwriting existing command: {command.name}") + + self._commands[command.name] = command + + # Register aliases + for alias in command.alt_names: + if alias in self._aliases: + logger.warning(f"Alias '{alias}' already registered") + self._aliases[alias] = command.name + + def unregister(self, name: str) -> bool: + """Unregister a command by name.""" + if name in self._commands: + cmd = self._commands.pop(name) + # Remove aliases + for alias in cmd.alt_names: + self._aliases.pop(alias, None) + return True + return False + + def get(self, name: str) -> Optional[SlashCommand]: + """Get a command by name or alias.""" + # Check primary name first + if name in self._commands: + return self._commands[name] + # Check aliases + if name in self._aliases: + return self._commands.get(self._aliases[name]) + return None + + def get_all(self) -> List[SlashCommand]: + """Get all registered commands.""" + return list(self._commands.values()) + + def get_names(self) -> List[str]: + """Get all command names (including aliases).""" + names = list(self._commands.keys()) + names.extend(self._aliases.keys()) + return sorted(set(names)) + + +class SlashCommandParser: + """ + Parser for slash commands. + + Handles parsing of user input into command and arguments. + """ + + def __init__(self, registry: SlashCommandRegistry): + self.registry = registry + + def is_slash_command(self, query: str) -> bool: + """Check if input is a slash command.""" + query = query.strip() + if not query.startswith('/'): + return False + # Exclude code comments + if query.startswith('//') or query.startswith('/*'): + return False + return True + + def parse(self, query: str) -> ParsedSlashCommand: + """ + Parse a slash command string. + + Args: + query: Raw input string (e.g., "/help model" or "/cost") + + Returns: + ParsedSlashCommand with resolved command and arguments + """ + if not self.is_slash_command(query): + return ParsedSlashCommand(error="Not a slash command") + + trimmed = query.strip() + parts = trimmed[1:].strip().split() # Remove leading / + + if not parts: + return ParsedSlashCommand(error="Empty command") + + command_path = parts + canonical_path: List[str] = [] + current_commands = self.registry.get_all() + command_to_execute: Optional[SlashCommand] = None + path_index = 0 + + for part in command_path: + # First pass: exact match on primary name + found_command = None + for cmd in current_commands: + if cmd.name == part: + found_command = cmd + break + + # Second pass: check aliases + if not found_command: + for cmd in current_commands: + if part in cmd.alt_names: + found_command = cmd + break + + if found_command: + command_to_execute = found_command + canonical_path.append(found_command.name) + path_index += 1 + if found_command.sub_commands: + current_commands = found_command.sub_commands + else: + break + else: + break + + args = ' '.join(parts[path_index:]) + + return ParsedSlashCommand( + command=command_to_execute, + args=args, + canonical_path=canonical_path + ) + + +# ============================================================================ +# Built-in Commands +# ============================================================================ + +def cmd_help(context: CommandContext, args: str) -> Dict[str, Any]: + """Show help for commands.""" + from rich.console import Console + from rich.table import Table + + console = Console() + + if args: + # Help for specific command + registry = context.config.get("command_registry") + if registry: + cmd = registry.get(args) + if cmd: + console.print(f"\n[bold cyan]/{cmd.name}[/bold cyan]") + console.print(f" {cmd.description}") + if cmd.alt_names: + console.print(f" Aliases: {', '.join(cmd.alt_names)}") + if cmd.sub_commands: + console.print("\n Sub-commands:") + for sub in cmd.sub_commands: + console.print(f" /{cmd.name} {sub.name} - {sub.description}") + return {"type": "help", "command": cmd.name} + else: + console.print(f"[red]Unknown command: {args}[/red]") + return {"type": "error", "message": f"Unknown command: {args}"} + + # General help + table = Table(title="Available Commands", show_header=True) + table.add_column("Command", style="cyan") + table.add_column("Description") + + registry = context.config.get("command_registry") + if registry: + for cmd in registry.get_all(): + aliases = f" ({', '.join(cmd.alt_names)})" if cmd.alt_names else "" + table.add_row(f"/{cmd.name}{aliases}", cmd.description) + + console.print(table) + return {"type": "help"} + + +def cmd_cost(context: CommandContext, args: str) -> Dict[str, Any]: + """Show session cost and token usage.""" + from rich.console import Console + from rich.panel import Panel + + console = Console() + + cost_info = f""" +[bold]Session Statistics[/bold] + +Total Tokens: {context.total_tokens:,} +Total Cost: ${context.total_cost:.4f} +Prompts: {context.prompt_count} +""" + + if context.session_start_time: + import time + duration = time.time() - context.session_start_time + minutes = int(duration // 60) + seconds = int(duration % 60) + cost_info += f"Duration: {minutes}m {seconds}s\n" + + console.print(Panel(cost_info, title="💰 Cost Tracker", border_style="green")) + + return { + "type": "stats", + "tokens": context.total_tokens, + "cost": context.total_cost, + "prompts": context.prompt_count + } + + +def cmd_clear(context: CommandContext, args: str) -> Dict[str, Any]: + """Clear the conversation history.""" + from rich.console import Console + console = Console() + + # Reset agent conversation if available + if context.agent and hasattr(context.agent, 'clear_history'): + context.agent.clear_history() + + console.print("[green]✓ Conversation cleared[/green]") + return {"type": "clear"} + + +def cmd_model(context: CommandContext, args: str) -> Dict[str, Any]: + """Show or change the current model.""" + from rich.console import Console + console = Console() + + if args: + # Change model + if context.agent and hasattr(context.agent, 'set_model'): + context.agent.set_model(args) + console.print(f"[green]✓ Model changed to: {args}[/green]") + return {"type": "model_change", "model": args} + else: + console.print(f"[yellow]Model change requested: {args}[/yellow]") + return {"type": "model_change", "model": args} + else: + # Show current model + current_model = "unknown" + if context.agent and hasattr(context.agent, 'llm'): + current_model = getattr(context.agent.llm, 'model', 'unknown') + console.print(f"[cyan]Current model: {current_model}[/cyan]") + return {"type": "model_info", "model": current_model} + + +def cmd_tokens(context: CommandContext, args: str) -> Dict[str, Any]: + """Show token usage breakdown.""" + from rich.console import Console + from rich.table import Table + + console = Console() + + table = Table(title="Token Usage", show_header=True) + table.add_column("Metric", style="cyan") + table.add_column("Value", justify="right") + + table.add_row("Total Tokens", f"{context.total_tokens:,}") + table.add_row("Prompts Sent", str(context.prompt_count)) + if context.prompt_count > 0: + avg = context.total_tokens / context.prompt_count + table.add_row("Avg Tokens/Prompt", f"{avg:.0f}") + + console.print(table) + return {"type": "tokens", "total": context.total_tokens} + + +def cmd_plan(context: CommandContext, args: str) -> Dict[str, Any]: + """Show or create an execution plan.""" + from rich.console import Console + console = Console() + + if args: + console.print(f"[cyan]Creating plan for: {args}[/cyan]") + # Trigger planning mode + return {"type": "submit_prompt", "content": f"Create a detailed plan for: {args}"} + else: + console.print("[yellow]Usage: /plan [/yellow]") + return {"type": "help", "command": "plan"} + + +def cmd_undo(context: CommandContext, args: str) -> Dict[str, Any]: + """Undo the last change.""" + from rich.console import Console + console = Console() + + # Check for git integration + try: + import subprocess + result = subprocess.run( + ["git", "diff", "--stat", "HEAD~1"], + capture_output=True, + text=True, + cwd="." + ) + if result.returncode == 0: + console.print("[yellow]Last commit changes:[/yellow]") + console.print(result.stdout) + console.print("\n[cyan]Run 'git reset --soft HEAD~1' to undo[/cyan]") + else: + console.print("[red]No git history available[/red]") + except Exception as e: + console.print(f"[red]Error: {e}[/red]") + + return {"type": "undo"} + + +def cmd_diff(context: CommandContext, args: str) -> Dict[str, Any]: + """Show git diff of changes.""" + from rich.console import Console + from rich.syntax import Syntax + + console = Console() + + try: + import subprocess + result = subprocess.run( + ["git", "diff"], + capture_output=True, + text=True, + cwd="." + ) + if result.stdout: + syntax = Syntax(result.stdout, "diff", theme="monokai") + console.print(syntax) + else: + console.print("[green]No uncommitted changes[/green]") + except Exception as e: + console.print(f"[red]Error: {e}[/red]") + + return {"type": "diff"} + + +def cmd_commit(context: CommandContext, args: str) -> Dict[str, Any]: + """Commit changes with AI-generated message.""" + from rich.console import Console + console = Console() + + console.print("[cyan]Generating commit message...[/cyan]") + + # This will trigger the commit command handler + return {"type": "commit", "message": args if args else None} + + +def cmd_exit(context: CommandContext, args: str) -> Dict[str, Any]: + """Exit the CLI.""" + from rich.console import Console + console = Console() + + console.print("[yellow]Goodbye![/yellow]") + return {"type": "exit"} + + +def cmd_settings(context: CommandContext, args: str) -> Dict[str, Any]: + """Show current settings.""" + from rich.console import Console + from rich.table import Table + + console = Console() + + table = Table(title="Current Settings", show_header=True) + table.add_column("Setting", style="cyan") + table.add_column("Value") + + for key, value in context.config.items(): + if key != "command_registry": # Skip internal objects + table.add_row(key, str(value)) + + console.print(table) + return {"type": "settings"} + + +def cmd_map(context: CommandContext, args: str) -> Dict[str, Any]: + """Show repository map.""" + from rich.console import Console + console = Console() + + console.print("[cyan]Generating repository map...[/cyan]") + # This will be implemented with RepoMap feature + return {"type": "map"} + + +# ============================================================================ +# Registry Setup +# ============================================================================ + +def create_default_registry() -> SlashCommandRegistry: + """Create registry with default built-in commands.""" + registry = SlashCommandRegistry() + + # Core commands + registry.register(SlashCommand( + name="help", + description="Show help for commands", + action=cmd_help, + alt_names=["h", "?"] + )) + + registry.register(SlashCommand( + name="cost", + description="Show session cost and token usage", + action=cmd_cost, + alt_names=["usage", "stats"] + )) + + registry.register(SlashCommand( + name="clear", + description="Clear conversation history", + action=cmd_clear, + alt_names=["reset"] + )) + + registry.register(SlashCommand( + name="model", + description="Show or change the current model", + action=cmd_model, + alt_names=["m"] + )) + + registry.register(SlashCommand( + name="tokens", + description="Show token usage breakdown", + action=cmd_tokens + )) + + registry.register(SlashCommand( + name="plan", + description="Create an execution plan for a task", + action=cmd_plan + )) + + registry.register(SlashCommand( + name="undo", + description="Undo the last change", + action=cmd_undo + )) + + registry.register(SlashCommand( + name="diff", + description="Show git diff of changes", + action=cmd_diff + )) + + registry.register(SlashCommand( + name="commit", + description="Commit changes with AI-generated message", + action=cmd_commit + )) + + registry.register(SlashCommand( + name="exit", + description="Exit the CLI", + action=cmd_exit, + alt_names=["quit", "q"] + )) + + registry.register(SlashCommand( + name="settings", + description="Show current settings", + action=cmd_settings + )) + + registry.register(SlashCommand( + name="map", + description="Show repository map", + action=cmd_map, + alt_names=["repo"] + )) + + return registry + + +# ============================================================================ +# Handler for CLI Integration +# ============================================================================ + +class SlashCommandHandler: + """ + Handler for integrating slash commands with PraisonAI CLI. + """ + + def __init__(self, verbose: bool = False): + self.verbose = verbose + self.registry = create_default_registry() + self.parser = SlashCommandParser(self.registry) + self._context: Optional[CommandContext] = None + + def set_context(self, context: CommandContext) -> None: + """Set the command context.""" + self._context = context + context.config["command_registry"] = self.registry + + def is_command(self, query: str) -> bool: + """Check if input is a slash command.""" + return self.parser.is_slash_command(query) + + def execute(self, query: str) -> Optional[Dict[str, Any]]: + """ + Execute a slash command. + + Returns: + Command result dict or None if not a command + """ + if not self.is_command(query): + return None + + parsed = self.parser.parse(query) + + if parsed.error: + if self.verbose: + print(f"[red]Error: {parsed.error}[/red]") + return {"type": "error", "message": parsed.error} + + if not parsed.command: + return {"type": "error", "message": "Unknown command"} + + if not parsed.command.action: + return {"type": "error", "message": f"Command '{parsed.command.name}' has no action"} + + # Create context if not set + if not self._context: + self._context = CommandContext() + self._context.config["command_registry"] = self.registry + + self._context.raw_input = query + self._context.command_name = parsed.command.name + self._context.args = parsed.args + + try: + result = parsed.command.action(self._context, parsed.args) + return result + except Exception as e: + logger.error(f"Error executing command: {e}") + return {"type": "error", "message": str(e)} + + def register_command(self, command: SlashCommand) -> None: + """Register a custom command.""" + self.registry.register(command) + + def get_completions(self, partial: str) -> List[str]: + """Get command completions for partial input.""" + if not partial.startswith('/'): + return [] + + partial_cmd = partial[1:].lower() + completions = [] + + for name in self.registry.get_names(): + if name.lower().startswith(partial_cmd): + completions.append(f"/{name}") + + return sorted(completions) diff --git a/src/praisonai/praisonai/deploy.py b/src/praisonai/praisonai/deploy.py index 41f98add0..dc6b81019 100644 --- a/src/praisonai/praisonai/deploy.py +++ b/src/praisonai/praisonai/deploy.py @@ -57,7 +57,7 @@ class CloudDeployer: file.write("FROM python:3.11-slim\n") file.write("WORKDIR /app\n") file.write("COPY . .\n") - file.write("RUN pip install flask praisonai==2.3.44 gunicorn markdown\n") + file.write("RUN pip install flask praisonai==2.3.45 gunicorn markdown\n") file.write("EXPOSE 8080\n") file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n') diff --git a/src/praisonai/praisonai/version.py b/src/praisonai/praisonai/version.py index b473487bd..d60085dd0 100644 --- a/src/praisonai/praisonai/version.py +++ b/src/praisonai/praisonai/version.py @@ -1 +1 @@ -__version__ = "2.3.44" \ No newline at end of file +__version__ = "2.3.45" \ No newline at end of file diff --git a/src/praisonai/test_routing_live.py b/src/praisonai/test_routing_live.py deleted file mode 100644 index 47655b3a2..000000000 --- a/src/praisonai/test_routing_live.py +++ /dev/null @@ -1,231 +0,0 @@ -#!/usr/bin/env python -""" -Live test for model routing with real API calls. -Run this script directly to test with your API keys. -""" - -import os -import sys - -sys.path.insert(0, os.path.dirname(__file__)) - -from praisonai.cli.features.router import RouterHandler - - -def test_router_basics(): - """Test basic router functionality.""" - print("\n" + "="*60) - print("TEST 1: Router Basics") - print("="*60) - - handler = RouterHandler(verbose=True) - - # Test complexity analysis - print("\n📊 Complexity Analysis:") - prompts = [ - ("What is 2+2?", "simple"), - ("Define AI", "simple"), - ("Analyze and research the comprehensive impact of AI on healthcare", "complex"), - ] - - for prompt, expected in prompts: - result = handler.analyze_complexity(prompt) - status = "✅" if result == expected else "❌" - print(f" {status} '{prompt[:40]}...' -> {result} (expected: {expected})") - - print("\n✅ Router basics test passed!") - - -def test_model_selection(): - """Test model selection with different providers.""" - print("\n" + "="*60) - print("TEST 2: Model Selection by Provider") - print("="*60) - - handler = RouterHandler(verbose=True) - - providers = ['openai', 'anthropic', 'google'] - - for provider in providers: - model = handler.select_model("Simple question", preferred_provider=provider) - print(f" {provider}: {model}") - - print("\n✅ Model selection test passed!") - - -def test_custom_models(): - """Test custom model configuration.""" - print("\n" + "="*60) - print("TEST 3: Custom Models from YAML Config") - print("="*60) - - handler = RouterHandler(verbose=True) - - # Simulate YAML config - models_config = { - 'cheap-model': { - 'provider': 'openai', - 'complexity': ['simple'], - 'cost_per_1k': 0.0001, - 'capabilities': ['text'] - }, - 'premium-model': { - 'provider': 'anthropic', - 'complexity': ['complex', 'very_complex'], - 'cost_per_1k': 0.015, - 'capabilities': ['text', 'vision'] - } - } - - handler.load_models_from_config(models_config) - - print(f"\n Available models: {handler.get_available_models()}") - - simple_model = handler.select_model("What is 2+2?") - print(f" Simple task -> {simple_model}") - assert simple_model == 'cheap-model', f"Expected cheap-model, got {simple_model}" - - complex_model = handler.select_model("Analyze and research comprehensive AI impact") - print(f" Complex task -> {complex_model}") - assert complex_model == 'premium-model', f"Expected premium-model, got {complex_model}" - - print("\n✅ Custom models test passed!") - - -def test_real_openai(): - """Test real OpenAI API call.""" - print("\n" + "="*60) - print("TEST 4: Real OpenAI API Call") - print("="*60) - - api_key = os.getenv('OPENAI_API_KEY') - if not api_key or api_key == 'test-key': - print(" ⚠️ OPENAI_API_KEY not set, skipping") - return - - print(f" API Key: {api_key[:20]}...") - - from praisonaiagents import Agent - - handler = RouterHandler(verbose=True) - model = handler.select_model("What is 2+2?", preferred_provider='openai') - print(f" Selected model: {model}") - - agent = Agent( - name="MathAgent", - role="Calculator", - goal="Answer math questions", - llm=model, - verbose=False - ) - - result = agent.start("What is 2+2? Reply with just the number.") - print(f" Result: {result}") - - if result and '4' in str(result): - print("\n✅ OpenAI API test passed!") - else: - print("\n❌ OpenAI API test failed!") - - -def test_real_anthropic(): - """Test real Anthropic API call.""" - print("\n" + "="*60) - print("TEST 5: Real Anthropic API Call") - print("="*60) - - api_key = os.getenv('ANTHROPIC_API_KEY') - if not api_key or api_key == 'test-key': - print(" ⚠️ ANTHROPIC_API_KEY not set, skipping") - return - - print(f" API Key: {api_key[:20]}...") - - from praisonaiagents import Agent - - handler = RouterHandler(verbose=True) - model = handler.select_model("What is the capital of France?", preferred_provider='anthropic') - print(f" Selected model: {model}") - - # Model now includes provider prefix (anthropic/...) - print(f" Using LLM: {model}") - - agent = Agent( - name="GeoAgent", - role="Geography Expert", - goal="Answer geography questions", - llm=model, - verbose=False - ) - - result = agent.start("What is the capital of France? Reply with just the city name.") - print(f" Result: {result}") - - if result and 'paris' in str(result).lower(): - print("\n✅ Anthropic API test passed!") - else: - print("\n❌ Anthropic API test failed!") - - -def test_real_gemini(): - """Test real Gemini API call.""" - print("\n" + "="*60) - print("TEST 6: Real Gemini API Call") - print("="*60) - - api_key = os.getenv('GEMINI_API_KEY') or os.getenv('GOOGLE_API_KEY') - if not api_key or api_key == 'test-key': - print(" ⚠️ GEMINI_API_KEY/GOOGLE_API_KEY not set, skipping") - return - - print(f" API Key: {api_key[:20]}...") - - from praisonaiagents import Agent - - handler = RouterHandler(verbose=True) - model = handler.select_model("What is 3 times 7?", preferred_provider='google') - print(f" Selected model: {model}") - - # Model should already have gemini/ prefix now - llm_model = model - print(f" Using LLM: {llm_model}") - - agent = Agent( - name="MathAgent", - role="Calculator", - goal="Answer math questions", - llm=llm_model, - verbose=False - ) - - result = agent.start("What is 3 times 7? Reply with just the number.") - print(f" Result: {result}") - - if result and '21' in str(result): - print("\n✅ Gemini API test passed!") - else: - print("\n❌ Gemini API test failed!") - - -if __name__ == '__main__': - print("\n" + "="*60) - print("MODEL ROUTING LIVE TESTS") - print("="*60) - - # Check environment variables - print("\n🔑 Environment Variables:") - print(f" OPENAI_API_KEY: {'✅ Set' if os.getenv('OPENAI_API_KEY') else '❌ Not set'}") - print(f" ANTHROPIC_API_KEY: {'✅ Set' if os.getenv('ANTHROPIC_API_KEY') else '❌ Not set'}") - print(f" GEMINI_API_KEY: {'✅ Set' if os.getenv('GEMINI_API_KEY') else '❌ Not set'}") - - # Run tests - test_router_basics() - test_model_selection() - test_custom_models() - test_real_openai() - test_real_anthropic() - test_real_gemini() - - print("\n" + "="*60) - print("ALL TESTS COMPLETED") - print("="*60 + "\n") diff --git a/src/praisonai/tests/e2e/test_cli_features_e2e.py b/src/praisonai/tests/e2e/test_cli_features_e2e.py new file mode 100644 index 000000000..d610e2713 --- /dev/null +++ b/src/praisonai/tests/e2e/test_cli_features_e2e.py @@ -0,0 +1,309 @@ +""" +End-to-End Tests for CLI Features with Real API Keys. + +These tests require actual API keys to be set in environment variables: +- OPENAI_API_KEY +- ANTHROPIC_API_KEY (optional) +- GEMINI_API_KEY (optional) + +Run with: pytest tests/e2e/test_cli_features_e2e.py -v --tb=short +""" + +import pytest +import os +import time +from unittest.mock import patch + +# Skip all tests if no API key is available +pytestmark = pytest.mark.skipif( + not os.environ.get("OPENAI_API_KEY"), + reason="OPENAI_API_KEY not set" +) + + +class TestSlashCommandsE2E: + """End-to-end tests for slash commands.""" + + def test_slash_command_handler_initialization(self): + """Test that slash command handler initializes correctly.""" + from praisonai.cli.features import SlashCommandHandler + + handler = SlashCommandHandler(verbose=True) + assert handler.registry is not None + assert len(handler.registry.get_all()) > 0 + + def test_help_command_execution(self): + """Test /help command execution.""" + from praisonai.cli.features import SlashCommandHandler + + handler = SlashCommandHandler() + + with patch('rich.console.Console.print'): + result = handler.execute("/help") + + assert result is not None + assert result["type"] == "help" + + def test_cost_command_with_context(self): + """Test /cost command with session context.""" + from praisonai.cli.features.slash_commands import ( + SlashCommandHandler, + CommandContext + ) + + handler = SlashCommandHandler() + ctx = CommandContext( + total_tokens=5000, + total_cost=0.015, + prompt_count=10, + session_start_time=time.time() - 300 + ) + handler.set_context(ctx) + + with patch('rich.console.Console.print'): + result = handler.execute("/cost") + + assert result["type"] == "stats" + assert result["tokens"] == 5000 + assert result["cost"] == 0.015 + + def test_command_completions(self): + """Test command auto-completion.""" + from praisonai.cli.features import SlashCommandHandler + + handler = SlashCommandHandler() + + completions = handler.get_completions("/he") + assert "/help" in completions + + completions = handler.get_completions("/mo") + assert "/model" in completions + + +class TestAutonomyModeE2E: + """End-to-end tests for autonomy modes.""" + + def test_autonomy_mode_initialization(self): + """Test autonomy mode handler initialization.""" + from praisonai.cli.features import AutonomyModeHandler + + handler = AutonomyModeHandler(verbose=True) + manager = handler.initialize(mode="suggest") + + assert manager is not None + assert handler.get_mode() == "suggest" + + def test_mode_transitions(self): + """Test transitioning between autonomy modes.""" + from praisonai.cli.features import AutonomyModeHandler + + handler = AutonomyModeHandler() + handler.initialize(mode="suggest") + + # Test all mode transitions + for mode in ["suggest", "auto_edit", "full_auto"]: + handler.set_mode(mode) + assert handler.get_mode() == mode + + def test_approval_flow_suggest_mode(self): + """Test approval flow in suggest mode.""" + from praisonai.cli.features.autonomy_mode import ( + AutonomyModeHandler, + ActionRequest, + ActionType, + ApprovalResult + ) + + # Mock approval callback + def mock_approval(action): + return ApprovalResult(approved=True) + + handler = AutonomyModeHandler() + handler.initialize(mode="suggest", approval_callback=mock_approval) + + # File read should be auto-approved + read_action = ActionRequest(ActionType.FILE_READ, "Read file") + result = handler.request_approval(read_action) + assert result.approved is True + + # File write should require approval (via callback) + write_action = ActionRequest(ActionType.FILE_WRITE, "Write file") + result = handler.request_approval(write_action) + assert result.approved is True + + def test_full_auto_mode(self): + """Test full auto mode approves everything.""" + from praisonai.cli.features.autonomy_mode import ( + AutonomyModeHandler, + ActionRequest, + ActionType + ) + + handler = AutonomyModeHandler() + handler.initialize(mode="full_auto") + + # All actions should be auto-approved + actions = [ + ActionRequest(ActionType.FILE_WRITE, "Write"), + ActionRequest(ActionType.SHELL_COMMAND, "Command"), + ActionRequest(ActionType.FILE_DELETE, "Delete"), + ] + + for action in actions: + result = handler.request_approval(action) + assert result.approved is True + + +class TestCostTrackerE2E: + """End-to-end tests for cost tracking.""" + + def test_cost_tracker_initialization(self): + """Test cost tracker initialization.""" + from praisonai.cli.features import CostTrackerHandler + + handler = CostTrackerHandler(verbose=True) + tracker = handler.initialize(session_id="e2e-test") + + assert tracker is not None + assert tracker.session_id == "e2e-test" + + def test_track_multiple_requests(self): + """Test tracking multiple requests.""" + from praisonai.cli.features import CostTrackerHandler + + handler = CostTrackerHandler() + handler.initialize() + + # Track several requests + models = ["gpt-4o", "gpt-4o-mini", "gpt-4o"] + for i, model in enumerate(models): + handler.track_request( + model=model, + input_tokens=1000 * (i + 1), + output_tokens=500 * (i + 1) + ) + + assert handler.get_tokens() == 9000 # (1500 + 3000 + 4500) + assert handler.get_cost() > 0 + + def test_cost_calculation_accuracy(self): + """Test cost calculation is accurate.""" + from praisonai.cli.features.cost_tracker import CostTracker + + tracker = CostTracker() + + # gpt-4o: $2.50/1M input, $10.00/1M output + stats = tracker.track_request( + model="gpt-4o", + input_tokens=100_000, + output_tokens=50_000 + ) + + # Expected: (100K/1M * 2.50) + (50K/1M * 10.00) = 0.25 + 0.50 = 0.75 + assert abs(stats.cost - 0.75) < 0.01 + + def test_session_export(self): + """Test session data export.""" + from praisonai.cli.features.cost_tracker import CostTracker + import json + + tracker = CostTracker(session_id="export-test") + tracker.track_request("gpt-4o", 1000, 500) + tracker.track_request("gpt-4o-mini", 2000, 1000) + + json_str = tracker.export_json() + data = json.loads(json_str) + + assert data["session"]["session_id"] == "export-test" + assert len(data["requests"]) == 2 + + +class TestIntegratedFeaturesE2E: + """End-to-end tests for integrated features.""" + + def test_all_features_together(self): + """Test all features working together.""" + from praisonai.cli.features import ( + SlashCommandHandler, + AutonomyModeHandler, + CostTrackerHandler + ) + from praisonai.cli.features.slash_commands import CommandContext + from praisonai.cli.features.autonomy_mode import ( + ActionRequest, + ActionType, + ) + + # Initialize all handlers + slash_handler = SlashCommandHandler() + autonomy_handler = AutonomyModeHandler() + cost_handler = CostTrackerHandler() + + # Set up autonomy mode + autonomy_handler.initialize(mode="auto_edit") + + # Initialize cost tracking + cost_handler.initialize(session_id="integrated-test") + + # Track some requests + cost_handler.track_request("gpt-4o", 1000, 500) + cost_handler.track_request("gpt-4o", 2000, 1000) + + # Set up slash command context with cost data + ctx = CommandContext( + total_tokens=cost_handler.get_tokens(), + total_cost=cost_handler.get_cost(), + prompt_count=2 + ) + slash_handler.set_context(ctx) + + # Execute cost command + with patch('rich.console.Console.print'): + result = slash_handler.execute("/cost") + + assert result["tokens"] == 4500 + assert result["cost"] > 0 + + # Test autonomy mode + write_action = ActionRequest(ActionType.FILE_WRITE, "Write file") + approval = autonomy_handler.request_approval(write_action) + assert approval.approved is True # Auto-approved in auto_edit mode + + +class TestRealAPIIntegration: + """Tests that use real API calls.""" + + @pytest.mark.skipif( + not os.environ.get("OPENAI_API_KEY"), + reason="OPENAI_API_KEY required for this test" + ) + def test_cost_tracking_with_real_agent(self): + """Test cost tracking with a real agent call.""" + from praisonai.cli.features import CostTrackerHandler + + handler = CostTrackerHandler() + handler.initialize() + + # Simulate what would happen after an agent call + # In real usage, this would be called after agent.start() + handler.track_request( + model="gpt-4o-mini", + input_tokens=150, + output_tokens=50, + duration_ms=500.0 + ) + + summary = handler.get_summary() + + assert summary["total_requests"] == 1 + assert summary["total_tokens"] == 200 + assert summary["total_cost"] > 0 + + # Print summary for verification + print("\nCost tracking test:") + print(f" Tokens: {summary['total_tokens']}") + print(f" Cost: ${summary['total_cost']:.6f}") + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--tb=short"]) diff --git a/src/praisonai/tests/unit/cli/test_autonomy_mode.py b/src/praisonai/tests/unit/cli/test_autonomy_mode.py new file mode 100644 index 000000000..5b423fc93 --- /dev/null +++ b/src/praisonai/tests/unit/cli/test_autonomy_mode.py @@ -0,0 +1,590 @@ +""" +Tests for Autonomy Mode System. + +Test-Driven Development approach for autonomy/approval modes. +""" + +import pytest +from unittest.mock import MagicMock + +from praisonai.cli.features.autonomy_mode import ( + AutonomyMode, + ActionType, + ActionRequest, + ApprovalResult, + AutonomyPolicy, + AutonomyManager, + AutonomyModeHandler, + create_file_write_action, + create_shell_command_action, + create_git_action, + create_tool_call_action, +) + + +# ============================================================================ +# AutonomyMode Tests +# ============================================================================ + +class TestAutonomyMode: + """Tests for AutonomyMode enum.""" + + def test_mode_values(self): + """Test mode enum values.""" + assert AutonomyMode.SUGGEST.value == "suggest" + assert AutonomyMode.AUTO_EDIT.value == "auto_edit" + assert AutonomyMode.FULL_AUTO.value == "full_auto" + + def test_from_string_valid(self): + """Test parsing valid mode strings.""" + assert AutonomyMode.from_string("suggest") == AutonomyMode.SUGGEST + assert AutonomyMode.from_string("auto_edit") == AutonomyMode.AUTO_EDIT + assert AutonomyMode.from_string("full_auto") == AutonomyMode.FULL_AUTO + + def test_from_string_with_hyphen(self): + """Test parsing mode strings with hyphens.""" + assert AutonomyMode.from_string("auto-edit") == AutonomyMode.AUTO_EDIT + assert AutonomyMode.from_string("full-auto") == AutonomyMode.FULL_AUTO + + def test_from_string_case_insensitive(self): + """Test case insensitive parsing.""" + assert AutonomyMode.from_string("SUGGEST") == AutonomyMode.SUGGEST + assert AutonomyMode.from_string("Auto_Edit") == AutonomyMode.AUTO_EDIT + + def test_from_string_invalid(self): + """Test parsing invalid mode string.""" + with pytest.raises(ValueError, match="Unknown autonomy mode"): + AutonomyMode.from_string("invalid_mode") + + +# ============================================================================ +# ActionType Tests +# ============================================================================ + +class TestActionType: + """Tests for ActionType enum.""" + + def test_action_types_exist(self): + """Test all action types exist.""" + assert ActionType.FILE_READ.value == "file_read" + assert ActionType.FILE_WRITE.value == "file_write" + assert ActionType.FILE_DELETE.value == "file_delete" + assert ActionType.SHELL_COMMAND.value == "shell_command" + assert ActionType.NETWORK_REQUEST.value == "network_request" + assert ActionType.CODE_EXECUTION.value == "code_execution" + assert ActionType.GIT_OPERATION.value == "git_operation" + assert ActionType.TOOL_CALL.value == "tool_call" + + +# ============================================================================ +# ActionRequest Tests +# ============================================================================ + +class TestActionRequest: + """Tests for ActionRequest dataclass.""" + + def test_create_basic_action(self): + """Test creating a basic action request.""" + action = ActionRequest( + action_type=ActionType.FILE_WRITE, + description="Write to config.py" + ) + assert action.action_type == ActionType.FILE_WRITE + assert action.description == "Write to config.py" + assert action.details == {} + assert action.risk_level == "medium" + assert action.reversible is True + + def test_create_action_with_details(self): + """Test creating action with details.""" + action = ActionRequest( + action_type=ActionType.SHELL_COMMAND, + description="Run tests", + details={"command": "pytest", "cwd": "/project"}, + risk_level="low", + reversible=False + ) + assert action.details["command"] == "pytest" + assert action.risk_level == "low" + assert action.reversible is False + + def test_action_str(self): + """Test action string representation.""" + action = ActionRequest( + action_type=ActionType.FILE_WRITE, + description="Write to file.py" + ) + assert "file_write" in str(action) + assert "Write to file.py" in str(action) + + +# ============================================================================ +# ApprovalResult Tests +# ============================================================================ + +class TestApprovalResult: + """Tests for ApprovalResult dataclass.""" + + def test_approved_result(self): + """Test approved result.""" + result = ApprovalResult(approved=True) + assert result.approved is True + assert result.reason is None + assert result.remember_choice is False + + def test_denied_result_with_reason(self): + """Test denied result with reason.""" + result = ApprovalResult( + approved=False, + reason="User denied the action" + ) + assert result.approved is False + assert result.reason == "User denied the action" + + def test_result_with_remember(self): + """Test result with remember choice.""" + result = ApprovalResult( + approved=True, + remember_choice=True + ) + assert result.remember_choice is True + + +# ============================================================================ +# AutonomyPolicy Tests +# ============================================================================ + +class TestAutonomyPolicy: + """Tests for AutonomyPolicy.""" + + def test_suggest_mode_policy(self): + """Test policy for suggest mode.""" + policy = AutonomyPolicy.for_mode(AutonomyMode.SUGGEST) + + assert policy.mode == AutonomyMode.SUGGEST + assert ActionType.FILE_READ in policy.auto_approve + assert ActionType.FILE_WRITE in policy.require_approval + assert ActionType.SHELL_COMMAND in policy.require_approval + assert policy.show_preview is True + + def test_auto_edit_mode_policy(self): + """Test policy for auto-edit mode.""" + policy = AutonomyPolicy.for_mode(AutonomyMode.AUTO_EDIT) + + assert policy.mode == AutonomyMode.AUTO_EDIT + assert ActionType.FILE_READ in policy.auto_approve + assert ActionType.FILE_WRITE in policy.auto_approve + assert ActionType.SHELL_COMMAND in policy.require_approval + assert ActionType.FILE_DELETE in policy.require_approval + + def test_full_auto_mode_policy(self): + """Test policy for full-auto mode.""" + policy = AutonomyPolicy.for_mode(AutonomyMode.FULL_AUTO) + + assert policy.mode == AutonomyMode.FULL_AUTO + assert ActionType.FILE_WRITE in policy.auto_approve + assert ActionType.SHELL_COMMAND in policy.auto_approve + assert ActionType.FILE_DELETE in policy.auto_approve + assert len(policy.require_approval) == 0 + assert policy.max_auto_actions == 100 + + def test_custom_policy(self): + """Test creating custom policy.""" + policy = AutonomyPolicy( + mode=AutonomyMode.SUGGEST, + auto_approve={ActionType.FILE_READ, ActionType.TOOL_CALL}, + blocked={ActionType.FILE_DELETE}, + trusted_paths={"/safe/path"} + ) + + assert ActionType.TOOL_CALL in policy.auto_approve + assert ActionType.FILE_DELETE in policy.blocked + assert "/safe/path" in policy.trusted_paths + + +# ============================================================================ +# AutonomyManager Tests +# ============================================================================ + +class TestAutonomyManager: + """Tests for AutonomyManager.""" + + def test_create_manager_default(self): + """Test creating manager with defaults.""" + manager = AutonomyManager() + + assert manager.mode == AutonomyMode.SUGGEST + assert manager.policy is not None + + def test_create_manager_with_mode(self): + """Test creating manager with specific mode.""" + manager = AutonomyManager(mode=AutonomyMode.FULL_AUTO) + + assert manager.mode == AutonomyMode.FULL_AUTO + + def test_set_mode(self): + """Test changing mode.""" + manager = AutonomyManager(mode=AutonomyMode.SUGGEST) + manager.set_mode(AutonomyMode.AUTO_EDIT) + + assert manager.mode == AutonomyMode.AUTO_EDIT + + def test_auto_approve_file_read(self): + """Test auto-approval of file read in suggest mode.""" + manager = AutonomyManager(mode=AutonomyMode.SUGGEST) + + action = ActionRequest( + action_type=ActionType.FILE_READ, + description="Read config.py" + ) + + result = manager.request_approval(action) + + assert result.approved is True + assert "Auto-approved" in result.reason + + def test_require_approval_file_write_suggest(self): + """Test file write requires approval in suggest mode.""" + # Mock the approval callback + mock_callback = MagicMock(return_value=ApprovalResult(approved=True)) + manager = AutonomyManager( + mode=AutonomyMode.SUGGEST, + approval_callback=mock_callback + ) + + action = ActionRequest( + action_type=ActionType.FILE_WRITE, + description="Write to config.py" + ) + + result = manager.request_approval(action) + + assert result.approved is True + mock_callback.assert_called_once() + + def test_auto_approve_file_write_auto_edit(self): + """Test file write is auto-approved in auto-edit mode.""" + manager = AutonomyManager(mode=AutonomyMode.AUTO_EDIT) + + action = ActionRequest( + action_type=ActionType.FILE_WRITE, + description="Write to config.py" + ) + + result = manager.request_approval(action) + + assert result.approved is True + assert "Auto-approved" in result.reason + + def test_auto_approve_everything_full_auto(self): + """Test everything is auto-approved in full-auto mode.""" + manager = AutonomyManager(mode=AutonomyMode.FULL_AUTO) + + actions = [ + ActionRequest(ActionType.FILE_WRITE, "Write file"), + ActionRequest(ActionType.SHELL_COMMAND, "Run command"), + ActionRequest(ActionType.FILE_DELETE, "Delete file"), + ] + + for action in actions: + result = manager.request_approval(action) + assert result.approved is True + + def test_blocked_action(self): + """Test blocked action is denied.""" + policy = AutonomyPolicy( + mode=AutonomyMode.SUGGEST, + blocked={ActionType.FILE_DELETE} + ) + manager = AutonomyManager(mode=AutonomyMode.SUGGEST, policy=policy) + + action = ActionRequest( + action_type=ActionType.FILE_DELETE, + description="Delete important file" + ) + + result = manager.request_approval(action) + + assert result.approved is False + assert "blocked" in result.reason.lower() + + def test_remember_approval(self): + """Test remembered approval decisions.""" + mock_callback = MagicMock(return_value=ApprovalResult( + approved=True, + remember_choice=True + )) + manager = AutonomyManager( + mode=AutonomyMode.SUGGEST, + approval_callback=mock_callback + ) + + action = ActionRequest( + action_type=ActionType.FILE_WRITE, + description="Write to config.py", + details={"path": "/project/config.py"} + ) + + # First request - should call callback + result1 = manager.request_approval(action) + assert result1.approved is True + assert mock_callback.call_count == 1 + + # Second request - should use remembered choice + result2 = manager.request_approval(action) + assert result2.approved is True + assert mock_callback.call_count == 1 # Not called again + + def test_max_auto_actions(self): + """Test max auto actions limit.""" + mock_callback = MagicMock(return_value=ApprovalResult(approved=True)) + policy = AutonomyPolicy( + mode=AutonomyMode.SUGGEST, + auto_approve={ActionType.FILE_READ}, + max_auto_actions=3 + ) + manager = AutonomyManager( + mode=AutonomyMode.SUGGEST, + policy=policy, + approval_callback=mock_callback + ) + + action = ActionRequest(ActionType.FILE_READ, "Read file") + + # First 3 should be auto-approved + for _ in range(3): + result = manager.request_approval(action) + assert result.approved is True + + # 4th should require approval (callback called) + result = manager.request_approval(action) + assert mock_callback.call_count == 1 + + def test_get_stats(self): + """Test getting approval statistics.""" + manager = AutonomyManager(mode=AutonomyMode.FULL_AUTO) + + # Perform some actions + manager.request_approval(ActionRequest(ActionType.FILE_READ, "Read")) + manager.request_approval(ActionRequest(ActionType.FILE_WRITE, "Write")) + + stats = manager.get_stats() + + assert stats["total_actions"] == 2 + assert stats["auto_approved"] == 2 + assert stats["denied"] == 0 + + def test_reset_stats(self): + """Test resetting statistics.""" + manager = AutonomyManager(mode=AutonomyMode.FULL_AUTO) + + manager.request_approval(ActionRequest(ActionType.FILE_READ, "Read")) + manager.reset_stats() + + stats = manager.get_stats() + assert stats["total_actions"] == 0 + + def test_clear_remembered(self): + """Test clearing remembered approvals.""" + mock_callback = MagicMock(return_value=ApprovalResult( + approved=True, + remember_choice=True + )) + manager = AutonomyManager( + mode=AutonomyMode.SUGGEST, + approval_callback=mock_callback + ) + + action = ActionRequest(ActionType.FILE_WRITE, "Write", details={"path": "/a"}) + + # Remember approval + manager.request_approval(action) + assert mock_callback.call_count == 1 + + # Clear remembered + manager.clear_remembered() + + # Should call callback again + manager.request_approval(action) + assert mock_callback.call_count == 2 + + +# ============================================================================ +# AutonomyModeHandler Tests +# ============================================================================ + +class TestAutonomyModeHandler: + """Tests for AutonomyModeHandler.""" + + def test_handler_creation(self): + """Test handler creation.""" + handler = AutonomyModeHandler() + assert handler.feature_name == "autonomy_mode" + + def test_initialize_default(self): + """Test initializing with defaults.""" + handler = AutonomyModeHandler() + manager = handler.initialize() + + assert manager is not None + assert manager.mode == AutonomyMode.SUGGEST + + def test_initialize_with_mode(self): + """Test initializing with specific mode.""" + handler = AutonomyModeHandler() + manager = handler.initialize(mode="full_auto") + + assert manager.mode == AutonomyMode.FULL_AUTO + + def test_initialize_invalid_mode(self): + """Test initializing with invalid mode falls back to suggest.""" + handler = AutonomyModeHandler() + manager = handler.initialize(mode="invalid_mode") + + assert manager.mode == AutonomyMode.SUGGEST + + def test_get_manager(self): + """Test getting manager.""" + handler = AutonomyModeHandler() + + assert handler.get_manager() is None + + handler.initialize() + assert handler.get_manager() is not None + + def test_request_approval(self): + """Test requesting approval through handler.""" + handler = AutonomyModeHandler() + handler.initialize(mode="full_auto") + + action = ActionRequest(ActionType.FILE_WRITE, "Write file") + result = handler.request_approval(action) + + assert result.approved is True + + def test_set_mode(self): + """Test setting mode through handler.""" + handler = AutonomyModeHandler() + handler.initialize(mode="suggest") + + handler.set_mode("auto_edit") + + assert handler.get_mode() == "auto_edit" + + def test_get_mode(self): + """Test getting mode through handler.""" + handler = AutonomyModeHandler() + + # Before initialization + assert handler.get_mode() == "suggest" + + handler.initialize(mode="full_auto") + assert handler.get_mode() == "full_auto" + + +# ============================================================================ +# Convenience Function Tests +# ============================================================================ + +class TestConvenienceFunctions: + """Tests for convenience functions.""" + + def test_create_file_write_action(self): + """Test creating file write action.""" + action = create_file_write_action("/path/to/file.py", "content preview") + + assert action.action_type == ActionType.FILE_WRITE + assert "/path/to/file.py" in action.description + assert action.details["path"] == "/path/to/file.py" + assert action.risk_level == "medium" + assert action.reversible is True + + def test_create_shell_command_action(self): + """Test creating shell command action.""" + action = create_shell_command_action("pytest tests/", "/project") + + assert action.action_type == ActionType.SHELL_COMMAND + assert action.details["command"] == "pytest tests/" + assert action.details["cwd"] == "/project" + assert action.risk_level == "medium" + assert action.reversible is False + + def test_create_shell_command_dangerous(self): + """Test dangerous shell command has high risk.""" + action = create_shell_command_action("rm -rf /") + + assert action.risk_level == "high" + + def test_create_git_action(self): + """Test creating git action.""" + action = create_git_action("commit", {"message": "Fix bug"}) + + assert action.action_type == ActionType.GIT_OPERATION + assert "commit" in action.description + assert action.details["message"] == "Fix bug" + + def test_create_git_action_status_low_risk(self): + """Test git status has low risk.""" + action = create_git_action("status") + + assert action.risk_level == "low" + + def test_create_tool_call_action(self): + """Test creating tool call action.""" + action = create_tool_call_action("web_search", {"query": "python"}) + + assert action.action_type == ActionType.TOOL_CALL + assert action.details["tool"] == "web_search" + assert action.details["args"]["query"] == "python" + + +# ============================================================================ +# Integration Tests +# ============================================================================ + +class TestAutonomyModeIntegration: + """Integration tests for autonomy mode.""" + + def test_full_workflow_suggest_mode(self): + """Test full workflow in suggest mode.""" + approvals = [] + + def track_approval(action): + approvals.append(action) + return ApprovalResult(approved=True) + + handler = AutonomyModeHandler() + handler.initialize(mode="suggest", approval_callback=track_approval) + + # File read should be auto-approved + read_action = ActionRequest(ActionType.FILE_READ, "Read file") + result = handler.request_approval(read_action) + assert result.approved is True + assert len(approvals) == 0 # No callback for auto-approved + + # File write should require approval + write_action = ActionRequest(ActionType.FILE_WRITE, "Write file") + result = handler.request_approval(write_action) + assert result.approved is True + assert len(approvals) == 1 # Callback was called + + def test_mode_transition(self): + """Test transitioning between modes.""" + handler = AutonomyModeHandler() + handler.initialize(mode="suggest") + + # In suggest mode, file write requires approval + mock_callback = MagicMock(return_value=ApprovalResult(approved=True)) + handler._manager.approval_callback = mock_callback + + write_action = ActionRequest(ActionType.FILE_WRITE, "Write") + handler.request_approval(write_action) + assert mock_callback.call_count == 1 + + # Switch to auto-edit mode + handler.set_mode("auto_edit") + mock_callback.reset_mock() + + # Now file write should be auto-approved + handler.request_approval(write_action) + assert mock_callback.call_count == 0 # Not called diff --git a/src/praisonai/tests/unit/cli/test_cost_tracker.py b/src/praisonai/tests/unit/cli/test_cost_tracker.py new file mode 100644 index 000000000..f0e796c30 --- /dev/null +++ b/src/praisonai/tests/unit/cli/test_cost_tracker.py @@ -0,0 +1,620 @@ +""" +Tests for Cost Tracking System. + +Test-Driven Development approach for cost and token tracking. +""" + +import pytest +from unittest.mock import MagicMock +from datetime import datetime, timedelta +import json + +from praisonai.cli.features.cost_tracker import ( + ModelPricing, + TokenUsage, + RequestStats, + SessionStats, + CostTracker, + CostTrackerHandler, + get_pricing, + DEFAULT_PRICING, +) + + +# ============================================================================ +# ModelPricing Tests +# ============================================================================ + +class TestModelPricing: + """Tests for ModelPricing dataclass.""" + + def test_create_pricing(self): + """Test creating model pricing.""" + pricing = ModelPricing( + model_name="gpt-4o", + input_price_per_1m=2.50, + output_price_per_1m=10.00 + ) + assert pricing.model_name == "gpt-4o" + assert pricing.input_price_per_1m == 2.50 + assert pricing.output_price_per_1m == 10.00 + assert pricing.context_window == 128000 + + def test_calculate_cost(self): + """Test cost calculation.""" + pricing = ModelPricing( + model_name="test", + input_price_per_1m=1.00, + output_price_per_1m=2.00 + ) + + # 1000 input tokens, 500 output tokens + cost = pricing.calculate_cost(1000, 500) + + # Expected: (1000/1M * 1.00) + (500/1M * 2.00) = 0.001 + 0.001 = 0.002 + assert cost == pytest.approx(0.002, rel=1e-6) + + def test_calculate_cost_large_tokens(self): + """Test cost calculation with large token counts.""" + pricing = ModelPricing( + model_name="test", + input_price_per_1m=2.50, + output_price_per_1m=10.00 + ) + + # 100K input, 50K output + cost = pricing.calculate_cost(100_000, 50_000) + + # Expected: (100K/1M * 2.50) + (50K/1M * 10.00) = 0.25 + 0.50 = 0.75 + assert cost == pytest.approx(0.75, rel=1e-6) + + +# ============================================================================ +# get_pricing Tests +# ============================================================================ + +class TestGetPricing: + """Tests for get_pricing function.""" + + def test_get_exact_match(self): + """Test getting pricing with exact model name.""" + pricing = get_pricing("gpt-4o") + assert pricing.model_name == "gpt-4o" + + def test_get_partial_match(self): + """Test getting pricing with partial model name.""" + pricing = get_pricing("gpt-4o-2024-08-06") + assert "gpt-4o" in pricing.model_name.lower() or pricing == DEFAULT_PRICING["default"] + + def test_get_unknown_model(self): + """Test getting pricing for unknown model returns default.""" + pricing = get_pricing("unknown-model-xyz") + assert pricing == DEFAULT_PRICING["default"] + + def test_default_pricing_exists(self): + """Test that default pricing models exist.""" + assert "gpt-4o" in DEFAULT_PRICING + assert "gpt-4o-mini" in DEFAULT_PRICING + assert "claude-3-5-sonnet" in DEFAULT_PRICING + assert "gemini-2.0-flash" in DEFAULT_PRICING + assert "default" in DEFAULT_PRICING + + +# ============================================================================ +# TokenUsage Tests +# ============================================================================ + +class TestTokenUsage: + """Tests for TokenUsage dataclass.""" + + def test_create_usage(self): + """Test creating token usage.""" + usage = TokenUsage( + input_tokens=100, + output_tokens=50 + ) + assert usage.input_tokens == 100 + assert usage.output_tokens == 50 + assert usage.total_tokens == 150 + assert usage.cached_tokens == 0 + + def test_usage_with_cached(self): + """Test usage with cached tokens.""" + usage = TokenUsage( + input_tokens=100, + output_tokens=50, + cached_tokens=30 + ) + assert usage.cached_tokens == 30 + + def test_total_tokens_auto_calculated(self): + """Test total tokens is auto-calculated.""" + usage = TokenUsage(input_tokens=200, output_tokens=100) + assert usage.total_tokens == 300 + + def test_total_tokens_explicit(self): + """Test explicit total tokens is preserved.""" + usage = TokenUsage( + input_tokens=100, + output_tokens=50, + total_tokens=200 # Explicit value + ) + assert usage.total_tokens == 200 + + +# ============================================================================ +# RequestStats Tests +# ============================================================================ + +class TestRequestStats: + """Tests for RequestStats dataclass.""" + + def test_create_request_stats(self): + """Test creating request stats.""" + usage = TokenUsage(input_tokens=100, output_tokens=50) + stats = RequestStats( + timestamp=datetime.now(), + model="gpt-4o", + usage=usage, + cost=0.001 + ) + assert stats.model == "gpt-4o" + assert stats.usage.total_tokens == 150 + assert stats.cost == 0.001 + + def test_to_dict(self): + """Test converting to dictionary.""" + usage = TokenUsage(input_tokens=100, output_tokens=50) + stats = RequestStats( + timestamp=datetime(2024, 1, 1, 12, 0, 0), + model="gpt-4o", + usage=usage, + cost=0.001, + duration_ms=500.0 + ) + + data = stats.to_dict() + + assert data["model"] == "gpt-4o" + assert data["input_tokens"] == 100 + assert data["output_tokens"] == 50 + assert data["cost"] == 0.001 + assert data["duration_ms"] == 500.0 + assert "timestamp" in data + + +# ============================================================================ +# SessionStats Tests +# ============================================================================ + +class TestSessionStats: + """Tests for SessionStats dataclass.""" + + def test_create_session_stats(self): + """Test creating session stats.""" + stats = SessionStats( + session_id="test-123", + start_time=datetime.now() + ) + assert stats.session_id == "test-123" + assert stats.total_requests == 0 + assert stats.total_cost == 0.0 + + def test_add_request(self): + """Test adding request to session.""" + session = SessionStats( + session_id="test", + start_time=datetime.now() + ) + + usage = TokenUsage(input_tokens=100, output_tokens=50) + request = RequestStats( + timestamp=datetime.now(), + model="gpt-4o", + usage=usage, + cost=0.001 + ) + + session.add_request(request) + + assert session.total_requests == 1 + assert session.total_input_tokens == 100 + assert session.total_output_tokens == 50 + assert session.total_tokens == 150 + assert session.total_cost == 0.001 + assert "gpt-4o" in session.models_used + + def test_add_multiple_requests(self): + """Test adding multiple requests.""" + session = SessionStats( + session_id="test", + start_time=datetime.now() + ) + + for i in range(3): + usage = TokenUsage(input_tokens=100, output_tokens=50) + request = RequestStats( + timestamp=datetime.now(), + model="gpt-4o", + usage=usage, + cost=0.001 + ) + session.add_request(request) + + assert session.total_requests == 3 + assert session.total_tokens == 450 + assert session.total_cost == pytest.approx(0.003) + + def test_avg_tokens_per_request(self): + """Test average tokens per request.""" + session = SessionStats( + session_id="test", + start_time=datetime.now() + ) + + for tokens in [100, 200, 300]: + usage = TokenUsage(input_tokens=tokens, output_tokens=0) + request = RequestStats( + timestamp=datetime.now(), + model="gpt-4o", + usage=usage, + cost=0.0 + ) + session.add_request(request) + + assert session.avg_tokens_per_request == 200.0 + + def test_avg_cost_per_request(self): + """Test average cost per request.""" + session = SessionStats( + session_id="test", + start_time=datetime.now() + ) + + for cost in [0.001, 0.002, 0.003]: + usage = TokenUsage(input_tokens=100, output_tokens=50) + request = RequestStats( + timestamp=datetime.now(), + model="gpt-4o", + usage=usage, + cost=cost + ) + session.add_request(request) + + assert session.avg_cost_per_request == pytest.approx(0.002) + + def test_duration_seconds(self): + """Test session duration calculation.""" + start = datetime.now() - timedelta(seconds=60) + session = SessionStats( + session_id="test", + start_time=start + ) + + assert session.duration_seconds >= 60 + + def test_to_dict(self): + """Test converting to dictionary.""" + session = SessionStats( + session_id="test-123", + start_time=datetime(2024, 1, 1, 12, 0, 0) + ) + + data = session.to_dict() + + assert data["session_id"] == "test-123" + assert data["total_requests"] == 0 + assert "start_time" in data + + +# ============================================================================ +# CostTracker Tests +# ============================================================================ + +class TestCostTracker: + """Tests for CostTracker.""" + + def test_create_tracker(self): + """Test creating cost tracker.""" + tracker = CostTracker() + + assert tracker.session_id is not None + assert len(tracker.session_id) == 8 + assert tracker.session_stats is not None + + def test_create_tracker_with_session_id(self): + """Test creating tracker with custom session ID.""" + tracker = CostTracker(session_id="custom-id") + + assert tracker.session_id == "custom-id" + + def test_track_request(self): + """Test tracking a request.""" + tracker = CostTracker() + + stats = tracker.track_request( + model="gpt-4o", + input_tokens=1000, + output_tokens=500 + ) + + assert stats is not None + assert stats.model == "gpt-4o" + assert stats.usage.input_tokens == 1000 + assert stats.usage.output_tokens == 500 + assert stats.cost > 0 + + def test_track_multiple_requests(self): + """Test tracking multiple requests.""" + tracker = CostTracker() + + tracker.track_request("gpt-4o", 1000, 500) + tracker.track_request("gpt-4o-mini", 2000, 1000) + tracker.track_request("gpt-4o", 500, 250) + + assert tracker.get_request_count() == 3 + assert tracker.get_total_tokens() == 5250 + assert tracker.get_total_cost() > 0 + + def test_track_with_cached_tokens(self): + """Test tracking with cached tokens.""" + tracker = CostTracker() + + stats = tracker.track_request( + model="gpt-4o", + input_tokens=1000, + output_tokens=500, + cached_tokens=200 + ) + + assert stats.usage.cached_tokens == 200 + assert tracker.session_stats.total_cached_tokens == 200 + + def test_track_with_duration(self): + """Test tracking with duration.""" + tracker = CostTracker() + + stats = tracker.track_request( + model="gpt-4o", + input_tokens=1000, + output_tokens=500, + duration_ms=1500.0 + ) + + assert stats.duration_ms == 1500.0 + + def test_get_current_stats(self): + """Test getting current stats.""" + tracker = CostTracker() + tracker.track_request("gpt-4o", 1000, 500) + + stats = tracker.get_current_stats() + + assert stats.total_requests == 1 + assert stats.total_tokens == 1500 + + def test_format_summary(self): + """Test formatting summary.""" + tracker = CostTracker(session_id="test-123") + tracker.track_request("gpt-4o", 1000, 500) + + summary = tracker.format_summary() + + assert "test-123" in summary + assert "Tokens:" in summary + assert "Cost:" in summary + + def test_export_json(self): + """Test exporting to JSON.""" + tracker = CostTracker(session_id="test-123") + tracker.track_request("gpt-4o", 1000, 500) + + json_str = tracker.export_json() + data = json.loads(json_str) + + assert "session" in data + assert "requests" in data + assert data["session"]["session_id"] == "test-123" + assert len(data["requests"]) == 1 + + def test_end_session(self): + """Test ending session.""" + tracker = CostTracker() + tracker.track_request("gpt-4o", 1000, 500) + + stats = tracker.end_session() + + assert stats.end_time is not None + + def test_models_used_tracking(self): + """Test tracking of models used.""" + tracker = CostTracker() + + tracker.track_request("gpt-4o", 1000, 500) + tracker.track_request("gpt-4o", 1000, 500) + tracker.track_request("gpt-4o-mini", 1000, 500) + + stats = tracker.get_current_stats() + + assert stats.models_used["gpt-4o"] == 2 + assert stats.models_used["gpt-4o-mini"] == 1 + + +# ============================================================================ +# CostTracker Response Tracking Tests +# ============================================================================ + +class TestCostTrackerResponseTracking: + """Tests for tracking from LLM responses.""" + + def test_track_from_openai_response(self): + """Test tracking from OpenAI-style response.""" + tracker = CostTracker() + + # Mock OpenAI response + mock_response = MagicMock() + mock_response.usage = MagicMock() + mock_response.usage.prompt_tokens = 100 + mock_response.usage.completion_tokens = 50 + mock_response.usage.cached_tokens = 0 + + stats = tracker.track_from_response("gpt-4o", mock_response) + + assert stats is not None + assert stats.usage.input_tokens == 100 + assert stats.usage.output_tokens == 50 + + def test_track_from_dict_response(self): + """Test tracking from dict response.""" + tracker = CostTracker() + + response = { + "usage": { + "prompt_tokens": 200, + "completion_tokens": 100 + } + } + + stats = tracker.track_from_response("gpt-4o", response) + + assert stats is not None + assert stats.usage.input_tokens == 200 + assert stats.usage.output_tokens == 100 + + def test_track_from_invalid_response(self): + """Test tracking from invalid response returns None.""" + tracker = CostTracker() + + stats = tracker.track_from_response("gpt-4o", "invalid response") + + assert stats is None + + +# ============================================================================ +# CostTrackerHandler Tests +# ============================================================================ + +class TestCostTrackerHandler: + """Tests for CostTrackerHandler.""" + + def test_handler_creation(self): + """Test handler creation.""" + handler = CostTrackerHandler() + assert handler.feature_name == "cost_tracker" + + def test_initialize(self): + """Test initializing tracker.""" + handler = CostTrackerHandler() + tracker = handler.initialize() + + assert tracker is not None + assert handler.get_tracker() is tracker + + def test_initialize_with_session_id(self): + """Test initializing with session ID.""" + handler = CostTrackerHandler() + tracker = handler.initialize(session_id="custom-123") + + assert tracker.session_id == "custom-123" + + def test_track_request(self): + """Test tracking request through handler.""" + handler = CostTrackerHandler() + + stats = handler.track_request( + model="gpt-4o", + input_tokens=1000, + output_tokens=500 + ) + + assert stats is not None + assert handler.get_cost() > 0 + assert handler.get_tokens() == 1500 + + def test_get_summary(self): + """Test getting summary.""" + handler = CostTrackerHandler() + handler.track_request("gpt-4o", 1000, 500) + + summary = handler.get_summary() + + assert "total_requests" in summary + assert summary["total_requests"] == 1 + + def test_get_cost_before_init(self): + """Test getting cost before initialization.""" + handler = CostTrackerHandler() + + assert handler.get_cost() == 0.0 + + def test_get_tokens_before_init(self): + """Test getting tokens before initialization.""" + handler = CostTrackerHandler() + + assert handler.get_tokens() == 0 + + +# ============================================================================ +# Integration Tests +# ============================================================================ + +class TestCostTrackerIntegration: + """Integration tests for cost tracking.""" + + def test_full_session_workflow(self): + """Test full session workflow.""" + handler = CostTrackerHandler(verbose=False) + handler.initialize(session_id="integration-test") + + # Simulate multiple requests + models = ["gpt-4o", "gpt-4o-mini", "gpt-4o"] + for i, model in enumerate(models): + handler.track_request( + model=model, + input_tokens=1000 * (i + 1), + output_tokens=500 * (i + 1), + duration_ms=100.0 * (i + 1) + ) + + # Verify totals + summary = handler.get_summary() + + assert summary["total_requests"] == 3 + assert summary["total_input_tokens"] == 6000 # 1000 + 2000 + 3000 + assert summary["total_output_tokens"] == 3000 # 500 + 1000 + 1500 + assert summary["total_tokens"] == 9000 + assert handler.get_cost() > 0 + + def test_cost_calculation_accuracy(self): + """Test cost calculation accuracy.""" + tracker = CostTracker() + + # Use known pricing + # gpt-4o: $2.50/1M input, $10.00/1M output + stats = tracker.track_request( + model="gpt-4o", + input_tokens=1_000_000, # 1M tokens + output_tokens=1_000_000 # 1M tokens + ) + + # Expected: $2.50 + $10.00 = $12.50 + assert stats.cost == pytest.approx(12.50, rel=0.01) + + def test_session_persistence(self): + """Test session data can be exported and contains all info.""" + tracker = CostTracker(session_id="persist-test") + + tracker.track_request("gpt-4o", 1000, 500) + tracker.track_request("gpt-4o-mini", 2000, 1000) + + json_str = tracker.export_json() + data = json.loads(json_str) + + # Verify structure + assert data["session"]["session_id"] == "persist-test" + assert len(data["requests"]) == 2 + assert data["session"]["total_requests"] == 2 + + # Verify request details preserved + assert data["requests"][0]["model"] == "gpt-4o" + assert data["requests"][1]["model"] == "gpt-4o-mini" diff --git a/src/praisonai/tests/unit/cli/test_git_integration.py b/src/praisonai/tests/unit/cli/test_git_integration.py new file mode 100644 index 000000000..09f028299 --- /dev/null +++ b/src/praisonai/tests/unit/cli/test_git_integration.py @@ -0,0 +1,368 @@ +""" +Tests for Git Integration System. + +Test-Driven Development approach for Git operations. +""" + +from pathlib import Path +import tempfile +import subprocess + +from praisonai.cli.features.git_integration import ( + GitStatus, + GitCommit, + GitManager, + CommitMessageGenerator, + DiffViewer, + GitIntegrationHandler, +) + + +# ============================================================================ +# GitStatus Tests +# ============================================================================ + +class TestGitStatus: + """Tests for GitStatus dataclass.""" + + def test_create_status(self): + """Test creating git status.""" + status = GitStatus() + assert status.branch == "" + assert status.is_clean is True + assert status.staged_files == [] + + def test_has_changes(self): + """Test has_changes property.""" + status = GitStatus() + assert status.has_changes is False + + status.modified_files = ["file.py"] + assert status.has_changes is True + + +# ============================================================================ +# GitCommit Tests +# ============================================================================ + +class TestGitCommit: + """Tests for GitCommit dataclass.""" + + def test_create_commit(self): + """Test creating git commit.""" + commit = GitCommit( + hash="abc123def456", + short_hash="abc123d", + message="Test commit", + author="Test User", + date="2024-01-01" + ) + assert commit.hash == "abc123def456" + assert commit.message == "Test commit" + + +# ============================================================================ +# GitManager Tests +# ============================================================================ + +class TestGitManager: + """Tests for GitManager.""" + + def test_create_manager(self): + """Test creating git manager.""" + manager = GitManager() + assert manager is not None + + def test_is_repo_false_for_non_repo(self): + """Test is_repo returns False for non-repo.""" + with tempfile.TemporaryDirectory() as tmpdir: + manager = GitManager(repo_path=tmpdir) + assert manager.is_repo is False + + def test_is_repo_true_for_repo(self): + """Test is_repo returns True for git repo.""" + with tempfile.TemporaryDirectory() as tmpdir: + # Initialize git repo + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + + manager = GitManager(repo_path=tmpdir) + assert manager.is_repo is True + + def test_get_status_empty_repo(self): + """Test getting status of empty repo.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + + manager = GitManager(repo_path=tmpdir) + status = manager.get_status() + + assert status.is_clean is True + + def test_get_status_with_changes(self): + """Test getting status with changes.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + + # Create a file + (Path(tmpdir) / "test.py").write_text("# test") + + manager = GitManager(repo_path=tmpdir) + status = manager.get_status() + + assert "test.py" in status.untracked_files + + def test_stage_files(self): + """Test staging files.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + (Path(tmpdir) / "test.py").write_text("# test") + + manager = GitManager(repo_path=tmpdir) + result = manager.stage_files(["test.py"]) + + assert result is True + + status = manager.get_status() + assert "test.py" in status.staged_files + + def test_commit(self): + """Test creating a commit.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + subprocess.run( + ["git", "config", "user.email", "test@test.com"], + cwd=tmpdir, capture_output=True + ) + subprocess.run( + ["git", "config", "user.name", "Test"], + cwd=tmpdir, capture_output=True + ) + + (Path(tmpdir) / "test.py").write_text("# test") + + manager = GitManager(repo_path=tmpdir) + manager.stage_files() + + commit = manager.commit("Test commit") + + assert commit is not None + assert commit.message == "Test commit" + + def test_get_log(self): + """Test getting commit log.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + subprocess.run( + ["git", "config", "user.email", "test@test.com"], + cwd=tmpdir, capture_output=True + ) + subprocess.run( + ["git", "config", "user.name", "Test"], + cwd=tmpdir, capture_output=True + ) + + (Path(tmpdir) / "test.py").write_text("# test") + + manager = GitManager(repo_path=tmpdir) + manager.stage_files() + manager.commit("First commit") + + commits = manager.get_log() + + assert len(commits) >= 1 + assert commits[0].message == "First commit" + + def test_undo_last_commit(self): + """Test undoing last commit.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + subprocess.run( + ["git", "config", "user.email", "test@test.com"], + cwd=tmpdir, capture_output=True + ) + subprocess.run( + ["git", "config", "user.name", "Test"], + cwd=tmpdir, capture_output=True + ) + + # Need at least 2 commits to undo one + (Path(tmpdir) / "test.py").write_text("# test") + + manager = GitManager(repo_path=tmpdir) + manager.stage_files() + manager.commit("First commit") + + (Path(tmpdir) / "test2.py").write_text("# test2") + manager.stage_files() + manager.commit("Second commit") + + result = manager.undo_last_commit(soft=True) + + assert result is True + + def test_get_branches(self): + """Test getting branches.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + subprocess.run( + ["git", "config", "user.email", "test@test.com"], + cwd=tmpdir, capture_output=True + ) + subprocess.run( + ["git", "config", "user.name", "Test"], + cwd=tmpdir, capture_output=True + ) + + (Path(tmpdir) / "test.py").write_text("# test") + + manager = GitManager(repo_path=tmpdir) + manager.stage_files() + manager.commit("Initial commit") + + branches = manager.get_branches() + + assert len(branches) >= 1 + + +# ============================================================================ +# CommitMessageGenerator Tests +# ============================================================================ + +class TestCommitMessageGenerator: + """Tests for CommitMessageGenerator.""" + + def test_create_generator(self): + """Test creating generator.""" + generator = CommitMessageGenerator() + assert generator is not None + + def test_generate_empty_diff(self): + """Test generating message for empty diff.""" + generator = CommitMessageGenerator(use_ai=False) + message = generator.generate("") + + assert message == "Empty commit" + + def test_generate_from_diff(self): + """Test generating message from diff.""" + generator = CommitMessageGenerator(use_ai=False) + + diff = """ +diff --git a/test.py b/test.py +--- a/test.py ++++ b/test.py +@@ -1 +1,5 @@ ++# New content ++def hello(): ++ pass +""" + + message = generator.generate(diff) + + assert len(message) > 0 + assert "test" in message.lower() or "feat" in message.lower() + + +# ============================================================================ +# DiffViewer Tests +# ============================================================================ + +class TestDiffViewer: + """Tests for DiffViewer.""" + + def test_create_viewer(self): + """Test creating viewer.""" + viewer = DiffViewer() + assert viewer is not None + + def test_display_diff_no_console(self): + """Test displaying diff without console.""" + viewer = DiffViewer() + viewer._console = None + + # Should not raise + viewer.display_diff("test diff") + + def test_display_status_no_console(self): + """Test displaying status without console.""" + viewer = DiffViewer() + viewer._console = None + + status = GitStatus(branch="main") + viewer.display_status(status) + + +# ============================================================================ +# GitIntegrationHandler Tests +# ============================================================================ + +class TestGitIntegrationHandler: + """Tests for GitIntegrationHandler.""" + + def test_handler_creation(self): + """Test handler creation.""" + handler = GitIntegrationHandler() + assert handler.feature_name == "git_integration" + + def test_initialize(self): + """Test initializing handler.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + + handler = GitIntegrationHandler() + git = handler.initialize(repo_path=tmpdir) + + assert git is not None + assert git.is_repo is True + + def test_show_status(self): + """Test showing status.""" + with tempfile.TemporaryDirectory() as tmpdir: + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + + handler = GitIntegrationHandler() + handler.initialize(repo_path=tmpdir) + + status = handler.show_status() + + assert status is not None + + +# ============================================================================ +# Integration Tests +# ============================================================================ + +class TestGitIntegrationIntegration: + """Integration tests for Git integration.""" + + def test_full_workflow(self): + """Test full git workflow.""" + with tempfile.TemporaryDirectory() as tmpdir: + # Initialize repo + subprocess.run(["git", "init"], cwd=tmpdir, capture_output=True) + subprocess.run( + ["git", "config", "user.email", "test@test.com"], + cwd=tmpdir, capture_output=True + ) + subprocess.run( + ["git", "config", "user.name", "Test"], + cwd=tmpdir, capture_output=True + ) + + handler = GitIntegrationHandler() + handler.initialize(repo_path=tmpdir) + + # Create file + (Path(tmpdir) / "main.py").write_text("# Main file\n") + + # Check status + status = handler.show_status() + assert "main.py" in status.untracked_files + + # Commit + commit = handler.commit(message="Add main.py") + assert commit is not None + + # Check log + commits = handler.show_log() + assert len(commits) >= 1 diff --git a/src/praisonai/tests/unit/cli/test_interactive_tui.py b/src/praisonai/tests/unit/cli/test_interactive_tui.py new file mode 100644 index 000000000..b26421ae8 --- /dev/null +++ b/src/praisonai/tests/unit/cli/test_interactive_tui.py @@ -0,0 +1,460 @@ +""" +Tests for Interactive TUI System. + +Test-Driven Development approach for interactive terminal interface. +""" + +from pathlib import Path +import tempfile + +from praisonai.cli.features.interactive_tui import ( + InteractiveConfig, + CommandCompleter, + HistoryManager, + StatusDisplay, + InteractiveSession, + InteractiveTUIHandler, +) + + +# ============================================================================ +# InteractiveConfig Tests +# ============================================================================ + +class TestInteractiveConfig: + """Tests for InteractiveConfig.""" + + def test_default_config(self): + """Test default configuration.""" + config = InteractiveConfig() + + assert config.prompt == ">>> " + assert config.multiline is True + assert config.enable_completions is True + assert config.vi_mode is False + + def test_custom_config(self): + """Test custom configuration.""" + config = InteractiveConfig( + prompt="$ ", + multiline=False, + vi_mode=True + ) + + assert config.prompt == "$ " + assert config.multiline is False + assert config.vi_mode is True + + +# ============================================================================ +# CommandCompleter Tests +# ============================================================================ + +class TestCommandCompleter: + """Tests for CommandCompleter.""" + + def test_create_completer(self): + """Test creating completer.""" + completer = CommandCompleter() + assert completer is not None + + def test_add_commands(self): + """Test adding commands.""" + completer = CommandCompleter() + completer.add_commands(["help", "exit", "cost"]) + + assert "help" in completer.commands + assert "exit" in completer.commands + + def test_complete_slash_command(self): + """Test completing slash commands.""" + completer = CommandCompleter(commands=["help", "history", "exit"]) + + completions = completer.get_completions("/he", 3) + + assert "/help" in completions + # /history should also match but may be filtered + assert "/exit" not in completions + + def test_complete_file_mention(self): + """Test completing file mentions.""" + completer = CommandCompleter() + completer._file_cache = ["main.py", "config.py", "utils.py"] + + completions = completer.get_completions("@main", 5) + + assert "@main.py" in completions + + def test_add_symbols(self): + """Test adding symbols.""" + completer = CommandCompleter() + completer.add_symbols(["MyClass", "my_function"]) + + assert "MyClass" in completer._symbol_cache + + def test_refresh_files(self): + """Test refreshing file cache.""" + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "test.py").write_text("# test") + (Path(tmpdir) / "main.py").write_text("# main") + + completer = CommandCompleter() + completer.refresh_files(Path(tmpdir)) + + assert len(completer._file_cache) >= 2 + + +# ============================================================================ +# HistoryManager Tests +# ============================================================================ + +class TestHistoryManager: + """Tests for HistoryManager.""" + + def test_create_manager(self): + """Test creating history manager.""" + manager = HistoryManager() + assert manager is not None + + def test_add_entry(self): + """Test adding history entry.""" + manager = HistoryManager() + manager.add("first command") + manager.add("second command") + + assert len(manager.get_all()) == 2 + + def test_no_duplicate_consecutive(self): + """Test no duplicate consecutive entries.""" + manager = HistoryManager() + manager.add("same command") + manager.add("same command") + + assert len(manager.get_all()) == 1 + + def test_get_previous(self): + """Test getting previous entry.""" + manager = HistoryManager() + manager.add("first") + manager.add("second") + manager.add("third") + + assert manager.get_previous() == "third" + assert manager.get_previous() == "second" + assert manager.get_previous() == "first" + + def test_get_next(self): + """Test getting next entry.""" + manager = HistoryManager() + manager.add("first") + manager.add("second") + + manager.get_previous() # second + manager.get_previous() # first + + assert manager.get_next() == "second" + + def test_search(self): + """Test searching history.""" + manager = HistoryManager() + manager.add("/help") + manager.add("/exit") + manager.add("/help model") + + results = manager.search("/help") + + assert len(results) == 2 + assert "/help" in results + assert "/help model" in results + + def test_clear(self): + """Test clearing history.""" + manager = HistoryManager() + manager.add("command") + manager.clear() + + assert len(manager.get_all()) == 0 + + def test_persistent_history(self): + """Test persistent history file.""" + with tempfile.TemporaryDirectory() as tmpdir: + history_file = Path(tmpdir) / "history.txt" + + # Create and add entries + manager1 = HistoryManager(history_file=str(history_file)) + manager1.add("command1") + manager1.add("command2") + + # Create new manager with same file + manager2 = HistoryManager(history_file=str(history_file)) + + assert len(manager2.get_all()) == 2 + + +# ============================================================================ +# StatusDisplay Tests +# ============================================================================ + +class TestStatusDisplay: + """Tests for StatusDisplay.""" + + def test_create_display(self): + """Test creating status display.""" + display = StatusDisplay() + assert display is not None + + def test_set_status(self): + """Test setting status.""" + display = StatusDisplay() + display.set_status("model", "gpt-4o") + + assert display._status_items["model"] == "gpt-4o" + + def test_clear_status(self): + """Test clearing status.""" + display = StatusDisplay() + display.set_status("key", "value") + display.clear_status("key") + + assert "key" not in display._status_items + + def test_print_welcome(self): + """Test printing welcome message.""" + display = StatusDisplay() + display._console = None # Force no console + + # Should not raise without console + display.print_welcome() + + def test_print_response(self): + """Test printing response.""" + display = StatusDisplay() + display._console = None # Force no console + + display.print_response("Test response") + + def test_print_error(self): + """Test printing error.""" + display = StatusDisplay() + display._console = None # Force no console + + display.print_error("Test error") + + +# ============================================================================ +# InteractiveSession Tests +# ============================================================================ + +class TestInteractiveSession: + """Tests for InteractiveSession.""" + + def test_create_session(self): + """Test creating session.""" + session = InteractiveSession() + assert session is not None + + def test_create_session_with_config(self): + """Test creating session with config.""" + config = InteractiveConfig(prompt="$ ") + session = InteractiveSession(config=config) + + assert session.config.prompt == "$ " + + def test_add_commands(self): + """Test adding commands.""" + session = InteractiveSession() + session.add_commands(["help", "exit"]) + + assert "help" in session.completer.commands + + def test_add_symbols(self): + """Test adding symbols.""" + session = InteractiveSession() + session.add_symbols(["MyClass"]) + + assert "MyClass" in session.completer._symbol_cache + + def test_process_empty_input(self): + """Test processing empty input.""" + session = InteractiveSession() + result = session.process_input("") + + assert result is None + + def test_process_slash_command(self): + """Test processing slash command.""" + command_called = [] + + def on_command(cmd): + command_called.append(cmd) + return {"type": "help"} + + session = InteractiveSession(on_command=on_command) + session.process_input("/help") + + assert len(command_called) == 1 + assert command_called[0] == "/help" + + def test_process_regular_input(self): + """Test processing regular input.""" + input_received = [] + + def on_input(text): + input_received.append(text) + return "Response" + + session = InteractiveSession(on_input=on_input) + result = session.process_input("Hello world") + + assert len(input_received) == 1 + assert result == "Response" + + def test_exit_command_stops_session(self): + """Test exit command stops session.""" + def on_command(cmd): + return {"type": "exit"} + + session = InteractiveSession(on_command=on_command) + session._running = True + + session.process_input("/exit") + + assert session._running is False + + def test_history_tracking(self): + """Test that input is added to history.""" + session = InteractiveSession() + session.process_input("test command") + + assert "test command" in session.history.get_all() + + def test_stop_session(self): + """Test stopping session.""" + session = InteractiveSession() + session._running = True + + session.stop() + + assert session._running is False + + +# ============================================================================ +# InteractiveTUIHandler Tests +# ============================================================================ + +class TestInteractiveTUIHandler: + """Tests for InteractiveTUIHandler.""" + + def test_handler_creation(self): + """Test handler creation.""" + handler = InteractiveTUIHandler() + assert handler.feature_name == "interactive_tui" + + def test_initialize(self): + """Test initializing handler.""" + handler = InteractiveTUIHandler() + session = handler.initialize() + + assert session is not None + assert handler.get_session() is session + + def test_initialize_with_config(self): + """Test initializing with config.""" + handler = InteractiveTUIHandler() + config = InteractiveConfig(prompt="$ ") + + session = handler.initialize(config=config) + + assert session.config.prompt == "$ " + + def test_initialize_with_callbacks(self): + """Test initializing with callbacks.""" + handler = InteractiveTUIHandler() + + def on_input(text): + return "response" + + def on_command(cmd): + return {"type": "help"} + + session = handler.initialize(on_input=on_input, on_command=on_command) + + assert session.on_input is on_input + assert session.on_command is on_command + + def test_stop(self): + """Test stopping handler.""" + handler = InteractiveTUIHandler() + handler.initialize() + handler._session._running = True + + handler.stop() + + assert handler._session._running is False + + +# ============================================================================ +# Integration Tests +# ============================================================================ + +class TestInteractiveTUIIntegration: + """Integration tests for Interactive TUI.""" + + def test_full_session_workflow(self): + """Test full session workflow.""" + + def on_input(text): + return f"Echo: {text}" + + def on_command(cmd): + if cmd == "/exit": + return {"type": "exit"} + return {"type": "command", "message": f"Executed: {cmd}"} + + handler = InteractiveTUIHandler() + session = handler.initialize(on_input=on_input, on_command=on_command) + + # Process various inputs + result1 = session.process_input("Hello") + assert result1 == "Echo: Hello" + + result2 = session.process_input("/help") + assert result2 == "Executed: /help" + + # Exit + session.process_input("/exit") + assert session._running is False + + def test_completion_integration(self): + """Test completion integration.""" + handler = InteractiveTUIHandler() + session = handler.initialize() + + # Add commands and symbols + session.add_commands(["help", "exit", "cost"]) + session.add_symbols(["MyClass", "my_function"]) + + # Test completions + completions = session.completer.get_completions("/he", 3) + assert "/help" in completions + + completions = session.completer.get_completions("My", 2) + assert "MyClass" in completions + + def test_history_integration(self): + """Test history integration.""" + with tempfile.TemporaryDirectory() as tmpdir: + history_file = Path(tmpdir) / "history.txt" + + config = InteractiveConfig(history_file=str(history_file)) + handler = InteractiveTUIHandler() + session = handler.initialize(config=config) + + # Add some history + session.process_input("command 1") + session.process_input("command 2") + + # Check history + history = session.history.get_all() + assert "command 1" in history + assert "command 2" in history diff --git a/src/praisonai/tests/unit/cli/test_repo_map.py b/src/praisonai/tests/unit/cli/test_repo_map.py new file mode 100644 index 000000000..82af72cf0 --- /dev/null +++ b/src/praisonai/tests/unit/cli/test_repo_map.py @@ -0,0 +1,611 @@ +""" +Tests for Repository Map System. + +Test-Driven Development approach for codebase mapping. +""" + +import pytest +from pathlib import Path +import tempfile + +from praisonai.cli.features.repo_map import ( + Symbol, + FileMap, + RepoMapConfig, + SymbolExtractor, + SymbolRanker, + RepoMap, + RepoMapHandler, +) + + +# ============================================================================ +# Symbol Tests +# ============================================================================ + +class TestSymbol: + """Tests for Symbol dataclass.""" + + def test_create_symbol(self): + """Test creating a symbol.""" + symbol = Symbol( + name="MyClass", + kind="class", + file_path="test.py", + line_number=10 + ) + assert symbol.name == "MyClass" + assert symbol.kind == "class" + assert symbol.file_path == "test.py" + assert symbol.line_number == 10 + + def test_symbol_with_signature(self): + """Test symbol with signature.""" + symbol = Symbol( + name="my_function", + kind="function", + file_path="test.py", + line_number=5, + signature="def my_function(x, y):" + ) + assert symbol.signature == "def my_function(x, y):" + + def test_symbol_with_parent(self): + """Test method symbol with parent class.""" + symbol = Symbol( + name="method", + kind="method", + file_path="test.py", + line_number=15, + parent="MyClass" + ) + assert symbol.parent == "MyClass" + assert symbol.qualified_name == "MyClass.method" + + def test_symbol_qualified_name_no_parent(self): + """Test qualified name without parent.""" + symbol = Symbol( + name="function", + kind="function", + file_path="test.py", + line_number=1 + ) + assert symbol.qualified_name == "function" + + def test_symbol_hash(self): + """Test symbol hashing for sets.""" + s1 = Symbol("test", "function", "file.py", 1) + s2 = Symbol("test", "function", "file.py", 1) + s3 = Symbol("test", "function", "file.py", 2) + + assert hash(s1) == hash(s2) + assert hash(s1) != hash(s3) + + +# ============================================================================ +# FileMap Tests +# ============================================================================ + +class TestFileMap: + """Tests for FileMap dataclass.""" + + def test_create_file_map(self): + """Test creating a file map.""" + file_map = FileMap(file_path="test.py") + assert file_map.file_path == "test.py" + assert file_map.symbols == [] + + def test_add_symbol(self): + """Test adding symbols to file map.""" + file_map = FileMap(file_path="test.py") + symbol = Symbol("test", "function", "test.py", 1) + + file_map.add_symbol(symbol) + + assert len(file_map.symbols) == 1 + assert file_map.symbols[0].name == "test" + + def test_get_summary(self): + """Test getting file summary.""" + file_map = FileMap(file_path="test.py") + file_map.add_symbol(Symbol( + "MyClass", "class", "test.py", 1, + signature="class MyClass:" + )) + file_map.add_symbol(Symbol( + "my_func", "function", "test.py", 10, + signature="def my_func(x):" + )) + + summary = file_map.get_summary() + + assert "test.py:" in summary + assert "class MyClass:" in summary or "MyClass" in summary + + +# ============================================================================ +# RepoMapConfig Tests +# ============================================================================ + +class TestRepoMapConfig: + """Tests for RepoMapConfig.""" + + def test_default_config(self): + """Test default configuration.""" + config = RepoMapConfig() + + assert config.max_tokens == 1024 + assert config.max_files == 50 + assert ".py" in config.file_extensions + assert "__pycache__" in config.exclude_patterns + + def test_custom_config(self): + """Test custom configuration.""" + config = RepoMapConfig( + max_tokens=2048, + max_files=100, + file_extensions={".py", ".js"} + ) + + assert config.max_tokens == 2048 + assert config.max_files == 100 + assert len(config.file_extensions) == 2 + + +# ============================================================================ +# SymbolExtractor Tests +# ============================================================================ + +class TestSymbolExtractor: + """Tests for SymbolExtractor.""" + + def test_create_extractor(self): + """Test creating extractor.""" + extractor = SymbolExtractor(use_tree_sitter=False) + assert extractor is not None + + def test_extract_python_class(self): + """Test extracting Python class.""" + extractor = SymbolExtractor(use_tree_sitter=False) + + content = """ +class MyClass: + def __init__(self): + pass + + def method(self): + pass +""" + + symbols = extractor.extract_symbols("test.py", content) + + class_symbols = [s for s in symbols if s.kind == "class"] + assert len(class_symbols) >= 1 + assert class_symbols[0].name == "MyClass" + + def test_extract_python_function(self): + """Test extracting Python function.""" + extractor = SymbolExtractor(use_tree_sitter=False) + + content = """ +def my_function(x, y): + return x + y + +def another_function(): + pass +""" + + symbols = extractor.extract_symbols("test.py", content) + + func_symbols = [s for s in symbols if s.kind == "function"] + assert len(func_symbols) >= 2 + names = [s.name for s in func_symbols] + assert "my_function" in names + assert "another_function" in names + + def test_extract_python_method(self): + """Test extracting Python method.""" + extractor = SymbolExtractor(use_tree_sitter=False) + + content = """ +class MyClass: + def method_one(self): + pass + + def method_two(self, x): + return x +""" + + symbols = extractor.extract_symbols("test.py", content) + + method_symbols = [s for s in symbols if s.kind == "method"] + assert len(method_symbols) >= 2 + + # Check parent is set + for method in method_symbols: + assert method.parent == "MyClass" + + def test_extract_javascript_class(self): + """Test extracting JavaScript class.""" + extractor = SymbolExtractor(use_tree_sitter=False) + + content = """ +class MyComponent { + constructor() {} +} + +export class AnotherClass {} +""" + + symbols = extractor.extract_symbols("test.js", content) + + class_symbols = [s for s in symbols if s.kind == "class"] + assert len(class_symbols) >= 1 + + def test_extract_javascript_function(self): + """Test extracting JavaScript function.""" + extractor = SymbolExtractor(use_tree_sitter=False) + + content = """ +function myFunction() {} + +const arrowFunc = () => {}; + +async function asyncFunc() {} +""" + + symbols = extractor.extract_symbols("test.js", content) + + func_symbols = [s for s in symbols if s.kind == "function"] + assert len(func_symbols) >= 1 + + def test_extract_go_symbols(self): + """Test extracting Go symbols.""" + extractor = SymbolExtractor(use_tree_sitter=False) + + content = """ +type MyStruct struct { + Name string +} + +func MyFunction() {} + +func (m *MyStruct) Method() {} +""" + + symbols = extractor.extract_symbols("test.go", content) + + assert len(symbols) >= 2 + + def test_extract_rust_symbols(self): + """Test extracting Rust symbols.""" + extractor = SymbolExtractor(use_tree_sitter=False) + + content = """ +pub struct MyStruct { + field: i32, +} + +fn my_function() {} + +pub async fn async_func() {} +""" + + symbols = extractor.extract_symbols("test.rs", content) + + assert len(symbols) >= 2 + + +# ============================================================================ +# SymbolRanker Tests +# ============================================================================ + +class TestSymbolRanker: + """Tests for SymbolRanker.""" + + def test_create_ranker(self): + """Test creating ranker.""" + ranker = SymbolRanker() + assert ranker is not None + + def test_analyze_references(self): + """Test analyzing references.""" + ranker = SymbolRanker() + + file_maps = { + "a.py": FileMap( + file_path="a.py", + symbols=[Symbol("MyClass", "class", "a.py", 1)] + ), + "b.py": FileMap( + file_path="b.py", + symbols=[Symbol("use_class", "function", "b.py", 1)] + ) + } + + all_content = { + "a.py": "class MyClass:\n pass", + "b.py": "from a import MyClass\nobj = MyClass()" + } + + ranker.analyze_references(file_maps, all_content) + + # MyClass should have references + my_class = file_maps["a.py"].symbols[0] + assert my_class.references > 0 + + def test_get_top_symbols(self): + """Test getting top symbols.""" + ranker = SymbolRanker() + + file_maps = { + "test.py": FileMap( + file_path="test.py", + symbols=[ + Symbol("HighRef", "class", "test.py", 1, references=10), + Symbol("LowRef", "function", "test.py", 10, references=1), + Symbol("MedRef", "function", "test.py", 20, references=5), + ] + ) + } + + # Manually set references + file_maps["test.py"].symbols[0].references = 10 + file_maps["test.py"].symbols[1].references = 1 + file_maps["test.py"].symbols[2].references = 5 + + top = ranker.get_top_symbols(file_maps, max_symbols=2) + + assert len(top) == 2 + assert top[0].name == "HighRef" + + +# ============================================================================ +# RepoMap Tests +# ============================================================================ + +class TestRepoMap: + """Tests for RepoMap.""" + + @pytest.fixture + def temp_repo(self): + """Create a temporary repository.""" + with tempfile.TemporaryDirectory() as tmpdir: + # Create some test files + (Path(tmpdir) / "main.py").write_text(""" +class Application: + def __init__(self): + self.config = Config() + + def run(self): + pass + +def main(): + app = Application() + app.run() +""") + + (Path(tmpdir) / "config.py").write_text(""" +class Config: + def __init__(self): + self.debug = False + + def load(self): + pass +""") + + (Path(tmpdir) / "utils.py").write_text(""" +def helper_function(): + pass + +def format_output(data): + return str(data) +""") + + yield tmpdir + + def test_create_repo_map(self, temp_repo): + """Test creating a repo map.""" + repo_map = RepoMap(root=temp_repo) + assert repo_map.root == Path(temp_repo) + + def test_scan_repository(self, temp_repo): + """Test scanning a repository.""" + repo_map = RepoMap(root=temp_repo) + repo_map.scan() + + assert len(repo_map._file_maps) > 0 + + def test_get_map(self, temp_repo): + """Test getting the repository map.""" + repo_map = RepoMap(root=temp_repo) + map_str = repo_map.get_map() + + assert len(map_str) > 0 + # Should contain file names + assert "main.py" in map_str or "config.py" in map_str + + def test_get_map_with_focus(self, temp_repo): + """Test getting map with focus files.""" + repo_map = RepoMap(root=temp_repo) + repo_map.scan() + + map_str = repo_map.get_map(focus_files=["main.py"]) + + assert "main.py" in map_str + + def test_get_file_symbols(self, temp_repo): + """Test getting symbols for a specific file.""" + repo_map = RepoMap(root=temp_repo) + repo_map.scan() + + symbols = repo_map.get_file_symbols("main.py") + + assert len(symbols) > 0 + + def test_get_symbol_context(self, temp_repo): + """Test getting context for a symbol.""" + repo_map = RepoMap(root=temp_repo) + repo_map.scan() + + context = repo_map.get_symbol_context("Application") + + if context: # May not find if symbol extraction differs + assert "Application" in context + + def test_refresh(self, temp_repo): + """Test refreshing the map.""" + repo_map = RepoMap(root=temp_repo) + repo_map.scan() + + initial_count = len(repo_map._file_maps) + + repo_map.refresh() + + assert len(repo_map._file_maps) == initial_count + + +# ============================================================================ +# RepoMapHandler Tests +# ============================================================================ + +class TestRepoMapHandler: + """Tests for RepoMapHandler.""" + + @pytest.fixture + def temp_repo(self): + """Create a temporary repository.""" + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "test.py").write_text(""" +class TestClass: + def test_method(self): + pass +""") + yield tmpdir + + def test_handler_creation(self): + """Test handler creation.""" + handler = RepoMapHandler() + assert handler.feature_name == "repo_map" + + def test_initialize(self, temp_repo): + """Test initializing handler.""" + handler = RepoMapHandler() + repo_map = handler.initialize(root=temp_repo) + + assert repo_map is not None + + def test_get_map(self, temp_repo): + """Test getting map through handler.""" + handler = RepoMapHandler() + handler.initialize(root=temp_repo) + + map_str = handler.get_map() + + assert len(map_str) > 0 + + def test_refresh(self, temp_repo): + """Test refreshing through handler.""" + handler = RepoMapHandler() + handler.initialize(root=temp_repo) + handler.get_map() + + handler.refresh() + + # Should not raise + map_str = handler.get_map() + assert len(map_str) > 0 + + +# ============================================================================ +# Integration Tests +# ============================================================================ + +class TestRepoMapIntegration: + """Integration tests for RepoMap.""" + + @pytest.fixture + def complex_repo(self): + """Create a more complex repository.""" + with tempfile.TemporaryDirectory() as tmpdir: + # Create directory structure + src_dir = Path(tmpdir) / "src" + src_dir.mkdir() + + (src_dir / "app.py").write_text(""" +from .models import User +from .utils import format_name + +class Application: + def __init__(self): + self.users = [] + + def add_user(self, name): + user = User(format_name(name)) + self.users.append(user) + return user + + def get_users(self): + return self.users +""") + + (src_dir / "models.py").write_text(""" +class User: + def __init__(self, name): + self.name = name + + def __repr__(self): + return f"User({self.name})" + +class Admin(User): + def __init__(self, name, role): + super().__init__(name) + self.role = role +""") + + (src_dir / "utils.py").write_text(""" +def format_name(name): + return name.strip().title() + +def validate_email(email): + return '@' in email +""") + + yield tmpdir + + def test_full_scan_and_map(self, complex_repo): + """Test full scan and map generation.""" + repo_map = RepoMap(root=complex_repo) + repo_map.scan() + + map_str = repo_map.get_map() + + # Should contain key symbols + assert len(map_str) > 0 + + def test_symbol_ranking(self, complex_repo): + """Test that frequently referenced symbols rank higher.""" + repo_map = RepoMap(root=complex_repo) + repo_map.scan() + + # User is referenced in multiple files + top_symbols = repo_map.ranker.get_top_symbols(repo_map._file_maps) + + # Should have some symbols + assert len(top_symbols) > 0 + + def test_handler_workflow(self, complex_repo): + """Test complete handler workflow.""" + handler = RepoMapHandler(verbose=False) + handler.initialize(root=complex_repo) + + # Get initial map + map1 = handler.get_map() + assert len(map1) > 0 + + # Refresh + handler.refresh() + + # Get map again + map2 = handler.get_map() + assert len(map2) > 0 diff --git a/src/praisonai/tests/unit/cli/test_sandbox_executor.py b/src/praisonai/tests/unit/cli/test_sandbox_executor.py new file mode 100644 index 000000000..d3dbf3886 --- /dev/null +++ b/src/praisonai/tests/unit/cli/test_sandbox_executor.py @@ -0,0 +1,373 @@ +""" +Tests for Sandbox Executor System. + +Test-Driven Development approach for sandboxed execution. +""" + +import tempfile + +from praisonai.cli.features.sandbox_executor import ( + SandboxMode, + SandboxPolicy, + ExecutionResult, + CommandValidator, + SubprocessSandbox, + SandboxExecutorHandler, +) + + +# ============================================================================ +# SandboxMode Tests +# ============================================================================ + +class TestSandboxMode: + """Tests for SandboxMode enum.""" + + def test_mode_values(self): + """Test mode enum values.""" + assert SandboxMode.DISABLED.value == "disabled" + assert SandboxMode.BASIC.value == "basic" + assert SandboxMode.STRICT.value == "strict" + assert SandboxMode.NETWORK_ISOLATED.value == "network_isolated" + + def test_from_string_valid(self): + """Test parsing valid mode strings.""" + assert SandboxMode.from_string("disabled") == SandboxMode.DISABLED + assert SandboxMode.from_string("basic") == SandboxMode.BASIC + assert SandboxMode.from_string("strict") == SandboxMode.STRICT + + def test_from_string_with_hyphen(self): + """Test parsing mode strings with hyphens.""" + assert SandboxMode.from_string("network-isolated") == SandboxMode.NETWORK_ISOLATED + + def test_from_string_case_insensitive(self): + """Test case insensitive parsing.""" + assert SandboxMode.from_string("BASIC") == SandboxMode.BASIC + assert SandboxMode.from_string("Strict") == SandboxMode.STRICT + + +# ============================================================================ +# SandboxPolicy Tests +# ============================================================================ + +class TestSandboxPolicy: + """Tests for SandboxPolicy.""" + + def test_default_policy(self): + """Test default policy.""" + policy = SandboxPolicy() + + assert policy.mode == SandboxMode.DISABLED + assert policy.max_memory_mb == 512 + assert policy.allow_network is True + + def test_for_mode_disabled(self): + """Test policy for disabled mode.""" + policy = SandboxPolicy.for_mode(SandboxMode.DISABLED) + + assert policy.mode == SandboxMode.DISABLED + + def test_for_mode_basic(self): + """Test policy for basic mode.""" + policy = SandboxPolicy.for_mode(SandboxMode.BASIC) + + assert policy.mode == SandboxMode.BASIC + assert policy.allow_network is True + + def test_for_mode_strict(self): + """Test policy for strict mode.""" + policy = SandboxPolicy.for_mode(SandboxMode.STRICT) + + assert policy.mode == SandboxMode.STRICT + assert policy.max_memory_mb == 256 + assert "curl" in policy.blocked_commands + + def test_for_mode_network_isolated(self): + """Test policy for network isolated mode.""" + policy = SandboxPolicy.for_mode(SandboxMode.NETWORK_ISOLATED) + + assert policy.mode == SandboxMode.NETWORK_ISOLATED + assert policy.allow_network is False + + +# ============================================================================ +# ExecutionResult Tests +# ============================================================================ + +class TestExecutionResult: + """Tests for ExecutionResult.""" + + def test_create_result(self): + """Test creating execution result.""" + result = ExecutionResult( + success=True, + exit_code=0, + stdout="output", + stderr="", + duration_ms=100.0, + was_sandboxed=True + ) + + assert result.success is True + assert result.exit_code == 0 + assert result.stdout == "output" + + def test_output_property(self): + """Test output property.""" + result = ExecutionResult( + success=True, + exit_code=0, + stdout="out", + stderr="err", + duration_ms=100.0, + was_sandboxed=True + ) + + assert result.output == "outerr" + + def test_policy_violations(self): + """Test policy violations.""" + result = ExecutionResult( + success=False, + exit_code=-1, + stdout="", + stderr="blocked", + duration_ms=0.0, + was_sandboxed=True, + policy_violations=["Command blocked"] + ) + + assert len(result.policy_violations) == 1 + + +# ============================================================================ +# CommandValidator Tests +# ============================================================================ + +class TestCommandValidator: + """Tests for CommandValidator.""" + + def test_validate_disabled_mode(self): + """Test validation in disabled mode.""" + policy = SandboxPolicy(mode=SandboxMode.DISABLED) + validator = CommandValidator(policy) + + violations = validator.validate("rm -rf /") + + assert len(violations) == 0 # No validation in disabled mode + + def test_validate_blocked_command(self): + """Test validation of blocked command.""" + policy = SandboxPolicy( + mode=SandboxMode.BASIC, + blocked_commands={"rm", "sudo"} + ) + validator = CommandValidator(policy) + + violations = validator.validate("rm file.txt") + + assert len(violations) > 0 + assert "rm" in violations[0] + + def test_validate_dangerous_pattern(self): + """Test validation of dangerous patterns.""" + policy = SandboxPolicy(mode=SandboxMode.BASIC) + validator = CommandValidator(policy) + + violations = validator.validate("echo test | sh") + + assert len(violations) > 0 + + def test_validate_blocked_path(self): + """Test validation of blocked paths.""" + policy = SandboxPolicy( + mode=SandboxMode.BASIC, + blocked_paths={"/etc", "/var"} + ) + validator = CommandValidator(policy) + + violations = validator.validate("cat /etc/passwd") + + assert len(violations) > 0 + + def test_is_allowed(self): + """Test is_allowed method.""" + policy = SandboxPolicy( + mode=SandboxMode.BASIC, + blocked_commands={"rm"} + ) + validator = CommandValidator(policy) + + assert validator.is_allowed("ls -la") is True + assert validator.is_allowed("rm file.txt") is False + + +# ============================================================================ +# SubprocessSandbox Tests +# ============================================================================ + +class TestSubprocessSandbox: + """Tests for SubprocessSandbox.""" + + def test_create_sandbox(self): + """Test creating sandbox.""" + sandbox = SubprocessSandbox() + assert sandbox is not None + + def test_execute_disabled_mode(self): + """Test execution in disabled mode.""" + policy = SandboxPolicy(mode=SandboxMode.DISABLED) + sandbox = SubprocessSandbox(policy=policy) + + result = sandbox.execute("echo hello") + + assert result.success is True + assert "hello" in result.stdout + assert result.was_sandboxed is False + + def test_execute_basic_mode(self): + """Test execution in basic mode.""" + policy = SandboxPolicy.for_mode(SandboxMode.BASIC) + sandbox = SubprocessSandbox(policy=policy) + + result = sandbox.execute("echo hello") + + assert result.success is True + assert "hello" in result.stdout + assert result.was_sandboxed is True + + def test_execute_blocked_command(self): + """Test execution of blocked command.""" + policy = SandboxPolicy( + mode=SandboxMode.BASIC, + blocked_commands={"rm"} + ) + sandbox = SubprocessSandbox(policy=policy) + + result = sandbox.execute("rm test.txt") + + assert result.success is False + assert len(result.policy_violations) > 0 + + def test_execute_timeout(self): + """Test execution timeout.""" + policy = SandboxPolicy(mode=SandboxMode.BASIC) + sandbox = SubprocessSandbox(policy=policy) + + result = sandbox.execute("sleep 10", timeout=0.1) + + assert result.success is False + assert "timed out" in result.stderr.lower() + + def test_execute_with_working_dir(self): + """Test execution with working directory.""" + with tempfile.TemporaryDirectory() as tmpdir: + policy = SandboxPolicy(mode=SandboxMode.BASIC) + sandbox = SubprocessSandbox(policy=policy, working_dir=tmpdir) + + result = sandbox.execute("pwd") + + assert result.success is True + assert tmpdir in result.stdout + + +# ============================================================================ +# SandboxExecutorHandler Tests +# ============================================================================ + +class TestSandboxExecutorHandler: + """Tests for SandboxExecutorHandler.""" + + def test_handler_creation(self): + """Test handler creation.""" + handler = SandboxExecutorHandler() + assert handler.feature_name == "sandbox_executor" + assert handler.is_enabled is False + + def test_initialize_disabled(self): + """Test initializing with disabled mode.""" + handler = SandboxExecutorHandler() + sandbox = handler.initialize(mode="disabled") + + assert sandbox is not None + assert handler.is_enabled is False + + def test_initialize_basic(self): + """Test initializing with basic mode.""" + handler = SandboxExecutorHandler() + sandbox = handler.initialize(mode="basic") + + assert sandbox is not None + assert handler.is_enabled is True + + def test_execute(self): + """Test executing command.""" + handler = SandboxExecutorHandler() + handler.initialize(mode="basic") + + result = handler.execute("echo test") + + assert result.success is True + assert "test" in result.stdout + + def test_validate_command(self): + """Test validating command.""" + handler = SandboxExecutorHandler() + handler.initialize(mode="basic") + + violations = handler.validate_command("echo hello") + assert len(violations) == 0 + + def test_get_mode(self): + """Test getting mode.""" + handler = SandboxExecutorHandler() + + assert handler.get_mode() == "disabled" + + handler.initialize(mode="strict") + assert handler.get_mode() == "strict" + + +# ============================================================================ +# Integration Tests +# ============================================================================ + +class TestSandboxIntegration: + """Integration tests for Sandbox.""" + + def test_full_workflow_disabled(self): + """Test full workflow with sandbox disabled.""" + handler = SandboxExecutorHandler() + handler.initialize(mode="disabled") + + # Should execute without sandboxing + result = handler.execute("echo hello") + + assert result.success is True + assert result.was_sandboxed is False + + def test_full_workflow_basic(self): + """Test full workflow with basic sandbox.""" + handler = SandboxExecutorHandler() + handler.initialize(mode="basic") + + # Safe command should work + result = handler.execute("echo hello") + assert result.success is True + assert result.was_sandboxed is True + + # Blocked command should fail + result = handler.execute("rm -rf /") + assert result.success is False + assert len(result.policy_violations) > 0 + + def test_strict_mode_isolation(self): + """Test strict mode isolation.""" + handler = SandboxExecutorHandler() + handler.initialize(mode="strict") + + # Should work in isolated environment + result = handler.execute("echo $HOME") + + assert result.success is True + assert result.was_sandboxed is True diff --git a/src/praisonai/tests/unit/cli/test_slash_commands.py b/src/praisonai/tests/unit/cli/test_slash_commands.py new file mode 100644 index 000000000..642926410 --- /dev/null +++ b/src/praisonai/tests/unit/cli/test_slash_commands.py @@ -0,0 +1,608 @@ +""" +Tests for Slash Commands System. + +Test-Driven Development approach: +1. Write tests first +2. Implement features to pass tests +3. Refactor while keeping tests green +""" + +import pytest +from unittest.mock import MagicMock, patch +import time + +from praisonai.cli.features.slash_commands import ( + SlashCommand, + CommandContext, + CommandKind, + SlashCommandRegistry, + SlashCommandParser, + SlashCommandHandler, + create_default_registry, + cmd_help, + cmd_cost, + cmd_clear, + cmd_model, + cmd_tokens, + cmd_plan, + cmd_exit, +) + + +# ============================================================================ +# SlashCommand Tests +# ============================================================================ + +class TestSlashCommand: + """Tests for SlashCommand dataclass.""" + + def test_create_basic_command(self): + """Test creating a basic slash command.""" + cmd = SlashCommand( + name="test", + description="A test command" + ) + assert cmd.name == "test" + assert cmd.description == "A test command" + assert cmd.action is None + assert cmd.alt_names == [] + assert cmd.sub_commands == [] + assert cmd.kind == CommandKind.BUILT_IN + assert cmd.auto_execute is True + + def test_create_command_with_action(self): + """Test creating a command with an action.""" + def my_action(ctx, args): + return {"result": "success"} + + cmd = SlashCommand( + name="mycommand", + description="My command", + action=my_action + ) + assert cmd.action is not None + result = cmd.action(None, "") + assert result == {"result": "success"} + + def test_create_command_with_aliases(self): + """Test creating a command with aliases.""" + cmd = SlashCommand( + name="help", + description="Show help", + alt_names=["h", "?"] + ) + assert "h" in cmd.alt_names + assert "?" in cmd.alt_names + + def test_create_command_with_subcommands(self): + """Test creating a command with sub-commands.""" + sub1 = SlashCommand(name="sub1", description="Sub 1") + sub2 = SlashCommand(name="sub2", description="Sub 2") + + cmd = SlashCommand( + name="parent", + description="Parent command", + sub_commands=[sub1, sub2] + ) + assert len(cmd.sub_commands) == 2 + assert cmd.sub_commands[0].name == "sub1" + + def test_empty_name_raises_error(self): + """Test that empty name raises ValueError.""" + with pytest.raises(ValueError, match="Command name cannot be empty"): + SlashCommand(name="", description="Empty name") + + def test_command_kinds(self): + """Test different command kinds.""" + builtin = SlashCommand(name="builtin", description="Built-in", kind=CommandKind.BUILT_IN) + custom = SlashCommand(name="custom", description="Custom", kind=CommandKind.CUSTOM) + mcp = SlashCommand(name="mcp", description="MCP", kind=CommandKind.MCP) + + assert builtin.kind == CommandKind.BUILT_IN + assert custom.kind == CommandKind.CUSTOM + assert mcp.kind == CommandKind.MCP + + +# ============================================================================ +# CommandContext Tests +# ============================================================================ + +class TestCommandContext: + """Tests for CommandContext dataclass.""" + + def test_create_default_context(self): + """Test creating a default context.""" + ctx = CommandContext() + assert ctx.agent is None + assert ctx.config == {} + assert ctx.session_id is None + assert ctx.total_tokens == 0 + assert ctx.total_cost == 0.0 + assert ctx.prompt_count == 0 + + def test_context_with_metrics(self): + """Test context with metrics.""" + ctx = CommandContext( + total_tokens=1500, + total_cost=0.0045, + prompt_count=3 + ) + assert ctx.total_tokens == 1500 + assert ctx.total_cost == 0.0045 + assert ctx.prompt_count == 3 + + def test_context_with_session(self): + """Test context with session info.""" + start_time = time.time() + ctx = CommandContext( + session_id="test-session-123", + session_start_time=start_time + ) + assert ctx.session_id == "test-session-123" + assert ctx.session_start_time == start_time + + +# ============================================================================ +# SlashCommandRegistry Tests +# ============================================================================ + +class TestSlashCommandRegistry: + """Tests for SlashCommandRegistry.""" + + def test_register_command(self): + """Test registering a command.""" + registry = SlashCommandRegistry() + cmd = SlashCommand(name="test", description="Test") + + registry.register(cmd) + + assert registry.get("test") is not None + assert registry.get("test").name == "test" + + def test_register_command_with_aliases(self): + """Test registering a command with aliases.""" + registry = SlashCommandRegistry() + cmd = SlashCommand( + name="help", + description="Help", + alt_names=["h", "?"] + ) + + registry.register(cmd) + + # Should be accessible by name and aliases + assert registry.get("help") is not None + assert registry.get("h") is not None + assert registry.get("?") is not None + + # All should return the same command + assert registry.get("help") == registry.get("h") + assert registry.get("h") == registry.get("?") + + def test_get_nonexistent_command(self): + """Test getting a non-existent command.""" + registry = SlashCommandRegistry() + assert registry.get("nonexistent") is None + + def test_unregister_command(self): + """Test unregistering a command.""" + registry = SlashCommandRegistry() + cmd = SlashCommand(name="test", description="Test", alt_names=["t"]) + + registry.register(cmd) + assert registry.get("test") is not None + + result = registry.unregister("test") + assert result is True + assert registry.get("test") is None + assert registry.get("t") is None # Alias should also be removed + + def test_unregister_nonexistent(self): + """Test unregistering a non-existent command.""" + registry = SlashCommandRegistry() + result = registry.unregister("nonexistent") + assert result is False + + def test_get_all_commands(self): + """Test getting all commands.""" + registry = SlashCommandRegistry() + registry.register(SlashCommand(name="cmd1", description="Cmd 1")) + registry.register(SlashCommand(name="cmd2", description="Cmd 2")) + registry.register(SlashCommand(name="cmd3", description="Cmd 3")) + + all_cmds = registry.get_all() + assert len(all_cmds) == 3 + names = [cmd.name for cmd in all_cmds] + assert "cmd1" in names + assert "cmd2" in names + assert "cmd3" in names + + def test_get_names(self): + """Test getting all command names including aliases.""" + registry = SlashCommandRegistry() + registry.register(SlashCommand(name="help", description="Help", alt_names=["h"])) + registry.register(SlashCommand(name="exit", description="Exit", alt_names=["q", "quit"])) + + names = registry.get_names() + assert "help" in names + assert "h" in names + assert "exit" in names + assert "q" in names + assert "quit" in names + + +# ============================================================================ +# SlashCommandParser Tests +# ============================================================================ + +class TestSlashCommandParser: + """Tests for SlashCommandParser.""" + + @pytest.fixture + def parser(self): + """Create a parser with test commands.""" + registry = SlashCommandRegistry() + registry.register(SlashCommand(name="help", description="Help", alt_names=["h"])) + registry.register(SlashCommand(name="cost", description="Cost")) + registry.register(SlashCommand( + name="memory", + description="Memory", + sub_commands=[ + SlashCommand(name="add", description="Add memory"), + SlashCommand(name="list", description="List memories"), + ] + )) + return SlashCommandParser(registry) + + def test_is_slash_command(self, parser): + """Test slash command detection.""" + assert parser.is_slash_command("/help") is True + assert parser.is_slash_command("/cost") is True + assert parser.is_slash_command("help") is False + assert parser.is_slash_command("") is False + + def test_is_not_code_comment(self, parser): + """Test that code comments are not detected as commands.""" + assert parser.is_slash_command("// This is a comment") is False + assert parser.is_slash_command("/* Block comment */") is False + + def test_parse_simple_command(self, parser): + """Test parsing a simple command.""" + result = parser.parse("/help") + + assert result.command is not None + assert result.command.name == "help" + assert result.args == "" + assert result.canonical_path == ["help"] + assert result.error is None + + def test_parse_command_with_args(self, parser): + """Test parsing a command with arguments.""" + result = parser.parse("/help model") + + assert result.command is not None + assert result.command.name == "help" + assert result.args == "model" + + def test_parse_command_by_alias(self, parser): + """Test parsing a command by alias.""" + result = parser.parse("/h") + + assert result.command is not None + assert result.command.name == "help" + + def test_parse_subcommand(self, parser): + """Test parsing a sub-command.""" + result = parser.parse("/memory add some data") + + assert result.command is not None + assert result.command.name == "add" + assert result.args == "some data" + assert result.canonical_path == ["memory", "add"] + + def test_parse_unknown_command(self, parser): + """Test parsing an unknown command.""" + result = parser.parse("/unknown") + + assert result.command is None + assert result.canonical_path == [] + + def test_parse_empty_command(self, parser): + """Test parsing empty command.""" + result = parser.parse("/") + + assert result.error == "Empty command" + + def test_parse_not_slash_command(self, parser): + """Test parsing non-slash command.""" + result = parser.parse("regular text") + + assert result.error == "Not a slash command" + + +# ============================================================================ +# Built-in Command Tests +# ============================================================================ + +class TestBuiltinCommands: + """Tests for built-in commands.""" + + @pytest.fixture + def context(self): + """Create a test context.""" + ctx = CommandContext( + total_tokens=1000, + total_cost=0.003, + prompt_count=5, + session_start_time=time.time() - 120 # 2 minutes ago + ) + ctx.config["command_registry"] = create_default_registry() + return ctx + + def test_cmd_help_general(self, context): + """Test general help command.""" + with patch('rich.console.Console.print'): + result = cmd_help(context, "") + + assert result["type"] == "help" + + def test_cmd_help_specific(self, context): + """Test help for specific command.""" + with patch('rich.console.Console.print'): + result = cmd_help(context, "cost") + + assert result["type"] == "help" + assert result["command"] == "cost" + + def test_cmd_help_unknown(self, context): + """Test help for unknown command.""" + with patch('rich.console.Console.print'): + result = cmd_help(context, "unknown_cmd") + + assert result["type"] == "error" + + def test_cmd_cost(self, context): + """Test cost command.""" + with patch('rich.console.Console.print'): + result = cmd_cost(context, "") + + assert result["type"] == "stats" + assert result["tokens"] == 1000 + assert result["cost"] == 0.003 + assert result["prompts"] == 5 + + def test_cmd_clear(self, context): + """Test clear command.""" + mock_agent = MagicMock() + context.agent = mock_agent + + with patch('rich.console.Console.print'): + result = cmd_clear(context, "") + + assert result["type"] == "clear" + mock_agent.clear_history.assert_called_once() + + def test_cmd_model_show(self, context): + """Test model command without args (show current).""" + mock_agent = MagicMock() + mock_agent.llm = MagicMock(model="gpt-4o") + context.agent = mock_agent + + with patch('rich.console.Console.print'): + result = cmd_model(context, "") + + assert result["type"] == "model_info" + assert result["model"] == "gpt-4o" + + def test_cmd_model_change(self, context): + """Test model command with args (change model).""" + mock_agent = MagicMock() + context.agent = mock_agent + + with patch('rich.console.Console.print'): + result = cmd_model(context, "gpt-4o-mini") + + assert result["type"] == "model_change" + assert result["model"] == "gpt-4o-mini" + + def test_cmd_tokens(self, context): + """Test tokens command.""" + with patch('rich.console.Console.print'): + result = cmd_tokens(context, "") + + assert result["type"] == "tokens" + assert result["total"] == 1000 + + def test_cmd_plan_with_args(self, context): + """Test plan command with task.""" + with patch('rich.console.Console.print'): + result = cmd_plan(context, "build a REST API") + + assert result["type"] == "submit_prompt" + assert "build a REST API" in result["content"] + + def test_cmd_plan_without_args(self, context): + """Test plan command without task.""" + with patch('rich.console.Console.print'): + result = cmd_plan(context, "") + + assert result["type"] == "help" + + def test_cmd_exit(self, context): + """Test exit command.""" + with patch('rich.console.Console.print'): + result = cmd_exit(context, "") + + assert result["type"] == "exit" + + +# ============================================================================ +# SlashCommandHandler Tests +# ============================================================================ + +class TestSlashCommandHandler: + """Tests for SlashCommandHandler.""" + + def test_handler_creation(self): + """Test handler creation.""" + handler = SlashCommandHandler() + assert handler.registry is not None + assert handler.parser is not None + + def test_is_command(self): + """Test command detection.""" + handler = SlashCommandHandler() + + assert handler.is_command("/help") is True + assert handler.is_command("/cost") is True + assert handler.is_command("regular text") is False + + def test_execute_help(self): + """Test executing help command.""" + handler = SlashCommandHandler() + + with patch('rich.console.Console.print'): + result = handler.execute("/help") + + assert result is not None + assert result["type"] == "help" + + def test_execute_unknown_command(self): + """Test executing unknown command.""" + handler = SlashCommandHandler() + + result = handler.execute("/unknown_command_xyz") + + assert result is not None + assert result["type"] == "error" + + def test_execute_non_command(self): + """Test executing non-command text.""" + handler = SlashCommandHandler() + + result = handler.execute("regular text") + + assert result is None + + def test_register_custom_command(self): + """Test registering a custom command.""" + handler = SlashCommandHandler() + + def custom_action(ctx, args): + return {"type": "custom", "args": args} + + handler.register_command(SlashCommand( + name="custom", + description="Custom command", + action=custom_action + )) + + result = handler.execute("/custom test args") + + assert result is not None + assert result["type"] == "custom" + assert result["args"] == "test args" + + def test_get_completions(self): + """Test command completions.""" + handler = SlashCommandHandler() + + completions = handler.get_completions("/he") + assert "/help" in completions + + completions = handler.get_completions("/co") + assert "/cost" in completions + assert "/commit" in completions + + completions = handler.get_completions("regular") + assert completions == [] + + def test_set_context(self): + """Test setting context.""" + handler = SlashCommandHandler() + ctx = CommandContext(total_tokens=500) + + handler.set_context(ctx) + + assert handler._context is ctx + assert "command_registry" in ctx.config + + +# ============================================================================ +# Default Registry Tests +# ============================================================================ + +class TestDefaultRegistry: + """Tests for default command registry.""" + + def test_default_registry_has_core_commands(self): + """Test that default registry has core commands.""" + registry = create_default_registry() + + # Core commands should exist + assert registry.get("help") is not None + assert registry.get("cost") is not None + assert registry.get("clear") is not None + assert registry.get("model") is not None + assert registry.get("tokens") is not None + assert registry.get("plan") is not None + assert registry.get("exit") is not None + + def test_default_registry_aliases(self): + """Test that default registry has aliases.""" + registry = create_default_registry() + + # Help aliases + assert registry.get("h") is not None + assert registry.get("?") is not None + + # Exit aliases + assert registry.get("q") is not None + assert registry.get("quit") is not None + + # Cost aliases + assert registry.get("usage") is not None + assert registry.get("stats") is not None + + +# ============================================================================ +# Integration Tests (require real execution) +# ============================================================================ + +class TestSlashCommandsIntegration: + """Integration tests for slash commands.""" + + def test_full_command_flow(self): + """Test full command execution flow.""" + handler = SlashCommandHandler(verbose=True) + + # Set up context + ctx = CommandContext( + total_tokens=2500, + total_cost=0.0075, + prompt_count=10, + session_start_time=time.time() - 300 + ) + handler.set_context(ctx) + + # Execute commands + with patch('rich.console.Console.print'): + help_result = handler.execute("/help") + cost_result = handler.execute("/cost") + tokens_result = handler.execute("/tokens") + + assert help_result["type"] == "help" + assert cost_result["type"] == "stats" + assert cost_result["tokens"] == 2500 + assert tokens_result["type"] == "tokens" + + def test_command_with_special_characters(self): + """Test commands with special characters in args.""" + handler = SlashCommandHandler() + + with patch('rich.console.Console.print'): + result = handler.execute("/plan Build a REST API with /users endpoint") + + assert result is not None + assert result["type"] == "submit_prompt" + assert "/users" in result["content"]