Compare commits

...
19 Commits
Author SHA1 Message Date
Mervin Praison deef65c4bb Release v2.3.47 2025-12-19 12:23:05 +00:00
Mervin Praison d51327effd Release v2.3.46 2025-12-19 12:22:29 +00:00
Mervin Praison 30ce8857e8 fix: suppress crewai.cli.config INFO log at startup 2025-12-19 12:18:23 +00:00
Mervin Praison 397874ccfc Release v2.3.45 2025-12-19 12:16:12 +00:00
Mervin Praison 2016902e9a docs: Add examples README with quick links to docs 2025-12-19 11:14:24 +00:00
Mervin Praison d2d19511da docs: Update README with Model Router info and fix links
- Rename 'Router Agent' to 'Model Router' for consistency
- Fix docs link to point to /features/model-router
- Add model-routing template CLI command
- Add link to Model Router documentation
2025-12-19 11:13:48 +00:00
Mervin Praison bc9579b497 Add Code category with individual examples for external CLI integrations
- Add Code category in Key Features table (like Workflows)
- Add individual examples for each external agent:
  - claude_code_example.py
  - gemini_cli_example.py
  - codex_cli_example.py
  - cursor_cli_example.py
  - external_agents_example.py (all integrations)
  - code_editing_example.py
- Link to individual docs pages for each integration
- Organize code-related features under Code category
2025-12-19 11:07:25 +00:00
Mervin Praison 3161389341 Add External Agents CLI documentation to README
- Add External Agents feature to Key Features table
- Update Code Editing docs link to /code/editing
- Add External Agents CLI section with examples
- Include Python API usage examples
- Add environment variables documentation
- Link to detailed docs at docs.praison.ai/code/external-agents
2025-12-19 11:01:24 +00:00
Mervin Praison b01eb2b8fa Release v2.3.44 2025-12-19 10:50:33 +00:00
Mervin Praison e609fb7425 Release v2.3.43 2025-12-19 10:32:33 +00:00
Mervin Praison 168ef98473 feat: Implement all remaining AutoAgents improvements (TDD)
TODO 2: Single-Agent Option
- Add SingleAgentStructure Pydantic model
- Add single_agent parameter to AutoGenerator and WorkflowAutoGenerator

TODO 3: LLM-Based Pattern Recommendation
- Add PatternRecommendation Pydantic model (pattern, reasoning, confidence)
- Add recommend_pattern_llm() method to WorkflowAutoGenerator

TODO 4: Validation Gates
- Add ValidationGate Pydantic model (criteria, pass_action, fail_action)
- Add optional gates field to WorkflowStructure

TODO 6: Pattern Support for AutoGenerator
- Add pattern parameter to AutoGenerator
- Add recommend_pattern() method to AutoGenerator
- Add pattern-specific guidance in get_user_content()

TODO 7: Merge Support for WorkflowAutoGenerator
- Add merge parameter to generate() method
- Add merge_with_existing_workflow() method

TODO 10: Framework Support for WorkflowAutoGenerator
- Add framework parameter to WorkflowAutoGenerator
- Default framework is 'praisonai'

Test Coverage:
- 51 tests pass (up from 34)
- 17 new tests added using TDD approach
- All tests verify new functionality
2025-12-19 10:29:40 +00:00
Mervin Praison ce2c2926f9 refactor: Add BaseAutoGenerator class and improve WorkflowAutoGenerator
- Create BaseAutoGenerator with shared functionality (DRY principle)
  - Lazy-loaded instructor client (LiteLLM or OpenAI fallback)
  - Environment variable handling for model/API configuration
  - analyze_complexity() static method for task complexity detection
  - get_available_tools() static method for tool list
  - AVAILABLE_TOOLS constant shared between generators

- Update AutoGenerator to inherit from BaseAutoGenerator
  - Removes ~40 lines of duplicate code
  - Maintains all existing functionality

- Update WorkflowAutoGenerator to inherit from BaseAutoGenerator
  - Removes ~40 lines of duplicate code
  - Add dynamic agent count based on task complexity
  - Add tools list to workflow prompts
  - Improved prompt with 3-step analysis

- Add 8 new unit tests for new functionality
  - TestBaseAutoGenerator: complexity analysis, tools list
  - TestWorkflowDynamicAgentCount: prompt improvements

All 34 tests pass with no performance impact.
2025-12-19 10:23:29 +00:00
Mervin Praison 0a569e2b46 docs: Add AutoAgents workflow patterns examples and update README
- Add autoagents_workflow_patterns.py with all 6 workflow patterns
- Update auto_agents_example.py with new features and CLI examples
- Add orchestrator-workers and evaluator-optimizer patterns to README CLI section
- Update pattern options table with all 6 patterns
- Add documentation links to examples
2025-12-19 09:54:05 +00:00
Mervin Praison 054aea8f66 Add MCP Registry badge and ignore MCP tokens 2025-12-18 22:18:33 +00:00
Mervin Praison 34f1293cfc Release v2.3.42 2025-12-18 22:13:08 +00:00
Mervin Praison 7c1e32afc0 upate 2025-12-18 22:11:02 +00:00
Mervin Praison 1125c835ec Release v2.3.41 2025-12-18 22:09:58 +00:00
Mervin Praison 889d04a843 Release v2.3.40 2025-12-18 22:07:54 +00:00
Mervin Praison fba811c477 Release v2.3.39 2025-12-18 21:54:19 +00:00
113 changed files with 23547 additions and 153 deletions

No files matched your search

+4 -1
View File
@@ -77,4 +77,7 @@ agents/praisonaiagents/praisonaiagents.egg-info
.qodo
CopilotKit*
.gemini
issue_details.json
issue_details.json
# MCP Registry tokens
.mcpregistry_*
+94 -4
View File
@@ -6,10 +6,13 @@
</picture>
</p>
<!-- mcp-name: io.github.MervinPraison/praisonai -->
<p align="center">
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://static.pepy.tech/badge/PraisonAI" alt="Total Downloads" /></a>
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://img.shields.io/github/v/release/MervinPraison/PraisonAI" alt="Latest Stable Version" /></a>
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="License" /></a>
<a href="https://registry.modelcontextprotocol.io/servers/io.github.MervinPraison/praisonai"><img src="https://img.shields.io/badge/MCP-Registry-blue" alt="MCP Registry" /></a>
</p>
<div align="center">
@@ -60,6 +63,7 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
- [Planning & Workflows](#planning--workflows)
- [Specialised Agents](#specialised-agents)
- [MCP Protocol](#mcp-protocol)
- [A2A & A2UI Protocols](#a2a--a2ui-protocols)
- [Safety & Control](#safety--control)
- [Developer Tools](#developer-tools)
- [📊 Process Types & Patterns](#-process-types--patterns)
@@ -187,7 +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) |
| **💻 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) |
@@ -205,6 +216,7 @@ npm install praisonai
| 🧠 Claude Memory Tool | [Example](#claude-memory-tool-cli) | [📖](https://docs.praison.ai/features/claude-memory-tool) |
| 💾 File-Based Memory | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
| 🔍 Built-in Search Tools | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/tools/tavily) |
| 🔎 Unified Web Search | [Example](src/praisonai-agents/examples/web_search_example.py) | [📖](https://docs.praison.ai/tools/web-search) |
| 📋 Planning Mode | [Example](examples/python/agents/planning-agent.py) | [📖](https://docs.praison.ai/features/planning-mode) |
| 🔧 Planning Tools | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
| 🧠 Planning Reasoning | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
@@ -224,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) |
@@ -246,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) |
---
@@ -1118,6 +1138,15 @@ praisonai workflow auto "Research AI trends" --pattern parallel
# Generate routing workflow (classifier routes to specialists)
praisonai workflow auto "Build a chatbot" --pattern routing
# Generate orchestrator-workers workflow (central orchestrator delegates)
praisonai workflow auto "Comprehensive market analysis" --pattern orchestrator-workers
# Generate evaluator-optimizer workflow (iterative refinement)
praisonai workflow auto "Write and refine article" --pattern evaluator-optimizer
# Specify output file
praisonai workflow auto "Build a chatbot" --pattern routing
# Specify output file
praisonai workflow auto "Research AI" --pattern sequential --output my_workflow.yaml
```
@@ -1127,7 +1156,7 @@ praisonai workflow auto "Research AI" --pattern sequential --output my_workflow.
|------|-------------|
| `--workflow-var key=value` | Set workflow variable (can be repeated) |
| `--var key=value` | Set variable for YAML workflows |
| `--pattern <pattern>` | Pattern for auto-generation (sequential, parallel, routing, loop) |
| `--pattern <pattern>` | Pattern for auto-generation (sequential, parallel, routing, loop, orchestrator-workers, evaluator-optimizer) |
| `--output <file>` | Output file for auto-generation |
| `--llm <model>` | LLM model (e.g., openai/gpt-4o-mini) |
| `--tools <tools>` | Tools (comma-separated, e.g., tavily) |
@@ -1359,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
@@ -1449,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
@@ -2251,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
@@ -2260,6 +2343,13 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
- 🔐 **MCP Security** - Origin validation, DNS rebinding prevention, secure sessions
- 🔄 **MCP Resumability** - SSE stream recovery via Last-Event-ID
### A2A & A2UI Protocols
- 🔗 **A2A Protocol** - Agent-to-Agent communication for inter-agent collaboration
- 🖼️ **A2UI Protocol** - Agent-to-User Interface for generating rich UIs from agents
- 📋 **UI Templates** - ChatTemplate, ListTemplate, FormTemplate, DashboardTemplate
- 🔧 **Surface Builder** - Fluent API for building declarative UIs
### Safety & Control
- 🤝 **Agent Handoffs** - Transfer context between specialised agents
+1 -1
View File
@@ -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.38" \
"praisonai>=2.3.47" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -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.38" \
"praisonai>=2.3.47" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -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.38" \
"praisonai>=2.3.47" \
"praisonai[ui]" \
"praisonai[crewai]"
+49
View File
@@ -0,0 +1,49 @@
# PraisonAI Examples
This folder contains examples for PraisonAI. For detailed documentation, visit [docs.praison.ai](https://docs.praison.ai).
## Structure
```
examples/
├── python/ # Python examples
│ ├── agents/ # Agent examples (single, multi, router, etc.)
│ ├── workflows/ # Workflow patterns (routing, parallel, loop)
│ ├── tools/ # Custom tools examples
│ ├── mcp/ # MCP protocol examples
│ ├── memory/ # Memory and sessions
│ ├── code/ # Code editing and external CLI tools
│ └── ...
├── yaml/ # YAML workflow examples
└── cookbooks/ # Complete use-case examples
```
## Quick Links
| Category | Examples | Docs |
|----------|----------|------|
| **Agents** | [python/agents/](python/agents/) | [📖](https://docs.praison.ai/concepts/agents) |
| **Workflows** | [python/workflows/](python/workflows/) | [📖](https://docs.praison.ai/features/workflows) |
| **Model Router** | [python/agents/router-agent-cost-optimization.py](python/agents/router-agent-cost-optimization.py) | [📖](https://docs.praison.ai/features/model-router) |
| **MCP** | [python/mcp/](python/mcp/) | [📖](https://docs.praison.ai/mcp) |
| **Memory** | [python/memory/](python/memory/) | [📖](https://docs.praison.ai/concepts/memory) |
| **Tools** | [python/tools/](python/tools/) | [📖](https://docs.praison.ai/tools) |
| **Code** | [python/code/](python/code/) | [📖](https://docs.praison.ai/code) |
| **YAML** | [yaml/](yaml/) | [📖](https://docs.praison.ai/features/yaml-workflows) |
## Running Examples
```bash
# Install PraisonAI
pip install praisonai
# Set API key
export OPENAI_API_KEY=your_key_here
# Run an example
python examples/python/agents/single-agent.py
```
## CLI Commands
See the main [README.md](../README.md#-cli--no-code-interface) for all CLI commands.
@@ -0,0 +1,132 @@
"""
AutoAgents Workflow Patterns Example
This example demonstrates the new workflow auto-generation features:
- Dynamic agent count based on task complexity
- Multiple workflow patterns (sequential, parallel, routing, orchestrator-workers, evaluator-optimizer)
- Pattern recommendation based on task keywords
Documentation: https://docs.praison.ai/features/autoagents
CLI Reference: https://docs.praison.ai/nocode/auto
"""
from praisonai.auto import AutoGenerator, WorkflowAutoGenerator
import os
import tempfile
# Ensure API key is set
if not os.environ.get("OPENAI_API_KEY"):
print("Please set OPENAI_API_KEY environment variable")
exit(1)
print("=" * 60)
print("AutoAgents Workflow Patterns Demo")
print("=" * 60)
# =============================================================================
# Example 1: Dynamic Agent Count (Simple vs Complex Tasks)
# =============================================================================
print("\n📊 Example 1: Dynamic Agent Count")
print("-" * 40)
# Simple task - should create 1-2 agents
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
simple_file = f.name
simple_generator = AutoGenerator(
topic="Write a haiku about spring",
agent_file=simple_file,
framework="praisonai"
)
print("Simple task: 'Write a haiku about spring'")
result = simple_generator.generate()
print(f"Generated: {result}")
# Complex task - should create 3-4 agents
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
complex_file = f.name
complex_generator = AutoGenerator(
topic="Research AI trends, analyze market data, write a comprehensive report with executive summary",
agent_file=complex_file,
framework="praisonai"
)
print("\nComplex task: 'Research AI trends, analyze market data...'")
result = complex_generator.generate()
print(f"Generated: {result}")
# =============================================================================
# Example 2: Pattern Recommendation
# =============================================================================
print("\n🎯 Example 2: Pattern Recommendation")
print("-" * 40)
test_tasks = [
("Write a blog post", "sequential"),
("Research from multiple sources simultaneously", "parallel"),
("Classify and route customer requests", "routing"),
("Comprehensive analysis and break down the problem", "orchestrator-workers"),
("Refine and improve content quality", "evaluator-optimizer"),
]
for task, expected in test_tasks:
generator = WorkflowAutoGenerator(topic=task)
recommended = generator.recommend_pattern()
status = "✅" if recommended == expected else "⚠️"
print(f"{status} Task: '{task[:40]}...'")
print(f" Recommended: {recommended}")
# =============================================================================
# Example 3: Generate Workflow with Specific Pattern
# =============================================================================
print("\n🔄 Example 3: Generate Workflows with Patterns")
print("-" * 40)
patterns = ["sequential", "parallel", "orchestrator-workers", "evaluator-optimizer"]
for pattern in patterns:
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
workflow_file = f.name
generator = WorkflowAutoGenerator(
topic="Create a market analysis report",
workflow_file=workflow_file
)
print(f"\n📋 Pattern: {pattern}")
result = generator.generate(pattern=pattern)
print(f" Generated: {result}")
# Read and show structure
import yaml
with open(result, 'r') as f:
data = yaml.safe_load(f)
print(f" Agents: {list(data.get('agents', {}).keys())}")
print(f" Steps: {len(data.get('steps', []))}")
# =============================================================================
# Example 4: CLI Usage (for reference)
# =============================================================================
print("\n💻 Example 4: CLI Commands (for reference)")
print("-" * 40)
print("""
# Auto-generate agents (dynamic count)
praisonai --auto "Write a haiku about spring"
# Auto-generate workflow with pattern
praisonai workflow auto "Research and write a report" --pattern sequential
praisonai workflow auto "Research from multiple sources" --pattern parallel
praisonai workflow auto "Comprehensive market analysis" --pattern orchestrator-workers
praisonai workflow auto "Refine content quality" --pattern evaluator-optimizer
# With output file
praisonai workflow auto "Customer support routing" --pattern routing --output support.yaml
""")
print("\n" + "=" * 60)
print("✅ AutoAgents Workflow Patterns Demo Complete!")
print("=" * 60)
print("\nDocumentation: https://docs.praison.ai/features/autoagents")
print("CLI Reference: https://docs.praison.ai/nocode/auto")
@@ -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}")
@@ -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()
@@ -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
@@ -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()
+24
View File
@@ -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")
@@ -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}")
@@ -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', '')}")
@@ -0,0 +1,34 @@
"""
Claude Code CLI Integration Example
Demonstrates how to use Claude Code CLI as an agent tool in PraisonAI.
"""
import asyncio
from praisonai.integrations import ClaudeCodeIntegration
async def main():
# Create Claude Code integration
claude = ClaudeCodeIntegration(
workspace=".",
output_format="json",
skip_permissions=True,
system_prompt="Be concise and follow best practices",
allowed_tools=["Read", "Write", "Bash"]
)
# Check availability
print(f"Claude CLI available: {claude.is_available}")
print(f"Claude SDK available: {claude.sdk_available}")
# Execute a coding task
result = await claude.execute("List all Python files in the current directory")
print(f"Result: {result}")
# Stream output
print("\nStreaming output:")
async for event in claude.stream("Explain the project structure"):
print(event)
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,60 @@
"""
Code Editing Tools Example
Demonstrates how to use PraisonAI's code editing tools with SEARCH/REPLACE diffs.
"""
from praisonai.code import (
set_workspace,
code_read_file,
code_write_file,
code_apply_diff,
code_list_files,
code_execute_command,
CODE_TOOLS
)
# Set workspace
set_workspace(".")
# List files
print("Files in current directory:")
files = code_list_files(".", recursive=False)
for f in files:
print(f" {f}")
# Read a file
print("\nReading example file:")
try:
content = code_read_file("README.md", start_line=1, end_line=10)
print(content)
except Exception as e:
print(f" File not found: {e}")
# Write a file
print("\nWriting test file:")
code_write_file("test_output.txt", "Hello from PraisonAI Code Tools!")
print(" Created test_output.txt")
# Apply a diff
print("\nApplying diff:")
diff = """
<<<<<<< SEARCH
Hello from PraisonAI Code Tools!
=======
Hello from PraisonAI Code Tools!
This line was added by a diff.
>>>>>>> REPLACE
"""
result = code_apply_diff("test_output.txt", diff)
print(f" {result}")
# Execute command
print("\nExecuting command:")
output = code_execute_command("cat test_output.txt")
print(f" {output}")
# Cleanup
import os
os.remove("test_output.txt")
print("\nCleanup complete.")
+36
View File
@@ -0,0 +1,36 @@
"""
Codex CLI Integration Example
Demonstrates how to use OpenAI Codex CLI as an agent tool in PraisonAI.
"""
import asyncio
from praisonai.integrations import CodexCLIIntegration
async def main():
# Create Codex CLI integration
codex = CodexCLIIntegration(
workspace=".",
full_auto=True,
sandbox="danger-full-access",
json_output=True
)
# Check availability
print(f"Codex CLI available: {codex.is_available}")
print(f"Full Auto: {codex.full_auto}")
print(f"Sandbox: {codex.sandbox}")
# Execute a coding task
result = await codex.execute("List files in the current directory")
print(f"Result: {result}")
# Stream JSON events
print("\nStreaming events:")
async for event in codex.stream("Explain the project"):
event_type = event.get("type", "")
if event_type == "item.completed":
print(f" {event}")
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,37 @@
"""
Cursor CLI Integration Example
Demonstrates how to use Cursor CLI as an agent tool in PraisonAI.
"""
import asyncio
from praisonai.integrations import CursorCLIIntegration
async def main():
# Create Cursor CLI integration
cursor = CursorCLIIntegration(
workspace=".",
output_format="json",
force=True,
model="gpt-5",
stream_partial=True
)
# Check availability
print(f"Cursor CLI available: {cursor.is_available}")
print(f"Force mode: {cursor.force}")
print(f"Model: {cursor.model}")
# Execute a coding task
result = await cursor.execute("List files in the current directory")
print(f"Result: {result}")
# Stream output
print("\nStreaming output:")
async for event in cursor.stream("Explain the project"):
content = event.get("content", "")
if content:
print(content, end="", flush=True)
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,58 @@
"""
External Agents Integration Example
Demonstrates how to use all external CLI integrations as agent tools.
"""
import asyncio
from praisonai.integrations import (
ClaudeCodeIntegration,
GeminiCLIIntegration,
CodexCLIIntegration,
CursorCLIIntegration,
get_available_integrations
)
async def main():
# Check availability of all integrations
print("=== External CLI Integrations ===\n")
availability = get_available_integrations()
print("Availability:")
for name, avail in availability.items():
status = "✅" if avail else "❌"
print(f" {status} {name}")
# Create integrations
print("\n--- Claude Code ---")
claude = ClaudeCodeIntegration(workspace=".", skip_permissions=True)
print(f"CLI: {claude.cli_command}")
print(f"Available: {claude.is_available}")
tool = claude.as_tool()
print(f"Tool: {tool.__name__}")
print("\n--- Gemini CLI ---")
gemini = GeminiCLIIntegration(workspace=".", model="gemini-2.5-flash")
print(f"CLI: {gemini.cli_command}")
print(f"Available: {gemini.is_available}")
tool = gemini.as_tool()
print(f"Tool: {tool.__name__}")
print("\n--- Codex CLI ---")
codex = CodexCLIIntegration(workspace=".", full_auto=True)
print(f"CLI: {codex.cli_command}")
print(f"Available: {codex.is_available}")
tool = codex.as_tool()
print(f"Tool: {tool.__name__}")
print("\n--- Cursor CLI ---")
cursor = CursorCLIIntegration(workspace=".", force=True)
print(f"CLI: {cursor.cli_command}")
print(f"Available: {cursor.is_available}")
tool = cursor.as_tool()
print(f"Tool: {tool.__name__}")
print("\n=== All integrations ready ===")
if __name__ == "__main__":
asyncio.run(main())
@@ -0,0 +1,33 @@
"""
Gemini CLI Integration Example
Demonstrates how to use Gemini CLI as an agent tool in PraisonAI.
"""
import asyncio
from praisonai.integrations import GeminiCLIIntegration
async def main():
# Create Gemini CLI integration
gemini = GeminiCLIIntegration(
workspace=".",
model="gemini-2.5-flash",
output_format="json",
include_directories=["../lib", "../docs"]
)
# Check availability
print(f"Gemini CLI available: {gemini.is_available}")
print(f"Model: {gemini.model}")
# Execute a coding task
result = await gemini.execute("Analyze the code structure")
print(f"Result: {result}")
# Execute with stats
result, stats = await gemini.execute_with_stats("Explain main.py")
print(f"Result: {result}")
print(f"Stats: {stats}")
if __name__ == "__main__":
asyncio.run(main())
+40 -3
View File
@@ -1,11 +1,48 @@
"""
AutoAgents Example
AutoAgents automatically creates and manages AI agents based on high-level instructions.
It features dynamic agent count based on task complexity.
Documentation: https://docs.praison.ai/features/autoagents
CLI Reference: https://docs.praison.ai/nocode/auto
New Features:
- Dynamic agent count (1-4 based on task complexity)
- Workflow patterns: sequential, parallel, routing, orchestrator-workers, evaluator-optimizer
- Pattern recommendation based on task keywords
- Tool preservation from LLM suggestions
"""
from praisonaiagents import AutoAgents
from praisonaiagents.tools import duckduckgo
# Basic usage - AutoAgents analyzes complexity and creates optimal agents
agents = AutoAgents(
instructions="Search for information about AI Agents",
tools=[duckduckgo],
process="sequential",
verbose=True
process="sequential", # or "hierarchical"
verbose=True,
max_agents=3 # Maximum number of agents to create
)
agents.start()
result = agents.start()
print(result)
# =============================================================================
# CLI Usage Examples (for reference)
# =============================================================================
#
# # Auto-generate agents (dynamic count based on complexity)
# praisonai --auto "Write a haiku about spring"
# praisonai --auto "Research AI trends, analyze data, write report"
#
# # Auto-generate workflow with specific pattern
# praisonai workflow auto "Research and write" --pattern sequential
# praisonai workflow auto "Research from multiple sources" --pattern parallel
# praisonai workflow auto "Comprehensive analysis" --pattern orchestrator-workers
# praisonai workflow auto "Refine content quality" --pattern evaluator-optimizer
#
# # With framework selection
# praisonai --framework crewai --auto "Create a movie script"
# praisonai --framework autogen --auto "Create a movie script"
@@ -0,0 +1,90 @@
"""
Unified Web Search Example
This example demonstrates the search_web tool that automatically tries
multiple search providers with fallback support.
Search Provider Priority:
1. Tavily (requires TAVILY_API_KEY + tavily-python)
2. Exa (requires EXA_API_KEY + exa_py)
3. You.com (requires YDC_API_KEY + youdotcom)
4. DuckDuckGo (requires duckduckgo_search package, no API key)
5. SearxNG (requires requests + running SearxNG instance)
Installation:
pip install praisonaiagents duckduckgo_search
Usage:
python web_search_example.py
"""
from praisonaiagents import Agent
from praisonaiagents.tools import search_web, get_available_providers
def main():
print("=" * 60)
print("Unified Web Search - Provider Status")
print("=" * 60)
# Check which providers are available
providers = get_available_providers()
for p in providers:
status = "✓ Available" if p["available"] else f"✗ {p['reason']}"
print(f" {p['name']:12} {status}")
print("\n" + "=" * 60)
print("Example 1: Basic Search")
print("=" * 60)
# Simple search - automatically uses best available provider
results = search_web("Python programming best practices", max_results=3)
if results and "error" not in results[0]:
print(f"\nFound {len(results)} results using {results[0].get('provider', 'unknown')}:\n")
for r in results:
print(f" Title: {r.get('title', 'N/A')[:60]}...")
print(f" URL: {r.get('url', 'N/A')}")
print()
else:
print(f"\nSearch failed: {results}")
print("=" * 60)
print("Example 2: Specify Providers")
print("=" * 60)
# Only try specific providers
results = search_web(
"machine learning tutorials",
max_results=3,
providers=["duckduckgo", "tavily"] # Try these in order
)
if results and "error" not in results[0]:
print(f"\nFound {len(results)} results using {results[0].get('provider', 'unknown')}:\n")
for r in results:
print(f" Title: {r.get('title', 'N/A')[:60]}...")
print(f" URL: {r.get('url', 'N/A')}")
print()
else:
print(f"\nSearch failed: {results}")
print("=" * 60)
print("Example 3: With PraisonAI Agent")
print("=" * 60)
# Use search_web as an agent tool
agent = Agent(
name="SearchAgent",
role="Web Researcher",
goal="Find and summarize information from the web",
instructions="Search the web and provide concise summaries of findings.",
tools=[search_web]
)
result = agent.start("What are the top 3 AI trends in 2025?")
print(f"\nAgent Response:\n{result}")
if __name__ == "__main__":
main()
@@ -29,7 +29,9 @@ SUPPRESSED_PATTERNS = [
"PydanticSerializationUnexpectedValue",
"Expected 9 fields but got 5 for type `Message`",
"Expected `StreamingChoices` but got `Choices`",
"serialized value may not be as expected"
"serialized value may not be as expected",
"Mixing V1 models and V2 models",
"Please upgrade `Settings` to V2"
]
@functools.wraps(_original_warn)
@@ -218,6 +218,10 @@ TOOL_MAPPINGS = {
'crawl4ai_extract_sync': ('.crawl4ai_tools', None),
'Crawl4AITools': ('.crawl4ai_tools', 'Crawl4AITools'),
'crawl4ai_tools': ('.crawl4ai_tools', None),
# Unified Web Search (auto-fallback across providers)
'search_web': ('.web_search', None),
'get_available_providers': ('.web_search', None),
}
_instances = {} # Cache for class instances
@@ -254,7 +258,8 @@ def __getattr__(name: str) -> Any:
'exa', 'exa_search', 'exa_search_contents', 'exa_find_similar', 'exa_answer',
'exa_search_async', 'exa_search_contents_async', 'exa_answer_async',
'crawl4ai', 'crawl4ai_many', 'crawl4ai_extract', 'crawl4ai_llm_extract',
'crawl4ai_sync', 'crawl4ai_extract_sync'
'crawl4ai_sync', 'crawl4ai_extract_sync',
'search_web', 'get_available_providers'
]:
return getattr(module, name)
if name in ['file_tools', 'pandas_tools', 'wikipedia_tools',
@@ -14,6 +14,7 @@ Usage:
context, cleaned_prompt = parser.process("@file:main.py explain this code")
"""
import os
import re
import logging
from pathlib import Path
@@ -43,10 +44,14 @@ class MentionsParser:
"url": re.compile(r'@url:(https?://[^\s]+)'),
}
# Default max file chars: 500K (~125K tokens) - fits GPT-4o (128K), Claude 3.5 (200K)
DEFAULT_MAX_FILE_CHARS = 500000
def __init__(
self,
workspace_path: Optional[str] = None,
verbose: int = 0
verbose: int = 0,
max_file_chars: Optional[int] = None
):
"""
Initialize MentionsParser.
@@ -54,10 +59,24 @@ class MentionsParser:
Args:
workspace_path: Path to workspace/project root
verbose: Verbosity level
max_file_chars: Maximum characters to include from files.
Default: 500000 (500K chars, ~125K tokens).
Set to 0 for no limit.
Can be overridden via PRAISON_MAX_FILE_CHARS env var.
"""
self.workspace_path = Path(workspace_path) if workspace_path else Path.cwd()
self.verbose = verbose
# Configure max file chars: constructor > env var > default
if max_file_chars is not None:
self.max_file_chars = max_file_chars
else:
env_val = os.environ.get('PRAISON_MAX_FILE_CHARS')
if env_val is not None:
self.max_file_chars = int(env_val)
else:
self.max_file_chars = self.DEFAULT_MAX_FILE_CHARS
# Lazy-loaded managers
self._docs_manager = None
self._rules_manager = None
@@ -158,9 +177,15 @@ class MentionsParser:
content = full_path.read_text(encoding="utf-8")
# Limit content size
if len(content) > 50000:
content = content[:50000] + "\n... (truncated)"
# Limit content size (0 means no limit)
if self.max_file_chars > 0 and len(content) > self.max_file_chars:
original_size = len(content)
content = content[:self.max_file_chars] + "\n... (truncated)"
self._log(
f"File {file_path} truncated from {original_size:,} to {self.max_file_chars:,} chars. "
f"Set PRAISON_MAX_FILE_CHARS=0 for no limit.",
logging.WARNING
)
# Detect language for code fence
ext = full_path.suffix.lower()
@@ -253,9 +278,9 @@ class MentionsParser:
with urllib.request.urlopen(req, timeout=10) as response:
content = response.read().decode('utf-8', errors='ignore')
# Limit content size
if len(content) > 50000:
content = content[:50000] + "\n... (truncated)"
# Limit content size (0 means no limit)
if self.max_file_chars > 0 and len(content) > self.max_file_chars:
content = content[:self.max_file_chars] + "\n... (truncated)"
# Try to extract text from HTML
if '<html' in content.lower():
@@ -0,0 +1,329 @@
"""Unified web search tool with automatic fallback across multiple providers.
This module provides a single `search_web` function that automatically tries
multiple search providers in order, falling back to the next if one fails.
Search Provider Priority:
1. Tavily (requires TAVILY_API_KEY + tavily-python)
2. Exa (requires EXA_API_KEY + exa_py)
3. You.com (requires YDC_API_KEY + youdotcom)
4. DuckDuckGo (requires duckduckgo_search package, no API key)
5. SearxNG (requires requests + running SearxNG instance)
Usage:
from praisonaiagents.tools import search_web
# Simple search - automatically uses best available provider
results = search_web("AI news 2024")
# With max results
results = search_web("Python tutorials", max_results=10)
"""
from typing import List, Dict, Any, Optional
import logging
import os
from importlib import util
logger = logging.getLogger(__name__)
def _check_tavily() -> tuple[bool, Optional[str]]:
"""Check if Tavily is available."""
if util.find_spec("tavily") is None:
return False, "tavily-python package not installed"
if not os.environ.get("TAVILY_API_KEY"):
return False, "TAVILY_API_KEY not set"
return True, None
def _check_exa() -> tuple[bool, Optional[str]]:
"""Check if Exa is available."""
if util.find_spec("exa_py") is None:
return False, "exa_py package not installed"
if not os.environ.get("EXA_API_KEY"):
return False, "EXA_API_KEY not set"
return True, None
def _check_youdotcom() -> tuple[bool, Optional[str]]:
"""Check if You.com is available."""
if util.find_spec("youdotcom") is None:
return False, "youdotcom package not installed"
if not os.environ.get("YDC_API_KEY"):
return False, "YDC_API_KEY not set"
return True, None
def _check_duckduckgo() -> tuple[bool, Optional[str]]:
"""Check if DuckDuckGo is available."""
if util.find_spec("duckduckgo_search") is None:
return False, "duckduckgo_search package not installed"
return True, None
def _check_searxng() -> tuple[bool, Optional[str]]:
"""Check if SearxNG is available."""
if util.find_spec("requests") is None:
return False, "requests package not installed"
return True, None
def _search_tavily(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
"""Search using Tavily."""
from tavily import TavilyClient
client = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))
response = client.search(query=query, max_results=max_results)
results = []
for r in response.get("results", []):
results.append({
"title": r.get("title", ""),
"url": r.get("url", ""),
"snippet": r.get("content", ""),
"provider": "tavily"
})
return results
def _search_exa(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
"""Search using Exa."""
from exa_py import Exa
client = Exa(os.environ.get("EXA_API_KEY"))
response = client.search(query=query, num_results=max_results)
results = []
for r in response.results:
results.append({
"title": getattr(r, "title", "") or "",
"url": r.url,
"snippet": getattr(r, "text", "") or "",
"provider": "exa"
})
return results
def _search_youdotcom(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
"""Search using You.com."""
from youdotcom import You
client = You(api_key_auth=os.environ.get("YDC_API_KEY"))
response = client.search.unified(query=query, count=max_results)
results = []
# Handle response object
if hasattr(response, 'results'):
web_results = getattr(response.results, 'web', []) or []
for r in web_results[:max_results]:
results.append({
"title": getattr(r, "title", "") or "",
"url": getattr(r, "url", "") or "",
"snippet": getattr(r, "description", "") or "",
"provider": "youdotcom"
})
elif isinstance(response, dict):
web_results = response.get("results", {}).get("web", [])
for r in web_results[:max_results]:
results.append({
"title": r.get("title", ""),
"url": r.get("url", ""),
"snippet": r.get("description", ""),
"provider": "youdotcom"
})
return results
def _search_duckduckgo(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
"""Search using DuckDuckGo."""
from duckduckgo_search import DDGS
results = []
ddgs = DDGS()
for result in ddgs.text(keywords=query, max_results=max_results):
results.append({
"title": result.get("title", ""),
"url": result.get("href", ""),
"snippet": result.get("body", ""),
"provider": "duckduckgo"
})
return results
def _search_searxng(query: str, max_results: int = 5, searxng_url: Optional[str] = None) -> List[Dict[str, Any]]:
"""Search using SearxNG."""
import requests
url = searxng_url or os.environ.get("SEARXNG_URL", "http://localhost:32768/search")
params = {
'q': query,
'format': 'json',
'engines': 'google,bing,duckduckgo',
'safesearch': '1'
}
response = requests.get(url, params=params, timeout=10)
response.raise_for_status()
raw_results = response.json().get('results', [])
results = []
for result in raw_results[:max_results]:
results.append({
"title": result.get("title", ""),
"url": result.get("url", ""),
"snippet": result.get("content", ""),
"provider": "searxng"
})
return results
# Provider configuration: (name, check_func, search_func)
SEARCH_PROVIDERS = [
("tavily", _check_tavily, _search_tavily),
("exa", _check_exa, _search_exa),
("youdotcom", _check_youdotcom, _search_youdotcom),
("duckduckgo", _check_duckduckgo, _search_duckduckgo),
("searxng", _check_searxng, _search_searxng),
]
def search_web(
query: str,
max_results: int = 5,
providers: Optional[List[str]] = None,
searxng_url: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Search the web using multiple providers with automatic fallback.
Tries each search provider in order until one succeeds. For each provider,
it first checks if the API key is set (if required) and if the package is
installed, then attempts the search. If the search fails, it moves to the
next provider.
Args:
query: The search query string
max_results: Maximum number of results to return (default: 5)
providers: Optional list of provider names to try, in order.
If not specified, uses default order:
["tavily", "exa", "youdotcom", "duckduckgo", "searxng"]
searxng_url: Optional custom SearxNG instance URL
Returns:
List of search results, each containing:
- title: Result title
- url: Result URL
- snippet: Result description/snippet
- provider: Name of the provider that returned the result
If all providers fail, returns a list with a single error dict.
Example:
# Use default provider order
results = search_web("AI news 2024")
# Specify providers to try
results = search_web("Python tutorials", providers=["duckduckgo", "tavily"])
# With custom SearxNG URL
results = search_web("tech news", searxng_url="http://my-searxng:8080/search")
"""
errors = []
# Determine which providers to try
if providers:
# Filter to only valid provider names
provider_order = []
for name in providers:
for pname, check_fn, search_fn in SEARCH_PROVIDERS:
if pname == name.lower():
provider_order.append((pname, check_fn, search_fn))
break
else:
provider_order = SEARCH_PROVIDERS
for provider_name, check_func, search_func in provider_order:
# Step 1: Check if provider is available (API key + package)
is_available, check_error = check_func()
if not is_available:
logger.debug(f"Skipping {provider_name}: {check_error}")
errors.append(f"{provider_name}: {check_error}")
continue
# Step 2: Try to search
try:
logger.debug(f"Trying {provider_name} for query: {query}")
if provider_name == "searxng":
results = search_func(query, max_results, searxng_url)
else:
results = search_func(query, max_results)
if results:
logger.info(f"search_web: Successfully used {provider_name}")
return results
else:
errors.append(f"{provider_name}: No results returned")
except Exception as e:
error_msg = f"{provider_name}: {str(e)}"
logger.debug(f"Search failed with {provider_name}: {e}")
errors.append(error_msg)
continue
# All providers failed
error_summary = "; ".join(errors) if errors else "No search providers available"
logger.error(f"search_web: All providers failed - {error_summary}")
return [{"error": f"All search providers failed: {error_summary}"}]
def get_available_providers() -> List[Dict[str, Any]]:
"""Get list of available search providers and their status.
Returns:
List of dicts with provider info:
- name: Provider name
- available: Whether the provider is ready to use
- reason: Reason if not available
Example:
providers = get_available_providers()
for p in providers:
print(f"{p['name']}: {'✓' if p['available'] else '✗'} {p.get('reason', '')}")
"""
result = []
for name, check_func, _ in SEARCH_PROVIDERS:
is_available, error = check_func()
result.append({
"name": name,
"available": is_available,
"reason": error if not is_available else None
})
return result
if __name__ == "__main__":
# Example usage and provider status check
print("=" * 60)
print("Web Search Tool - Provider Status")
print("=" * 60)
providers = get_available_providers()
for p in providers:
status = "✓ Available" if p["available"] else f"✗ {p['reason']}"
print(f" {p['name']:12} {status}")
print("\n" + "=" * 60)
print("Testing search_web...")
print("=" * 60)
results = search_web("Python programming", max_results=3)
if results and "error" not in results[0]:
print(f"\nFound {len(results)} results using {results[0].get('provider', 'unknown')}:")
for r in results:
print(f"\n Title: {r.get('title', 'N/A')[:50]}...")
print(f" URL: {r.get('url', 'N/A')}")
else:
print(f"\nSearch failed: {results}")
@@ -0,0 +1,126 @@
"""
A2UI (Agent-to-User Interface) Protocol Integration for PraisonAI Agents
This module provides A2UI protocol support, enabling PraisonAI Agents
to generate rich, interactive UIs declaratively via JSON.
A2UI is Google's open standard that allows agents to "speak UI" by sending
declarative JSON describing UI components, which clients then render natively.
Usage:
from praisonaiagents.ui.a2ui import (
# Message types
CreateSurfaceMessage,
UpdateComponentsMessage,
UpdateDataModelMessage,
DeleteSurfaceMessage,
# Component types
TextComponent,
ImageComponent,
ButtonComponent,
CardComponent,
RowComponent,
ColumnComponent,
# Extension helpers
create_a2ui_part,
is_a2ui_part,
)
"""
from praisonaiagents.ui.a2ui.types import (
# Data binding
PathBinding,
resolve_string_or_path,
# Action types
ActionContext,
Action,
# Children types
ChildrenTemplate,
# Component types
TextComponent,
ImageComponent,
IconComponent,
ButtonComponent,
CardComponent,
RowComponent,
ColumnComponent,
ListComponent,
TextFieldComponent,
CheckBoxComponent,
SliderComponent,
DateTimeInputComponent,
# Message types
CreateSurfaceMessage,
UpdateComponentsMessage,
UpdateDataModelMessage,
DeleteSurfaceMessage,
# Data part
A2UIDataPart,
)
from praisonaiagents.ui.a2ui.extension import (
A2UI_MIME_TYPE,
A2UI_EXTENSION_URI,
STANDARD_CATALOG_ID,
create_a2ui_part,
is_a2ui_part,
get_a2ui_agent_extension,
AgentExtension,
)
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.agent import A2UIAgent
from praisonaiagents.ui.a2ui.templates import (
ChatTemplate,
ListTemplate,
FormTemplate,
DashboardTemplate,
)
__all__ = [
# Data binding
"PathBinding",
"resolve_string_or_path",
# Action types
"ActionContext",
"Action",
# Children types
"ChildrenTemplate",
# Component types
"TextComponent",
"ImageComponent",
"IconComponent",
"ButtonComponent",
"CardComponent",
"RowComponent",
"ColumnComponent",
"ListComponent",
"TextFieldComponent",
"CheckBoxComponent",
"SliderComponent",
"DateTimeInputComponent",
# Message types
"CreateSurfaceMessage",
"UpdateComponentsMessage",
"UpdateDataModelMessage",
"DeleteSurfaceMessage",
# Data part
"A2UIDataPart",
# Extension
"A2UI_MIME_TYPE",
"A2UI_EXTENSION_URI",
"STANDARD_CATALOG_ID",
"create_a2ui_part",
"is_a2ui_part",
"get_a2ui_agent_extension",
"AgentExtension",
# Surface builder
"Surface",
# Agent wrapper
"A2UIAgent",
# Templates
"ChatTemplate",
"ListTemplate",
"FormTemplate",
"DashboardTemplate",
]
@@ -0,0 +1,358 @@
"""
A2UI Agent Integration
Provides A2UIAgent wrapper for generating A2UI responses from PraisonAI agents.
"""
import uuid
from typing import Any, Dict, List, Optional, TYPE_CHECKING
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import (
PathBinding,
A2UIDataPart,
)
from praisonaiagents.ui.a2ui.extension import (
create_a2ui_part,
STANDARD_CATALOG_ID,
)
if TYPE_CHECKING:
from praisonaiagents import Agent
class A2UIAgent:
"""
A2UI wrapper for PraisonAI agents.
Provides methods to generate A2UI responses from agent interactions.
Example:
>>> from praisonaiagents import Agent
>>> from praisonaiagents.ui.a2ui import A2UIAgent
>>>
>>> agent = Agent(name="Assistant", role="Helper", goal="Help users")
>>> a2ui_agent = A2UIAgent(agent=agent)
>>>
>>> # Render a simple text response
>>> messages = a2ui_agent.render_text("Hello, World!")
>>>
>>> # Render a list of items
>>> messages = a2ui_agent.render_list(
... title="Results",
... items=[{"name": "Item 1"}, {"name": "Item 2"}]
... )
"""
def __init__(
self,
agent: "Agent",
surface_id: Optional[str] = None,
catalog_id: str = STANDARD_CATALOG_ID,
):
"""
Initialize A2UIAgent wrapper.
Args:
agent: The PraisonAI Agent to wrap
surface_id: Unique ID for the UI surface (auto-generated if not provided)
catalog_id: The component catalog to use
"""
self.agent = agent
self.surface_id = surface_id or f"surface-{uuid.uuid4().hex[:8]}"
self.catalog_id = catalog_id
self._surface: Optional[Surface] = None
def get_surface(self) -> Surface:
"""
Get the underlying Surface object.
Creates a new Surface if one doesn't exist.
Returns:
The Surface object
"""
if self._surface is None:
self._surface = Surface(
surface_id=self.surface_id,
catalog_id=self.catalog_id,
)
return self._surface
def _reset_surface(self) -> Surface:
"""Reset and return a new Surface."""
self._surface = Surface(
surface_id=self.surface_id,
catalog_id=self.catalog_id,
)
return self._surface
# =========================================================================
# Rendering methods
# =========================================================================
def render_text(
self,
text: str,
title: Optional[str] = None,
usage_hint: str = "body",
) -> List[Dict[str, Any]]:
"""
Render a simple text response as A2UI.
Args:
text: The text content to display
title: Optional title heading
usage_hint: Text style hint (h1, h2, body, caption)
Returns:
List of A2UI message dictionaries
"""
surface = self._reset_surface()
children = []
if title:
surface.text(id="title", text=title, usage_hint="h1")
children.append("title")
surface.text(id="content", text=text, usage_hint=usage_hint)
children.append("content")
# Add root column
surface.column(id="root", children=children)
return surface.to_messages()
def render_list(
self,
title: str,
items: List[Dict[str, Any]],
item_title_key: str = "title",
item_description_key: str = "description",
item_image_key: Optional[str] = None,
item_action: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""
Render a list of items as A2UI cards.
Args:
title: List title
items: List of item dictionaries
item_title_key: Key for item title in each dict
item_description_key: Key for item description
item_image_key: Optional key for item image URL
item_action: Optional action name for item buttons
Returns:
List of A2UI message dictionaries
"""
surface = self._reset_surface()
# Title
surface.text(id="list-title", text=title, usage_hint="h1")
# Create card for each item
card_ids = []
for i, item in enumerate(items):
card_id = f"card-{i}"
content_id = f"card-content-{i}"
title_id = f"item-title-{i}"
desc_id = f"item-desc-{i}"
# Item title
surface.text(
id=title_id,
text=item.get(item_title_key, f"Item {i+1}"),
usage_hint="h3"
)
# Item description
surface.text(
id=desc_id,
text=item.get(item_description_key, ""),
usage_hint="body"
)
content_children = [title_id, desc_id]
# Optional action button
if item_action:
btn_id = f"btn-{i}"
btn_text_id = f"btn-text-{i}"
surface.text(id=btn_text_id, text="Select")
surface.button(
id=btn_id,
child=btn_text_id,
action_name=item_action,
action_context=[{"key": "index", "value": i}],
primary=True
)
content_children.append(btn_id)
# Card content column
surface.column(id=content_id, children=content_children)
# Card
surface.card(id=card_id, child=content_id)
card_ids.append(card_id)
# List column
surface.column(id="item-list", children=card_ids)
# Root column
surface.column(id="root", children=["list-title", "item-list"])
# Set data
surface.set_data("items", items)
return surface.to_messages()
def render_form(
self,
title: str,
fields: List[Dict[str, str]],
submit_action: str = "submit",
submit_label: str = "Submit",
) -> List[Dict[str, Any]]:
"""
Render a form as A2UI.
Args:
title: Form title
fields: List of field definitions with id, label, type
submit_action: Action name for submit button
submit_label: Label for submit button
Returns:
List of A2UI message dictionaries
"""
surface = self._reset_surface()
# Title
surface.text(id="form-title", text=title, usage_hint="h1")
# Fields
field_ids = []
form_data = {}
for field in fields:
field_id = field.get("id", f"field-{len(field_ids)}")
label = field.get("label", field_id)
field_type = field.get("type", "text")
surface.text_field(
id=field_id,
label=label,
text=PathBinding(path=f"/{field_id}"),
field_type=field_type,
)
field_ids.append(field_id)
form_data[field_id] = field.get("default", "")
# Submit button
surface.text(id="submit-text", text=submit_label)
# Build action context from fields
action_context = [
{"key": f["id"], "value": {"path": f"/{f['id']}"}}
for f in fields
]
surface.button(
id="submit-btn",
child="submit-text",
action_name=submit_action,
action_context=action_context,
primary=True,
)
field_ids.append("submit-btn")
# Form column
surface.column(id="form-fields", children=field_ids)
# Root column
surface.column(id="root", children=["form-title", "form-fields"])
# Set initial form data
for key, value in form_data.items():
surface.set_data(key, value)
return surface.to_messages()
# =========================================================================
# Agent interaction methods
# =========================================================================
async def chat_async(self, message: str) -> str:
"""
Send a message to the agent and get a response.
Args:
message: The message to send
Returns:
The agent's response text
"""
# Use the agent's chat method
if hasattr(self.agent, 'chat'):
response = self.agent.chat(message)
return response
elif hasattr(self.agent, 'run'):
response = self.agent.run(message)
return response
else:
raise NotImplementedError("Agent does not have chat or run method")
def chat(self, message: str) -> str:
"""
Send a message to the agent and get a response (sync).
Args:
message: The message to send
Returns:
The agent's response text
"""
if hasattr(self.agent, 'chat'):
return self.agent.chat(message)
elif hasattr(self.agent, 'run'):
return self.agent.run(message)
else:
raise NotImplementedError("Agent does not have chat or run method")
# =========================================================================
# Export methods
# =========================================================================
def to_messages(self) -> List[Dict[str, Any]]:
"""
Get the current surface as A2UI messages.
Returns:
List of A2UI message dictionaries
"""
return self.get_surface().to_messages()
def to_json(self, indent: Optional[int] = None) -> str:
"""
Get the current surface as JSON string.
Args:
indent: JSON indentation level
Returns:
JSON string of A2UI messages
"""
return self.get_surface().to_json(indent=indent)
def to_a2a_part(self) -> A2UIDataPart:
"""
Get the current surface as an A2A DataPart.
Returns:
A2UIDataPart with A2UI MIME type
"""
messages = self.to_messages()
# For A2A, we send messages as a list
return create_a2ui_part({"messages": messages})
@@ -0,0 +1,115 @@
"""
A2UI Extension for A2A Protocol
Provides helpers for integrating A2UI with the A2A (Agent-to-Agent) protocol.
"""
from dataclasses import dataclass
from typing import Any, Dict, Optional
from praisonaiagents.ui.a2ui.types import A2UIDataPart
# =============================================================================
# Constants
# =============================================================================
A2UI_EXTENSION_URI = "https://a2ui.org/a2a-extension/a2ui/v0.9"
A2UI_MIME_TYPE = "application/json+a2ui"
STANDARD_CATALOG_ID = "https://raw.githubusercontent.com/google/A2UI/refs/heads/main/specification/0.9/json/standard_catalog_definition.json"
# =============================================================================
# Agent Extension Type
# =============================================================================
@dataclass
class AgentExtension:
"""
A2A Agent Extension configuration.
Describes an extension that an agent supports.
"""
uri: str
description: str = ""
params: Optional[Dict[str, Any]] = None
# =============================================================================
# Helper Functions
# =============================================================================
def create_a2ui_part(a2ui_data: Dict[str, Any]) -> A2UIDataPart:
"""
Create an A2UI data part for A2A integration.
Wraps A2UI data in an A2UIDataPart with the appropriate MIME type.
Args:
a2ui_data: The A2UI data dictionary (e.g., createSurface, updateComponents)
Returns:
An A2UIDataPart with the A2UI MIME type in metadata
Example:
>>> part = create_a2ui_part({"createSurface": {"surfaceId": "main", "catalogId": "..."}})
>>> part.metadata["mimeType"]
'application/json+a2ui'
"""
return A2UIDataPart(
data=a2ui_data,
metadata={"mimeType": A2UI_MIME_TYPE}
)
def is_a2ui_part(part: A2UIDataPart) -> bool:
"""
Check if a data part contains A2UI data.
Args:
part: The data part to check
Returns:
True if the part contains A2UI data (has A2UI MIME type), False otherwise
Example:
>>> part = create_a2ui_part({"createSurface": {...}})
>>> is_a2ui_part(part)
True
"""
if not isinstance(part, A2UIDataPart):
return False
if part.metadata is None:
return False
return part.metadata.get("mimeType") == A2UI_MIME_TYPE
def get_a2ui_agent_extension(
accepts_inline_custom_catalog: bool = False,
) -> AgentExtension:
"""
Create the A2UI AgentExtension configuration.
Returns an AgentExtension that can be added to an A2A Agent Card
to indicate A2UI support.
Args:
accepts_inline_custom_catalog: Whether the agent accepts inline custom catalogs
Returns:
The configured A2UI AgentExtension
Example:
>>> ext = get_a2ui_agent_extension()
>>> ext.uri
'https://a2ui.org/a2a-extension/a2ui/v0.9'
"""
params: Optional[Dict[str, Any]] = None
if accepts_inline_custom_catalog:
params = {"acceptsInlineCustomCatalog": True}
return AgentExtension(
uri=A2UI_EXTENSION_URI,
description="Provides agent-driven UI using the A2UI JSON format.",
params=params
)
@@ -0,0 +1,364 @@
"""
A2UI Surface Builder
Provides a fluent API for building A2UI surfaces with components and data.
"""
import json
from typing import Any, Dict, List, Optional
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
from praisonaiagents.ui.a2ui.types import (
Component,
TextComponent,
ImageComponent,
ButtonComponent,
CardComponent,
RowComponent,
ColumnComponent,
ListComponent,
TextFieldComponent,
Action,
ActionContext,
PathBinding,
CreateSurfaceMessage,
UpdateComponentsMessage,
UpdateDataModelMessage,
StringOrPath,
ChildrenType,
)
class Surface:
"""
A2UI Surface builder for creating rich agent-generated UIs.
Provides a fluent API for adding components and data, then
generating the A2UI JSON messages.
Example:
>>> surface = Surface(surface_id="main")
>>> surface.text(id="title", text="Hello World", usage_hint="h1")
>>> surface.button(id="submit", child="submit-text", action_name="submit")
>>> surface.set_data("title", "Hello World")
>>> messages = surface.to_messages()
"""
def __init__(
self,
surface_id: str,
catalog_id: str = STANDARD_CATALOG_ID,
):
"""
Initialize a new Surface.
Args:
surface_id: Unique identifier for this surface
catalog_id: The component catalog to use (defaults to standard catalog)
"""
self.surface_id = surface_id
self.catalog_id = catalog_id
self.components: List[Component] = []
self.data: Dict[str, Any] = {}
def add(self, component: Component) -> "Surface":
"""
Add a component to the surface.
Args:
component: The component to add
Returns:
self for method chaining
"""
self.components.append(component)
return self
def set_data(self, key: str, value: Any) -> "Surface":
"""
Set a value in the data model.
Args:
key: The key in the data model
value: The value to set
Returns:
self for method chaining
"""
self.data[key] = value
return self
# =========================================================================
# Helper methods for common components
# =========================================================================
def text(
self,
id: str,
text: StringOrPath,
usage_hint: Optional[str] = None,
weight: Optional[float] = None,
) -> "Surface":
"""
Add a Text component.
Args:
id: Component ID
text: Text content or path binding
usage_hint: Style hint (h1, h2, h3, body, caption)
weight: Flex weight
Returns:
self for method chaining
"""
return self.add(TextComponent(
id=id,
text=text,
usage_hint=usage_hint,
weight=weight,
))
def image(
self,
id: str,
url: StringOrPath,
fit: Optional[str] = None,
usage_hint: Optional[str] = None,
weight: Optional[float] = None,
) -> "Surface":
"""
Add an Image component.
Args:
id: Component ID
url: Image URL or path binding
fit: Image fit (contain, cover, fill, none, scale-down)
usage_hint: Size hint (icon, avatar, smallFeature, etc.)
weight: Flex weight
Returns:
self for method chaining
"""
return self.add(ImageComponent(
id=id,
url=url,
fit=fit,
usage_hint=usage_hint,
weight=weight,
))
def button(
self,
id: str,
child: str,
action_name: Optional[str] = None,
action_context: Optional[List[Dict[str, Any]]] = None,
primary: bool = False,
weight: Optional[float] = None,
) -> "Surface":
"""
Add a Button component.
Args:
id: Component ID
child: Child component ID (usually a Text component)
action_name: Name of the action to trigger
action_context: List of {key, value} dicts for action context
primary: Whether this is a primary button
weight: Flex weight
Returns:
self for method chaining
"""
action = None
if action_name:
context = []
if action_context:
for ctx in action_context:
value = ctx.get("value")
if isinstance(value, dict) and "path" in value:
value = PathBinding(path=value["path"])
context.append(ActionContext(key=ctx["key"], value=value))
action = Action(name=action_name, context=context)
return self.add(ButtonComponent(
id=id,
child=child,
action=action,
primary=primary,
weight=weight,
))
def card(
self,
id: str,
child: str,
weight: Optional[float] = None,
) -> "Surface":
"""
Add a Card component.
Args:
id: Component ID
child: Child component ID
weight: Flex weight
Returns:
self for method chaining
"""
return self.add(CardComponent(
id=id,
child=child,
weight=weight,
))
def row(
self,
id: str,
children: ChildrenType,
weight: Optional[float] = None,
) -> "Surface":
"""
Add a Row (horizontal layout) component.
Args:
id: Component ID
children: List of child IDs or ChildrenTemplate
weight: Flex weight
Returns:
self for method chaining
"""
return self.add(RowComponent(
id=id,
children=children,
weight=weight,
))
def column(
self,
id: str,
children: ChildrenType,
weight: Optional[float] = None,
) -> "Surface":
"""
Add a Column (vertical layout) component.
Args:
id: Component ID
children: List of child IDs or ChildrenTemplate
weight: Flex weight
Returns:
self for method chaining
"""
return self.add(ColumnComponent(
id=id,
children=children,
weight=weight,
))
def list(
self,
id: str,
children: ChildrenType,
direction: str = "vertical",
weight: Optional[float] = None,
) -> "Surface":
"""
Add a List component.
Args:
id: Component ID
children: List of child IDs or ChildrenTemplate for dynamic lists
direction: "vertical" or "horizontal"
weight: Flex weight
Returns:
self for method chaining
"""
return self.add(ListComponent(
id=id,
children=children,
direction=direction,
weight=weight,
))
def text_field(
self,
id: str,
label: StringOrPath,
text: StringOrPath,
field_type: str = "text",
weight: Optional[float] = None,
) -> "Surface":
"""
Add a TextField input component.
Args:
id: Component ID
label: Field label
text: Text value or path binding
field_type: Input type (text, number, email, etc.)
weight: Flex weight
Returns:
self for method chaining
"""
return self.add(TextFieldComponent(
id=id,
label=label,
text=text,
field_type=field_type,
weight=weight,
))
# =========================================================================
# Message generation
# =========================================================================
def to_messages(self) -> List[Dict[str, Any]]:
"""
Generate A2UI messages for this surface.
Returns:
List of A2UI message dictionaries
"""
messages: List[Dict[str, Any]] = []
# 1. CreateSurface message
create_msg = CreateSurfaceMessage(
surface_id=self.surface_id,
catalog_id=self.catalog_id,
)
messages.append(create_msg.to_dict())
# 2. UpdateComponents message (if we have components)
if self.components:
update_msg = UpdateComponentsMessage(
surface_id=self.surface_id,
components=self.components,
)
messages.append(update_msg.to_dict())
# 3. UpdateDataModel message (if we have data)
if self.data:
data_msg = UpdateDataModelMessage(
surface_id=self.surface_id,
value=self.data,
)
messages.append(data_msg.to_dict())
return messages
def to_json(self, indent: Optional[int] = None) -> str:
"""
Generate A2UI messages as a JSON string.
Args:
indent: JSON indentation level (None for compact)
Returns:
JSON string of A2UI messages
"""
return json.dumps(self.to_messages(), indent=indent)
@@ -0,0 +1,17 @@
"""
A2UI Templates
Pre-built UI templates for common use cases.
"""
from praisonaiagents.ui.a2ui.templates.chat import ChatTemplate
from praisonaiagents.ui.a2ui.templates.list import ListTemplate
from praisonaiagents.ui.a2ui.templates.form import FormTemplate
from praisonaiagents.ui.a2ui.templates.dashboard import DashboardTemplate
__all__ = [
"ChatTemplate",
"ListTemplate",
"FormTemplate",
"DashboardTemplate",
]
@@ -0,0 +1,122 @@
"""
Chat Template for A2UI
Provides a conversational chat UI template.
"""
from typing import Any, Dict, List, Optional
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
class ChatTemplate:
"""
Chat UI template for conversational interfaces.
Creates a chat-style UI with user and agent messages.
Example:
>>> template = ChatTemplate(surface_id="chat")
>>> template.add_user_message("Hello!")
>>> template.add_agent_message("Hi there! How can I help?")
>>> messages = template.to_messages()
"""
def __init__(
self,
surface_id: str = "chat",
catalog_id: str = STANDARD_CATALOG_ID,
title: Optional[str] = None,
):
"""
Initialize ChatTemplate.
Args:
surface_id: Unique ID for the surface
catalog_id: Component catalog to use
title: Optional title for the chat
"""
self.surface_id = surface_id
self.catalog_id = catalog_id
self.title = title
self.messages: List[Dict[str, Any]] = []
def add_user_message(self, content: str) -> "ChatTemplate":
"""
Add a user message.
Args:
content: Message content
Returns:
self for chaining
"""
self.messages.append({
"role": "user",
"content": content,
})
return self
def add_agent_message(self, content: str) -> "ChatTemplate":
"""
Add an agent message.
Args:
content: Message content
Returns:
self for chaining
"""
self.messages.append({
"role": "agent",
"content": content,
})
return self
def to_messages(self) -> List[Dict[str, Any]]:
"""
Generate A2UI messages for the chat.
Returns:
List of A2UI message dictionaries
"""
surface = Surface(
surface_id=self.surface_id,
catalog_id=self.catalog_id,
)
children = []
# Title if provided
if self.title:
surface.text(id="chat-title", text=self.title, usage_hint="h1")
children.append("chat-title")
# Message components
msg_ids = []
for i, msg in enumerate(self.messages):
msg_id = f"msg-{i}"
role = msg["role"]
content = msg["content"]
# Style based on role
if role == "user":
surface.text(id=msg_id, text=f"You: {content}", usage_hint="body")
else:
surface.text(id=msg_id, text=f"Agent: {content}", usage_hint="body")
msg_ids.append(msg_id)
# Messages column
if msg_ids:
surface.column(id="messages", children=msg_ids)
children.append("messages")
# Root column
surface.column(id="root", children=children)
# Set data
surface.set_data("messages", self.messages)
return surface.to_messages()
@@ -0,0 +1,128 @@
"""
Dashboard Template for A2UI
Provides a dashboard UI template with multiple panels.
"""
from typing import Any, Dict, List, Optional
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
class DashboardTemplate:
"""
Dashboard UI template with multiple panels.
Creates a dashboard layout with cards for each panel.
Example:
>>> template = DashboardTemplate(surface_id="dashboard", title="My Dashboard")
>>> template.add_panel(id="stats", title="Statistics", content="100 items")
>>> template.add_panel(id="chart", title="Chart", content="[Chart here]")
>>> messages = template.to_messages()
"""
def __init__(
self,
surface_id: str = "dashboard",
title: str = "Dashboard",
catalog_id: str = STANDARD_CATALOG_ID,
columns: int = 2,
):
"""
Initialize DashboardTemplate.
Args:
surface_id: Unique ID for the surface
title: Dashboard title
catalog_id: Component catalog to use
columns: Number of columns in the grid
"""
self.surface_id = surface_id
self.title = title
self.catalog_id = catalog_id
self.columns = columns
self.panels: List[Dict[str, Any]] = []
def add_panel(
self,
id: str,
title: str,
content: str,
weight: Optional[float] = None,
) -> "DashboardTemplate":
"""
Add a panel to the dashboard.
Args:
id: Panel ID
title: Panel title
content: Panel content text
weight: Optional flex weight
Returns:
self for chaining
"""
self.panels.append({
"id": id,
"title": title,
"content": content,
"weight": weight,
})
return self
def to_messages(self) -> List[Dict[str, Any]]:
"""
Generate A2UI messages for the dashboard.
Returns:
List of A2UI message dictionaries
"""
surface = Surface(
surface_id=self.surface_id,
catalog_id=self.catalog_id,
)
# Title
surface.text(id="dashboard-title", text=self.title, usage_hint="h1")
# Create cards for each panel
card_ids = []
for panel in self.panels:
panel_id = panel["id"]
content_id = f"{panel_id}-content"
title_id = f"{panel_id}-title"
text_id = f"{panel_id}-text"
# Panel title
surface.text(id=title_id, text=panel["title"], usage_hint="h3")
# Panel content
surface.text(id=text_id, text=panel["content"], usage_hint="body")
# Panel content column
surface.column(id=content_id, children=[title_id, text_id])
# Panel card
surface.card(id=panel_id, child=content_id, weight=panel.get("weight"))
card_ids.append(panel_id)
# Arrange panels in rows based on columns setting
row_ids = []
for i in range(0, len(card_ids), self.columns):
row_cards = card_ids[i:i + self.columns]
row_id = f"row-{i // self.columns}"
surface.row(id=row_id, children=row_cards)
row_ids.append(row_id)
# Panels container
surface.column(id="panels", children=row_ids)
# Root column
surface.column(id="root", children=["dashboard-title", "panels"])
# Set data
surface.set_data("panels", self.panels)
return surface.to_messages()
@@ -0,0 +1,205 @@
"""
Form Template for A2UI
Provides a form UI template for data collection.
"""
from typing import Any, Dict, List
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import PathBinding
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
class FormTemplate:
"""
Form UI template for data collection.
Creates a form with various input fields and a submit button.
Example:
>>> template = FormTemplate(surface_id="form", title="Contact")
>>> template.add_text_field(id="name", label="Name")
>>> template.add_email_field(id="email", label="Email")
>>> template.set_submit_action("submit_form", "Send")
>>> messages = template.to_messages()
"""
def __init__(
self,
surface_id: str = "form",
title: str = "Form",
catalog_id: str = STANDARD_CATALOG_ID,
):
"""
Initialize FormTemplate.
Args:
surface_id: Unique ID for the surface
title: Form title
catalog_id: Component catalog to use
"""
self.surface_id = surface_id
self.title = title
self.catalog_id = catalog_id
self.fields: List[Dict[str, Any]] = []
self.submit_action = "submit"
self.submit_label = "Submit"
def add_text_field(
self,
id: str,
label: str,
default: str = "",
required: bool = False,
) -> "FormTemplate":
"""
Add a text input field.
Args:
id: Field ID
label: Field label
default: Default value
required: Whether field is required
Returns:
self for chaining
"""
self.fields.append({
"id": id,
"label": label,
"type": "text",
"default": default,
"required": required,
})
return self
def add_number_field(
self,
id: str,
label: str,
default: float = 0,
required: bool = False,
) -> "FormTemplate":
"""
Add a number input field.
Args:
id: Field ID
label: Field label
default: Default value
required: Whether field is required
Returns:
self for chaining
"""
self.fields.append({
"id": id,
"label": label,
"type": "number",
"default": default,
"required": required,
})
return self
def add_email_field(
self,
id: str,
label: str,
default: str = "",
required: bool = False,
) -> "FormTemplate":
"""
Add an email input field.
Args:
id: Field ID
label: Field label
default: Default value
required: Whether field is required
Returns:
self for chaining
"""
self.fields.append({
"id": id,
"label": label,
"type": "email",
"default": default,
"required": required,
})
return self
def set_submit_action(
self,
action: str,
label: str = "Submit",
) -> "FormTemplate":
"""
Set the submit button action.
Args:
action: Action name
label: Button label
Returns:
self for chaining
"""
self.submit_action = action
self.submit_label = label
return self
def to_messages(self) -> List[Dict[str, Any]]:
"""
Generate A2UI messages for the form.
Returns:
List of A2UI message dictionaries
"""
surface = Surface(
surface_id=self.surface_id,
catalog_id=self.catalog_id,
)
# Title
surface.text(id="form-title", text=self.title, usage_hint="h1")
# Fields
field_ids = []
for field in self.fields:
field_id = field["id"]
surface.text_field(
id=field_id,
label=field["label"],
text=PathBinding(path=f"/{field_id}"),
field_type=field["type"],
)
field_ids.append(field_id)
# Set default value
surface.set_data(field_id, field.get("default", ""))
# Submit button
surface.text(id="submit-text", text=self.submit_label)
action_context = [
{"key": f["id"], "value": {"path": f"/{f['id']}"}}
for f in self.fields
]
surface.button(
id="submit-btn",
child="submit-text",
action_name=self.submit_action,
action_context=action_context,
primary=True,
)
field_ids.append("submit-btn")
# Form fields column
surface.column(id="fields", children=field_ids)
# Root column
surface.column(id="root", children=["form-title", "fields"])
return surface.to_messages()
@@ -0,0 +1,147 @@
"""
List Template for A2UI
Provides a list UI template for displaying items.
"""
from typing import Any, Dict, List, Optional
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
class ListTemplate:
"""
List UI template for displaying items.
Creates a list of cards with title, description, and optional image.
Example:
>>> template = ListTemplate(surface_id="list", title="Results")
>>> template.add_item(title="Item 1", description="First item")
>>> template.add_item(title="Item 2", description="Second item")
>>> messages = template.to_messages()
"""
def __init__(
self,
surface_id: str = "list",
title: str = "Items",
catalog_id: str = STANDARD_CATALOG_ID,
):
"""
Initialize ListTemplate.
Args:
surface_id: Unique ID for the surface
title: List title
catalog_id: Component catalog to use
"""
self.surface_id = surface_id
self.title = title
self.catalog_id = catalog_id
self.items: List[Dict[str, Any]] = []
def add_item(
self,
title: str,
description: str = "",
image_url: Optional[str] = None,
action: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None,
) -> "ListTemplate":
"""
Add an item to the list.
Args:
title: Item title
description: Item description
image_url: Optional image URL
action: Optional action name for item button
metadata: Optional additional metadata
Returns:
self for chaining
"""
item = {
"title": title,
"description": description,
}
if image_url:
item["image_url"] = image_url
if action:
item["action"] = action
if metadata:
item.update(metadata)
self.items.append(item)
return self
def to_messages(self) -> List[Dict[str, Any]]:
"""
Generate A2UI messages for the list.
Returns:
List of A2UI message dictionaries
"""
surface = Surface(
surface_id=self.surface_id,
catalog_id=self.catalog_id,
)
# Title
surface.text(id="list-title", text=self.title, usage_hint="h1")
# Create cards for each item
card_ids = []
for i, item in enumerate(self.items):
card_id = f"card-{i}"
content_id = f"content-{i}"
title_id = f"title-{i}"
desc_id = f"desc-{i}"
# Item title
surface.text(id=title_id, text=item["title"], usage_hint="h3")
# Item description
surface.text(id=desc_id, text=item.get("description", ""), usage_hint="body")
content_children = [title_id, desc_id]
# Optional image
if item.get("image_url"):
img_id = f"img-{i}"
surface.image(id=img_id, url=item["image_url"], usage_hint="smallFeature")
content_children.insert(0, img_id)
# Optional action button
if item.get("action"):
btn_id = f"btn-{i}"
btn_text_id = f"btn-text-{i}"
surface.text(id=btn_text_id, text="Select")
surface.button(
id=btn_id,
child=btn_text_id,
action_name=item["action"],
action_context=[{"key": "index", "value": i}],
primary=True,
)
content_children.append(btn_id)
# Card content
surface.column(id=content_id, children=content_children)
# Card
surface.card(id=card_id, child=content_id)
card_ids.append(card_id)
# Items column
surface.column(id="items", children=card_ids)
# Root column
surface.column(id="root", children=["list-title", "items"])
# Set data
surface.set_data("items", self.items)
return surface.to_messages()
@@ -0,0 +1,523 @@
"""
A2UI Protocol Types
Pydantic-style classes for the A2UI (Agent-to-User Interface) protocol.
Based on the A2UI Specification v0.9.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union
from enum import Enum
# =============================================================================
# Data Binding Types
# =============================================================================
@dataclass
class PathBinding:
"""
A reference to a value in the data model.
Used for data binding in A2UI components.
"""
path: str
def to_dict(self) -> Dict[str, str]:
return {"path": self.path}
# Type alias for values that can be literal or path-bound
StringOrPath = Union[str, PathBinding]
NumberOrPath = Union[int, float, PathBinding]
BooleanOrPath = Union[bool, PathBinding]
def resolve_string_or_path(value: StringOrPath) -> Union[str, Dict[str, str]]:
"""Resolve a StringOrPath to its serialized form."""
if isinstance(value, PathBinding):
return value.to_dict()
return value
def resolve_number_or_path(value: NumberOrPath) -> Union[int, float, Dict[str, str]]:
"""Resolve a NumberOrPath to its serialized form."""
if isinstance(value, PathBinding):
return value.to_dict()
return value
def resolve_boolean_or_path(value: BooleanOrPath) -> Union[bool, Dict[str, str]]:
"""Resolve a BooleanOrPath to its serialized form."""
if isinstance(value, PathBinding):
return value.to_dict()
return value
# =============================================================================
# Action Types
# =============================================================================
@dataclass
class ActionContext:
"""
A key-value pair for action context.
Used to pass data from UI to action handlers.
"""
key: str
value: Union[str, int, float, bool, PathBinding]
def to_dict(self) -> Dict[str, Any]:
if isinstance(self.value, PathBinding):
return {"key": self.key, "value": self.value.to_dict()}
return {"key": self.key, "value": self.value}
@dataclass
class Action:
"""
An action triggered by a UI component (e.g., button click).
"""
name: str
context: List[ActionContext] = field(default_factory=list)
def to_dict(self) -> Dict[str, Any]:
return {
"name": self.name,
"context": [c.to_dict() for c in self.context]
}
# =============================================================================
# Children Types
# =============================================================================
@dataclass
class ChildrenTemplate:
"""
A template for generating dynamic children from a data model list.
"""
component_id: str
path: str
def to_dict(self) -> Dict[str, str]:
return {
"componentId": self.component_id,
"path": self.path
}
# Type alias for children - either explicit list or template
ChildrenType = Union[List[str], ChildrenTemplate]
def resolve_children(children: ChildrenType) -> Union[List[str], Dict[str, str]]:
"""Resolve children to serialized form."""
if isinstance(children, ChildrenTemplate):
return children.to_dict()
return children
# =============================================================================
# Component Types
# =============================================================================
class UsageHint(str, Enum):
"""Text usage hints for styling."""
H1 = "h1"
H2 = "h2"
H3 = "h3"
H4 = "h4"
H5 = "h5"
CAPTION = "caption"
BODY = "body"
class ImageFit(str, Enum):
"""Image fit options."""
CONTAIN = "contain"
COVER = "cover"
FILL = "fill"
NONE = "none"
SCALE_DOWN = "scale-down"
class ImageUsageHint(str, Enum):
"""Image usage hints."""
ICON = "icon"
AVATAR = "avatar"
SMALL_FEATURE = "smallFeature"
MEDIUM_FEATURE = "mediumFeature"
LARGE_FEATURE = "largeFeature"
HEADER = "header"
@dataclass
class TextComponent:
"""Text display component."""
id: str
text: StringOrPath
usage_hint: Optional[str] = None
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Text",
"text": resolve_string_or_path(self.text)
}
if self.usage_hint:
result["usageHint"] = self.usage_hint
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class ImageComponent:
"""Image display component."""
id: str
url: StringOrPath
fit: Optional[str] = None
usage_hint: Optional[str] = None
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Image",
"url": resolve_string_or_path(self.url)
}
if self.fit:
result["fit"] = self.fit
if self.usage_hint:
result["usageHint"] = self.usage_hint
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class IconComponent:
"""Icon display component."""
id: str
name: StringOrPath
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Icon",
"name": resolve_string_or_path(self.name)
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class ButtonComponent:
"""Button component with action."""
id: str
child: str
action: Optional[Action] = None
primary: bool = False
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Button",
"child": self.child
}
if self.primary:
result["primary"] = True
if self.action:
result["action"] = self.action.to_dict()
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class CardComponent:
"""Card container component."""
id: str
child: str
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Card",
"child": self.child
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class RowComponent:
"""Horizontal layout component."""
id: str
children: ChildrenType
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Row",
"children": resolve_children(self.children)
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class ColumnComponent:
"""Vertical layout component."""
id: str
children: ChildrenType
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Column",
"children": resolve_children(self.children)
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class ListComponent:
"""List component for repeating items."""
id: str
children: ChildrenType
direction: str = "vertical"
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "List",
"direction": self.direction,
"children": resolve_children(self.children)
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class TextFieldComponent:
"""Text input field component."""
id: str
label: StringOrPath
text: StringOrPath
field_type: str = "text"
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "TextField",
"label": resolve_string_or_path(self.label),
"text": resolve_string_or_path(self.text),
"type": self.field_type
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class CheckBoxComponent:
"""Checkbox input component."""
id: str
label: StringOrPath
checked: BooleanOrPath = False
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "CheckBox",
"label": resolve_string_or_path(self.label),
"checked": resolve_boolean_or_path(self.checked)
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class SliderComponent:
"""Slider input component."""
id: str
min_value: NumberOrPath
max_value: NumberOrPath
value: NumberOrPath
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "Slider",
"min": resolve_number_or_path(self.min_value),
"max": resolve_number_or_path(self.max_value),
"value": resolve_number_or_path(self.value)
}
if self.weight is not None:
result["weight"] = self.weight
return result
@dataclass
class DateTimeInputComponent:
"""Date/time input component."""
id: str
label: StringOrPath
value: StringOrPath
enable_date: bool = True
enable_time: bool = False
weight: Optional[float] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"id": self.id,
"component": "DateTimeInput",
"label": resolve_string_or_path(self.label),
"value": resolve_string_or_path(self.value),
"enableDate": self.enable_date,
"enableTime": self.enable_time
}
if self.weight is not None:
result["weight"] = self.weight
return result
# Type alias for any component
Component = Union[
TextComponent,
ImageComponent,
IconComponent,
ButtonComponent,
CardComponent,
RowComponent,
ColumnComponent,
ListComponent,
TextFieldComponent,
CheckBoxComponent,
SliderComponent,
DateTimeInputComponent,
]
# =============================================================================
# Message Types
# =============================================================================
@dataclass
class CreateSurfaceMessage:
"""
Message to create a new UI surface.
Signals the client to create a new surface and begin rendering it.
"""
surface_id: str
catalog_id: str
def to_dict(self) -> Dict[str, Any]:
return {
"createSurface": {
"surfaceId": self.surface_id,
"catalogId": self.catalog_id
}
}
@dataclass
class UpdateComponentsMessage:
"""
Message to update surface components.
Updates a surface with a new set of components.
"""
surface_id: str
components: List[Component]
def to_dict(self) -> Dict[str, Any]:
return {
"updateComponents": {
"surfaceId": self.surface_id,
"components": [c.to_dict() for c in self.components]
}
}
@dataclass
class UpdateDataModelMessage:
"""
Message to update the data model.
Updates the data model for an existing surface.
"""
surface_id: str
value: Any
path: str = "/"
op: str = "replace"
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {
"updateDataModel": {
"surfaceId": self.surface_id,
"path": self.path,
"op": self.op,
"value": self.value
}
}
return result
@dataclass
class DeleteSurfaceMessage:
"""
Message to delete a surface.
Signals the client to delete the surface.
"""
surface_id: str
def to_dict(self) -> Dict[str, Any]:
return {
"deleteSurface": {
"surfaceId": self.surface_id
}
}
# =============================================================================
# Data Part Type
# =============================================================================
@dataclass
class A2UIDataPart:
"""
A2UI data part for A2A integration.
Wraps A2UI data with the appropriate MIME type.
"""
data: Dict[str, Any]
metadata: Optional[Dict[str, Any]] = None
def to_dict(self) -> Dict[str, Any]:
result: Dict[str, Any] = {"data": self.data}
if self.metadata:
result["metadata"] = self.metadata
return result
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "praisonaiagents"
version = "0.1.16"
version = "0.1.17"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
readme = "README.md"
requires-python = ">=3.10"
@@ -0,0 +1 @@
"""A2UI Tests Package."""
@@ -0,0 +1,201 @@
"""
Tests for A2UI Agent Integration
TDD: Write tests first, then implement the A2UIAgent class.
"""
import os
import pytest
# Skip tests if no API key
pytestmark = pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY"),
reason="OPENAI_API_KEY not set"
)
class TestA2UIAgent:
"""Test A2UIAgent wrapper class."""
def test_a2ui_agent_creation(self):
"""Test A2UIAgent creation with an agent."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
agent = Agent(
name="TestAgent",
role="Test assistant",
goal="Help with testing"
)
a2ui_agent = A2UIAgent(agent=agent)
assert a2ui_agent.agent is agent
assert a2ui_agent.surface_id is not None
def test_a2ui_agent_with_custom_surface_id(self):
"""Test A2UIAgent with custom surface ID."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
agent = Agent(name="TestAgent", role="Test", goal="Test")
a2ui_agent = A2UIAgent(agent=agent, surface_id="custom-surface")
assert a2ui_agent.surface_id == "custom-surface"
def test_render_text_response(self):
"""Test rendering a simple text response as A2UI."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
agent = Agent(name="TestAgent", role="Test", goal="Test")
a2ui_agent = A2UIAgent(agent=agent)
# Render a simple text response
messages = a2ui_agent.render_text("Hello, World!")
assert isinstance(messages, list)
assert len(messages) >= 2 # createSurface + updateComponents
# Check createSurface
assert "createSurface" in messages[0]
# Check updateComponents has a text component
update_msg = next((m for m in messages if "updateComponents" in m), None)
assert update_msg is not None
components = update_msg["updateComponents"]["components"]
assert any(c.get("component") == "Text" for c in components)
def test_render_list_response(self):
"""Test rendering a list of items as A2UI cards."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
agent = Agent(name="TestAgent", role="Test", goal="Test")
a2ui_agent = A2UIAgent(agent=agent)
items = [
{"title": "Item 1", "description": "First item"},
{"title": "Item 2", "description": "Second item"},
]
messages = a2ui_agent.render_list(
title="My Items",
items=items,
item_title_key="title",
item_description_key="description"
)
assert isinstance(messages, list)
assert len(messages) >= 2
def test_render_form(self):
"""Test rendering a form as A2UI."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
agent = Agent(name="TestAgent", role="Test", goal="Test")
a2ui_agent = A2UIAgent(agent=agent)
fields = [
{"id": "name", "label": "Name", "type": "text"},
{"id": "email", "label": "Email", "type": "text"},
]
messages = a2ui_agent.render_form(
title="Contact Form",
fields=fields,
submit_action="submit_form"
)
assert isinstance(messages, list)
assert len(messages) >= 2
# Check for TextField components
update_msg = next((m for m in messages if "updateComponents" in m), None)
assert update_msg is not None
components = update_msg["updateComponents"]["components"]
text_fields = [c for c in components if c.get("component") == "TextField"]
assert len(text_fields) >= 2
def test_get_surface(self):
"""Test getting the underlying Surface object."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
from praisonaiagents.ui.a2ui.surface import Surface
agent = Agent(name="TestAgent", role="Test", goal="Test")
a2ui_agent = A2UIAgent(agent=agent)
surface = a2ui_agent.get_surface()
assert isinstance(surface, Surface)
assert surface.surface_id == a2ui_agent.surface_id
class TestA2UIAgentWithLLM:
"""Test A2UIAgent with actual LLM calls."""
@pytest.mark.asyncio
async def test_chat_with_ui(self):
"""Test chat that returns A2UI response."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
agent = Agent(
name="UIAgent",
role="UI Generator",
goal="Generate helpful UI responses",
llm="gpt-4o-mini"
)
a2ui_agent = A2UIAgent(agent=agent)
# This should work with the agent's chat
response = await a2ui_agent.chat_async("Say hello")
assert response is not None
# Response should contain text
assert isinstance(response, str) or isinstance(response, dict)
class TestA2UIAgentIntegration:
"""Integration tests for A2UIAgent with A2A."""
def test_get_a2a_data_part(self):
"""Test getting A2UI as A2A DataPart."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
from praisonaiagents.ui.a2ui.extension import is_a2ui_part
agent = Agent(name="TestAgent", role="Test", goal="Test")
a2ui_agent = A2UIAgent(agent=agent)
# Render something
a2ui_agent.render_text("Hello")
# Get as A2A DataPart
part = a2ui_agent.to_a2a_part()
assert part is not None
assert is_a2ui_part(part)
def test_to_json(self):
"""Test getting A2UI as JSON string."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui.agent import A2UIAgent
import json
agent = Agent(name="TestAgent", role="Test", goal="Test")
a2ui_agent = A2UIAgent(agent=agent)
# Render something
a2ui_agent.render_text("Hello")
# Get as JSON
json_str = a2ui_agent.to_json()
# Should be valid JSON
parsed = json.loads(json_str)
assert isinstance(parsed, list)
@@ -0,0 +1,115 @@
"""
Tests for A2UI Extension Module
TDD: Write tests first, then implement the extension.
"""
class TestA2UIExtension:
"""Test A2UI extension helpers for A2A integration."""
def test_a2ui_mime_type(self):
"""Test A2UI MIME type constant."""
from praisonaiagents.ui.a2ui.extension import A2UI_MIME_TYPE
assert A2UI_MIME_TYPE == "application/json+a2ui"
def test_a2ui_extension_uri(self):
"""Test A2UI extension URI constant."""
from praisonaiagents.ui.a2ui.extension import A2UI_EXTENSION_URI
assert "a2ui" in A2UI_EXTENSION_URI.lower()
def test_create_a2ui_part(self):
"""Test create_a2ui_part wraps A2UI data in A2A DataPart."""
from praisonaiagents.ui.a2ui.extension import create_a2ui_part, A2UI_MIME_TYPE
a2ui_data = {
"createSurface": {
"surfaceId": "main",
"catalogId": "https://a2ui.dev/standard"
}
}
part = create_a2ui_part(a2ui_data)
# Should return a Part with DataPart root
assert part is not None
assert hasattr(part, 'data')
assert part.data == a2ui_data
assert part.metadata is not None
assert part.metadata.get("mimeType") == A2UI_MIME_TYPE
def test_is_a2ui_part_true(self):
"""Test is_a2ui_part returns True for A2UI parts."""
from praisonaiagents.ui.a2ui.extension import create_a2ui_part, is_a2ui_part
a2ui_data = {"createSurface": {"surfaceId": "main", "catalogId": "test"}}
part = create_a2ui_part(a2ui_data)
assert is_a2ui_part(part) is True
def test_is_a2ui_part_false(self):
"""Test is_a2ui_part returns False for non-A2UI parts."""
from praisonaiagents.ui.a2ui.extension import is_a2ui_part
from praisonaiagents.ui.a2ui.types import A2UIDataPart
# Create a regular data part without A2UI mime type
part = A2UIDataPart(data={"foo": "bar"}, metadata={"mimeType": "application/json"})
assert is_a2ui_part(part) is False
def test_get_a2ui_agent_extension(self):
"""Test get_a2ui_agent_extension returns AgentExtension config."""
from praisonaiagents.ui.a2ui.extension import get_a2ui_agent_extension
ext = get_a2ui_agent_extension()
assert ext is not None
assert hasattr(ext, 'uri')
assert "a2ui" in ext.uri.lower()
assert hasattr(ext, 'description')
def test_get_a2ui_agent_extension_with_inline_catalog(self):
"""Test get_a2ui_agent_extension with inline catalog support."""
from praisonaiagents.ui.a2ui.extension import get_a2ui_agent_extension
ext = get_a2ui_agent_extension(accepts_inline_custom_catalog=True)
assert ext.params is not None
assert ext.params.get("acceptsInlineCustomCatalog") is True
def test_standard_catalog_id(self):
"""Test standard catalog ID constant."""
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
assert "standard_catalog" in STANDARD_CATALOG_ID.lower() or "a2ui" in STANDARD_CATALOG_ID.lower()
class TestA2UIDataPart:
"""Test A2UI DataPart type."""
def test_a2ui_data_part_creation(self):
"""Test A2UIDataPart creation."""
from praisonaiagents.ui.a2ui.types import A2UIDataPart
part = A2UIDataPart(
data={"test": "value"},
metadata={"mimeType": "application/json+a2ui"}
)
assert part.data == {"test": "value"}
assert part.metadata["mimeType"] == "application/json+a2ui"
def test_a2ui_data_part_to_dict(self):
"""Test A2UIDataPart serialization."""
from praisonaiagents.ui.a2ui.types import A2UIDataPart
part = A2UIDataPart(
data={"createSurface": {"surfaceId": "main"}},
metadata={"mimeType": "application/json+a2ui"}
)
data = part.to_dict()
assert "data" in data
assert "metadata" in data
@@ -0,0 +1,225 @@
"""
End-to-End Integration Tests for A2UI
Tests the complete A2UI workflow from agent to JSON output.
"""
import json
import os
import pytest
class TestA2UIEndToEnd:
"""End-to-end integration tests."""
def test_complete_surface_workflow(self):
"""Test complete workflow: create surface, add components, generate JSON."""
from praisonaiagents.ui.a2ui import (
Surface,
PathBinding,
)
# Create surface
surface = Surface(surface_id="test-surface")
# Add components using fluent API
surface.text(id="title", text="Welcome", usage_hint="h1")
surface.text(id="subtitle", text=PathBinding(path="/subtitle"), usage_hint="h2")
surface.button(
id="action-btn",
child="btn-text",
action_name="do_action",
action_context=[{"key": "id", "value": "123"}],
primary=True
)
surface.text(id="btn-text", text="Click Me")
surface.column(id="root", children=["title", "subtitle", "action-btn"])
# Set data
surface.set_data("subtitle", "Hello World")
# Generate messages
messages = surface.to_messages()
# Verify structure
assert len(messages) == 3 # createSurface, updateComponents, updateDataModel
# Verify createSurface
assert "createSurface" in messages[0]
assert messages[0]["createSurface"]["surfaceId"] == "test-surface"
# Verify updateComponents
assert "updateComponents" in messages[1]
components = messages[1]["updateComponents"]["components"]
assert len(components) == 5 # title, subtitle, btn, btn-text, root
# Verify updateDataModel
assert "updateDataModel" in messages[2]
assert messages[2]["updateDataModel"]["value"]["subtitle"] == "Hello World"
# Verify JSON output
json_str = surface.to_json()
parsed = json.loads(json_str)
assert len(parsed) == 3
def test_a2ui_agent_render_workflow(self):
"""Test A2UIAgent rendering workflow."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui import A2UIAgent
agent = Agent(
name="TestAgent",
role="Test assistant",
goal="Help with testing"
)
a2ui_agent = A2UIAgent(agent=agent, surface_id="agent-surface")
# Render text
messages = a2ui_agent.render_text(
text="This is a test response",
title="Test Title"
)
# Verify structure
assert len(messages) >= 2
assert "createSurface" in messages[0]
# Verify JSON export
json_str = a2ui_agent.to_json()
parsed = json.loads(json_str)
assert isinstance(parsed, list)
def test_template_workflow(self):
"""Test template-based workflow."""
from praisonaiagents.ui.a2ui import (
ChatTemplate,
ListTemplate,
FormTemplate,
DashboardTemplate,
)
# Chat template
chat = ChatTemplate(surface_id="chat")
chat.add_user_message("Hello!")
chat.add_agent_message("Hi there!")
chat_messages = chat.to_messages()
assert len(chat_messages) >= 2
# List template
list_template = ListTemplate(surface_id="list", title="Items")
list_template.add_item(title="Item 1", description="First")
list_template.add_item(title="Item 2", description="Second")
list_messages = list_template.to_messages()
assert len(list_messages) >= 2
# Form template
form = FormTemplate(surface_id="form", title="Contact")
form.add_text_field(id="name", label="Name")
form.add_email_field(id="email", label="Email")
form.set_submit_action("submit", "Send")
form_messages = form.to_messages()
assert len(form_messages) >= 2
# Dashboard template
dashboard = DashboardTemplate(surface_id="dashboard", title="Dashboard")
dashboard.add_panel(id="stats", title="Stats", content="100 items")
dashboard.add_panel(id="chart", title="Chart", content="[Chart]")
dashboard_messages = dashboard.to_messages()
assert len(dashboard_messages) >= 2
def test_a2a_integration(self):
"""Test A2A protocol integration."""
from praisonaiagents.ui.a2ui import (
Surface,
create_a2ui_part,
is_a2ui_part,
A2UI_MIME_TYPE,
)
# Create surface
surface = Surface(surface_id="a2a-test")
surface.text(id="msg", text="Hello from A2A")
surface.column(id="root", children=["msg"])
# Get messages
messages = surface.to_messages()
# Wrap in A2A DataPart
part = create_a2ui_part({"messages": messages})
# Verify A2UI part
assert is_a2ui_part(part)
assert part.metadata["mimeType"] == A2UI_MIME_TYPE
assert "messages" in part.data
assert len(part.data["messages"]) >= 2
def test_data_binding_workflow(self):
"""Test data binding with path references."""
from praisonaiagents.ui.a2ui import (
Surface,
PathBinding,
ChildrenTemplate,
)
surface = Surface(surface_id="data-binding-test")
# Add components with data binding
surface.text(id="title", text=PathBinding(path="/title"), usage_hint="h1")
surface.text(id="item-template", text=PathBinding(path="/name"))
# Add list with template
surface.list(
id="items",
children=ChildrenTemplate(component_id="item-template", path="/items"),
direction="vertical"
)
surface.column(id="root", children=["title", "items"])
# Set data
surface.set_data("title", "My List")
surface.set_data("items", [
{"name": "Item 1"},
{"name": "Item 2"},
{"name": "Item 3"},
])
messages = surface.to_messages()
# Verify data model
data_msg = next((m for m in messages if "updateDataModel" in m), None)
assert data_msg is not None
assert data_msg["updateDataModel"]["value"]["title"] == "My List"
assert len(data_msg["updateDataModel"]["value"]["items"]) == 3
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY"),
reason="OPENAI_API_KEY not set"
)
class TestA2UIWithLLM:
"""Tests that require LLM API calls."""
def test_agent_chat_and_render(self):
"""Test agent chat followed by UI rendering."""
from praisonaiagents import Agent
from praisonaiagents.ui.a2ui import A2UIAgent
agent = Agent(
name="UIAgent",
role="UI Generator",
goal="Generate helpful responses",
llm="gpt-4o-mini"
)
a2ui_agent = A2UIAgent(agent=agent)
# Chat with agent
response = a2ui_agent.chat("What is 2+2?")
# Render response as UI
messages = a2ui_agent.render_text(response, title="Answer")
assert len(messages) >= 2
assert "createSurface" in messages[0]
@@ -0,0 +1,228 @@
"""
Tests for A2UI Surface Builder
TDD: Write tests first, then implement the Surface class.
"""
class TestSurface:
"""Test Surface builder class."""
def test_surface_creation(self):
"""Test Surface creation with ID and catalog."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
assert surface.surface_id == "main"
assert surface.catalog_id is not None # Should have default
def test_surface_with_custom_catalog(self):
"""Test Surface with custom catalog ID."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(
surface_id="main",
catalog_id="https://custom.catalog/v1"
)
assert surface.catalog_id == "https://custom.catalog/v1"
def test_add_component(self):
"""Test adding a component to the surface."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent
surface = Surface(surface_id="main")
text = TextComponent(id="title", text="Hello World", usage_hint="h1")
surface.add(text)
assert len(surface.components) == 1
assert surface.components[0].id == "title"
def test_add_multiple_components(self):
"""Test adding multiple components."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent, ButtonComponent
surface = Surface(surface_id="main")
surface.add(TextComponent(id="title", text="Hello"))
surface.add(ButtonComponent(id="btn", child="btn-text"))
assert len(surface.components) == 2
def test_set_data(self):
"""Test setting data model values."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
surface.set_data("title", "Hello World")
surface.set_data("count", 42)
assert surface.data["title"] == "Hello World"
assert surface.data["count"] == 42
def test_set_nested_data(self):
"""Test setting nested data model values."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
surface.set_data("user", {"name": "John", "age": 30})
assert surface.data["user"]["name"] == "John"
def test_to_messages_creates_surface(self):
"""Test to_messages generates CreateSurface message."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent
surface = Surface(surface_id="main")
surface.add(TextComponent(id="title", text="Hello"))
messages = surface.to_messages()
assert len(messages) >= 1
assert "createSurface" in messages[0]
assert messages[0]["createSurface"]["surfaceId"] == "main"
def test_to_messages_includes_components(self):
"""Test to_messages includes UpdateComponents message."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent
surface = Surface(surface_id="main")
surface.add(TextComponent(id="title", text="Hello"))
messages = surface.to_messages()
# Should have createSurface and updateComponents
assert len(messages) >= 2
update_msg = next((m for m in messages if "updateComponents" in m), None)
assert update_msg is not None
assert len(update_msg["updateComponents"]["components"]) == 1
def test_to_messages_includes_data_model(self):
"""Test to_messages includes UpdateDataModel message when data is set."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent
surface = Surface(surface_id="main")
surface.add(TextComponent(id="title", text="Hello"))
surface.set_data("title", "Hello World")
messages = surface.to_messages()
# Should have createSurface, updateComponents, and updateDataModel
assert len(messages) == 3
data_msg = next((m for m in messages if "updateDataModel" in m), None)
assert data_msg is not None
assert data_msg["updateDataModel"]["value"]["title"] == "Hello World"
def test_to_json(self):
"""Test to_json returns JSON string."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent
import json
surface = Surface(surface_id="main")
surface.add(TextComponent(id="title", text="Hello"))
json_str = surface.to_json()
# Should be valid JSON
parsed = json.loads(json_str)
assert isinstance(parsed, list)
assert len(parsed) >= 2
class TestSurfaceChaining:
"""Test Surface method chaining."""
def test_add_returns_self(self):
"""Test add() returns self for chaining."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent
surface = Surface(surface_id="main")
result = surface.add(TextComponent(id="title", text="Hello"))
assert result is surface
def test_set_data_returns_self(self):
"""Test set_data() returns self for chaining."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
result = surface.set_data("key", "value")
assert result is surface
def test_fluent_api(self):
"""Test fluent API for building surfaces."""
from praisonaiagents.ui.a2ui.surface import Surface
from praisonaiagents.ui.a2ui.types import TextComponent, ButtonComponent
surface = (
Surface(surface_id="main")
.add(TextComponent(id="title", text="Hello"))
.add(ButtonComponent(id="btn", child="btn-text"))
.set_data("title", "Hello World")
)
assert len(surface.components) == 2
assert surface.data["title"] == "Hello World"
class TestSurfaceHelpers:
"""Test Surface helper methods."""
def test_text_helper(self):
"""Test text() helper method."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
surface.text(id="title", text="Hello", usage_hint="h1")
assert len(surface.components) == 1
assert surface.components[0].id == "title"
def test_button_helper(self):
"""Test button() helper method."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
surface.button(id="submit", child="submit-text", action_name="submit")
assert len(surface.components) == 1
assert surface.components[0].id == "submit"
def test_card_helper(self):
"""Test card() helper method."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
surface.card(id="card-1", child="card-content")
assert len(surface.components) == 1
assert surface.components[0].id == "card-1"
def test_row_helper(self):
"""Test row() helper method."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
surface.row(id="row-1", children=["col-1", "col-2"])
assert len(surface.components) == 1
assert surface.components[0].id == "row-1"
def test_column_helper(self):
"""Test column() helper method."""
from praisonaiagents.ui.a2ui.surface import Surface
surface = Surface(surface_id="main")
surface.column(id="col-1", children=["item-1", "item-2"])
assert len(surface.components) == 1
assert surface.components[0].id == "col-1"
@@ -0,0 +1,210 @@
"""
Tests for A2UI Templates
TDD: Write tests first, then implement the templates.
"""
class TestChatTemplate:
"""Test ChatTemplate for conversational UIs."""
def test_chat_template_creation(self):
"""Test ChatTemplate creation."""
from praisonaiagents.ui.a2ui.templates import ChatTemplate
template = ChatTemplate(surface_id="chat")
assert template.surface_id == "chat"
def test_add_user_message(self):
"""Test adding a user message."""
from praisonaiagents.ui.a2ui.templates import ChatTemplate
template = ChatTemplate(surface_id="chat")
template.add_user_message("Hello!")
assert len(template.messages) == 1
assert template.messages[0]["role"] == "user"
assert template.messages[0]["content"] == "Hello!"
def test_add_agent_message(self):
"""Test adding an agent message."""
from praisonaiagents.ui.a2ui.templates import ChatTemplate
template = ChatTemplate(surface_id="chat")
template.add_agent_message("Hi there!")
assert len(template.messages) == 1
assert template.messages[0]["role"] == "agent"
def test_to_messages(self):
"""Test generating A2UI messages."""
from praisonaiagents.ui.a2ui.templates import ChatTemplate
template = ChatTemplate(surface_id="chat")
template.add_user_message("Hello!")
template.add_agent_message("Hi there!")
messages = template.to_messages()
assert isinstance(messages, list)
assert len(messages) >= 2 # createSurface + updateComponents
class TestListTemplate:
"""Test ListTemplate for displaying lists of items."""
def test_list_template_creation(self):
"""Test ListTemplate creation."""
from praisonaiagents.ui.a2ui.templates import ListTemplate
template = ListTemplate(
surface_id="list",
title="My Items"
)
assert template.surface_id == "list"
assert template.title == "My Items"
def test_add_item(self):
"""Test adding an item to the list."""
from praisonaiagents.ui.a2ui.templates import ListTemplate
template = ListTemplate(surface_id="list", title="Items")
template.add_item(title="Item 1", description="First item")
assert len(template.items) == 1
assert template.items[0]["title"] == "Item 1"
def test_add_item_with_image(self):
"""Test adding an item with an image."""
from praisonaiagents.ui.a2ui.templates import ListTemplate
template = ListTemplate(surface_id="list", title="Items")
template.add_item(
title="Item 1",
description="First item",
image_url="https://example.com/image.jpg"
)
assert template.items[0]["image_url"] == "https://example.com/image.jpg"
def test_to_messages(self):
"""Test generating A2UI messages."""
from praisonaiagents.ui.a2ui.templates import ListTemplate
template = ListTemplate(surface_id="list", title="Items")
template.add_item(title="Item 1", description="First")
template.add_item(title="Item 2", description="Second")
messages = template.to_messages()
assert isinstance(messages, list)
assert len(messages) >= 2
class TestFormTemplate:
"""Test FormTemplate for data collection forms."""
def test_form_template_creation(self):
"""Test FormTemplate creation."""
from praisonaiagents.ui.a2ui.templates import FormTemplate
template = FormTemplate(
surface_id="form",
title="Contact Form"
)
assert template.surface_id == "form"
assert template.title == "Contact Form"
def test_add_text_field(self):
"""Test adding a text field."""
from praisonaiagents.ui.a2ui.templates import FormTemplate
template = FormTemplate(surface_id="form", title="Form")
template.add_text_field(id="name", label="Name")
assert len(template.fields) == 1
assert template.fields[0]["id"] == "name"
assert template.fields[0]["type"] == "text"
def test_add_number_field(self):
"""Test adding a number field."""
from praisonaiagents.ui.a2ui.templates import FormTemplate
template = FormTemplate(surface_id="form", title="Form")
template.add_number_field(id="age", label="Age")
assert template.fields[0]["type"] == "number"
def test_add_email_field(self):
"""Test adding an email field."""
from praisonaiagents.ui.a2ui.templates import FormTemplate
template = FormTemplate(surface_id="form", title="Form")
template.add_email_field(id="email", label="Email")
assert template.fields[0]["type"] == "email"
def test_set_submit_action(self):
"""Test setting submit action."""
from praisonaiagents.ui.a2ui.templates import FormTemplate
template = FormTemplate(surface_id="form", title="Form")
template.set_submit_action("submit_form", "Send")
assert template.submit_action == "submit_form"
assert template.submit_label == "Send"
def test_to_messages(self):
"""Test generating A2UI messages."""
from praisonaiagents.ui.a2ui.templates import FormTemplate
template = FormTemplate(surface_id="form", title="Form")
template.add_text_field(id="name", label="Name")
template.add_email_field(id="email", label="Email")
messages = template.to_messages()
assert isinstance(messages, list)
assert len(messages) >= 2
class TestDashboardTemplate:
"""Test DashboardTemplate for multi-panel dashboards."""
def test_dashboard_template_creation(self):
"""Test DashboardTemplate creation."""
from praisonaiagents.ui.a2ui.templates import DashboardTemplate
template = DashboardTemplate(
surface_id="dashboard",
title="My Dashboard"
)
assert template.surface_id == "dashboard"
assert template.title == "My Dashboard"
def test_add_panel(self):
"""Test adding a panel."""
from praisonaiagents.ui.a2ui.templates import DashboardTemplate
template = DashboardTemplate(surface_id="dashboard", title="Dashboard")
template.add_panel(id="panel-1", title="Panel 1", content="Some content")
assert len(template.panels) == 1
assert template.panels[0]["id"] == "panel-1"
def test_to_messages(self):
"""Test generating A2UI messages."""
from praisonaiagents.ui.a2ui.templates import DashboardTemplate
template = DashboardTemplate(surface_id="dashboard", title="Dashboard")
template.add_panel(id="panel-1", title="Panel 1", content="Content 1")
template.add_panel(id="panel-2", title="Panel 2", content="Content 2")
messages = template.to_messages()
assert isinstance(messages, list)
assert len(messages) >= 2
@@ -0,0 +1,318 @@
"""
Tests for A2UI Types Module
TDD: Write tests first, then implement the types.
"""
class TestA2UIMessageTypes:
"""Test A2UI message types: CreateSurface, UpdateComponents, UpdateDataModel, DeleteSurface."""
def test_create_surface_message(self):
"""Test CreateSurface message creation."""
from praisonaiagents.ui.a2ui.types import CreateSurfaceMessage
msg = CreateSurfaceMessage(
surface_id="main",
catalog_id="https://a2ui.dev/standard"
)
assert msg.surface_id == "main"
assert msg.catalog_id == "https://a2ui.dev/standard"
# Test serialization
data = msg.to_dict()
assert "createSurface" in data
assert data["createSurface"]["surfaceId"] == "main"
assert data["createSurface"]["catalogId"] == "https://a2ui.dev/standard"
def test_update_components_message(self):
"""Test UpdateComponents message creation."""
from praisonaiagents.ui.a2ui.types import UpdateComponentsMessage, TextComponent
text = TextComponent(id="title", text="Hello World", usage_hint="h1")
msg = UpdateComponentsMessage(
surface_id="main",
components=[text]
)
assert msg.surface_id == "main"
assert len(msg.components) == 1
# Test serialization
data = msg.to_dict()
assert "updateComponents" in data
assert data["updateComponents"]["surfaceId"] == "main"
assert len(data["updateComponents"]["components"]) == 1
def test_update_data_model_message(self):
"""Test UpdateDataModel message creation."""
from praisonaiagents.ui.a2ui.types import UpdateDataModelMessage
msg = UpdateDataModelMessage(
surface_id="main",
path="/",
op="replace",
value={"title": "Hello", "items": []}
)
assert msg.surface_id == "main"
assert msg.path == "/"
assert msg.op == "replace"
assert msg.value == {"title": "Hello", "items": []}
# Test serialization
data = msg.to_dict()
assert "updateDataModel" in data
assert data["updateDataModel"]["surfaceId"] == "main"
def test_delete_surface_message(self):
"""Test DeleteSurface message creation."""
from praisonaiagents.ui.a2ui.types import DeleteSurfaceMessage
msg = DeleteSurfaceMessage(surface_id="main")
assert msg.surface_id == "main"
# Test serialization
data = msg.to_dict()
assert "deleteSurface" in data
assert data["deleteSurface"]["surfaceId"] == "main"
class TestA2UIComponentTypes:
"""Test A2UI component types."""
def test_text_component(self):
"""Test Text component creation."""
from praisonaiagents.ui.a2ui.types import TextComponent
text = TextComponent(id="title", text="Hello World", usage_hint="h1")
assert text.id == "title"
assert text.text == "Hello World"
assert text.usage_hint == "h1"
# Test serialization
data = text.to_dict()
assert data["id"] == "title"
assert data["component"] == "Text"
assert data["text"] == "Hello World"
assert data["usageHint"] == "h1"
def test_text_component_with_path(self):
"""Test Text component with data binding path."""
from praisonaiagents.ui.a2ui.types import TextComponent, PathBinding
text = TextComponent(id="title", text=PathBinding(path="/title"))
data = text.to_dict()
assert data["text"] == {"path": "/title"}
def test_image_component(self):
"""Test Image component creation."""
from praisonaiagents.ui.a2ui.types import ImageComponent
img = ImageComponent(
id="hero",
url="https://example.com/image.jpg",
fit="cover",
usage_hint="header"
)
assert img.id == "hero"
assert img.url == "https://example.com/image.jpg"
assert img.fit == "cover"
data = img.to_dict()
assert data["component"] == "Image"
assert data["url"] == "https://example.com/image.jpg"
def test_button_component(self):
"""Test Button component creation."""
from praisonaiagents.ui.a2ui.types import ButtonComponent, Action, ActionContext
action = Action(
name="submit",
context=[
ActionContext(key="value", value="test")
]
)
btn = ButtonComponent(
id="submit-btn",
child="submit-text",
primary=True,
action=action
)
assert btn.id == "submit-btn"
assert btn.child == "submit-text"
assert btn.primary is True
data = btn.to_dict()
assert data["component"] == "Button"
assert data["action"]["name"] == "submit"
def test_card_component(self):
"""Test Card component creation."""
from praisonaiagents.ui.a2ui.types import CardComponent
card = CardComponent(id="card-1", child="card-content")
assert card.id == "card-1"
assert card.child == "card-content"
data = card.to_dict()
assert data["component"] == "Card"
assert data["child"] == "card-content"
def test_row_component(self):
"""Test Row component creation."""
from praisonaiagents.ui.a2ui.types import RowComponent
row = RowComponent(id="row-1", children=["col-1", "col-2"])
assert row.id == "row-1"
assert row.children == ["col-1", "col-2"]
data = row.to_dict()
assert data["component"] == "Row"
assert data["children"] == ["col-1", "col-2"]
def test_column_component(self):
"""Test Column component creation."""
from praisonaiagents.ui.a2ui.types import ColumnComponent
col = ColumnComponent(id="col-1", children=["item-1", "item-2"])
assert col.id == "col-1"
assert col.children == ["item-1", "item-2"]
data = col.to_dict()
assert data["component"] == "Column"
def test_list_component(self):
"""Test List component creation."""
from praisonaiagents.ui.a2ui.types import ListComponent, ChildrenTemplate
template = ChildrenTemplate(component_id="item-template", path="/items")
lst = ListComponent(id="item-list", direction="vertical", children=template)
assert lst.id == "item-list"
assert lst.direction == "vertical"
data = lst.to_dict()
assert data["component"] == "List"
assert data["children"]["componentId"] == "item-template"
assert data["children"]["path"] == "/items"
def test_text_field_component(self):
"""Test TextField component creation."""
from praisonaiagents.ui.a2ui.types import TextFieldComponent, PathBinding
field = TextFieldComponent(
id="name-field",
label="Name",
text=PathBinding(path="/name"),
field_type="text"
)
assert field.id == "name-field"
assert field.label == "Name"
data = field.to_dict()
assert data["component"] == "TextField"
assert data["label"] == "Name"
assert data["text"] == {"path": "/name"}
def test_checkbox_component(self):
"""Test CheckBox component creation."""
from praisonaiagents.ui.a2ui.types import CheckBoxComponent, PathBinding
cb = CheckBoxComponent(
id="agree-cb",
label="I agree",
checked=PathBinding(path="/agreed")
)
data = cb.to_dict()
assert data["component"] == "CheckBox"
assert data["label"] == "I agree"
def test_slider_component(self):
"""Test Slider component creation."""
from praisonaiagents.ui.a2ui.types import SliderComponent
slider = SliderComponent(
id="volume",
min_value=0,
max_value=100,
value=50
)
data = slider.to_dict()
assert data["component"] == "Slider"
assert data["min"] == 0
assert data["max"] == 100
class TestDataBindingTypes:
"""Test data binding types."""
def test_path_binding(self):
"""Test PathBinding for data model references."""
from praisonaiagents.ui.a2ui.types import PathBinding
path = PathBinding(path="/user/name")
assert path.path == "/user/name"
assert path.to_dict() == {"path": "/user/name"}
def test_string_or_path_literal(self):
"""Test StringOrPath with literal string."""
from praisonaiagents.ui.a2ui.types import resolve_string_or_path
result = resolve_string_or_path("Hello")
assert result == "Hello"
def test_string_or_path_binding(self):
"""Test StringOrPath with path binding."""
from praisonaiagents.ui.a2ui.types import resolve_string_or_path, PathBinding
result = resolve_string_or_path(PathBinding(path="/title"))
assert result == {"path": "/title"}
class TestActionTypes:
"""Test action types for button events."""
def test_action_creation(self):
"""Test Action creation."""
from praisonaiagents.ui.a2ui.types import Action, ActionContext
action = Action(
name="book_restaurant",
context=[
ActionContext(key="restaurantName", value="The Fancy Place"),
ActionContext(key="address", value="123 Main St")
]
)
assert action.name == "book_restaurant"
assert len(action.context) == 2
data = action.to_dict()
assert data["name"] == "book_restaurant"
assert len(data["context"]) == 2
assert data["context"][0]["key"] == "restaurantName"
def test_action_context_with_path(self):
"""Test ActionContext with path binding."""
from praisonaiagents.ui.a2ui.types import ActionContext, PathBinding
ctx = ActionContext(key="name", value=PathBinding(path="/restaurant/name"))
data = ctx.to_dict()
assert data["key"] == "name"
assert data["value"] == {"path": "/restaurant/name"}
@@ -0,0 +1,217 @@
"""
Unit tests for configurable max_file_chars in MentionsParser.
TDD approach: Write tests first, then implement.
"""
import os
import pytest
import tempfile
from pathlib import Path
from unittest.mock import patch
class TestMentionsParserConfigurableLimit:
"""Test suite for configurable file truncation limit."""
@pytest.fixture
def temp_workspace(self):
"""Create a temporary workspace with test files."""
with tempfile.TemporaryDirectory() as tmpdir:
yield Path(tmpdir)
@pytest.fixture
def small_file(self, temp_workspace):
"""Create a small test file (1KB)."""
file_path = temp_workspace / "small.py"
content = "# Small file\n" + "x = 1\n" * 100 # ~700 chars
file_path.write_text(content)
return file_path
@pytest.fixture
def large_file(self, temp_workspace):
"""Create a large test file (100KB)."""
file_path = temp_workspace / "large.py"
content = "# Large file\n" + "x = 'a' * 100\n" * 7000 # ~100KB
file_path.write_text(content)
return file_path
@pytest.fixture
def huge_file(self, temp_workspace):
"""Create a huge test file (600KB)."""
file_path = temp_workspace / "huge.py"
content = "# Huge file\n" + "x = 'a' * 100\n" * 42000 # ~600KB
file_path.write_text(content)
return file_path
# ===== Test 1.1: Default value (500K) =====
def test_default_max_file_chars_is_500k(self, temp_workspace):
"""Default max_file_chars should be 500,000."""
from praisonaiagents.tools.mentions import MentionsParser
parser = MentionsParser(workspace_path=str(temp_workspace))
assert parser.max_file_chars == 500000
# ===== Test 1.2: Custom value via constructor =====
def test_custom_max_file_chars_via_constructor(self, temp_workspace):
"""max_file_chars should be configurable via constructor."""
from praisonaiagents.tools.mentions import MentionsParser
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=100000
)
assert parser.max_file_chars == 100000
def test_custom_max_file_chars_zero_means_no_limit(self, temp_workspace):
"""max_file_chars=0 should mean no truncation."""
from praisonaiagents.tools.mentions import MentionsParser
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=0
)
assert parser.max_file_chars == 0
# ===== Test 1.3: Environment variable override =====
def test_env_var_override(self, temp_workspace):
"""PRAISON_MAX_FILE_CHARS env var should override default."""
from praisonaiagents.tools.mentions import MentionsParser
with patch.dict(os.environ, {'PRAISON_MAX_FILE_CHARS': '200000'}):
# Need to reimport to pick up env var
parser = MentionsParser(workspace_path=str(temp_workspace))
assert parser.max_file_chars == 200000
def test_constructor_takes_precedence_over_env_var(self, temp_workspace):
"""Constructor value should take precedence over env var."""
from praisonaiagents.tools.mentions import MentionsParser
with patch.dict(os.environ, {'PRAISON_MAX_FILE_CHARS': '200000'}):
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=100000
)
assert parser.max_file_chars == 100000
# ===== Test 1.4: No-limit mode =====
def test_no_truncation_when_limit_is_zero(self, temp_workspace, huge_file):
"""File should not be truncated when max_file_chars=0."""
from praisonaiagents.tools.mentions import MentionsParser
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=0
)
context, _ = parser.process(f"@file:{huge_file.name} explain")
# Should contain full content, no truncation message
assert "(truncated)" not in context
# Should have the full file content
original_size = huge_file.read_text().__len__()
# Context includes header and code fence, so check it's close
assert len(context) >= original_size
# ===== Test 1.5: Truncation behavior =====
def test_truncation_at_custom_limit(self, temp_workspace, large_file):
"""File should be truncated at custom limit."""
from praisonaiagents.tools.mentions import MentionsParser
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=10000 # 10KB limit
)
context, _ = parser.process(f"@file:{large_file.name} explain")
# Should be truncated
assert "(truncated)" in context
# Content should be around the limit (plus header/fence overhead)
assert len(context) < 15000 # Some overhead for markdown
def test_no_truncation_when_file_under_limit(self, temp_workspace, small_file):
"""Small files should not be truncated."""
from praisonaiagents.tools.mentions import MentionsParser
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=500000
)
context, _ = parser.process(f"@file:{small_file.name} explain")
# Should NOT be truncated
assert "(truncated)" not in context
def test_truncation_warning_logged(self, temp_workspace, large_file, caplog):
"""Warning should be logged when file is truncated."""
from praisonaiagents.tools.mentions import MentionsParser
import logging
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=10000,
verbose=1
)
with caplog.at_level(logging.WARNING):
parser.process(f"@file:{large_file.name} explain")
# Should log a warning about truncation
assert any("truncat" in record.message.lower() for record in caplog.records)
# ===== Test: URL truncation also uses configurable limit =====
def test_url_truncation_uses_same_limit(self, temp_workspace):
"""URL content should also respect max_file_chars."""
from praisonaiagents.tools.mentions import MentionsParser
parser = MentionsParser(
workspace_path=str(temp_workspace),
max_file_chars=100000
)
# URL truncation should use the same limit
assert parser.max_file_chars == 100000
# ===== Test: Backward compatibility =====
def test_backward_compatible_no_args(self, temp_workspace):
"""Parser should work without max_file_chars argument."""
from praisonaiagents.tools.mentions import MentionsParser
# Should not raise any errors
parser = MentionsParser(workspace_path=str(temp_workspace))
assert hasattr(parser, 'max_file_chars')
# ===== Test: Performance - no impact on small files =====
def test_performance_small_file(self, temp_workspace, small_file):
"""Processing small files should be fast."""
from praisonaiagents.tools.mentions import MentionsParser
import time
parser = MentionsParser(workspace_path=str(temp_workspace))
start = time.time()
for _ in range(100):
parser.process(f"@file:{small_file.name} explain")
elapsed = time.time() - start
# 100 iterations should complete in under 1 second
assert elapsed < 1.0, f"Performance regression: {elapsed}s for 100 iterations"
class TestProcessMentionsFunction:
"""Test the convenience function."""
def test_process_mentions_uses_default_limit(self):
"""process_mentions function should use default limit."""
from praisonaiagents.tools.mentions import process_mentions
# The function should create a parser with default settings
# Just verify it works without errors
context, cleaned = process_mentions("hello world")
assert cleaned == "hello world"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
+1 -1
View File
@@ -2422,7 +2422,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.1.16"
version = "0.1.17"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },
+94 -4
View File
@@ -6,10 +6,13 @@
</picture>
</p>
<!-- mcp-name: io.github.MervinPraison/praisonai -->
<p align="center">
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://static.pepy.tech/badge/PraisonAI" alt="Total Downloads" /></a>
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://img.shields.io/github/v/release/MervinPraison/PraisonAI" alt="Latest Stable Version" /></a>
<a href="https://github.com/MervinPraison/PraisonAI"><img src="https://img.shields.io/badge/License-MIT-yellow.svg" alt="License" /></a>
<a href="https://registry.modelcontextprotocol.io/servers/io.github.MervinPraison/praisonai"><img src="https://img.shields.io/badge/MCP-Registry-blue" alt="MCP Registry" /></a>
</p>
<div align="center">
@@ -60,6 +63,7 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
- [Planning & Workflows](#planning--workflows)
- [Specialised Agents](#specialised-agents)
- [MCP Protocol](#mcp-protocol)
- [A2A & A2UI Protocols](#a2a--a2ui-protocols)
- [Safety & Control](#safety--control)
- [Developer Tools](#developer-tools)
- [📊 Process Types & Patterns](#-process-types--patterns)
@@ -187,7 +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) |
| **💻 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) |
@@ -205,6 +216,7 @@ npm install praisonai
| 🧠 Claude Memory Tool | [Example](#claude-memory-tool-cli) | [📖](https://docs.praison.ai/features/claude-memory-tool) |
| 💾 File-Based Memory | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
| 🔍 Built-in Search Tools | [Example](examples/python/agents/websearch-agent.py) | [📖](https://docs.praison.ai/tools/tavily) |
| 🔎 Unified Web Search | [Example](src/praisonai-agents/examples/web_search_example.py) | [📖](https://docs.praison.ai/tools/web-search) |
| 📋 Planning Mode | [Example](examples/python/agents/planning-agent.py) | [📖](https://docs.praison.ai/features/planning-mode) |
| 🔧 Planning Tools | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
| 🧠 Planning Reasoning | [Example](#3-agent-with-planning-mode) | [📖](https://docs.praison.ai/features/planning-mode) |
@@ -224,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) |
@@ -246,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) |
---
@@ -1118,6 +1138,15 @@ praisonai workflow auto "Research AI trends" --pattern parallel
# Generate routing workflow (classifier routes to specialists)
praisonai workflow auto "Build a chatbot" --pattern routing
# Generate orchestrator-workers workflow (central orchestrator delegates)
praisonai workflow auto "Comprehensive market analysis" --pattern orchestrator-workers
# Generate evaluator-optimizer workflow (iterative refinement)
praisonai workflow auto "Write and refine article" --pattern evaluator-optimizer
# Specify output file
praisonai workflow auto "Build a chatbot" --pattern routing
# Specify output file
praisonai workflow auto "Research AI" --pattern sequential --output my_workflow.yaml
```
@@ -1127,7 +1156,7 @@ praisonai workflow auto "Research AI" --pattern sequential --output my_workflow.
|------|-------------|
| `--workflow-var key=value` | Set workflow variable (can be repeated) |
| `--var key=value` | Set variable for YAML workflows |
| `--pattern <pattern>` | Pattern for auto-generation (sequential, parallel, routing, loop) |
| `--pattern <pattern>` | Pattern for auto-generation (sequential, parallel, routing, loop, orchestrator-workers, evaluator-optimizer) |
| `--output <file>` | Output file for auto-generation |
| `--llm <model>` | LLM model (e.g., openai/gpt-4o-mini) |
| `--tools <tools>` | Tools (comma-separated, e.g., tavily) |
@@ -1359,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
@@ -1449,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
@@ -2251,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
@@ -2260,6 +2343,13 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
- 🔐 **MCP Security** - Origin validation, DNS rebinding prevention, secure sessions
- 🔄 **MCP Resumability** - SSE stream recovery via Last-Event-ID
### A2A & A2UI Protocols
- 🔗 **A2A Protocol** - Agent-to-Agent communication for inter-agent collaboration
- 🖼️ **A2UI Protocol** - Agent-to-User Interface for generating rich UIs from agents
- 📋 **UI Templates** - ChatTemplate, ListTemplate, FormTemplate, DashboardTemplate
- 🔧 **Surface Builder** - Fluent API for building declarative UIs
### Safety & Control
- 🤝 **Agent Handoffs** - Transfer context between specialised agents
+2 -2
View File
@@ -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.38.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.38.tar.gz | shasum -a 256`.split.first
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.47.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.47.tar.gz | shasum -a 256`.split.first
license "MIT"
depends_on "[email protected]"
+4
View File
@@ -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"
+505 -105
View File
@@ -51,14 +51,140 @@ try:
except ImportError:
PRAISONAI_TOOLS_AVAILABLE = False
# LiteLLM availability check for multi-provider support
LITELLM_AVAILABLE = False
try:
import litellm # noqa: F401 - imported for availability check
LITELLM_AVAILABLE = True
except ImportError:
pass
logging.basicConfig(level=os.environ.get('LOGLEVEL', 'INFO').upper(), format='%(asctime)s - %(levelname)s - %(message)s')
# Define Pydantic models outside of the generate method
# =============================================================================
# Available Tools List (shared between generators)
# =============================================================================
AVAILABLE_TOOLS = [
"CodeDocsSearchTool", "CSVSearchTool", "DirectorySearchTool", "DOCXSearchTool",
"DirectoryReadTool", "FileReadTool", "TXTSearchTool", "JSONSearchTool",
"MDXSearchTool", "PDFSearchTool", "RagTool", "ScrapeElementFromWebsiteTool",
"ScrapeWebsiteTool", "WebsiteSearchTool", "XMLSearchTool",
"YoutubeChannelSearchTool", "YoutubeVideoSearchTool"
]
# =============================================================================
# Base Generator Class (DRY - shared functionality)
# =============================================================================
class BaseAutoGenerator:
"""
Base class for auto-generators with shared functionality.
Provides:
- Lazy-loaded instructor client (LiteLLM or OpenAI fallback)
- Environment variable handling for model/API configuration
- Config list management
"""
def __init__(self, config_list: Optional[List[Dict]] = None):
"""
Initialize base generator with LLM configuration.
Args:
config_list: Optional LLM configuration list
"""
# Support multiple environment variable patterns for better compatibility
model_name = os.environ.get("MODEL_NAME") or os.environ.get("OPENAI_MODEL_NAME", "gpt-4o-mini")
base_url = (
os.environ.get("OPENAI_BASE_URL") or
os.environ.get("OPENAI_API_BASE") or
os.environ.get("OLLAMA_API_BASE", "https://api.openai.com/v1")
)
self.config_list = config_list or [
{
'model': model_name,
'base_url': base_url,
'api_key': os.environ.get("OPENAI_API_KEY")
}
]
self._client = None # Lazy loading for performance
@property
def client(self):
"""Lazy load the instructor client to avoid performance impact.
Uses LiteLLM via instructor.from_provider if available for multi-provider support,
otherwise falls back to direct OpenAI SDK.
"""
if self._client is None:
if LITELLM_AVAILABLE:
# Use LiteLLM for multi-provider support (100+ LLMs)
model_name = self.config_list[0]['model']
self._client = instructor.from_provider(
f"litellm/{model_name}",
mode=instructor.Mode.JSON
)
else:
# Fallback to direct OpenAI SDK
self._client = instructor.patch(
OpenAI(
base_url=self.config_list[0]['base_url'],
api_key=self.config_list[0]['api_key'],
),
mode=instructor.Mode.JSON,
)
return self._client
@staticmethod
def get_available_tools() -> List[str]:
"""Return list of available tools for agent assignment."""
return AVAILABLE_TOOLS.copy()
@staticmethod
def analyze_complexity(topic: str) -> str:
"""
Analyze task complexity based on keywords.
Args:
topic: The task description
Returns:
str: Complexity level - 'simple', 'moderate', or 'complex'
"""
topic_lower = topic.lower()
# Complex task indicators
complex_keywords = [
'comprehensive', 'multi-step', 'analyze and', 'research and write',
'multiple', 'coordinate', 'complex', 'detailed analysis',
'full report', 'in-depth', 'thorough'
]
# Simple task indicators
simple_keywords = [
'write a', 'create a', 'simple', 'quick', 'brief',
'haiku', 'poem', 'summary', 'list', 'single'
]
if any(kw in topic_lower for kw in complex_keywords):
return 'complex'
elif any(kw in topic_lower for kw in simple_keywords):
return 'simple'
else:
return 'moderate'
# =============================================================================
# Pydantic Models for Structured Output
# =============================================================================
class TaskDetails(BaseModel):
"""Details for a single task."""
description: str
expected_output: str
class RoleDetails(BaseModel):
"""Details for a single role/agent."""
role: str
goal: str
backstory: str
@@ -66,14 +192,62 @@ class RoleDetails(BaseModel):
tools: List[str]
class TeamStructure(BaseModel):
"""Structure for multi-agent team."""
roles: Dict[str, RoleDetails]
class AutoGenerator:
def __init__(self, topic="Movie Story writing about AI", agent_file="test.yaml", framework="crewai", config_list: Optional[List[Dict]] = None):
class SingleAgentStructure(BaseModel):
"""Structure for single-agent generation (Anthropic's 'start simple' principle)."""
name: str
role: str
goal: str
backstory: str
instructions: str
tools: List[str] = []
task_description: str
expected_output: str
class PatternRecommendation(BaseModel):
"""LLM-based pattern recommendation with reasoning."""
pattern: str # sequential, parallel, routing, orchestrator-workers, evaluator-optimizer
reasoning: str # Why this pattern was chosen
confidence: float # 0.0 to 1.0 confidence score
class ValidationGate(BaseModel):
"""Validation gate for prompt chaining workflows."""
criteria: str # What to validate
pass_action: str # Action if validation passes (e.g., "continue", "next_step")
fail_action: str # Action if validation fails (e.g., "retry", "escalate", "abort")
class AutoGenerator(BaseAutoGenerator):
"""
Auto-generates agents.yaml files from a topic description.
Inherits from BaseAutoGenerator for shared LLM client functionality.
Usage:
generator = AutoGenerator(framework="crewai", topic="Create a movie script")
path = generator.generate()
"""
def __init__(self, topic="Movie Story writing about AI", agent_file="test.yaml",
framework="crewai", config_list: Optional[List[Dict]] = None,
pattern: str = "sequential", single_agent: bool = False):
"""
Initialize the AutoGenerator class with the specified topic, agent file, and framework.
Args:
topic: The task/topic for agent generation
agent_file: Output YAML file name
framework: Framework to use (crewai, autogen, praisonai)
config_list: Optional LLM configuration
pattern: Workflow pattern (sequential, parallel, routing, orchestrator-workers, evaluator-optimizer)
single_agent: If True, generate a single agent instead of a team
Note: autogen framework is different from this AutoGenerator class.
"""
# Initialize base class first (handles config_list and client)
super().__init__(config_list=config_list)
# Validate framework availability and show framework-specific messages
if framework == "crewai" and not CREWAI_AVAILABLE:
raise ImportError("""
@@ -106,35 +280,42 @@ Tools are not available for {framework}. To use tools, install:
pip install "praisonai[{framework}]"
""")
# Support multiple environment variable patterns for better compatibility
# Priority order: MODEL_NAME > OPENAI_MODEL_NAME for model selection
model_name = os.environ.get("MODEL_NAME") or os.environ.get("OPENAI_MODEL_NAME", "gpt-5-nano")
# Priority order for base_url: OPENAI_BASE_URL > OPENAI_API_BASE > OLLAMA_API_BASE
# OPENAI_BASE_URL is the standard OpenAI SDK environment variable
base_url = (
os.environ.get("OPENAI_BASE_URL") or
os.environ.get("OPENAI_API_BASE") or
os.environ.get("OLLAMA_API_BASE", "https://api.openai.com/v1")
)
self.config_list = config_list or [
{
'model': model_name,
'base_url': base_url,
'api_key': os.environ.get("OPENAI_API_KEY")
}
]
self.topic = topic
self.agent_file = agent_file
self.framework = framework or "praisonai"
self.client = instructor.patch(
OpenAI(
base_url=self.config_list[0]['base_url'],
api_key=self.config_list[0]['api_key'],
),
mode=instructor.Mode.JSON,
)
self.pattern = pattern
self.single_agent = single_agent
def recommend_pattern(self, topic: str = None) -> str:
"""
Recommend the best workflow pattern based on task characteristics.
Args:
topic: The task description (uses self.topic if not provided)
Returns:
str: Recommended pattern name
"""
task = topic or self.topic
task_lower = task.lower()
# Keywords that suggest specific patterns
parallel_keywords = ['multiple', 'concurrent', 'parallel', 'simultaneously', 'different sources', 'compare', 'various']
routing_keywords = ['classify', 'categorize', 'route', 'different types', 'depending on', 'if...then']
orchestrator_keywords = ['complex', 'comprehensive', 'multi-step', 'coordinate', 'delegate', 'break down', 'analyze and']
evaluator_keywords = ['refine', 'improve', 'iterate', 'quality', 'review', 'feedback', 'polish', 'optimize']
# Check for pattern indicators
if any(kw in task_lower for kw in evaluator_keywords):
return "evaluator-optimizer"
elif any(kw in task_lower for kw in orchestrator_keywords):
return "orchestrator-workers"
elif any(kw in task_lower for kw in routing_keywords):
return "routing"
elif any(kw in task_lower for kw in parallel_keywords):
return "parallel"
else:
return "sequential"
def generate(self, merge=False):
"""
@@ -157,7 +338,8 @@ Tools are not available for {framework}. To use tools, install:
response = self.client.chat.completions.create(
model=self.config_list[0]['model'],
response_model=TeamStructure,
max_retries=10,
max_retries=5,
timeout=120.0, # 2 minute timeout for complex generations
messages=[
{"role": "system", "content": "You are a helpful assistant designed to output complex team structures."},
{"role": "user", "content": self.get_user_content()}
@@ -194,8 +376,7 @@ Tools are not available for {framework}. To use tools, install:
"goal": role_details['goal'],
"role": role_details['role'],
"tasks": {},
# "tools": role_details.get('tools', []),
"tools": ['']
"tools": role_details.get('tools', [])
}
for task_id, task_details in role_details['tasks'].items():
@@ -265,7 +446,7 @@ Tools are not available for {framework}. To use tools, install:
"goal": role_details['goal'],
"role": role_details['role'],
"tasks": {},
"tools": ['']
"tools": role_details.get('tools', [])
}
# Add tasks for this role
@@ -292,45 +473,77 @@ Tools are not available for {framework}. To use tools, install:
prompt = generator.get_user_content()
print(prompt)
"""
user_content = """Generate a team structure for \"""" + self.topic + """\" task.
No Input data will be provided to the team.
The team will work in sequence. First role will pass the output to the next role, and so on.
The last role will generate the final output.
Think step by step.
With maximum 3 roles, each with 1 task. Include role goals, backstories, task descriptions, and expected outputs.
List of Available Tools: CodeDocsSearchTool, CSVSearchTool, DirectorySearchTool, DOCXSearchTool, DirectoryReadTool, FileReadTool, TXTSearchTool, JSONSearchTool, MDXSearchTool, PDFSearchTool, RagTool, ScrapeElementFromWebsiteTool, ScrapeWebsiteTool, WebsiteSearchTool, XMLSearchTool, YoutubeChannelSearchTool, YoutubeVideoSearchTool.
Only use Available Tools. Do Not use any other tools.
Example Below:
Use below example to understand the structure of the output.
The final role you create should satisfy the provided task: """ + self.topic + """.
{
"roles": {
"narrative_designer": {
"role": "Narrative Designer",
"goal": "Create AI storylines",
"backstory": "Skilled in narrative development for AI, with a focus on story resonance.",
"tools": ["ScrapeWebsiteTool"],
"tasks": {
"story_concept_development": {
"description": "Craft a unique AI story concept with depth and engagement using concept from this page the content https://www.asthebirdfliesblog.com/posts/how-to-write-book-story-development .",
"expected_output": "Document with narrative arcs, character bios, and settings."
}
}
},
"scriptwriter": {
"role": "Scriptwriter",
"goal": "Write scripts from AI concepts",
"backstory": "Expert in dialogue and script structure, translating concepts into scripts.",
"tasks": {
"scriptwriting_task": {
"description": "Turn narrative concepts into scripts, including dialogue and scenes.",
"expected_output": "Production-ready script with dialogue and scene details."
}
}
}
}
}
"""
# Pattern-specific guidance
pattern_guidance = {
"sequential": "The team will work in sequence. Each role passes output to the next.",
"parallel": "The team will work in parallel on independent subtasks, then combine results.",
"routing": "A classifier agent will route requests to specialized agents based on input type.",
"orchestrator-workers": "A central orchestrator will dynamically delegate tasks to specialized workers.",
"evaluator-optimizer": "One agent generates content, another evaluates it in a loop until quality criteria are met."
}
workflow_guidance = pattern_guidance.get(self.pattern, pattern_guidance["sequential"])
user_content = f"""Analyze and generate a team structure for: "{self.topic}"
STEP 1: TASK COMPLEXITY ANALYSIS
First, analyze the task complexity:
- Is this a simple task that 1-2 agents could handle?
- Does it require multiple specialized perspectives?
- Are there clear sequential dependencies between steps?
STEP 2: DETERMINE OPTIMAL TEAM SIZE
Based on your analysis, create the minimum number of agents needed:
- Simple tasks: 1-2 agents
- Moderate tasks: 2-3 agents
- Complex tasks: 3-4 agents (maximum)
IMPORTANT: Only create agents that provide meaningful specialization. Avoid unnecessary complexity.
STEP 3: DESIGN THE TEAM (Pattern: {self.pattern})
{workflow_guidance}
Each agent should have:
- A clear, distinct role
- A specific goal
- Relevant backstory
- 1 focused task with clear description and expected output
- Appropriate tools (only if needed)
Available Tools: CodeDocsSearchTool, CSVSearchTool, DirectorySearchTool, DOCXSearchTool, DirectoryReadTool, FileReadTool, TXTSearchTool, JSONSearchTool, MDXSearchTool, PDFSearchTool, RagTool, ScrapeElementFromWebsiteTool, ScrapeWebsiteTool, WebsiteSearchTool, XMLSearchTool, YoutubeChannelSearchTool, YoutubeVideoSearchTool.
Only use these tools if the task requires them. Use empty list [] if no tools needed.
Example structure (2 agents for a writing task):
{{
"roles": {{
"researcher": {{
"role": "Research Analyst",
"goal": "Gather comprehensive information on the topic",
"backstory": "Expert researcher skilled at finding and synthesizing information.",
"tools": ["WebsiteSearchTool"],
"tasks": {{
"research_task": {{
"description": "Research key information about the topic and compile findings.",
"expected_output": "Comprehensive research notes with key facts and insights."
}}
}}
}},
"writer": {{
"role": "Content Writer",
"goal": "Create polished final content",
"backstory": "Skilled writer who transforms research into engaging content.",
"tools": [],
"tasks": {{
"writing_task": {{
"description": "Write the final content based on research findings.",
"expected_output": "Polished, well-structured final document."
}}
}}
}}
}}
}}
Now generate the optimal team structure for: {self.topic}
"""
return user_content
@@ -372,13 +585,14 @@ class WorkflowStructure(BaseModel):
description: str
agents: Dict[str, WorkflowAgentDetails]
steps: List[Dict] # Can be agent steps, route, parallel, etc.
gates: Optional[List[ValidationGate]] = None # Optional validation gates
class WorkflowAutoGenerator:
class WorkflowAutoGenerator(BaseAutoGenerator):
"""
Auto-generates workflow.yaml files from a topic description.
Uses lazy loading to avoid performance impact when not used.
Inherits from BaseAutoGenerator for shared LLM client functionality.
Usage:
generator = WorkflowAutoGenerator(topic="Research AI trends and write a report")
@@ -387,7 +601,9 @@ class WorkflowAutoGenerator:
def __init__(self, topic: str = "Research and write about AI",
workflow_file: str = "workflow.yaml",
config_list: Optional[List[Dict]] = None):
config_list: Optional[List[Dict]] = None,
framework: str = "praisonai",
single_agent: bool = False):
"""
Initialize the WorkflowAutoGenerator.
@@ -395,45 +611,105 @@ class WorkflowAutoGenerator:
topic: The task/topic for the workflow
workflow_file: Output file name
config_list: Optional LLM configuration
framework: Framework to use (praisonai, crewai, autogen)
single_agent: If True, generate a single agent workflow
"""
# Support multiple environment variable patterns
model_name = os.environ.get("MODEL_NAME") or os.environ.get("OPENAI_MODEL_NAME", "gpt-4o-mini")
base_url = (
os.environ.get("OPENAI_BASE_URL") or
os.environ.get("OPENAI_API_BASE") or
os.environ.get("OLLAMA_API_BASE", "https://api.openai.com/v1")
)
# Initialize base class (handles config_list and client)
super().__init__(config_list=config_list)
self.config_list = config_list or [
{
'model': model_name,
'base_url': base_url,
'api_key': os.environ.get("OPENAI_API_KEY")
}
]
self.topic = topic
self.workflow_file = workflow_file
self._client = None # Lazy loading
self.framework = framework
self.single_agent = single_agent
@property
def client(self):
"""Lazy load the OpenAI client to avoid performance impact."""
if self._client is None:
self._client = instructor.patch(
OpenAI(
base_url=self.config_list[0]['base_url'],
api_key=self.config_list[0]['api_key'],
),
mode=instructor.Mode.JSON,
)
return self._client
def recommend_pattern(self, topic: str = None) -> str:
"""
Recommend the best workflow pattern based on task characteristics.
Args:
topic: The task description (uses self.topic if not provided)
Returns:
str: Recommended pattern name
Pattern recommendations based on Anthropic's best practices:
- sequential: Clear step-by-step dependencies
- parallel: Independent subtasks that can run concurrently
- routing: Different input types need different handling
- orchestrator-workers: Complex tasks needing dynamic decomposition
- evaluator-optimizer: Tasks requiring iterative refinement
"""
task = topic or self.topic
task_lower = task.lower()
# Keywords that suggest specific patterns
parallel_keywords = ['multiple', 'concurrent', 'parallel', 'simultaneously', 'different sources', 'compare', 'various']
routing_keywords = ['classify', 'categorize', 'route', 'different types', 'depending on', 'if...then']
orchestrator_keywords = ['complex', 'comprehensive', 'multi-step', 'coordinate', 'delegate', 'break down', 'analyze and']
evaluator_keywords = ['refine', 'improve', 'iterate', 'quality', 'review', 'feedback', 'polish', 'optimize']
# Check for pattern indicators
if any(kw in task_lower for kw in evaluator_keywords):
return "evaluator-optimizer"
elif any(kw in task_lower for kw in orchestrator_keywords):
return "orchestrator-workers"
elif any(kw in task_lower for kw in routing_keywords):
return "routing"
elif any(kw in task_lower for kw in parallel_keywords):
return "parallel"
else:
return "sequential"
def generate(self, pattern: str = "sequential") -> str:
def recommend_pattern_llm(self, topic: str = None) -> PatternRecommendation:
"""
Use LLM to recommend the best workflow pattern with reasoning.
Args:
topic: The task description (uses self.topic if not provided)
Returns:
PatternRecommendation: Pattern with reasoning and confidence score
"""
task = topic or self.topic
prompt = f"""Analyze this task and recommend the best workflow pattern:
Task: "{task}"
Available patterns:
1. sequential - Agents work one after another, passing output to the next
2. parallel - Multiple agents work concurrently on independent subtasks
3. routing - A classifier routes requests to specialized agents based on input type
4. orchestrator-workers - Central orchestrator dynamically delegates to specialized workers
5. evaluator-optimizer - Generator creates content, evaluator reviews in a loop until quality met
Respond with:
- pattern: The recommended pattern name
- reasoning: Why this pattern is best for this task
- confidence: Your confidence score (0.0 to 1.0)
"""
response = self.client.chat.completions.create(
model=self.config_list[0]['model'],
response_model=PatternRecommendation,
max_retries=3,
timeout=60.0,
messages=[
{"role": "system", "content": "You are an expert at designing AI agent workflows."},
{"role": "user", "content": prompt}
]
)
return response
def generate(self, pattern: str = "sequential", merge: bool = False) -> str:
"""
Generate a workflow YAML file.
Args:
pattern: Workflow pattern - "sequential", "routing", "parallel", "loop"
pattern: Workflow pattern - "sequential", "routing", "parallel", "loop",
"orchestrator-workers", "evaluator-optimizer"
merge: If True, merge with existing workflow file instead of overwriting
Returns:
Path to the generated workflow file
@@ -442,6 +718,7 @@ class WorkflowAutoGenerator:
model=self.config_list[0]['model'],
response_model=WorkflowStructure,
max_retries=5,
timeout=120.0, # 2 minute timeout for complex generations
messages=[
{"role": "system", "content": "You are a helpful assistant that designs workflow structures."},
{"role": "user", "content": self._get_prompt(pattern)}
@@ -449,16 +726,86 @@ class WorkflowAutoGenerator:
)
json_data = json.loads(response.model_dump_json())
if merge and os.path.exists(self.workflow_file):
return self._save_workflow(self.merge_with_existing_workflow(json_data), pattern)
return self._save_workflow(json_data, pattern)
def merge_with_existing_workflow(self, new_data: Dict) -> Dict:
"""
Merge new workflow data with existing workflow file.
Args:
new_data: The new workflow data to merge
Returns:
Dict: Merged workflow data
"""
try:
with open(self.workflow_file, 'r') as f:
existing_data = yaml.safe_load(f)
if not existing_data:
return new_data
except (yaml.YAMLError, FileNotFoundError) as e:
logging.warning(f"Could not load existing workflow file {self.workflow_file}: {e}")
return new_data
# Merge agents (avoid duplicates)
merged_agents = existing_data.get('agents', {}).copy()
for agent_id, agent_data in new_data.get('agents', {}).items():
# Rename if conflict
final_id = agent_id
counter = 1
while final_id in merged_agents:
final_id = f"{agent_id}_auto_{counter}"
counter += 1
merged_agents[final_id] = agent_data
# Merge steps (append new steps)
merged_steps = existing_data.get('steps', []) + new_data.get('steps', [])
# Create merged structure
merged = {
'name': existing_data.get('name', new_data.get('name', 'Merged Workflow')),
'description': f"{existing_data.get('description', '')} + {new_data.get('description', '')}",
'agents': merged_agents,
'steps': merged_steps
}
return merged
def _get_prompt(self, pattern: str) -> str:
"""Generate the prompt based on the workflow pattern."""
# Analyze complexity to determine agent count
complexity = self.analyze_complexity(self.topic)
if complexity == 'simple':
agent_guidance = "Create 1-2 agents (simple task detected)."
elif complexity == 'complex':
agent_guidance = "Create 3-4 agents (complex task detected)."
else:
agent_guidance = "Create 2-3 agents (moderate task detected)."
# Get available tools
tools_list = ", ".join(self.get_available_tools())
base_prompt = f"""Generate a workflow structure for: "{self.topic}"
STEP 1: ANALYZE TASK COMPLEXITY
- Is this a simple task (1-2 agents)?
- Does it require multiple specialists (2-3 agents)?
- Is it complex with many dependencies (3-4 agents)?
STEP 2: DESIGN WORKFLOW
The workflow should use the "{pattern}" pattern.
Create 2-4 agents with clear roles and instructions.
{agent_guidance}
Each agent should have clear roles and instructions.
Each step should have a clear action.
STEP 3: ASSIGN TOOLS (if needed)
Available Tools: {tools_list}
Only assign tools if the task requires them. Use empty list or null if no tools needed.
"""
if pattern == "routing":
@@ -501,6 +848,59 @@ Example structure:
{"agent": "aggregator", "action": "Combine all findings"}
]
}
"""
elif pattern == "orchestrator-workers":
base_prompt += """
Create an orchestrator-workers workflow where a central orchestrator dynamically delegates tasks to specialized workers.
The orchestrator analyzes the input, decides which workers are needed, and synthesizes results.
Example structure:
{
"name": "Orchestrator-Workers Workflow",
"description": "Central orchestrator delegates to specialized workers",
"agents": {
"orchestrator": {"name": "Orchestrator", "role": "Task Coordinator", "goal": "Analyze tasks and delegate to appropriate workers", "instructions": "Break down the task, identify required specialists, and coordinate their work. Output a JSON with 'subtasks' array listing which workers to invoke."},
"researcher": {"name": "Researcher", "role": "Research Specialist", "goal": "Gather information", "instructions": "Research and provide factual information"},
"analyst": {"name": "Analyst", "role": "Data Analyst", "goal": "Analyze data and patterns", "instructions": "Analyze information and identify insights"},
"writer": {"name": "Writer", "role": "Content Writer", "goal": "Create written content", "instructions": "Write clear, engaging content"},
"synthesizer": {"name": "Synthesizer", "role": "Results Synthesizer", "goal": "Combine all worker outputs", "instructions": "Synthesize all worker outputs into a coherent final result"}
},
"steps": [
{"agent": "orchestrator", "action": "Analyze task and determine required workers: {{input}}"},
{"name": "worker_dispatch", "parallel": [
{"agent": "researcher", "action": "Research: {{input}}"},
{"agent": "analyst", "action": "Analyze: {{input}}"},
{"agent": "writer", "action": "Draft content for: {{input}}"}
]},
{"agent": "synthesizer", "action": "Combine all worker outputs into final result"}
]
}
"""
elif pattern == "evaluator-optimizer":
base_prompt += """
Create an evaluator-optimizer workflow where one agent generates content and another evaluates it in a loop.
The generator improves based on evaluator feedback until quality criteria are met.
Example structure:
{
"name": "Evaluator-Optimizer Workflow",
"description": "Iterative refinement through generation and evaluation",
"agents": {
"generator": {"name": "Generator", "role": "Content Generator", "goal": "Generate high-quality content", "instructions": "Create content based on the input. If feedback is provided, improve the content accordingly."},
"evaluator": {"name": "Evaluator", "role": "Quality Evaluator", "goal": "Evaluate content quality", "instructions": "Evaluate the content on: clarity, accuracy, completeness, and relevance. Score 1-10 for each. If average score < 7, provide specific improvement feedback. If score >= 7, respond with 'APPROVED'."}
},
"steps": [
{"agent": "generator", "action": "Generate initial content for: {{input}}"},
{"name": "evaluation_loop", "loop": {
"agent": "evaluator",
"action": "Evaluate the generated content",
"condition": "output does not contain 'APPROVED'",
"max_iterations": 3,
"feedback_to": "generator"
}},
{"agent": "generator", "action": "Finalize content based on all feedback"}
]
}
"""
else: # sequential
base_prompt += """
@@ -21,6 +21,10 @@ Features:
- flow_display: Visual workflow tracking
- 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
@@ -43,6 +47,15 @@ __all__ = [
'FlowDisplayHandler',
'WorkflowHandler',
'N8nHandler',
'ExternalAgentsHandler',
# New CLI enhancement features
'SlashCommandHandler',
'AutonomyModeHandler',
'CostTrackerHandler',
'RepoMapHandler',
'InteractiveTUIHandler',
'GitIntegrationHandler',
'SandboxExecutorHandler',
]
def __getattr__(name):
@@ -95,4 +108,29 @@ def __getattr__(name):
elif name == 'N8nHandler':
from .n8n import N8nHandler
return N8nHandler
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}")
@@ -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
)
@@ -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
@@ -0,0 +1,231 @@
"""
External Agents Handler for CLI.
Provides integration with external AI coding CLI tools:
- Claude Code CLI
- Gemini CLI
- OpenAI Codex CLI
- Cursor CLI
Usage:
praisonai "Refactor the code" --external-agent claude
praisonai "Analyze codebase" --external-agent gemini
praisonai "Fix bugs" --external-agent codex
praisonai "Add tests" --external-agent cursor
"""
from typing import Dict, Any, List, Optional, Tuple
from .base import FlagHandler
class ExternalAgentsHandler(FlagHandler):
"""
Handler for --external-agent flag.
Integrates external AI coding CLI tools as agent tools.
Example:
praisonai "Refactor the auth module" --external-agent claude
praisonai "Analyze this codebase" --external-agent gemini
"""
# Mapping of integration names to their module paths
INTEGRATIONS = {
"claude": "claude_code.ClaudeCodeIntegration",
"gemini": "gemini_cli.GeminiCLIIntegration",
"codex": "codex_cli.CodexCLIIntegration",
"cursor": "cursor_cli.CursorCLIIntegration",
}
def __init__(self, verbose: bool = False):
"""Initialize the handler."""
super().__init__(verbose)
self._integration_cache: Dict[str, Any] = {}
@property
def feature_name(self) -> str:
"""Return the feature name."""
return "external_agents"
@property
def flag_name(self) -> str:
"""Return the flag name."""
return "external-agent"
@property
def flag_help(self) -> str:
"""Return the flag help text."""
return "External AI CLI tool to use (claude, gemini, codex, cursor)"
def check_dependencies(self) -> Tuple[bool, str]:
"""Check if integrations module is available."""
try:
import importlib.util
if importlib.util.find_spec("praisonai.integrations") is not None:
return True, ""
return False, "Integrations module not available"
except ImportError:
return False, "Integrations module not available"
def get_integration(self, name: str, **options):
"""
Get an integration instance by name.
Args:
name: Integration name (claude, gemini, codex, cursor)
**options: Options to pass to the integration
Returns:
Integration instance
Raises:
ValueError: If integration name is invalid
"""
if name not in self.INTEGRATIONS:
raise ValueError(f"Unknown integration: {name}. Available: {list(self.INTEGRATIONS.keys())}")
# Check cache first
cache_key = f"{name}_{hash(frozenset(options.items()))}"
if cache_key in self._integration_cache:
return self._integration_cache[cache_key]
# Lazy import the integration
module_path = self.INTEGRATIONS[name]
module_name, class_name = module_path.rsplit(".", 1)
import importlib
module = importlib.import_module(f"praisonai.integrations.{module_name}")
integration_class = getattr(module, class_name)
# Create instance with options
integration = integration_class(**options)
# Cache it
self._integration_cache[cache_key] = integration
return integration
def list_integrations(self) -> List[str]:
"""
List all available integration names.
Returns:
List of integration names
"""
return list(self.INTEGRATIONS.keys())
def check_availability(self) -> Dict[str, bool]:
"""
Check availability of all integrations.
Returns:
Dict mapping integration name to availability
"""
availability = {}
for name in self.INTEGRATIONS:
try:
integration = self.get_integration(name)
availability[name] = integration.is_available
except Exception:
availability[name] = False
return availability
def apply_to_agent_config(self, config: Dict[str, Any], flag_value: Any) -> Dict[str, Any]:
"""
Apply external agent configuration to agent config.
Args:
config: Agent configuration dictionary
flag_value: Integration name or dict with name and options
Returns:
Modified configuration with external agent tool
"""
if not flag_value:
return config
# Handle both string and dict formats
if isinstance(flag_value, str):
name = flag_value
options = {}
elif isinstance(flag_value, dict):
name = flag_value.get("name", "")
options = {k: v for k, v in flag_value.items() if k != "name"}
else:
return config
try:
integration = self.get_integration(name, **options)
if integration.is_available:
# Add tool to existing tools
existing_tools = config.get("tools", [])
if isinstance(existing_tools, list):
existing_tools.append(integration.as_tool())
else:
existing_tools = [integration.as_tool()]
config["tools"] = existing_tools
self.print_status(f"🔌 External agent connected: {name}", "success")
else:
self.print_status(f"⚠️ External agent not available: {name}", "warning")
except ValueError as e:
self.print_status(str(e), "error")
except Exception as e:
self.print_status(f"Failed to connect external agent: {e}", "error")
return config
def execute(
self,
integration_name: str = None,
**options
) -> Optional[Any]:
"""
Execute external agent setup.
Args:
integration_name: Name of the integration
**options: Options for the integration
Returns:
Integration instance or None
"""
if not integration_name:
# List available integrations
self.print_status("\n🔧 Available External Agents:", "info")
self.print_status("-" * 40, "info")
availability = self.check_availability()
for name, available in availability.items():
status = "✅ Available" if available else "❌ Not installed"
self.print_status(f" {name}: {status}", "info")
return None
try:
integration = self.get_integration(integration_name, **options)
if integration.is_available:
self.print_status(f"🔌 External agent ready: {integration_name}", "success")
return integration
else:
self.print_status(f"⚠️ External agent not installed: {integration_name}", "warning")
self.print_status(f" Install with: {self._get_install_instructions(integration_name)}", "info")
return None
except ValueError as e:
self.print_status(str(e), "error")
return None
def _get_install_instructions(self, name: str) -> str:
"""Get installation instructions for an integration."""
instructions = {
"claude": "npm install -g @anthropic-ai/claude-code",
"gemini": "npm install -g @anthropic-ai/gemini-cli",
"codex": "npm install -g @openai/codex",
"cursor": "Download from cursor.com",
}
return instructions.get(name, "See documentation")
@@ -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
@@ -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()
@@ -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"
))
+316 -10
View File
@@ -3,12 +3,45 @@ Router Handler for CLI.
Provides smart model selection based on task complexity.
Usage: praisonai "Complex task" --router
Supports custom model configuration from agents.yaml/workflow.yaml:
models:
gpt-4o-mini:
provider: openai
complexity: [simple, moderate]
cost_per_1k: 0.00075
capabilities: [text, function-calling]
"""
from typing import Any, Dict, Tuple
from typing import Any, Dict, List, Optional, Tuple
from dataclasses import dataclass
from enum import IntEnum
from .base import FlagHandler
class TaskComplexity(IntEnum):
"""Enum for task complexity levels."""
SIMPLE = 1
MODERATE = 2
COMPLEX = 3
VERY_COMPLEX = 4
@dataclass
class ModelProfile:
"""Profile for an LLM model with its characteristics."""
name: str
provider: str
complexity_range: Tuple[TaskComplexity, TaskComplexity]
cost_per_1k_tokens: float
capabilities: List[str]
context_window: int = 128000
supports_tools: bool = True
supports_streaming: bool = True
strengths: Optional[List[str]] = None
class RouterHandler(FlagHandler):
"""
Handler for --router flag.
@@ -20,13 +53,82 @@ class RouterHandler(FlagHandler):
praisonai "Complex analysis" --router
"""
# Model tiers for routing
MODEL_TIERS = {
'simple': ['gpt-4o-mini', 'claude-3-haiku', 'gemini-1.5-flash'],
'medium': ['gpt-4o', 'claude-3-sonnet', 'gemini-1.5-pro'],
'complex': ['gpt-4-turbo', 'claude-3-opus', 'o1-preview'],
# Default model tiers for routing (used when no custom models configured)
# Note: Use full model names for API compatibility
DEFAULT_MODEL_TIERS = {
'simple': ['gpt-4o-mini', 'anthropic/claude-3-haiku-20240307', 'gemini/gemini-1.5-flash'],
'medium': ['gpt-4o', 'anthropic/claude-3-5-sonnet-20241022', 'gemini/gemini-1.5-pro'],
'complex': ['gpt-4-turbo', 'anthropic/claude-3-opus-20240229', 'o1-preview'],
}
# Default model profiles
DEFAULT_MODELS: List[ModelProfile] = [
ModelProfile(
name="gpt-4o-mini",
provider="openai",
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE),
cost_per_1k_tokens=0.00075,
capabilities=["text", "function-calling"],
context_window=128000
),
ModelProfile(
name="gpt-4o",
provider="openai",
complexity_range=(TaskComplexity.MODERATE, TaskComplexity.COMPLEX),
cost_per_1k_tokens=0.0075,
capabilities=["text", "vision", "function-calling"],
context_window=128000
),
ModelProfile(
name="anthropic/claude-3-haiku-20240307",
provider="anthropic",
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE),
cost_per_1k_tokens=0.0008,
capabilities=["text", "function-calling"],
context_window=200000
),
ModelProfile(
name="anthropic/claude-3-5-sonnet-20241022",
provider="anthropic",
complexity_range=(TaskComplexity.MODERATE, TaskComplexity.VERY_COMPLEX),
cost_per_1k_tokens=0.009,
capabilities=["text", "vision", "function-calling"],
context_window=200000,
strengths=["reasoning", "code-generation", "analysis"]
),
ModelProfile(
name="gemini/gemini-1.5-flash",
provider="google",
complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE),
cost_per_1k_tokens=0.000125,
capabilities=["text", "vision", "function-calling"],
context_window=1048576
),
ModelProfile(
name="gemini/gemini-1.5-pro",
provider="google",
complexity_range=(TaskComplexity.COMPLEX, TaskComplexity.VERY_COMPLEX),
cost_per_1k_tokens=0.00625,
capabilities=["text", "vision", "function-calling"],
context_window=2097152
),
]
def __init__(self, verbose: bool = False):
"""Initialize RouterHandler with model routing support."""
super().__init__(verbose=verbose)
self._model_router: Optional[List[ModelProfile]] = None
self._custom_models: Dict[str, ModelProfile] = {}
self._cost_threshold: Optional[float] = None
self._model_tiers = self.DEFAULT_MODEL_TIERS.copy()
@property
def model_router(self) -> List[ModelProfile]:
"""Get the list of available model profiles."""
if self._model_router is None:
return self.DEFAULT_MODELS.copy()
return self._model_router
@property
def feature_name(self) -> str:
return "router"
@@ -94,17 +196,29 @@ class RouterHandler(FlagHandler):
Selected model name
"""
complexity = self.analyze_complexity(prompt)
models = self.MODEL_TIERS.get(complexity, self.MODEL_TIERS['medium'])
# If custom models are loaded, use ModelProfile-based selection
if self._custom_models:
return self._select_from_profiles(prompt, complexity, preferred_provider)
# Otherwise use tier-based selection
models = self._model_tiers.get(complexity, self._model_tiers['medium'])
# Filter by provider if specified
if preferred_provider:
provider_prefixes = {
'openai': ['gpt-', 'o1-'],
'anthropic': ['claude-'],
'google': ['gemini-']
'anthropic': ['anthropic/', 'claude-'],
'google': ['gemini/', 'gemini-']
}
prefixes = provider_prefixes.get(preferred_provider.lower(), [])
filtered = [m for m in models if any(m.startswith(p) for p in prefixes)]
filtered = [m for m in models if any(p in m for p in prefixes)]
if filtered:
models = filtered
# Apply cost threshold if set
if self._cost_threshold:
filtered = [m for m in models if self.get_model_cost(m) <= self._cost_threshold]
if filtered:
models = filtered
@@ -115,6 +229,48 @@ class RouterHandler(FlagHandler):
return selected
def _select_from_profiles(
self,
prompt: str,
complexity: str,
preferred_provider: str = None
) -> str:
"""Select model from custom ModelProfile objects."""
# Map string complexity to TaskComplexity enum
complexity_map = {
'simple': TaskComplexity.SIMPLE,
'medium': TaskComplexity.MODERATE,
'complex': TaskComplexity.COMPLEX,
}
task_complexity = complexity_map.get(complexity, TaskComplexity.MODERATE)
# Get all models that support this complexity level
candidates = []
for name, profile in self._custom_models.items():
min_c, max_c = profile.complexity_range
if min_c.value <= task_complexity.value <= max_c.value:
candidates.append(profile)
if not candidates:
return 'gpt-4o-mini'
# Filter by provider if specified
if preferred_provider:
filtered = [m for m in candidates if m.provider == preferred_provider]
if filtered:
candidates = filtered
# Filter by cost threshold if set
if self._cost_threshold:
filtered = [m for m in candidates if m.cost_per_1k_tokens <= self._cost_threshold]
if filtered:
candidates = filtered
# Sort by cost (cheapest first)
candidates.sort(key=lambda m: m.cost_per_1k_tokens)
return candidates[0].name if candidates else 'gpt-4o-mini'
def apply_to_agent_config(self, config: Dict[str, Any], flag_value: Any) -> Dict[str, Any]:
"""
Apply router configuration.
@@ -158,3 +314,153 @@ class RouterHandler(FlagHandler):
return 'gpt-4o-mini' # Default
return self.select_model(prompt, provider)
# =========================================================================
# Custom Model Configuration Methods
# =========================================================================
def yaml_to_model_profile(self, name: str, config: Dict[str, Any]) -> ModelProfile:
"""
Convert YAML model configuration to ModelProfile object.
Args:
name: Model name
config: Model configuration from YAML
Returns:
ModelProfile object
"""
# Parse complexity strings to TaskComplexity enum
complexity_list = config.get('complexity', ['moderate'])
complexity_map = {
'simple': TaskComplexity.SIMPLE,
'moderate': TaskComplexity.MODERATE,
'complex': TaskComplexity.COMPLEX,
'very_complex': TaskComplexity.VERY_COMPLEX,
}
complexities = [complexity_map.get(c.lower(), TaskComplexity.MODERATE) for c in complexity_list]
min_complexity = min(complexities, key=lambda x: x.value)
max_complexity = max(complexities, key=lambda x: x.value)
return ModelProfile(
name=name,
provider=config.get('provider', 'unknown'),
complexity_range=(min_complexity, max_complexity),
cost_per_1k_tokens=config.get('cost_per_1k', 0.001),
capabilities=config.get('capabilities', ['text']),
context_window=config.get('context_window', 128000),
supports_tools=config.get('supports_tools', True),
supports_streaming=config.get('supports_streaming', True),
strengths=config.get('strengths', None)
)
def load_models_from_config(
self,
models_config: Dict[str, Dict[str, Any]],
merge_with_defaults: bool = False
) -> None:
"""
Load custom models from YAML configuration.
Args:
models_config: Dictionary of model configurations from YAML
merge_with_defaults: If True, merge with default models
"""
# Start with defaults if merging
if merge_with_defaults:
for profile in self.DEFAULT_MODELS:
self._custom_models[profile.name] = profile
# Add custom models
for name, config in models_config.items():
profile = self.yaml_to_model_profile(name, config)
self._custom_models[name] = profile
# Rebuild model tiers from profiles
self._rebuild_model_tiers()
def _rebuild_model_tiers(self) -> None:
"""Rebuild MODEL_TIERS from custom model profiles."""
self._model_tiers = {
'simple': [],
'medium': [],
'complex': [],
}
for name, profile in self._custom_models.items():
min_c, max_c = profile.complexity_range
if min_c.value <= TaskComplexity.SIMPLE.value <= max_c.value:
self._model_tiers['simple'].append(name)
if min_c.value <= TaskComplexity.MODERATE.value <= max_c.value:
self._model_tiers['medium'].append(name)
if min_c.value <= TaskComplexity.COMPLEX.value <= max_c.value:
self._model_tiers['complex'].append(name)
# Sort by cost (cheapest first)
for tier in self._model_tiers:
self._model_tiers[tier].sort(
key=lambda n: self._custom_models[n].cost_per_1k_tokens
)
def get_available_models(self) -> List[str]:
"""
Get list of available model names.
Returns:
List of model names
"""
if self._custom_models:
return list(self._custom_models.keys())
return [p.name for p in self.DEFAULT_MODELS]
def get_model_cost(self, model_name: str) -> float:
"""
Get cost per 1k tokens for a model.
Args:
model_name: Name of the model
Returns:
Cost per 1k tokens
"""
# Check custom models first
if model_name in self._custom_models:
return self._custom_models[model_name].cost_per_1k_tokens
# Check default models
for profile in self.DEFAULT_MODELS:
if profile.name == model_name:
return profile.cost_per_1k_tokens
# Default cost for unknown models
return 0.01
def set_cost_threshold(self, threshold: float) -> None:
"""
Set maximum cost per 1k tokens for model selection.
Args:
threshold: Maximum cost per 1k tokens
"""
self._cost_threshold = threshold
def get_model_profile(self, model_name: str) -> Optional[ModelProfile]:
"""
Get ModelProfile for a model.
Args:
model_name: Name of the model
Returns:
ModelProfile or None if not found
"""
if model_name in self._custom_models:
return self._custom_models[model_name]
for profile in self.DEFAULT_MODELS:
if profile.name == model_name:
return profile
return None
@@ -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
@@ -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 <task description>[/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)
@@ -199,6 +199,73 @@ steps:
repeat:
until: "approved"
max_iterations: 3
""",
"model-routing": """# Model Routing Workflow
name: Cost-Optimized Model Routing Workflow
description: Uses custom models with automatic routing based on task complexity
# Custom models configuration - define your own models with costs and capabilities
models:
cheap-fast:
provider: openai
complexity: [simple]
cost_per_1k: 0.0001
capabilities: [text]
context_window: 16000
balanced:
provider: openai
complexity: [moderate]
cost_per_1k: 0.001
capabilities: [text, function-calling]
context_window: 128000
premium:
provider: anthropic
complexity: [complex, very_complex]
cost_per_1k: 0.015
capabilities: [text, vision, function-calling]
context_window: 200000
strengths: [reasoning, analysis, code-generation]
workflow:
verbose: true
router: true # Enable model routing
routing_strategy: cost-optimized # Options: auto, cost-optimized, performance-optimized
agents:
classifier:
name: Classifier
role: Request Classifier
goal: Classify incoming requests by complexity
instructions: "Classify the request as 'simple', 'moderate', or 'complex'."
llm: cheap-fast # Always use cheap model for classification
researcher:
name: Researcher
role: Research Analyst
goal: Research topics thoroughly
instructions: "Research the topic and provide detailed findings."
llm_routing: auto # Auto-select based on task complexity
llm_models: [balanced, premium] # Models to choose from
writer:
name: Writer
role: Content Writer
goal: Write high-quality content
instructions: "Write clear, engaging content based on research."
llm: premium # Always use premium for quality writing
steps:
- agent: classifier
action: "Classify complexity of: {{input}}"
- name: routing
route:
simple: [researcher]
moderate: [researcher]
complex: [researcher, writer]
default: [researcher]
"""
}
@@ -224,7 +291,7 @@ Workflow Commands:
praisonai workflow create --template <name> - Create from template
praisonai workflow auto "topic" --pattern <pattern> - Auto-generate workflow
Templates: simple, routing, parallel, loop, evaluator-optimizer
Templates: simple, routing, parallel, loop, evaluator-optimizer, model-routing
Patterns: sequential, routing, parallel
Example:
@@ -610,7 +677,7 @@ Example:
return None
# Validate pattern
valid_patterns = ["sequential", "routing", "parallel", "loop"]
valid_patterns = ["sequential", "routing", "parallel", "loop", "orchestrator-workers", "evaluator-optimizer"]
if pattern not in valid_patterns:
self.print_status(f"Unknown pattern: {pattern}", "error")
self.print_status(f"Valid patterns: {', '.join(valid_patterns)}", "info")
+84 -1
View File
@@ -209,6 +209,11 @@ class PraisonAI:
invocation_cmd = "praisonai"
version_string = f"PraisonAI version {__version__}"
# Handle --interactive flag - start interactive TUI mode
if getattr(args, 'interactive', False):
self._start_interactive_mode(args)
return None
self.framework = args.framework or self.framework
# Update config_list model if --model flag is provided
@@ -690,6 +695,15 @@ class PraisonAI:
parser.add_argument("--port", type=int, default=8005, help="Server port (default: 8005)")
parser.add_argument("--host", type=str, default="127.0.0.1", help="Server host (default: 127.0.0.1)")
# Interactive TUI - terminal interface with slash commands
parser.add_argument("--interactive", "-i", action="store_true", help="Start interactive terminal mode with slash commands")
# Autonomy Mode - control AI action approval
parser.add_argument("--autonomy", type=str, choices=["suggest", "auto_edit", "full_auto"], help="Set autonomy mode for AI actions")
# Sandbox Execution - secure command execution
parser.add_argument("--sandbox", type=str, choices=["off", "basic", "strict"], help="Enable sandboxed command execution")
# If we're in a test environment, parse with empty args to avoid pytest interference
if in_test_env:
args, unknown_args = parser.parse_known_args([])
@@ -2002,7 +2016,7 @@ class PraisonAI:
return
# Validate pattern
valid_patterns = ["sequential", "routing", "parallel", "loop"]
valid_patterns = ["sequential", "routing", "parallel", "loop", "orchestrator-workers", "evaluator-optimizer"]
if pattern not in valid_patterns:
print(f"[red]Unknown pattern: {pattern}[/red]")
print(f"[cyan]Valid patterns: {', '.join(valid_patterns)}[/cyan]")
@@ -3605,6 +3619,75 @@ Now, {final_instruction.lower()}:"""
print(f"[red]ERROR: Research failed: {e}[/red]")
sys.exit(1)
def _start_interactive_mode(self, args):
"""Start interactive TUI mode with slash commands."""
try:
from praisonai.cli.features import InteractiveTUIHandler, SlashCommandHandler, CostTrackerHandler
from praisonai.cli.features.interactive_tui import InteractiveConfig
# Initialize handlers
slash_handler = SlashCommandHandler()
cost_handler = CostTrackerHandler()
cost_handler.initialize()
# Configure interactive session
config = InteractiveConfig(
prompt="praisonai> ",
multiline=False,
enable_completions=True,
show_status_bar=True
)
# Define input handler
def on_input(text):
"""Handle regular input - send to agent."""
if not text.strip():
return ""
# Track cost if enabled
result = self.handle_direct_prompt(text)
return result
# Define command handler
def on_command(cmd):
"""Handle slash commands."""
result = slash_handler.execute(cmd)
if result.get("type") == "exit":
return result
# Handle special commands
if result.get("type") == "cost":
return {"type": "info", "message": cost_handler.get_summary()}
return result
# Initialize TUI handler
tui_handler = InteractiveTUIHandler()
session = tui_handler.initialize(
config=config,
on_input=on_input,
on_command=on_command
)
# Add slash commands for completion
session.add_commands(["help", "exit", "quit", "cost", "tokens", "model", "diff", "commit", "undo", "plan", "map", "clear", "settings"])
print("[bold green]PraisonAI Interactive Mode[/bold green]")
print("Type your prompt or use /help for commands. /exit to quit.\n")
# Run the interactive session
tui_handler.run()
except ImportError as e:
print(f"[red]ERROR: Interactive mode requires additional dependencies: {e}[/red]")
print("Install with: pip install prompt_toolkit")
sys.exit(1)
except KeyboardInterrupt:
print("\n[yellow]Exiting interactive mode.[/yellow]")
except Exception as e:
print(f"[red]ERROR: Interactive mode failed: {e}[/red]")
sys.exit(1)
if __name__ == "__main__":
praison_ai = PraisonAI()
praison_ai.main()
+93
View File
@@ -0,0 +1,93 @@
"""
PraisonAI Code - AI-powered code editing tools.
This module provides tools for AI agents to read, write, and modify code files
with precision using search/replace diff strategies similar to Kilo Code.
Features:
- read_file: Read file contents with optional line ranges
- write_file: Create or overwrite files
- list_files: List directory contents with filtering
- apply_diff: Apply SEARCH/REPLACE diffs with fuzzy matching
- search_replace: Multiple search/replace operations
- execute_command: Run shell commands safely
Example:
from praisonai.code import read_file, write_file, apply_diff
# Read a file
content = read_file("path/to/file.py")
# Write a file
write_file("path/to/new_file.py", "print('hello')")
# Apply a diff
result = apply_diff("path/to/file.py", diff_content)
"""
from .tools import (
read_file,
write_file,
list_files,
apply_diff,
search_replace,
execute_command,
)
from .diff import (
DiffResult,
apply_search_replace_diff,
parse_diff_blocks,
)
from .utils import (
FileAccessController,
add_line_numbers,
strip_line_numbers,
)
from .agent_tools import (
set_workspace,
get_workspace,
code_read_file,
code_write_file,
code_list_files,
code_apply_diff,
code_search_replace,
code_execute_command,
CODE_TOOLS,
create_diff_block,
create_multi_diff,
)
__all__ = [
# Main tools (low-level)
"read_file",
"write_file",
"list_files",
"apply_diff",
"search_replace",
"execute_command",
# Agent tools (high-level, string returns)
"set_workspace",
"get_workspace",
"code_read_file",
"code_write_file",
"code_list_files",
"code_apply_diff",
"code_search_replace",
"code_execute_command",
"CODE_TOOLS",
"create_diff_block",
"create_multi_diff",
# Diff utilities
"DiffResult",
"apply_search_replace_diff",
"parse_diff_blocks",
# Utilities
"FileAccessController",
"add_line_numbers",
"strip_line_numbers",
]
__version__ = "0.1.0"
+344
View File
@@ -0,0 +1,344 @@
"""
Agent-compatible tool wrappers for PraisonAI Code.
These functions are designed to be used as tools with PraisonAI agents.
They return string results suitable for agent consumption.
"""
import os
from typing import Optional
from .tools import (
read_file as _read_file,
write_file as _write_file,
list_files as _list_files,
apply_diff as _apply_diff,
search_replace as _search_replace,
execute_command as _execute_command,
)
from .tools.apply_diff import create_diff_block, create_multi_diff
# Global workspace setting for tools
_workspace_root: Optional[str] = None
def set_workspace(workspace_path: str) -> None:
"""
Set the workspace root for all code tools.
Args:
workspace_path: Absolute path to the workspace root
"""
global _workspace_root
_workspace_root = os.path.abspath(workspace_path)
def get_workspace() -> Optional[str]:
"""Get the current workspace root."""
return _workspace_root
def code_read_file(
path: str,
start_line: Optional[int] = None,
end_line: Optional[int] = None,
) -> str:
"""
Read the contents of a file.
Use this tool to read source code files. The content is returned with
line numbers for easy reference when making edits.
Args:
path: Path to the file (relative to workspace or absolute)
start_line: First line to read (1-indexed, optional)
end_line: Last line to read (1-indexed, optional)
Returns:
File content with line numbers, or error message
Example:
>>> content = code_read_file("src/main.py")
>>> content = code_read_file("src/main.py", start_line=10, end_line=50)
"""
result = _read_file(
path=path,
start_line=start_line,
end_line=end_line,
add_line_nums=True,
workspace=_workspace_root,
)
if result['success']:
header = f"File: {path}"
if start_line or end_line:
header += f" (lines {result['start_line']}-{result['end_line']} of {result['total_lines']})"
else:
header += f" ({result['total_lines']} lines)"
return f"{header}\n\n{result['content']}"
else:
return f"Error reading {path}: {result['error']}"
def code_write_file(path: str, content: str) -> str:
"""
Write content to a file, creating it if it doesn't exist.
Use this tool to create new files or completely replace existing files.
For partial modifications, use code_apply_diff instead.
Args:
path: Path to the file (relative to workspace or absolute)
content: Content to write to the file
Returns:
Success message or error
Example:
>>> result = code_write_file("src/new_module.py", "def hello():\\n print('Hello!')")
"""
result = _write_file(
path=path,
content=content,
workspace=_workspace_root,
create_directories=True,
strip_code_fences=True,
)
if result['success']:
action = "Created" if result['created'] else "Updated"
return f"{action} file: {path} ({result['bytes_written']} bytes)"
else:
return f"Error writing {path}: {result['error']}"
def code_list_files(
path: str = ".",
recursive: bool = False,
extensions: Optional[str] = None,
) -> str:
"""
List files in a directory.
Use this tool to explore the codebase structure.
Args:
path: Directory path (relative to workspace or absolute)
recursive: Whether to list recursively
extensions: Comma-separated list of extensions to filter (e.g., "py,js,ts")
Returns:
Formatted list of files and directories
Example:
>>> files = code_list_files("src", recursive=True, extensions="py")
"""
ext_list = None
if extensions:
ext_list = [e.strip() for e in extensions.split(',')]
result = _list_files(
path=path,
recursive=recursive,
workspace=_workspace_root,
extensions=ext_list,
)
if result['success']:
lines = [f"Contents of {path}:"]
# Add directories
for d in result['directories']:
lines.append(f" 📁 {d['path']}/")
# Add files
for f in result['files']:
size = f.get('size', 0)
if size < 1024:
size_str = f"{size}B"
elif size < 1024 * 1024:
size_str = f"{size / 1024:.1f}KB"
else:
size_str = f"{size / (1024 * 1024):.1f}MB"
lines.append(f" 📄 {f['path']} ({size_str})")
if result.get('truncated'):
lines.append(f"\n ... (showing {result['total_count']} items, more available)")
return '\n'.join(lines)
else:
return f"Error listing {path}: {result['error']}"
def code_apply_diff(path: str, diff: str) -> str:
"""
Apply a SEARCH/REPLACE diff to modify a file.
Use this tool for precise, surgical edits to existing files.
The SEARCH section must match existing content exactly.
Diff Format:
<<<<<<< SEARCH
:start_line:N
-------
[exact content to find]
=======
[new content to replace with]
>>>>>>> REPLACE
Args:
path: Path to the file to modify
diff: The SEARCH/REPLACE diff content
Returns:
Success message with number of changes, or error details
Example:
>>> diff = '''<<<<<<< SEARCH
... :start_line:10
... -------
... def old_function():
... pass
... =======
... def new_function():
... return True
... >>>>>>> REPLACE'''
>>> result = code_apply_diff("src/main.py", diff)
"""
result = _apply_diff(
path=path,
diff=diff,
workspace=_workspace_root,
fuzzy_threshold=0.95, # Allow slight fuzzy matching
)
if result['success']:
msg = f"Successfully applied {result['applied_count']} change(s) to {path}"
if result.get('failed_blocks'):
msg += f"\nWarning: {len(result['failed_blocks'])} block(s) failed to apply"
return msg
else:
error_msg = f"Error applying diff to {path}: {result['error']}"
if result.get('failed_blocks'):
for fb in result['failed_blocks'][:3]: # Show first 3 failures
error_msg += f"\n- {fb.get('error', 'Unknown error')}"
return error_msg
def code_search_replace(
path: str,
search: str,
replace: str,
is_regex: bool = False,
) -> str:
"""
Perform a search and replace operation on a file.
Use this for simple text replacements. For complex multi-part edits,
use code_apply_diff instead.
Args:
path: Path to the file
search: Text to search for (or regex pattern if is_regex=True)
replace: Text to replace with
is_regex: Whether to treat search as a regex pattern
Returns:
Success message with replacement count, or error
Example:
>>> result = code_search_replace("src/main.py", "old_name", "new_name")
"""
result = _search_replace(
path=path,
operations=[{
'search': search,
'replace': replace,
'is_regex': is_regex,
}],
workspace=_workspace_root,
)
if result['success']:
return f"Replaced {result['total_replacements']} occurrence(s) in {path}"
else:
return f"Error: {result['error']}"
def code_execute_command(
command: str,
cwd: Optional[str] = None,
) -> str:
"""
Execute a shell command.
Use this tool to run commands like tests, linters, or build tools.
Args:
command: The command to execute
cwd: Working directory (relative to workspace or absolute)
Returns:
Command output (stdout and stderr)
Example:
>>> result = code_execute_command("python -m pytest tests/")
>>> result = code_execute_command("npm test", cwd="frontend")
"""
work_dir = cwd
if work_dir and _workspace_root and not os.path.isabs(work_dir):
work_dir = os.path.join(_workspace_root, work_dir)
elif not work_dir and _workspace_root:
work_dir = _workspace_root
result = _execute_command(
command=command,
cwd=work_dir,
timeout=120,
)
output_parts = []
if result['stdout']:
output_parts.append(f"stdout:\n{result['stdout']}")
if result['stderr']:
output_parts.append(f"stderr:\n{result['stderr']}")
if result['success']:
status = f"Command completed successfully (exit code: {result['exit_code']})"
else:
if result.get('error'):
status = f"Command failed: {result['error']}"
else:
status = f"Command failed (exit code: {result['exit_code']})"
output_parts.insert(0, status)
return '\n\n'.join(output_parts)
# Export all agent tools
CODE_TOOLS = [
code_read_file,
code_write_file,
code_list_files,
code_apply_diff,
code_search_replace,
code_execute_command,
]
__all__ = [
'set_workspace',
'get_workspace',
'code_read_file',
'code_write_file',
'code_list_files',
'code_apply_diff',
'code_search_replace',
'code_execute_command',
'CODE_TOOLS',
'create_diff_block',
'create_multi_diff',
]
@@ -0,0 +1,21 @@
"""
Diff handling for PraisonAI Code module.
Provides SEARCH/REPLACE diff parsing and application with fuzzy matching.
"""
from .diff_strategy import (
DiffResult,
DiffBlock,
apply_search_replace_diff,
parse_diff_blocks,
validate_diff_format,
)
__all__ = [
"DiffResult",
"DiffBlock",
"apply_search_replace_diff",
"parse_diff_blocks",
"validate_diff_format",
]
@@ -0,0 +1,432 @@
"""
Multi-Search-Replace Diff Strategy for PraisonAI Code.
This module implements a diff strategy that applies SEARCH/REPLACE blocks
to file content with fuzzy matching support, similar to Kilo Code's approach.
Diff Format:
<<<<<<< SEARCH
:start_line:N
-------
[exact content to find]
=======
[new content to replace with]
>>>>>>> REPLACE
"""
import re
from dataclasses import dataclass, field
from typing import List, Optional, Tuple
from ..utils.text_utils import (
get_similarity,
fuzzy_search,
get_indentation,
)
from ..utils.file_utils import (
strip_line_numbers,
every_line_has_line_numbers,
add_line_numbers,
detect_line_ending,
)
# Default buffer lines for fuzzy search around start_line hint
BUFFER_LINES = 40
@dataclass
class DiffBlock:
"""
Represents a single SEARCH/REPLACE block.
Attributes:
search_content: The content to search for
replace_content: The content to replace with
start_line: Optional line number hint (1-indexed)
end_line: Optional end line hint (1-indexed)
"""
search_content: str
replace_content: str
start_line: Optional[int] = None
end_line: Optional[int] = None
@dataclass
class DiffResult:
"""
Result of applying a diff operation.
Attributes:
success: Whether the diff was applied successfully
content: The resulting content (if successful)
error: Error message (if failed)
applied_count: Number of blocks successfully applied
failed_blocks: List of failed block results
"""
success: bool
content: Optional[str] = None
error: Optional[str] = None
applied_count: int = 0
failed_blocks: List[dict] = field(default_factory=list)
def _unescape_markers(content: str) -> str:
"""
Unescape escaped diff markers in content.
Args:
content: Content with potentially escaped markers
Returns:
Content with markers unescaped
"""
replacements = [
(r'^\\\<\<\<\<\<\<\<', '<<<<<<<'),
(r'^\\=======', '======='),
(r'^\\\>\>\>\>\>\>\>', '>>>>>>>'),
(r'^\\-------', '-------'),
(r'^\\:end_line:', ':end_line:'),
(r'^\\:start_line:', ':start_line:'),
]
result = content
for pattern, replacement in replacements:
result = re.sub(pattern, replacement, result, flags=re.MULTILINE)
return result
def validate_diff_format(diff_content: str) -> Tuple[bool, Optional[str]]:
"""
Validate that a diff string has correct marker sequencing.
Args:
diff_content: The diff content to validate
Returns:
Tuple of (is_valid, error_message)
"""
# State machine states
STATE_START = 0
STATE_AFTER_SEARCH = 1
STATE_AFTER_SEPARATOR = 2
state = STATE_START
line_num = 0
SEARCH_PATTERN = re.compile(r'^<<<<<<< SEARCH>?$')
SEP = '======='
REPLACE = '>>>>>>> REPLACE'
lines = diff_content.split('\n')
for line in lines:
line_num += 1
marker = line.strip()
# Check for line markers in REPLACE sections
if state == STATE_AFTER_SEPARATOR:
if marker.startswith(':start_line:') and not line.strip().startswith('\\:start_line:'):
return False, f"ERROR: Line marker ':start_line:' found in REPLACE section at line {line_num}"
if marker.startswith(':end_line:') and not line.strip().startswith('\\:end_line:'):
return False, f"ERROR: Line marker ':end_line:' found in REPLACE section at line {line_num}"
if state == STATE_START:
if marker == SEP:
return False, f"ERROR: Unexpected '=======' at line {line_num}, expected '<<<<<<< SEARCH'"
if marker == REPLACE:
return False, f"ERROR: Unexpected '>>>>>>> REPLACE' at line {line_num}"
if SEARCH_PATTERN.match(marker):
state = STATE_AFTER_SEARCH
elif state == STATE_AFTER_SEARCH:
if SEARCH_PATTERN.match(marker):
return False, f"ERROR: Duplicate '<<<<<<< SEARCH' at line {line_num}"
if marker == REPLACE:
return False, f"ERROR: Missing '=======' before '>>>>>>> REPLACE' at line {line_num}"
if marker == SEP:
state = STATE_AFTER_SEPARATOR
elif state == STATE_AFTER_SEPARATOR:
if SEARCH_PATTERN.match(marker):
return False, f"ERROR: Unexpected '<<<<<<< SEARCH' at line {line_num}, expected '>>>>>>> REPLACE'"
if marker == SEP:
return False, f"ERROR: Duplicate '=======' at line {line_num}"
if marker == REPLACE:
state = STATE_START
if state != STATE_START:
expected = "'======='" if state == STATE_AFTER_SEARCH else "'>>>>>>> REPLACE'"
return False, f"ERROR: Unexpected end of diff, expected {expected}"
return True, None
def parse_diff_blocks(diff_content: str) -> Tuple[List[DiffBlock], Optional[str]]:
"""
Parse SEARCH/REPLACE blocks from diff content.
Args:
diff_content: The diff content to parse
Returns:
Tuple of (list of DiffBlock, error message if any)
"""
# Validate format first
is_valid, error = validate_diff_format(diff_content)
if not is_valid:
return [], error
# Regex to match SEARCH/REPLACE blocks
# Groups: 1=start_line line, 2=start_line number, 3=end_line line, 4=end_line number,
# 5=separator line, 6=search content, 7=replace content
pattern = re.compile(
r'(?:^|\n)(?<!\\)<<<<<<< SEARCH>?\s*\n'
r'((?::start_line:\s*(\d+)\s*\n))?'
r'((?::end_line:\s*(\d+)\s*\n))?'
r'((?<!\\)-------\s*\n)?'
r'([\s\S]*?)(?:\n)?'
r'(?:(?<=\n)(?<!\\)=======\s*\n)'
r'([\s\S]*?)(?:\n)?'
r'(?:(?<=\n)(?<!\\)>>>>>>> REPLACE)(?=\n|$)',
re.MULTILINE
)
matches = list(pattern.finditer(diff_content))
if not matches:
return [], "Invalid diff format - no valid SEARCH/REPLACE blocks found"
blocks = []
for match in matches:
start_line = int(match.group(2)) if match.group(2) else None
end_line = int(match.group(4)) if match.group(4) else None
search_content = match.group(6) or ""
replace_content = match.group(7) or ""
blocks.append(DiffBlock(
search_content=search_content,
replace_content=replace_content,
start_line=start_line,
end_line=end_line,
))
# Sort by start_line (blocks without start_line go last)
blocks.sort(key=lambda b: b.start_line if b.start_line else float('inf'))
return blocks, None
def apply_search_replace_diff(
original_content: str,
diff_content: str,
fuzzy_threshold: float = 1.0,
buffer_lines: int = BUFFER_LINES,
) -> DiffResult:
"""
Apply SEARCH/REPLACE diff blocks to original content.
Args:
original_content: The original file content
diff_content: The diff content with SEARCH/REPLACE blocks
fuzzy_threshold: Similarity threshold (0.0-1.0, 1.0 = exact match)
buffer_lines: Number of lines to search around start_line hint
Returns:
DiffResult with success status and resulting content
"""
# Parse diff blocks
blocks, parse_error = parse_diff_blocks(diff_content)
if parse_error:
return DiffResult(success=False, error=parse_error)
if not blocks:
return DiffResult(
success=False,
error="No valid SEARCH/REPLACE blocks found in diff"
)
# Detect line ending from original content
line_ending = detect_line_ending(original_content)
result_lines = original_content.split('\n')
if original_content.endswith('\n'):
# Handle trailing newline
if result_lines and result_lines[-1] == '':
result_lines = result_lines[:-1]
delta = 0 # Track line number changes from previous replacements
applied_count = 0
failed_blocks = []
for block in blocks:
search_content = _unescape_markers(block.search_content)
replace_content = _unescape_markers(block.replace_content)
# Handle line numbers in content
has_all_line_numbers = (
(every_line_has_line_numbers(search_content) and every_line_has_line_numbers(replace_content)) or
(every_line_has_line_numbers(search_content) and not replace_content.strip())
)
start_line = block.start_line
if has_all_line_numbers and not start_line:
# Extract start line from first line number in search content
first_line = search_content.split('\n')[0]
match = re.match(r'\s*(\d+)', first_line)
if match:
start_line = int(match.group(1))
if has_all_line_numbers:
search_content = strip_line_numbers(search_content)
replace_content = strip_line_numbers(replace_content)
# Validate search content
if search_content == replace_content:
failed_blocks.append({
'error': 'Search and replace content are identical - no changes would be made',
'search': search_content[:100],
})
continue
search_lines = search_content.split('\n') if search_content else []
replace_lines = replace_content.split('\n') if replace_content else []
if not search_lines or (len(search_lines) == 1 and not search_lines[0]):
failed_blocks.append({
'error': 'Empty search content is not allowed',
})
continue
# Apply delta to start_line
adjusted_start = start_line + delta if start_line else None
# Initialize search
match_index = -1
best_score = 0.0
best_match_content = ""
search_chunk = '\n'.join(search_lines)
# Determine search bounds
search_start_index = 0
search_end_index = len(result_lines)
if adjusted_start:
# Try exact match at start_line first
exact_start_index = adjusted_start - 1 # Convert to 0-indexed
exact_end_index = exact_start_index + len(search_lines)
if 0 <= exact_start_index < len(result_lines):
original_chunk = '\n'.join(result_lines[exact_start_index:exact_end_index])
similarity = get_similarity(original_chunk, search_chunk)
if similarity >= fuzzy_threshold:
match_index = exact_start_index
best_score = similarity
best_match_content = original_chunk
else:
# Set bounds for buffered search
search_start_index = max(0, adjusted_start - buffer_lines - 1)
search_end_index = min(len(result_lines), adjusted_start + len(search_lines) + buffer_lines)
# If no match found yet, try fuzzy search
if match_index == -1:
score, idx, content = fuzzy_search(
result_lines, search_chunk, search_start_index, search_end_index
)
match_index = idx
best_score = score
best_match_content = content
# Check if match meets threshold
if match_index == -1 or best_score < fuzzy_threshold:
# Try aggressive line number stripping as fallback
aggressive_search = strip_line_numbers(search_content, aggressive=True)
aggressive_replace = strip_line_numbers(replace_content, aggressive=True)
if aggressive_search != search_content:
aggressive_lines = aggressive_search.split('\n') if aggressive_search else []
aggressive_chunk = '\n'.join(aggressive_lines)
score, idx, content = fuzzy_search(
result_lines, aggressive_chunk, search_start_index, search_end_index
)
if idx != -1 and score >= fuzzy_threshold:
match_index = idx
best_score = score
best_match_content = content
search_content = aggressive_search
replace_content = aggressive_replace
search_lines = aggressive_lines
replace_lines = replace_content.split('\n') if replace_content else []
# Still no match - report error
if match_index == -1 or best_score < fuzzy_threshold:
line_info = f" at line {start_line}" if start_line else ""
# Build context for error message
context_start = max(0, (adjusted_start or 1) - 1 - buffer_lines)
context_end = min(len(result_lines), (adjusted_start or 1) + len(search_lines) + buffer_lines)
original_context = '\n'.join(result_lines[context_start:context_end])
failed_blocks.append({
'error': f"No sufficiently similar match found{line_info} ({int(best_score * 100)}% similar, needs {int(fuzzy_threshold * 100)}%)",
'search_content': search_chunk[:200],
'best_match': best_match_content[:200] if best_match_content else None,
'similarity': best_score,
'context': add_line_numbers(original_context, context_start + 1)[:500],
})
continue
# Apply the replacement with indentation preservation
matched_lines = result_lines[match_index:match_index + len(search_lines)]
# Get indentation from original
original_indents = [get_indentation(line) for line in matched_lines]
search_indents = [get_indentation(line) for line in search_lines]
# Apply replacement with preserved indentation
indented_replace_lines = []
for i, line in enumerate(replace_lines):
matched_indent = original_indents[0] if original_indents else ""
current_indent = get_indentation(line)
search_base_indent = search_indents[0] if search_indents else ""
# Calculate relative indentation
search_base_level = len(search_base_indent)
current_level = len(current_indent)
relative_level = current_level - search_base_level
if relative_level < 0:
final_indent = matched_indent[:max(0, len(matched_indent) + relative_level)]
else:
final_indent = matched_indent + current_indent[search_base_level:]
indented_replace_lines.append(final_indent + line.lstrip())
# Construct final content
before_match = result_lines[:match_index]
after_match = result_lines[match_index + len(search_lines):]
result_lines = before_match + indented_replace_lines + after_match
# Update delta
delta += len(replace_lines) - len(matched_lines)
applied_count += 1
# Build final content
final_content = line_ending.join(result_lines)
if applied_count == 0:
return DiffResult(
success=False,
error="No blocks were successfully applied",
failed_blocks=failed_blocks,
)
return DiffResult(
success=True,
content=final_content,
applied_count=applied_count,
failed_blocks=failed_blocks,
)
@@ -0,0 +1,27 @@
"""
Code editing tools for PraisonAI agents.
Provides tools for:
- read_file: Read file contents with optional line ranges
- write_file: Create or overwrite files
- list_files: List directory contents
- apply_diff: Apply SEARCH/REPLACE diffs
- search_replace: Multiple search/replace operations
- execute_command: Run shell commands
"""
from .read_file import read_file
from .write_file import write_file
from .list_files import list_files
from .apply_diff import apply_diff
from .search_replace import search_replace
from .execute_command import execute_command
__all__ = [
"read_file",
"write_file",
"list_files",
"apply_diff",
"search_replace",
"execute_command",
]
@@ -0,0 +1,221 @@
"""
Apply Diff Tool for PraisonAI Code.
Provides functionality to apply SEARCH/REPLACE diffs to files
with fuzzy matching support.
"""
import os
from typing import Optional, Dict, Any
from ..diff.diff_strategy import apply_search_replace_diff
from ..utils.file_utils import (
file_exists,
is_path_within_directory,
)
def apply_diff(
path: str,
diff: str,
workspace: Optional[str] = None,
fuzzy_threshold: float = 1.0,
buffer_lines: int = 40,
backup: bool = False,
encoding: str = 'utf-8',
) -> Dict[str, Any]:
"""
Apply a SEARCH/REPLACE diff to a file.
This tool applies precise, targeted modifications to an existing file
by searching for specific sections of content and replacing them.
Diff Format:
<<<<<<< SEARCH
:start_line:N
-------
[exact content to find]
=======
[new content to replace with]
>>>>>>> REPLACE
Args:
path: Path to the file (absolute or relative to workspace)
diff: The SEARCH/REPLACE diff content
workspace: Workspace root directory (for relative paths)
fuzzy_threshold: Similarity threshold (0.0-1.0, 1.0 = exact match)
buffer_lines: Lines to search around start_line hint
backup: Whether to create a backup before modifying
encoding: File encoding (default: utf-8)
Returns:
Dictionary with:
- success: bool
- path: str
- applied_count: int (number of blocks applied)
- failed_blocks: list (details of failed blocks)
- error: str (if success is False)
Example:
>>> diff = '''
... <<<<<<< SEARCH
... :start_line:1
... -------
... def old_function():
... pass
... =======
... def new_function():
... return True
... >>>>>>> REPLACE
... '''
>>> result = apply_diff("src/main.py", diff)
"""
# Resolve path
if workspace and not os.path.isabs(path):
abs_path = os.path.abspath(os.path.join(workspace, path))
else:
abs_path = os.path.abspath(path)
# Security check
if workspace:
if not is_path_within_directory(abs_path, workspace):
return {
'success': False,
'error': f"Path '{path}' is outside the workspace",
'path': path,
'applied_count': 0,
}
# Check if file exists
if not file_exists(abs_path):
return {
'success': False,
'error': f"File not found: {path}",
'path': path,
'applied_count': 0,
}
try:
# Read original content
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
original_content = f.read()
# Apply the diff
result = apply_search_replace_diff(
original_content=original_content,
diff_content=diff,
fuzzy_threshold=fuzzy_threshold,
buffer_lines=buffer_lines,
)
if not result.success:
return {
'success': False,
'error': result.error or "Failed to apply diff",
'path': path,
'applied_count': result.applied_count,
'failed_blocks': result.failed_blocks,
}
# Create backup if requested
backup_path = None
if backup:
import time
timestamp = int(time.time())
backup_path = f"{abs_path}.backup.{timestamp}"
with open(backup_path, 'w', encoding=encoding) as f:
f.write(original_content)
# Write the modified content
with open(abs_path, 'w', encoding=encoding) as f:
f.write(result.content)
return {
'success': True,
'path': path,
'absolute_path': abs_path,
'applied_count': result.applied_count,
'failed_blocks': result.failed_blocks if result.failed_blocks else None,
'backup_path': backup_path,
}
except PermissionError:
return {
'success': False,
'error': f"Permission denied: {path}",
'path': path,
'applied_count': 0,
}
except Exception as e:
return {
'success': False,
'error': f"Error applying diff to {path}: {str(e)}",
'path': path,
'applied_count': 0,
}
def create_diff_block(
search_content: str,
replace_content: str,
start_line: Optional[int] = None,
) -> str:
"""
Create a properly formatted SEARCH/REPLACE diff block.
Args:
search_content: The content to search for
replace_content: The content to replace with
start_line: Optional line number hint
Returns:
Formatted diff block string
Example:
>>> diff = create_diff_block(
... "def old():\\n pass",
... "def new():\\n return True",
... start_line=10
... )
"""
lines = ["<<<<<<< SEARCH"]
if start_line:
lines.append(f":start_line:{start_line}")
lines.append("-------")
lines.append(search_content)
lines.append("=======")
lines.append(replace_content)
lines.append(">>>>>>> REPLACE")
return '\n'.join(lines)
def create_multi_diff(blocks: list) -> str:
"""
Create a diff with multiple SEARCH/REPLACE blocks.
Args:
blocks: List of dicts with 'search', 'replace', and optional 'start_line'
Returns:
Combined diff string with all blocks
Example:
>>> diff = create_multi_diff([
... {'search': 'old1', 'replace': 'new1', 'start_line': 5},
... {'search': 'old2', 'replace': 'new2', 'start_line': 20},
... ])
"""
diff_blocks = []
for block in blocks:
diff_block = create_diff_block(
search_content=block.get('search', ''),
replace_content=block.get('replace', ''),
start_line=block.get('start_line'),
)
diff_blocks.append(diff_block)
return '\n\n'.join(diff_blocks)
@@ -0,0 +1,275 @@
"""
Execute Command Tool for PraisonAI Code.
Provides functionality to execute shell commands safely.
"""
import os
import subprocess
import shlex
from typing import Optional, Dict, Any
# Commands that are generally safe to auto-run
SAFE_COMMANDS = {
'ls', 'dir', 'pwd', 'echo', 'cat', 'head', 'tail', 'wc',
'grep', 'find', 'which', 'whereis', 'type', 'file',
'date', 'whoami', 'hostname', 'uname',
'python', 'python3', 'node', 'npm', 'npx', 'pip', 'pip3',
'git', 'cargo', 'go', 'ruby', 'java', 'javac',
}
# Commands that should never be auto-run
DANGEROUS_COMMANDS = {
'rm', 'rmdir', 'del', 'format', 'mkfs',
'dd', 'shred', 'chmod', 'chown', 'chgrp',
'kill', 'killall', 'pkill',
'shutdown', 'reboot', 'halt', 'poweroff',
'sudo', 'su', 'doas',
'curl', 'wget', 'ssh', 'scp', 'rsync',
'mv', 'cp', # Can be destructive
}
def execute_command(
command: str,
cwd: Optional[str] = None,
workspace: Optional[str] = None,
timeout: int = 120,
capture_output: bool = True,
shell: bool = True,
env: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""
Execute a shell command and return the result.
This tool executes shell commands with safety checks and output capture.
Args:
command: The command to execute
cwd: Working directory for the command
workspace: Workspace root (for security validation)
timeout: Command timeout in seconds
capture_output: Whether to capture stdout/stderr
shell: Whether to run through shell
env: Additional environment variables
Returns:
Dictionary with:
- success: bool
- exit_code: int
- stdout: str
- stderr: str
- command: str
- error: str (if success is False)
Example:
>>> result = execute_command("python --version")
>>> print(result['stdout'])
"""
if not command or not command.strip():
return {
'success': False,
'error': "Empty command",
'command': command,
'exit_code': -1,
'stdout': '',
'stderr': '',
}
# Resolve working directory
if cwd:
if workspace and not os.path.isabs(cwd):
work_dir = os.path.abspath(os.path.join(workspace, cwd))
else:
work_dir = os.path.abspath(cwd)
elif workspace:
work_dir = workspace
else:
work_dir = os.getcwd()
# Validate working directory exists
if not os.path.isdir(work_dir):
return {
'success': False,
'error': f"Working directory not found: {cwd}",
'command': command,
'exit_code': -1,
'stdout': '',
'stderr': '',
}
# Prepare environment
cmd_env = os.environ.copy()
if env:
cmd_env.update(env)
# Set PAGER to cat to avoid interactive pagers
cmd_env['PAGER'] = 'cat'
try:
# Execute command
if shell:
result = subprocess.run(
command,
shell=True,
cwd=work_dir,
capture_output=capture_output,
timeout=timeout,
env=cmd_env,
text=True,
)
else:
# Parse command for non-shell execution
args = shlex.split(command)
result = subprocess.run(
args,
cwd=work_dir,
capture_output=capture_output,
timeout=timeout,
env=cmd_env,
text=True,
)
return {
'success': result.returncode == 0,
'exit_code': result.returncode,
'stdout': result.stdout if capture_output else '',
'stderr': result.stderr if capture_output else '',
'command': command,
'cwd': work_dir,
}
except subprocess.TimeoutExpired:
return {
'success': False,
'error': f"Command timed out after {timeout} seconds",
'command': command,
'exit_code': -1,
'stdout': '',
'stderr': '',
}
except FileNotFoundError as e:
return {
'success': False,
'error': f"Command not found: {str(e)}",
'command': command,
'exit_code': -1,
'stdout': '',
'stderr': '',
}
except PermissionError:
return {
'success': False,
'error': "Permission denied",
'command': command,
'exit_code': -1,
'stdout': '',
'stderr': '',
}
except Exception as e:
return {
'success': False,
'error': f"Error executing command: {str(e)}",
'command': command,
'exit_code': -1,
'stdout': '',
'stderr': '',
}
def is_safe_command(command: str) -> bool:
"""
Check if a command is considered safe to auto-run.
Args:
command: The command to check
Returns:
True if the command is considered safe
"""
if not command:
return False
# Get the base command (first word)
parts = shlex.split(command)
if not parts:
return False
base_cmd = os.path.basename(parts[0]).lower()
# Check against dangerous commands
if base_cmd in DANGEROUS_COMMANDS:
return False
# Check against safe commands
if base_cmd in SAFE_COMMANDS:
return True
# Default to unsafe
return False
def get_command_info(command: str) -> Dict[str, Any]:
"""
Get information about a command without executing it.
Args:
command: The command to analyze
Returns:
Dictionary with command analysis
"""
if not command:
return {
'valid': False,
'error': "Empty command",
}
try:
parts = shlex.split(command)
base_cmd = parts[0] if parts else ''
return {
'valid': True,
'base_command': base_cmd,
'arguments': parts[1:] if len(parts) > 1 else [],
'is_safe': is_safe_command(command),
'is_dangerous': os.path.basename(base_cmd).lower() in DANGEROUS_COMMANDS,
}
except ValueError as e:
return {
'valid': False,
'error': f"Invalid command syntax: {str(e)}",
}
def run_python(
code: str,
cwd: Optional[str] = None,
timeout: int = 60,
) -> Dict[str, Any]:
"""
Execute Python code and return the result.
Args:
code: Python code to execute
cwd: Working directory
timeout: Execution timeout
Returns:
Dictionary with execution result
"""
# Create a temporary command to run the code
import sys
python_cmd = sys.executable
# Escape the code for command line
escaped_code = code.replace('\\', '\\\\').replace('"', '\\"')
command = f'{python_cmd} -c "{escaped_code}"'
return execute_command(
command=command,
cwd=cwd,
timeout=timeout,
)
@@ -0,0 +1,274 @@
"""
List Files Tool for PraisonAI Code.
Provides functionality to list directory contents with filtering options.
"""
import os
from typing import Optional, List, Dict, Any
from ..utils.ignore_utils import (
should_ignore_path,
load_gitignore_patterns,
)
def list_files(
path: str,
recursive: bool = False,
workspace: Optional[str] = None,
max_files: int = 200,
include_hidden: bool = False,
extensions: Optional[List[str]] = None,
respect_gitignore: bool = True,
) -> Dict[str, Any]:
"""
List files and directories in a given path.
This tool lists the contents of a directory, optionally recursively,
with support for filtering and gitignore patterns.
Args:
path: Path to the directory (absolute or relative to workspace)
recursive: Whether to list recursively
workspace: Workspace root directory (for relative paths)
max_files: Maximum number of files to return
include_hidden: Whether to include hidden files (starting with .)
extensions: List of file extensions to include (e.g., ['py', 'js'])
respect_gitignore: Whether to respect .gitignore patterns
Returns:
Dictionary with:
- success: bool
- files: List of file entries with path, type, size
- directories: List of directory entries
- total_count: int (total items found)
- truncated: bool (True if max_files was reached)
- error: str (if success is False)
Example:
>>> result = list_files("src", recursive=True, extensions=['py'])
>>> for f in result['files']:
... print(f['path'])
"""
# Resolve path
if workspace and not os.path.isabs(path):
abs_path = os.path.abspath(os.path.join(workspace, path))
else:
abs_path = os.path.abspath(path)
# Check if directory exists
if not os.path.isdir(abs_path):
return {
'success': False,
'error': f"Directory not found: {path}",
'files': [],
'directories': [],
'total_count': 0,
}
# Load gitignore patterns
ignore_patterns = []
if respect_gitignore:
# Load from workspace root if available
if workspace:
ignore_patterns = load_gitignore_patterns(workspace)
# Also load from the target directory
ignore_patterns.extend(load_gitignore_patterns(abs_path))
files = []
directories = []
total_count = 0
truncated = False
try:
if recursive:
# Walk directory tree
for root, dirs, filenames in os.walk(abs_path):
# Filter hidden directories
if not include_hidden:
dirs[:] = [d for d in dirs if not d.startswith('.')]
# Filter ignored directories
if respect_gitignore:
base_dir = workspace or abs_path
dirs[:] = [
d for d in dirs
if not should_ignore_path(
os.path.join(root, d),
ignore_patterns,
base_dir,
is_directory=True
)
]
# Process files
for filename in filenames:
if total_count >= max_files:
truncated = True
break
# Skip hidden files
if not include_hidden and filename.startswith('.'):
continue
file_path = os.path.join(root, filename)
rel_path = os.path.relpath(file_path, abs_path)
# Check gitignore
if respect_gitignore:
base_dir = workspace or abs_path
if should_ignore_path(file_path, ignore_patterns, base_dir):
continue
# Filter by extension
if extensions:
ext = os.path.splitext(filename)[1].lstrip('.')
if ext.lower() not in [e.lower() for e in extensions]:
continue
# Get file info
try:
stat = os.stat(file_path)
files.append({
'path': rel_path,
'name': filename,
'size': stat.st_size,
'modified': stat.st_mtime,
})
total_count += 1
except OSError:
continue
if truncated:
break
# Add directories
for dirname in dirs:
if total_count >= max_files:
truncated = True
break
dir_path = os.path.join(root, dirname)
rel_path = os.path.relpath(dir_path, abs_path)
directories.append({
'path': rel_path,
'name': dirname,
})
total_count += 1
else:
# List only top level
for entry in os.listdir(abs_path):
if total_count >= max_files:
truncated = True
break
# Skip hidden files
if not include_hidden and entry.startswith('.'):
continue
entry_path = os.path.join(abs_path, entry)
# Check gitignore
if respect_gitignore:
base_dir = workspace or abs_path
is_dir = os.path.isdir(entry_path)
if should_ignore_path(entry_path, ignore_patterns, base_dir, is_dir):
continue
if os.path.isfile(entry_path):
# Filter by extension
if extensions:
ext = os.path.splitext(entry)[1].lstrip('.')
if ext.lower() not in [e.lower() for e in extensions]:
continue
try:
stat = os.stat(entry_path)
files.append({
'path': entry,
'name': entry,
'size': stat.st_size,
'modified': stat.st_mtime,
})
total_count += 1
except OSError:
continue
elif os.path.isdir(entry_path):
directories.append({
'path': entry,
'name': entry,
})
total_count += 1
# Sort results
files.sort(key=lambda x: x['path'])
directories.sort(key=lambda x: x['path'])
return {
'success': True,
'files': files,
'directories': directories,
'total_count': total_count,
'truncated': truncated,
'path': path,
}
except PermissionError:
return {
'success': False,
'error': f"Permission denied: {path}",
'files': [],
'directories': [],
'total_count': 0,
}
except Exception as e:
return {
'success': False,
'error': f"Error listing {path}: {str(e)}",
'files': [],
'directories': [],
'total_count': 0,
}
def format_file_list(result: Dict[str, Any], show_size: bool = True) -> str:
"""
Format a list_files result as a readable string.
Args:
result: Result from list_files()
show_size: Whether to show file sizes
Returns:
Formatted string representation
"""
if not result['success']:
return f"Error: {result.get('error', 'Unknown error')}"
lines = []
# Add directories first
for d in result['directories']:
lines.append(f"📁 {d['path']}/")
# Add files
for f in result['files']:
if show_size:
size = f.get('size', 0)
if size < 1024:
size_str = f"{size}B"
elif size < 1024 * 1024:
size_str = f"{size / 1024:.1f}KB"
else:
size_str = f"{size / (1024 * 1024):.1f}MB"
lines.append(f"📄 {f['path']} ({size_str})")
else:
lines.append(f"📄 {f['path']}")
if result.get('truncated'):
lines.append(f"\n... (truncated, showing {result['total_count']} of more items)")
return '\n'.join(lines)
@@ -0,0 +1,206 @@
"""
Read File Tool for PraisonAI Code.
Provides functionality to read file contents with optional line ranges
and line number annotations.
"""
import os
from typing import Optional, List, Dict, Any
from ..utils.file_utils import (
add_line_numbers,
is_binary_file,
file_exists,
is_path_within_directory,
)
def read_file(
path: str,
start_line: Optional[int] = None,
end_line: Optional[int] = None,
add_line_nums: bool = True,
workspace: Optional[str] = None,
encoding: str = 'utf-8',
) -> Dict[str, Any]:
"""
Read file contents with optional line range and line numbers.
This tool reads the contents of a file, optionally limiting to specific
line ranges and adding line number annotations for AI context.
Args:
path: Path to the file (absolute or relative to workspace)
start_line: First line to read (1-indexed, inclusive)
end_line: Last line to read (1-indexed, inclusive)
add_line_nums: Whether to add line numbers to output
workspace: Workspace root directory (for relative paths)
encoding: File encoding (default: utf-8)
Returns:
Dictionary with:
- success: bool
- content: str (file content, possibly with line numbers)
- total_lines: int (total lines in file)
- start_line: int (actual start line read)
- end_line: int (actual end line read)
- error: str (if success is False)
Example:
>>> result = read_file("src/main.py", start_line=1, end_line=50)
>>> if result['success']:
... print(result['content'])
"""
# Resolve path
if workspace and not os.path.isabs(path):
abs_path = os.path.abspath(os.path.join(workspace, path))
else:
abs_path = os.path.abspath(path)
# Security check - ensure path is within workspace if specified
if workspace:
if not is_path_within_directory(abs_path, workspace):
return {
'success': False,
'error': f"Path '{path}' is outside the workspace",
'content': None,
'total_lines': 0,
}
# Check if file exists
if not file_exists(abs_path):
return {
'success': False,
'error': f"File not found: {path}",
'content': None,
'total_lines': 0,
}
# Check if binary
if is_binary_file(abs_path):
return {
'success': False,
'error': f"Cannot read binary file: {path}",
'content': None,
'total_lines': 0,
}
try:
# Read file content
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
lines = f.readlines()
total_lines = len(lines)
# Determine line range
actual_start = start_line if start_line and start_line > 0 else 1
actual_end = end_line if end_line and end_line > 0 else total_lines
# Clamp to valid range
actual_start = max(1, min(actual_start, total_lines))
actual_end = max(actual_start, min(actual_end, total_lines))
# Extract lines (convert to 0-indexed)
selected_lines = lines[actual_start - 1:actual_end]
content = ''.join(selected_lines)
# Remove trailing newline if present
if content.endswith('\n'):
content = content[:-1]
# Add line numbers if requested
if add_line_nums:
content = add_line_numbers(content, actual_start)
return {
'success': True,
'content': content,
'total_lines': total_lines,
'start_line': actual_start,
'end_line': actual_end,
'path': path,
}
except PermissionError:
return {
'success': False,
'error': f"Permission denied: {path}",
'content': None,
'total_lines': 0,
}
except UnicodeDecodeError as e:
return {
'success': False,
'error': f"Encoding error reading {path}: {e}",
'content': None,
'total_lines': 0,
}
except Exception as e:
return {
'success': False,
'error': f"Error reading {path}: {str(e)}",
'content': None,
'total_lines': 0,
}
def read_multiple_files(
files: List[Dict[str, Any]],
workspace: Optional[str] = None,
add_line_nums: bool = True,
) -> Dict[str, Any]:
"""
Read multiple files in a single operation.
Args:
files: List of file specifications, each with:
- path: str (required)
- start_line: int (optional)
- end_line: int (optional)
workspace: Workspace root directory
add_line_nums: Whether to add line numbers
Returns:
Dictionary with:
- success: bool (True if all files read successfully)
- results: List of individual file results
- errors: List of error messages
Example:
>>> files = [
... {'path': 'src/main.py', 'start_line': 1, 'end_line': 50},
... {'path': 'src/utils.py'},
... ]
>>> result = read_multiple_files(files, workspace='/project')
"""
results = []
errors = []
all_success = True
for file_spec in files:
path = file_spec.get('path')
if not path:
errors.append("Missing 'path' in file specification")
all_success = False
continue
result = read_file(
path=path,
start_line=file_spec.get('start_line'),
end_line=file_spec.get('end_line'),
add_line_nums=add_line_nums,
workspace=workspace,
)
results.append(result)
if not result['success']:
all_success = False
errors.append(result.get('error', f"Failed to read {path}"))
return {
'success': all_success,
'results': results,
'errors': errors if errors else None,
}
@@ -0,0 +1,248 @@
"""
Search and Replace Tool for PraisonAI Code.
Provides functionality to perform multiple search/replace operations
on a file in a single operation.
"""
import os
import re
from typing import Optional, List, Dict, Any
from ..utils.file_utils import (
file_exists,
is_path_within_directory,
)
def search_replace(
path: str,
operations: List[Dict[str, Any]],
workspace: Optional[str] = None,
backup: bool = False,
encoding: str = 'utf-8',
) -> Dict[str, Any]:
"""
Perform multiple search/replace operations on a file.
This tool allows precise modifications to file content by performing
multiple search and replace operations in sequence.
Args:
path: Path to the file (absolute or relative to workspace)
operations: List of operations, each with:
- search: str (text to find)
- replace: str (text to replace with)
- is_regex: bool (optional, treat search as regex)
- count: int (optional, max replacements, -1 for all)
workspace: Workspace root directory
backup: Whether to create a backup before modifying
encoding: File encoding
Returns:
Dictionary with:
- success: bool
- path: str
- operations_applied: int
- total_replacements: int
- failed_operations: list
- error: str (if success is False)
Example:
>>> operations = [
... {'search': 'old_name', 'replace': 'new_name'},
... {'search': r'def (\\w+)\\(', 'replace': r'def renamed_\\\\1(', 'is_regex': True},
... ]
>>> result = search_replace("src/main.py", operations)
"""
# Resolve path
if workspace and not os.path.isabs(path):
abs_path = os.path.abspath(os.path.join(workspace, path))
else:
abs_path = os.path.abspath(path)
# Security check
if workspace:
if not is_path_within_directory(abs_path, workspace):
return {
'success': False,
'error': f"Path '{path}' is outside the workspace",
'path': path,
'operations_applied': 0,
'total_replacements': 0,
}
# Check if file exists
if not file_exists(abs_path):
return {
'success': False,
'error': f"File not found: {path}",
'path': path,
'operations_applied': 0,
'total_replacements': 0,
}
if not operations:
return {
'success': False,
'error': "No operations provided",
'path': path,
'operations_applied': 0,
'total_replacements': 0,
}
try:
# Read original content
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
content = f.read()
original_content = content
operations_applied = 0
total_replacements = 0
failed_operations = []
# Apply each operation
for i, op in enumerate(operations):
search = op.get('search')
replace = op.get('replace', '')
is_regex = op.get('is_regex', False)
count = op.get('count', -1) # -1 means replace all
if not search:
failed_operations.append({
'index': i,
'error': "Missing 'search' field",
})
continue
try:
if is_regex:
# Regex replacement
if count == -1:
new_content, num_subs = re.subn(search, replace, content)
else:
new_content, num_subs = re.subn(search, replace, content, count=count)
else:
# Literal replacement
if count == -1:
num_subs = content.count(search)
new_content = content.replace(search, replace)
else:
num_subs = min(content.count(search), count)
new_content = content
for _ in range(count):
if search in new_content:
new_content = new_content.replace(search, replace, 1)
else:
break
if num_subs > 0:
content = new_content
operations_applied += 1
total_replacements += num_subs
else:
failed_operations.append({
'index': i,
'search': search[:50],
'error': "No matches found",
})
except re.error as e:
failed_operations.append({
'index': i,
'search': search[:50],
'error': f"Invalid regex: {str(e)}",
})
# Check if any changes were made
if content == original_content:
return {
'success': False,
'error': "No changes were made - no matches found",
'path': path,
'operations_applied': 0,
'total_replacements': 0,
'failed_operations': failed_operations,
}
# Create backup if requested
backup_path = None
if backup:
import time
timestamp = int(time.time())
backup_path = f"{abs_path}.backup.{timestamp}"
with open(backup_path, 'w', encoding=encoding) as f:
f.write(original_content)
# Write the modified content
with open(abs_path, 'w', encoding=encoding) as f:
f.write(content)
return {
'success': True,
'path': path,
'absolute_path': abs_path,
'operations_applied': operations_applied,
'total_replacements': total_replacements,
'failed_operations': failed_operations if failed_operations else None,
'backup_path': backup_path,
}
except PermissionError:
return {
'success': False,
'error': f"Permission denied: {path}",
'path': path,
'operations_applied': 0,
'total_replacements': 0,
}
except Exception as e:
return {
'success': False,
'error': f"Error performing search/replace on {path}: {str(e)}",
'path': path,
'operations_applied': 0,
'total_replacements': 0,
}
def simple_replace(
path: str,
search: str,
replace: str,
workspace: Optional[str] = None,
is_regex: bool = False,
count: int = -1,
backup: bool = False,
encoding: str = 'utf-8',
) -> Dict[str, Any]:
"""
Perform a single search/replace operation on a file.
Convenience wrapper around search_replace for single operations.
Args:
path: Path to the file
search: Text to find
replace: Text to replace with
workspace: Workspace root directory
is_regex: Whether search is a regex pattern
count: Max replacements (-1 for all)
backup: Whether to create a backup
encoding: File encoding
Returns:
Dictionary with success status and details
"""
return search_replace(
path=path,
operations=[{
'search': search,
'replace': replace,
'is_regex': is_regex,
'count': count,
}],
workspace=workspace,
backup=backup,
encoding=encoding,
)
@@ -0,0 +1,217 @@
"""
Write File Tool for PraisonAI Code.
Provides functionality to create or overwrite files with content.
"""
import os
from typing import Optional, Dict, Any
from ..utils.file_utils import (
file_exists,
create_directories_for_file,
is_path_within_directory,
detect_line_ending,
normalize_line_endings,
)
from ..utils.text_utils import (
unescape_html_entities,
strip_markdown_code_fences,
)
def write_file(
path: str,
content: str,
workspace: Optional[str] = None,
create_directories: bool = True,
backup: bool = False,
strip_code_fences: bool = True,
encoding: str = 'utf-8',
) -> Dict[str, Any]:
"""
Write content to a file, creating it if it doesn't exist.
This tool writes content to a file, optionally creating parent
directories and handling common AI output artifacts like code fences.
Args:
path: Path to the file (absolute or relative to workspace)
content: Content to write to the file
workspace: Workspace root directory (for relative paths)
create_directories: Whether to create parent directories
backup: Whether to create a backup of existing files
strip_code_fences: Whether to strip markdown code fences
encoding: File encoding (default: utf-8)
Returns:
Dictionary with:
- success: bool
- path: str (the path written to)
- created: bool (True if file was created, False if overwritten)
- backup_path: str (if backup was created)
- error: str (if success is False)
Example:
>>> result = write_file("src/new_file.py", "print('hello')")
>>> if result['success']:
... print(f"Wrote to {result['path']}")
"""
# Resolve path
if workspace and not os.path.isabs(path):
abs_path = os.path.abspath(os.path.join(workspace, path))
else:
abs_path = os.path.abspath(path)
# Security check - ensure path is within workspace if specified
if workspace:
if not is_path_within_directory(abs_path, workspace):
return {
'success': False,
'error': f"Path '{path}' is outside the workspace",
'path': path,
}
# Process content
processed_content = content
# Strip markdown code fences if requested
if strip_code_fences:
processed_content = strip_markdown_code_fences(processed_content)
# Unescape HTML entities (common in AI output)
processed_content = unescape_html_entities(processed_content)
# Check if file exists
file_existed = file_exists(abs_path)
backup_path = None
try:
# Create backup if requested and file exists
if backup and file_existed:
import time
timestamp = int(time.time())
backup_path = f"{abs_path}.backup.{timestamp}"
with open(abs_path, 'r', encoding=encoding, errors='replace') as src:
with open(backup_path, 'w', encoding=encoding) as dst:
dst.write(src.read())
# Create parent directories if needed
if create_directories:
if not create_directories_for_file(abs_path):
return {
'success': False,
'error': f"Failed to create directories for {path}",
'path': path,
}
# Preserve line endings if file exists
if file_existed:
try:
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
original_content = f.read()
original_line_ending = detect_line_ending(original_content)
processed_content = normalize_line_endings(processed_content, original_line_ending)
except Exception:
pass # If we can't read, just use the content as-is
# Write the file
with open(abs_path, 'w', encoding=encoding) as f:
f.write(processed_content)
return {
'success': True,
'path': path,
'absolute_path': abs_path,
'created': not file_existed,
'backup_path': backup_path,
'bytes_written': len(processed_content.encode(encoding)),
}
except PermissionError:
return {
'success': False,
'error': f"Permission denied: {path}",
'path': path,
}
except OSError as e:
return {
'success': False,
'error': f"OS error writing {path}: {str(e)}",
'path': path,
}
except Exception as e:
return {
'success': False,
'error': f"Error writing {path}: {str(e)}",
'path': path,
}
def append_to_file(
path: str,
content: str,
workspace: Optional[str] = None,
create_if_missing: bool = True,
encoding: str = 'utf-8',
) -> Dict[str, Any]:
"""
Append content to an existing file.
Args:
path: Path to the file
content: Content to append
workspace: Workspace root directory
create_if_missing: Create file if it doesn't exist
encoding: File encoding
Returns:
Dictionary with success status and details
"""
# Resolve path
if workspace and not os.path.isabs(path):
abs_path = os.path.abspath(os.path.join(workspace, path))
else:
abs_path = os.path.abspath(path)
# Security check
if workspace:
if not is_path_within_directory(abs_path, workspace):
return {
'success': False,
'error': f"Path '{path}' is outside the workspace",
'path': path,
}
file_existed = file_exists(abs_path)
if not file_existed and not create_if_missing:
return {
'success': False,
'error': f"File not found: {path}",
'path': path,
}
try:
# Create directories if needed
if not file_existed:
create_directories_for_file(abs_path)
# Append to file
with open(abs_path, 'a', encoding=encoding) as f:
f.write(content)
return {
'success': True,
'path': path,
'absolute_path': abs_path,
'created': not file_existed,
'bytes_appended': len(content.encode(encoding)),
}
except Exception as e:
return {
'success': False,
'error': f"Error appending to {path}: {str(e)}",
'path': path,
}
@@ -0,0 +1,46 @@
"""
Utility functions for PraisonAI Code module.
"""
from .file_utils import (
add_line_numbers,
strip_line_numbers,
every_line_has_line_numbers,
normalize_line_endings,
get_file_extension,
is_binary_file,
create_directories_for_file,
file_exists,
)
from .text_utils import (
normalize_string,
unescape_html_entities,
get_similarity,
)
from .ignore_utils import (
FileAccessController,
load_gitignore_patterns,
should_ignore_path,
)
__all__ = [
# File utilities
"add_line_numbers",
"strip_line_numbers",
"every_line_has_line_numbers",
"normalize_line_endings",
"get_file_extension",
"is_binary_file",
"create_directories_for_file",
"file_exists",
# Text utilities
"normalize_string",
"unescape_html_entities",
"get_similarity",
# Ignore utilities
"FileAccessController",
"load_gitignore_patterns",
"should_ignore_path",
]
@@ -0,0 +1,307 @@
"""
File utility functions for PraisonAI Code module.
Provides utilities for file operations including:
- Line number handling (add/strip)
- File existence and type checking
- Directory creation
- Line ending normalization
"""
import os
import re
from typing import Optional, Tuple
def add_line_numbers(content: str, start_line: int = 1) -> str:
"""
Add line numbers to content in the format: " N | content"
Args:
content: The text content to add line numbers to
start_line: The starting line number (1-indexed)
Returns:
Content with line numbers prefixed to each line
Example:
>>> add_line_numbers("hello\\nworld", 1)
' 1 | hello\\n 2 | world'
"""
if not content:
return ""
lines = content.split('\n')
max_line_num = start_line + len(lines) - 1
width = len(str(max_line_num))
numbered_lines = []
for i, line in enumerate(lines):
line_num = start_line + i
numbered_lines.append(f"{line_num:>{width}} | {line}")
return '\n'.join(numbered_lines)
def strip_line_numbers(content: str, aggressive: bool = False) -> str:
"""
Strip line numbers from content.
Handles formats like:
- " 1 | content"
- "1| content"
- " 1\tcontent"
Args:
content: Content with line numbers
aggressive: If True, also strip formats like "1:" or just leading numbers
Returns:
Content with line numbers removed
"""
if not content:
return ""
lines = content.split('\n')
stripped_lines = []
# Pattern for standard format: optional spaces, digits, optional spaces, pipe/tab, content
standard_pattern = re.compile(r'^\s*\d+\s*[|\t]\s?(.*)$')
# Aggressive pattern: digits followed by colon or just leading digits with space
aggressive_pattern = re.compile(r'^\s*\d+[:\s]\s?(.*)$')
for line in lines:
match = standard_pattern.match(line)
if match:
stripped_lines.append(match.group(1))
elif aggressive:
match = aggressive_pattern.match(line)
if match:
stripped_lines.append(match.group(1))
else:
stripped_lines.append(line)
else:
stripped_lines.append(line)
return '\n'.join(stripped_lines)
def every_line_has_line_numbers(content: str) -> bool:
"""
Check if every non-empty line in content has line numbers.
Args:
content: The content to check
Returns:
True if all non-empty lines have line numbers
"""
if not content or not content.strip():
return False
lines = content.split('\n')
pattern = re.compile(r'^\s*\d+\s*[|\t]')
for line in lines:
if line.strip(): # Only check non-empty lines
if not pattern.match(line):
return False
return True
def normalize_line_endings(content: str, target: str = '\n') -> str:
"""
Normalize line endings to a consistent format.
Args:
content: The content to normalize
target: The target line ending ('\\n' or '\\r\\n')
Returns:
Content with normalized line endings
"""
# First normalize all to \n, then convert to target
normalized = content.replace('\r\n', '\n').replace('\r', '\n')
if target == '\r\n':
normalized = normalized.replace('\n', '\r\n')
return normalized
def detect_line_ending(content: str) -> str:
"""
Detect the predominant line ending in content.
Args:
content: The content to analyze
Returns:
'\\r\\n' if Windows-style, '\\n' otherwise
"""
if '\r\n' in content:
return '\r\n'
return '\n'
def get_file_extension(file_path: str) -> str:
"""
Get the file extension from a path.
Args:
file_path: Path to the file
Returns:
File extension without the dot, or empty string
"""
ext = os.path.splitext(file_path)[1]
return ext[1:] if ext else ""
def is_binary_file(file_path: str, sample_size: int = 8192) -> bool:
"""
Check if a file is binary by reading a sample.
Args:
file_path: Path to the file
sample_size: Number of bytes to sample
Returns:
True if the file appears to be binary
"""
try:
with open(file_path, 'rb') as f:
sample = f.read(sample_size)
# Check for null bytes (common in binary files)
if b'\x00' in sample:
return True
# Check the ratio of non-text characters
text_chars = bytearray({7, 8, 9, 10, 12, 13, 27} | set(range(0x20, 0x100)) - {0x7f})
non_text = sum(1 for byte in sample if byte not in text_chars)
# If more than 30% non-text characters, consider it binary
return len(sample) > 0 and (non_text / len(sample)) > 0.30
except (IOError, OSError):
return False
def create_directories_for_file(file_path: str) -> bool:
"""
Create parent directories for a file path if they don't exist.
Args:
file_path: Path to the file
Returns:
True if directories were created or already exist
"""
try:
parent_dir = os.path.dirname(file_path)
if parent_dir:
os.makedirs(parent_dir, exist_ok=True)
return True
except OSError:
return False
def file_exists(file_path: str) -> bool:
"""
Check if a file exists.
Args:
file_path: Path to the file
Returns:
True if the file exists and is a file (not directory)
"""
return os.path.isfile(file_path)
def get_relative_path(file_path: str, base_path: str) -> str:
"""
Get the relative path from base_path to file_path.
Args:
file_path: The target file path
base_path: The base directory path
Returns:
Relative path string
"""
try:
return os.path.relpath(file_path, base_path)
except ValueError:
# On Windows, relpath fails for paths on different drives
return file_path
def is_path_within_directory(file_path: str, directory: str) -> bool:
"""
Check if a file path is within a directory (prevents path traversal).
Args:
file_path: The file path to check
directory: The directory that should contain the file
Returns:
True if file_path is within directory
"""
# Resolve to absolute paths
abs_file = os.path.abspath(file_path)
abs_dir = os.path.abspath(directory)
# Ensure directory ends with separator for proper prefix matching
if not abs_dir.endswith(os.sep):
abs_dir += os.sep
return abs_file.startswith(abs_dir) or abs_file == abs_dir.rstrip(os.sep)
def read_file_lines(file_path: str, start_line: int = 1, end_line: Optional[int] = None) -> Tuple[str, int]:
"""
Read specific lines from a file.
Args:
file_path: Path to the file
start_line: First line to read (1-indexed)
end_line: Last line to read (inclusive), None for all remaining
Returns:
Tuple of (content, total_lines)
"""
with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
lines = f.readlines()
total_lines = len(lines)
# Convert to 0-indexed
start_idx = max(0, start_line - 1)
end_idx = end_line if end_line else total_lines
selected_lines = lines[start_idx:end_idx]
content = ''.join(selected_lines)
# Remove trailing newline if present
if content.endswith('\n'):
content = content[:-1]
return content, total_lines
def count_file_lines(file_path: str) -> int:
"""
Count the number of lines in a file efficiently.
Args:
file_path: Path to the file
Returns:
Number of lines in the file
"""
count = 0
with open(file_path, 'rb') as f:
for _ in f:
count += 1
return count
@@ -0,0 +1,308 @@
"""
File access control utilities for PraisonAI Code module.
Provides utilities for:
- Loading and parsing .gitignore patterns
- Checking if paths should be ignored
- Controlling file access within workspaces
"""
import os
import fnmatch
from typing import List, Optional, Set
def load_gitignore_patterns(directory: str) -> List[str]:
"""
Load .gitignore patterns from a directory.
Args:
directory: The directory to search for .gitignore
Returns:
List of gitignore patterns
"""
patterns = []
gitignore_path = os.path.join(directory, '.gitignore')
if os.path.isfile(gitignore_path):
try:
with open(gitignore_path, 'r', encoding='utf-8') as f:
for line in f:
line = line.strip()
# Skip empty lines and comments
if line and not line.startswith('#'):
patterns.append(line)
except (IOError, OSError):
pass
return patterns
def _pattern_to_regex_parts(pattern: str) -> tuple:
"""
Convert a gitignore pattern to matching components.
Args:
pattern: A gitignore pattern
Returns:
Tuple of (is_negation, is_directory_only, cleaned_pattern)
"""
is_negation = pattern.startswith('!')
if is_negation:
pattern = pattern[1:]
is_directory_only = pattern.endswith('/')
if is_directory_only:
pattern = pattern[:-1]
return is_negation, is_directory_only, pattern
def should_ignore_path(
path: str,
patterns: List[str],
base_dir: str,
is_directory: bool = False
) -> bool:
"""
Check if a path should be ignored based on gitignore patterns.
Args:
path: The path to check (relative or absolute)
patterns: List of gitignore patterns
base_dir: The base directory for relative path calculation
is_directory: Whether the path is a directory
Returns:
True if the path should be ignored
"""
# Get relative path
if os.path.isabs(path):
try:
rel_path = os.path.relpath(path, base_dir)
except ValueError:
rel_path = path
else:
rel_path = path
# Normalize path separators
rel_path = rel_path.replace(os.sep, '/')
# Check each pattern
ignored = False
for pattern in patterns:
is_negation, is_directory_only, clean_pattern = _pattern_to_regex_parts(pattern)
# Skip directory-only patterns for files
if is_directory_only and not is_directory:
continue
# Check if pattern matches
matches = False
# Handle patterns with /
if '/' in clean_pattern:
# Pattern with / matches from root
if clean_pattern.startswith('/'):
clean_pattern = clean_pattern[1:]
matches = fnmatch.fnmatch(rel_path, clean_pattern)
else:
# Pattern without / matches any path component
path_parts = rel_path.split('/')
for part in path_parts:
if fnmatch.fnmatch(part, clean_pattern):
matches = True
break
# Also check full path for ** patterns
if not matches and '**' in clean_pattern:
matches = fnmatch.fnmatch(rel_path, clean_pattern)
if matches:
ignored = not is_negation
return ignored
class FileAccessController:
"""
Controls file access within a workspace.
Respects .gitignore patterns and can be configured with
additional ignore/allow patterns.
Attributes:
workspace_root: The root directory of the workspace
ignore_patterns: Patterns for files to ignore
protected_patterns: Patterns for files that are read-only
"""
def __init__(
self,
workspace_root: str,
load_gitignore: bool = True,
additional_ignore_patterns: Optional[List[str]] = None,
protected_patterns: Optional[List[str]] = None
):
"""
Initialize the FileAccessController.
Args:
workspace_root: The root directory of the workspace
load_gitignore: Whether to load .gitignore patterns
additional_ignore_patterns: Extra patterns to ignore
protected_patterns: Patterns for read-only files
"""
self.workspace_root = os.path.abspath(workspace_root)
self.ignore_patterns: List[str] = []
self.protected_patterns: List[str] = protected_patterns or []
self._ignored_cache: Set[str] = set()
# Load .gitignore patterns
if load_gitignore:
self.ignore_patterns.extend(load_gitignore_patterns(self.workspace_root))
# Add additional patterns
if additional_ignore_patterns:
self.ignore_patterns.extend(additional_ignore_patterns)
# Always ignore common patterns
default_ignores = [
'.git',
'.git/**',
'__pycache__',
'__pycache__/**',
'*.pyc',
'.DS_Store',
'node_modules',
'node_modules/**',
]
for pattern in default_ignores:
if pattern not in self.ignore_patterns:
self.ignore_patterns.append(pattern)
def is_path_allowed(self, path: str) -> bool:
"""
Check if a path is allowed for access.
Args:
path: The path to check
Returns:
True if the path is allowed
"""
# Resolve to absolute path
if not os.path.isabs(path):
abs_path = os.path.abspath(os.path.join(self.workspace_root, path))
else:
abs_path = os.path.abspath(path)
# Check if within workspace
if not self._is_within_workspace(abs_path):
return False
# Check if ignored
is_dir = os.path.isdir(abs_path)
if should_ignore_path(abs_path, self.ignore_patterns, self.workspace_root, is_dir):
return False
return True
def is_write_protected(self, path: str) -> bool:
"""
Check if a path is write-protected.
Args:
path: The path to check
Returns:
True if the path is write-protected
"""
if not self.protected_patterns:
return False
# Get relative path
if os.path.isabs(path):
try:
rel_path = os.path.relpath(path, self.workspace_root)
except ValueError:
rel_path = path
else:
rel_path = path
rel_path = rel_path.replace(os.sep, '/')
for pattern in self.protected_patterns:
if fnmatch.fnmatch(rel_path, pattern):
return True
return False
def validate_access(self, path: str, write: bool = False) -> bool:
"""
Validate if access to a path is allowed.
Args:
path: The path to validate
write: Whether write access is needed
Returns:
True if access is allowed
"""
if not self.is_path_allowed(path):
return False
if write and self.is_write_protected(path):
return False
return True
def _is_within_workspace(self, abs_path: str) -> bool:
"""
Check if an absolute path is within the workspace.
Args:
abs_path: The absolute path to check
Returns:
True if within workspace
"""
# Normalize paths
workspace = os.path.normpath(self.workspace_root)
target = os.path.normpath(abs_path)
# Check if target starts with workspace path
return target.startswith(workspace + os.sep) or target == workspace
def get_relative_path(self, path: str) -> str:
"""
Get the relative path from workspace root.
Args:
path: The path to convert
Returns:
Relative path string
"""
if os.path.isabs(path):
try:
return os.path.relpath(path, self.workspace_root)
except ValueError:
return path
return path
def get_absolute_path(self, path: str) -> str:
"""
Get the absolute path for a relative path.
Args:
path: The relative path
Returns:
Absolute path string
"""
if os.path.isabs(path):
return path
return os.path.abspath(os.path.join(self.workspace_root, path))
@@ -0,0 +1,276 @@
"""
Text utility functions for PraisonAI Code module.
Provides utilities for text processing including:
- String normalization (smart quotes, unicode)
- HTML entity unescaping
- Similarity calculation using Levenshtein distance
"""
import html
import re
import unicodedata
def normalize_string(text: str) -> str:
"""
Normalize a string for comparison.
Handles:
- Smart quotes to regular quotes
- Unicode normalization
- Consistent whitespace
Args:
text: The text to normalize
Returns:
Normalized text string
"""
if not text:
return ""
# Unicode normalization (NFC form)
normalized = unicodedata.normalize('NFC', text)
# Replace smart quotes with regular quotes
quote_replacements = {
'\u2018': "'", # Left single quote
'\u2019': "'", # Right single quote
'\u201C': '"', # Left double quote
'\u201D': '"', # Right double quote
'\u2013': '-', # En dash
'\u2014': '-', # Em dash
'\u2026': '...', # Ellipsis
}
for smart, regular in quote_replacements.items():
normalized = normalized.replace(smart, regular)
return normalized
def unescape_html_entities(text: str) -> str:
"""
Unescape HTML entities in text.
Args:
text: Text with HTML entities
Returns:
Text with entities converted to characters
Example:
>>> unescape_html_entities("&lt;div&gt;")
'<div>'
"""
if not text:
return ""
return html.unescape(text)
def levenshtein_distance(s1: str, s2: str) -> int:
"""
Calculate the Levenshtein distance between two strings.
This is the minimum number of single-character edits (insertions,
deletions, or substitutions) required to change one string into the other.
Args:
s1: First string
s2: Second string
Returns:
The edit distance between the strings
"""
if len(s1) < len(s2):
return levenshtein_distance(s2, s1)
if len(s2) == 0:
return len(s1)
previous_row = range(len(s2) + 1)
for i, c1 in enumerate(s1):
current_row = [i + 1]
for j, c2 in enumerate(s2):
# Cost is 0 if characters match, 1 otherwise
insertions = previous_row[j + 1] + 1
deletions = current_row[j] + 1
substitutions = previous_row[j] + (c1 != c2)
current_row.append(min(insertions, deletions, substitutions))
previous_row = current_row
return previous_row[-1]
def get_similarity(original: str, search: str) -> float:
"""
Calculate similarity ratio between two strings using Levenshtein distance.
Args:
original: The original string
search: The string to compare against
Returns:
Similarity ratio from 0.0 (completely different) to 1.0 (identical)
"""
if not search:
return 0.0
# Normalize both strings for comparison
normalized_original = normalize_string(original)
normalized_search = normalize_string(search)
if normalized_original == normalized_search:
return 1.0
# Calculate Levenshtein distance
dist = levenshtein_distance(normalized_original, normalized_search)
# Calculate similarity ratio
max_length = max(len(normalized_original), len(normalized_search))
if max_length == 0:
return 1.0
return 1.0 - (dist / max_length)
def fuzzy_search(
lines: list,
search_chunk: str,
start_index: int,
end_index: int
) -> tuple:
"""
Perform a middle-out fuzzy search to find the best match for search_chunk.
Searches outward from the middle of the range to find the slice of lines
that is most similar to the search chunk.
Args:
lines: List of lines to search in
search_chunk: The text to search for
start_index: Start of search range (inclusive)
end_index: End of search range (exclusive)
Returns:
Tuple of (best_score, best_match_index, best_match_content)
"""
best_score = 0.0
best_match_index = -1
best_match_content = ""
search_lines = search_chunk.split('\n')
search_len = len(search_lines)
if search_len == 0:
return (0.0, -1, "")
# Middle-out search from the midpoint
mid_point = (start_index + end_index) // 2
left_index = mid_point
right_index = mid_point + 1
while left_index >= start_index or right_index <= end_index - search_len:
# Check left side
if left_index >= start_index:
original_chunk = '\n'.join(lines[left_index:left_index + search_len])
similarity = get_similarity(original_chunk, search_chunk)
if similarity > best_score:
best_score = similarity
best_match_index = left_index
best_match_content = original_chunk
left_index -= 1
# Check right side
if right_index <= end_index - search_len:
original_chunk = '\n'.join(lines[right_index:right_index + search_len])
similarity = get_similarity(original_chunk, search_chunk)
if similarity > best_score:
best_score = similarity
best_match_index = right_index
best_match_content = original_chunk
right_index += 1
return (best_score, best_match_index, best_match_content)
def strip_markdown_code_fences(content: str) -> str:
"""
Strip markdown code fences from content.
Handles content that starts with ```language and ends with ```.
Args:
content: Content potentially wrapped in code fences
Returns:
Content with code fences removed
"""
if not content:
return ""
lines = content.split('\n')
# Check if starts with code fence
if lines and lines[0].strip().startswith('```'):
lines = lines[1:]
# Check if ends with code fence
if lines and lines[-1].strip() == '```':
lines = lines[:-1]
return '\n'.join(lines)
def get_indentation(line: str) -> str:
"""
Extract the leading whitespace (indentation) from a line.
Args:
line: The line to extract indentation from
Returns:
The leading whitespace string (tabs and/or spaces)
"""
match = re.match(r'^[\t ]*', line)
return match.group(0) if match else ""
def preserve_indentation(original_lines: list, replacement_lines: list) -> list:
"""
Apply replacement lines while preserving the indentation from original.
Args:
original_lines: The original lines being replaced
replacement_lines: The new lines to insert
Returns:
Replacement lines with adjusted indentation
"""
if not original_lines or not replacement_lines:
return replacement_lines
# Get base indentation from first original line
original_base_indent = get_indentation(original_lines[0])
# Get base indentation from first replacement line
replacement_base_indent = get_indentation(replacement_lines[0])
result = []
for line in replacement_lines:
current_indent = get_indentation(line)
content = line.lstrip()
# Calculate relative indentation
if len(current_indent) >= len(replacement_base_indent):
relative_indent = current_indent[len(replacement_base_indent):]
new_indent = original_base_indent + relative_indent
else:
# Line has less indentation than base, adjust accordingly
new_indent = original_base_indent[:max(0, len(original_base_indent) - (len(replacement_base_indent) - len(current_indent)))]
result.append(new_indent + content)
return result
+1 -1
View File
@@ -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.38 gunicorn markdown\n")
file.write("RUN pip install flask praisonai==2.3.47 gunicorn markdown\n")
file.write("EXPOSE 8080\n")
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
@@ -0,0 +1,56 @@
"""
PraisonAI Integrations - External CLI tool integrations.
This module provides integrations with external AI coding CLI tools:
- Claude Code CLI
- Gemini CLI
- OpenAI Codex CLI
- Cursor CLI
All integrations use lazy loading to avoid performance impact.
Usage:
from praisonai.integrations import ClaudeCodeIntegration, GeminiCLIIntegration
# Create integration
claude = ClaudeCodeIntegration(workspace="/path/to/project")
# Use as agent tool
tool = claude.as_tool()
# Or execute directly
result = await claude.execute("Refactor this code")
"""
# Lazy imports to avoid performance impact
__all__ = [
'BaseCLIIntegration',
'ClaudeCodeIntegration',
'GeminiCLIIntegration',
'CodexCLIIntegration',
'CursorCLIIntegration',
'get_available_integrations',
]
def __getattr__(name):
"""Lazy load integrations to minimize import overhead."""
if name == 'BaseCLIIntegration':
from .base import BaseCLIIntegration
return BaseCLIIntegration
elif name == 'ClaudeCodeIntegration':
from .claude_code import ClaudeCodeIntegration
return ClaudeCodeIntegration
elif name == 'GeminiCLIIntegration':
from .gemini_cli import GeminiCLIIntegration
return GeminiCLIIntegration
elif name == 'CodexCLIIntegration':
from .codex_cli import CodexCLIIntegration
return CodexCLIIntegration
elif name == 'CursorCLIIntegration':
from .cursor_cli import CursorCLIIntegration
return CursorCLIIntegration
elif name == 'get_available_integrations':
from .base import get_available_integrations
return get_available_integrations
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -0,0 +1,300 @@
"""
Base class for external CLI tool integrations.
Provides a consistent interface for integrating with AI coding CLI tools
like Claude Code, Gemini CLI, Codex CLI, and Cursor CLI.
Features:
- Lazy availability checking with caching
- Async subprocess execution
- Tool wrapper for agent integration
- Timeout handling
"""
from abc import ABC, abstractmethod
from typing import AsyncIterator, Optional, Dict, Any, Tuple, List
import asyncio
import shutil
import os
class BaseCLIIntegration(ABC):
"""
Abstract base class for external CLI tool integrations.
Provides:
- Lazy availability checking with class-level caching
- Async subprocess execution with timeout handling
- Tool wrapper for PraisonAI agent integration
Subclasses must implement:
- cli_command property: The CLI command name
- execute method: Execute the CLI and return result
- stream method: Stream output from the CLI
Example:
class MyIntegration(BaseCLIIntegration):
@property
def cli_command(self) -> str:
return "my-cli"
async def execute(self, prompt: str, **options) -> str:
return await self.execute_async(["my-cli", "-p", prompt])
async def stream(self, prompt: str, **options):
async for line in self.stream_async(["my-cli", "-p", prompt]):
yield line
"""
# Class-level cache for availability checks (shared across instances)
_availability_cache: Dict[str, bool] = {}
def __init__(self, workspace: str = ".", timeout: int = 300):
"""
Initialize the CLI integration.
Args:
workspace: Working directory for CLI execution (default: current dir)
timeout: Default timeout in seconds for CLI execution (default: 300)
"""
self.workspace = workspace
self.timeout = timeout
@property
@abstractmethod
def cli_command(self) -> str:
"""
Return the CLI command name.
This is used for availability checking and tool naming.
Returns:
str: The CLI command (e.g., "claude", "gemini", "codex", "cursor-agent")
"""
pass
@property
def is_available(self) -> bool:
"""
Check if the CLI tool is installed and available.
Uses class-level caching to avoid repeated filesystem checks.
Returns:
bool: True if the CLI is available, False otherwise
"""
if self.cli_command not in self._availability_cache:
self._availability_cache[self.cli_command] = shutil.which(self.cli_command) is not None
return self._availability_cache[self.cli_command]
@abstractmethod
async def execute(self, prompt: str, **options) -> str:
"""
Execute the CLI tool and return the result.
Args:
prompt: The prompt/query to send to the CLI
**options: Additional options for the CLI
Returns:
str: The CLI output
"""
pass
@abstractmethod
async def stream(self, prompt: str, **options) -> AsyncIterator[Dict[str, Any]]:
"""
Stream output from the CLI tool.
Args:
prompt: The prompt/query to send to the CLI
**options: Additional options for the CLI
Yields:
dict: Parsed output events from the CLI
"""
pass
async def execute_async(self, cmd: List[str], timeout: Optional[int] = None) -> str:
"""
Execute a command asynchronously and return stdout.
Args:
cmd: Command and arguments as a list
timeout: Timeout in seconds (uses self.timeout if not specified)
Returns:
str: The command's stdout
Raises:
TimeoutError: If the command times out
"""
timeout = timeout or self.timeout
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=self.workspace
)
try:
stdout, stderr = await asyncio.wait_for(
proc.communicate(),
timeout=timeout
)
return stdout.decode()
except asyncio.TimeoutError:
proc.kill()
await proc.wait()
raise TimeoutError(f"Command timed out after {timeout}s: {' '.join(cmd)}")
async def execute_async_with_stderr(
self,
cmd: List[str],
timeout: Optional[int] = None
) -> Tuple[str, str]:
"""
Execute a command asynchronously and return both stdout and stderr.
Args:
cmd: Command and arguments as a list
timeout: Timeout in seconds (uses self.timeout if not specified)
Returns:
Tuple[str, str]: (stdout, stderr)
Raises:
TimeoutError: If the command times out
"""
timeout = timeout or self.timeout
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=self.workspace
)
try:
stdout, stderr = await asyncio.wait_for(
proc.communicate(),
timeout=timeout
)
return stdout.decode(), stderr.decode()
except asyncio.TimeoutError:
proc.kill()
await proc.wait()
raise TimeoutError(f"Command timed out after {timeout}s: {' '.join(cmd)}")
async def stream_async(
self,
cmd: List[str],
timeout: Optional[int] = None
) -> AsyncIterator[str]:
"""
Stream stdout lines from a command asynchronously.
Args:
cmd: Command and arguments as a list
timeout: Timeout in seconds for the entire operation
Yields:
str: Each line of output
"""
timeout = timeout or self.timeout
proc = await asyncio.create_subprocess_exec(
*cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=self.workspace
)
try:
async def read_lines():
while True:
line = await proc.stdout.readline()
if not line:
break
yield line.decode().rstrip('\n')
async for line in read_lines():
yield line
except asyncio.TimeoutError:
proc.kill()
await proc.wait()
raise TimeoutError(f"Stream timed out after {timeout}s")
finally:
if proc.returncode is None:
proc.kill()
await proc.wait()
def as_tool(self) -> callable:
"""
Return a callable suitable for use as a PraisonAI agent tool.
The returned function is synchronous and wraps the async execute method.
Uses asyncio.run() for proper event loop handling (Python 3.7+).
Returns:
callable: A function that can be used as an agent tool
"""
# Capture self for closure
integration = self
def tool_func(query: str) -> str:
"""Execute the CLI tool with the given query."""
# Use asyncio.run() for clean event loop management
# This creates a new event loop, runs the coroutine, and closes it
try:
# Check if we're already in an async context
asyncio.get_running_loop()
# If we're in an async context, use ThreadPoolExecutor to avoid nested loop
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(asyncio.run, integration.execute(query))
return future.result()
except RuntimeError:
# No running loop, safe to use asyncio.run()
return asyncio.run(integration.execute(query))
# Set function metadata for agent tool registration
tool_func.__name__ = f"{self.cli_command}_tool"
tool_func.__doc__ = f"Execute {self.cli_command} for coding tasks."
return tool_func
def get_env(self) -> Dict[str, str]:
"""
Get environment variables for CLI execution.
Subclasses can override this to add tool-specific environment variables.
Returns:
dict: Environment variables to pass to the CLI
"""
return dict(os.environ)
def get_available_integrations() -> Dict[str, bool]:
"""
Get a dictionary of all integrations and their availability status.
Returns:
dict: Mapping of integration name to availability (True/False)
"""
from .claude_code import ClaudeCodeIntegration
from .gemini_cli import GeminiCLIIntegration
from .codex_cli import CodexCLIIntegration
from .cursor_cli import CursorCLIIntegration
integrations = {
'claude': ClaudeCodeIntegration(),
'gemini': GeminiCLIIntegration(),
'codex': CodexCLIIntegration(),
'cursor': CursorCLIIntegration(),
}
return {name: integration.is_available for name, integration in integrations.items()}
@@ -0,0 +1,270 @@
"""
Claude Code CLI Integration.
Provides integration with Claude Code CLI for AI-powered coding tasks.
Features:
- Headless mode execution with JSON output
- Session continuation support
- System prompt customization
- Tool restrictions
- Optional SDK integration (when claude-agent-sdk is installed)
Usage:
from praisonai.integrations import ClaudeCodeIntegration
# Create integration
claude = ClaudeCodeIntegration(workspace="/path/to/project")
# Execute a coding task
result = await claude.execute("Refactor the auth module")
# Stream output
async for event in claude.stream("Add error handling"):
print(event)
# Use as agent tool
tool = claude.as_tool()
"""
import json
import os
from typing import AsyncIterator, Dict, Any, Optional, List
from .base import BaseCLIIntegration
# Check if Claude Code SDK is available
CLAUDE_SDK_AVAILABLE = False
try:
from claude_agent_sdk import query as claude_query, ClaudeAgentOptions
CLAUDE_SDK_AVAILABLE = True
except ImportError:
pass
class ClaudeCodeIntegration(BaseCLIIntegration):
"""
Integration with Claude Code CLI.
Supports both subprocess-based execution and SDK-based execution
(when claude-agent-sdk is installed).
Attributes:
output_format: Output format ("json", "text", "stream-json")
skip_permissions: Whether to skip permission prompts
system_prompt: Custom system prompt to append
allowed_tools: List of allowed tools (e.g., ["Read", "Write", "Bash"])
disallowed_tools: List of disallowed tools
use_sdk: Whether to use the SDK instead of subprocess
"""
def __init__(
self,
workspace: str = ".",
timeout: int = 300,
output_format: str = "json",
skip_permissions: bool = True,
system_prompt: Optional[str] = None,
allowed_tools: Optional[List[str]] = None,
disallowed_tools: Optional[List[str]] = None,
use_sdk: bool = False,
model: Optional[str] = None,
):
"""
Initialize Claude Code integration.
Args:
workspace: Working directory for CLI execution
timeout: Timeout in seconds for CLI execution
output_format: Output format ("json", "text", "stream-json")
skip_permissions: Whether to skip permission prompts (--dangerously-skip-permissions)
system_prompt: Custom system prompt to append
allowed_tools: List of allowed tools
disallowed_tools: List of disallowed tools
use_sdk: Whether to use the SDK instead of subprocess
model: Model to use (e.g., "sonnet", "opus")
"""
super().__init__(workspace=workspace, timeout=timeout)
self.output_format = output_format
self.skip_permissions = skip_permissions
self.system_prompt = system_prompt
self.allowed_tools = allowed_tools
self.disallowed_tools = disallowed_tools
self.use_sdk = use_sdk and CLAUDE_SDK_AVAILABLE
self.model = model
# Session management
self._session_active = False
@property
def cli_command(self) -> str:
"""Return the CLI command name."""
return "claude"
@property
def sdk_available(self) -> bool:
"""Check if the Claude Code SDK is available."""
return CLAUDE_SDK_AVAILABLE
def _build_command(
self,
prompt: str,
continue_session: bool = False,
**options
) -> List[str]:
"""
Build the Claude CLI command.
Args:
prompt: The prompt to send
continue_session: Whether to continue a previous session
**options: Additional options
Returns:
List of command arguments
"""
cmd = ["claude"]
# Add continue flag if needed
if continue_session or self._session_active:
cmd.append("--continue")
# Add skip permissions flag
if self.skip_permissions:
cmd.append("--dangerously-skip-permissions")
# Add print mode flag
cmd.append("-p")
# Add prompt
cmd.append(prompt)
# Add output format
cmd.extend(["--output-format", self.output_format])
# Add model if specified
if self.model:
cmd.extend(["--model", self.model])
# Add system prompt if specified
if self.system_prompt:
cmd.extend(["--append-system-prompt", self.system_prompt])
# Add allowed tools if specified
if self.allowed_tools:
cmd.extend(["--allowedTools", ",".join(self.allowed_tools)])
# Add disallowed tools if specified
if self.disallowed_tools:
cmd.extend(["--disallowedTools", ",".join(self.disallowed_tools)])
return cmd
async def execute(self, prompt: str, **options) -> str:
"""
Execute Claude Code CLI and return the result.
Args:
prompt: The prompt/query to send
**options: Additional options (continue_session, etc.)
Returns:
str: The CLI output (parsed from JSON if output_format is "json")
"""
if self.use_sdk:
return await self._execute_sdk(prompt, **options)
return await self._execute_subprocess(prompt, **options)
async def _execute_subprocess(self, prompt: str, **options) -> str:
"""Execute using subprocess."""
cmd = self._build_command(prompt, **options)
output = await self.execute_async(cmd)
# Mark session as active for continuation
self._session_active = True
# Parse JSON output if applicable
if self.output_format == "json":
try:
data = json.loads(output)
# Extract the main result
if isinstance(data, dict):
return data.get("result", data.get("content", str(data)))
return str(data)
except json.JSONDecodeError:
return output
return output
async def _execute_sdk(self, prompt: str, **options) -> str:
"""Execute using the Claude Code SDK."""
if not CLAUDE_SDK_AVAILABLE:
raise RuntimeError("Claude Code SDK not available. Install with: pip install claude-agent-sdk")
sdk_options = ClaudeAgentOptions(
cwd=self.workspace,
system_prompt=self.system_prompt,
)
if self.allowed_tools:
sdk_options.allowed_tools = self.allowed_tools
if self.skip_permissions:
sdk_options.permission_mode = 'acceptEdits'
result_parts = []
async for message in claude_query(prompt=prompt, options=sdk_options):
if hasattr(message, 'content'):
for block in message.content:
if hasattr(block, 'text'):
result_parts.append(block.text)
return '\n'.join(result_parts)
async def stream(self, prompt: str, **options) -> AsyncIterator[Dict[str, Any]]:
"""
Stream output from Claude Code CLI.
Args:
prompt: The prompt/query to send
**options: Additional options
Yields:
dict: Parsed JSON events from the CLI
"""
# Use stream-json format for streaming
original_format = self.output_format
self.output_format = "stream-json"
try:
cmd = self._build_command(prompt, **options)
async for line in self.stream_async(cmd):
if line.strip():
try:
event = json.loads(line)
yield event
except json.JSONDecodeError:
yield {"type": "text", "content": line}
finally:
self.output_format = original_format
def reset_session(self):
"""Reset the session state."""
self._session_active = False
def get_env(self) -> Dict[str, str]:
"""Get environment variables for CLI execution."""
env = super().get_env()
# Add Anthropic API key if available
if "ANTHROPIC_API_KEY" in os.environ:
env["ANTHROPIC_API_KEY"] = os.environ["ANTHROPIC_API_KEY"]
elif "CLAUDE_API_KEY" in os.environ:
env["ANTHROPIC_API_KEY"] = os.environ["CLAUDE_API_KEY"]
return env
@@ -0,0 +1,252 @@
"""
OpenAI Codex CLI Integration.
Provides integration with OpenAI's Codex CLI for AI-powered coding tasks.
Features:
- Non-interactive execution with `codex exec`
- Full auto mode for file modifications
- Sandbox modes for security
- JSON streaming output
- Structured output with schemas
Usage:
from praisonai.integrations import CodexCLIIntegration
# Create integration
codex = CodexCLIIntegration(workspace="/path/to/project")
# Execute a coding task
result = await codex.execute("Fix the authentication bug")
# Execute with full auto (allows file modifications)
codex_auto = CodexCLIIntegration(full_auto=True)
result = await codex_auto.execute("Refactor the module")
# Use as agent tool
tool = codex.as_tool()
"""
import json
import os
from typing import AsyncIterator, Dict, Any, Optional, List
from .base import BaseCLIIntegration
class CodexCLIIntegration(BaseCLIIntegration):
"""
Integration with OpenAI's Codex CLI.
Attributes:
full_auto: Whether to allow file modifications
sandbox: Sandbox mode ("default", "danger-full-access")
json_output: Whether to use JSON streaming output
output_schema: Path to JSON schema for structured output
"""
def __init__(
self,
workspace: str = ".",
timeout: int = 300,
full_auto: bool = False,
sandbox: str = "default",
json_output: bool = False,
output_schema: Optional[str] = None,
output_file: Optional[str] = None,
):
"""
Initialize Codex CLI integration.
Args:
workspace: Working directory for CLI execution
timeout: Timeout in seconds for CLI execution
full_auto: Whether to allow file modifications (--full-auto)
sandbox: Sandbox mode ("default", "danger-full-access")
json_output: Whether to use JSON streaming output (--json)
output_schema: Path to JSON schema for structured output
output_file: Path to save the final output (-o)
"""
super().__init__(workspace=workspace, timeout=timeout)
self.full_auto = full_auto
self.sandbox = sandbox
self.json_output = json_output
self.output_schema = output_schema
self.output_file = output_file
@property
def cli_command(self) -> str:
"""Return the CLI command name."""
return "codex"
def _build_command(self, task: str, **options) -> List[str]:
"""
Build the Codex CLI command.
Args:
task: The task to execute
**options: Additional options
Returns:
List of command arguments
"""
cmd = ["codex", "exec"]
# Add task
cmd.append(task)
# Add full auto flag if enabled
if self.full_auto:
cmd.append("--full-auto")
# Add sandbox mode if not default
if self.sandbox and self.sandbox != "default":
cmd.extend(["--sandbox", self.sandbox])
# Add JSON output flag if enabled
if self.json_output:
cmd.append("--json")
# Add output schema if specified
if self.output_schema:
cmd.extend(["--output-schema", self.output_schema])
# Add output file if specified
if self.output_file:
cmd.extend(["-o", self.output_file])
return cmd
async def execute(self, prompt: str, **options) -> str:
"""
Execute Codex CLI and return the result.
Args:
prompt: The task/prompt to execute
**options: Additional options
Returns:
str: The CLI output
"""
cmd = self._build_command(prompt, **options)
output = await self.execute_async(cmd)
# Parse JSON Lines output if json_output is enabled
if self.json_output:
return self._parse_json_events(output)
return output
def _parse_json_events(self, output: str) -> str:
"""
Parse JSON Lines output and extract the final result.
Args:
output: Raw JSON Lines output
Returns:
str: The extracted result
"""
result_parts = []
for line in output.strip().split('\n'):
if not line.strip():
continue
try:
event = json.loads(line)
event_type = event.get("type", "")
# Extract agent messages
if event_type == "item.completed":
item = event.get("item", {})
if item.get("type") == "agent_message":
text = item.get("text", "")
if text:
result_parts.append(text)
except json.JSONDecodeError:
continue
return '\n'.join(result_parts) if result_parts else output
async def stream(self, prompt: str, **options) -> AsyncIterator[Dict[str, Any]]:
"""
Stream output from Codex CLI.
Args:
prompt: The task/prompt to execute
**options: Additional options
Yields:
dict: Parsed JSON events from the CLI
"""
# Ensure JSON output is enabled for streaming
original_json = self.json_output
self.json_output = True
try:
cmd = self._build_command(prompt, **options)
async for line in self.stream_async(cmd):
if line.strip():
try:
event = json.loads(line)
yield event
except json.JSONDecodeError:
yield {"type": "text", "content": line}
finally:
self.json_output = original_json
async def execute_with_schema(
self,
prompt: str,
schema_path: str,
output_path: Optional[str] = None
) -> Dict[str, Any]:
"""
Execute Codex with structured output schema.
Args:
prompt: The task/prompt to execute
schema_path: Path to the JSON schema file
output_path: Optional path to save the output
Returns:
dict: Parsed structured output
"""
original_schema = self.output_schema
original_output = self.output_file
self.output_schema = schema_path
if output_path:
self.output_file = output_path
try:
cmd = self._build_command(prompt)
output = await self.execute_async(cmd)
# If output file was specified, read from it
if output_path and os.path.exists(output_path):
with open(output_path, 'r') as f:
return json.load(f)
# Otherwise parse the output
try:
return json.loads(output)
except json.JSONDecodeError:
return {"result": output}
finally:
self.output_schema = original_schema
self.output_file = original_output
def get_env(self) -> Dict[str, str]:
"""Get environment variables for CLI execution."""
env = super().get_env()
# Codex uses ChatGPT authentication or API key
# The CLI handles authentication internally
return env
@@ -0,0 +1,192 @@
"""
Cursor CLI Integration.
Provides integration with Cursor's CLI (cursor-agent) for AI-powered coding tasks.
Features:
- Headless mode execution with JSON output
- Force mode for file modifications
- Model selection
- Session resume support
- Streaming output with partial deltas
Usage:
from praisonai.integrations import CursorCLIIntegration
# Create integration
cursor = CursorCLIIntegration(workspace="/path/to/project")
# Execute a coding task
result = await cursor.execute("Fix the authentication bug")
# Execute with force mode (allows file modifications)
cursor_force = CursorCLIIntegration(force=True)
result = await cursor_force.execute("Refactor the module")
# Use as agent tool
tool = cursor.as_tool()
"""
import json
import os
from typing import AsyncIterator, Dict, Any, Optional, List
from .base import BaseCLIIntegration
class CursorCLIIntegration(BaseCLIIntegration):
"""
Integration with Cursor CLI (cursor-agent).
Attributes:
output_format: Output format ("json", "text", "stream-json")
force: Whether to allow file modifications
model: Model to use (e.g., "gpt-5")
stream_partial: Whether to stream partial output
resume_session: Session ID to resume
"""
def __init__(
self,
workspace: str = ".",
timeout: int = 300,
output_format: str = "json",
force: bool = False,
model: Optional[str] = None,
stream_partial: bool = False,
resume_session: Optional[str] = None,
):
"""
Initialize Cursor CLI integration.
Args:
workspace: Working directory for CLI execution
timeout: Timeout in seconds for CLI execution
output_format: Output format ("json", "text", "stream-json")
force: Whether to allow file modifications (--force)
model: Model to use (e.g., "gpt-5")
stream_partial: Whether to stream partial output (--stream-partial-output)
resume_session: Session ID to resume (--resume)
"""
super().__init__(workspace=workspace, timeout=timeout)
self.output_format = output_format
self.force = force
self.model = model
self.stream_partial = stream_partial
self.resume_session = resume_session
@property
def cli_command(self) -> str:
"""Return the CLI command name."""
return "cursor-agent"
def _build_command(self, prompt: str, **options) -> List[str]:
"""
Build the Cursor CLI command.
Args:
prompt: The prompt to send
**options: Additional options
Returns:
List of command arguments
"""
cmd = ["cursor-agent"]
# Add print mode flag
cmd.append("-p")
# Add force flag if enabled
if self.force:
cmd.append("--force")
# Add model if specified
if self.model:
cmd.extend(["-m", self.model])
# Add output format
cmd.extend(["--output-format", self.output_format])
# Add stream partial flag if enabled
if self.stream_partial:
cmd.append("--stream-partial-output")
# Add resume session if specified
if self.resume_session:
cmd.extend(["--resume", self.resume_session])
# Add prompt (must be last)
cmd.append(prompt)
return cmd
async def execute(self, prompt: str, **options) -> str:
"""
Execute Cursor CLI and return the result.
Args:
prompt: The prompt/query to send
**options: Additional options
Returns:
str: The CLI output (parsed from JSON if output_format is "json")
"""
cmd = self._build_command(prompt, **options)
output = await self.execute_async(cmd)
# Parse JSON output if applicable
if self.output_format == "json":
try:
data = json.loads(output)
# Extract the main result
if isinstance(data, dict):
return data.get("result", data.get("content", str(data)))
return str(data)
except json.JSONDecodeError:
return output
return output
async def stream(self, prompt: str, **options) -> AsyncIterator[Dict[str, Any]]:
"""
Stream output from Cursor CLI.
Args:
prompt: The prompt/query to send
**options: Additional options
Yields:
dict: Parsed JSON events from the CLI
"""
# Use stream-json format for streaming
original_format = self.output_format
original_partial = self.stream_partial
self.output_format = "stream-json"
self.stream_partial = True
try:
cmd = self._build_command(prompt, **options)
async for line in self.stream_async(cmd):
if line.strip():
try:
event = json.loads(line)
yield event
except json.JSONDecodeError:
yield {"type": "text", "content": line}
finally:
self.output_format = original_format
self.stream_partial = original_partial
def get_env(self) -> Dict[str, str]:
"""Get environment variables for CLI execution."""
env = super().get_env()
# Add Cursor API key if available
if "CURSOR_API_KEY" in os.environ:
env["CURSOR_API_KEY"] = os.environ["CURSOR_API_KEY"]
return env
@@ -0,0 +1,222 @@
"""
Gemini CLI Integration.
Provides integration with Google's Gemini CLI for AI-powered coding tasks.
Features:
- Headless mode execution with JSON output
- Model selection (gemini-2.5-pro, gemini-2.5-flash, etc.)
- Multi-directory context support
- Google Search grounding
- Usage statistics
Usage:
from praisonai.integrations import GeminiCLIIntegration
# Create integration
gemini = GeminiCLIIntegration(workspace="/path/to/project")
# Execute a coding task
result = await gemini.execute("Analyze this codebase")
# Get result with stats
result, stats = await gemini.execute_with_stats("Explain the architecture")
# Use as agent tool
tool = gemini.as_tool()
"""
import json
import os
from typing import AsyncIterator, Dict, Any, Optional, List, Tuple
from .base import BaseCLIIntegration
class GeminiCLIIntegration(BaseCLIIntegration):
"""
Integration with Google's Gemini CLI.
Attributes:
output_format: Output format ("json", "text", "stream-json")
model: Gemini model to use
include_directories: Additional directories to include in context
sandbox: Whether to run in sandbox mode
"""
def __init__(
self,
workspace: str = ".",
timeout: int = 300,
output_format: str = "json",
model: str = "gemini-2.5-pro",
include_directories: Optional[List[str]] = None,
sandbox: bool = False,
):
"""
Initialize Gemini CLI integration.
Args:
workspace: Working directory for CLI execution
timeout: Timeout in seconds for CLI execution
output_format: Output format ("json", "text", "stream-json")
model: Gemini model to use (e.g., "gemini-2.5-pro", "gemini-2.5-flash")
include_directories: Additional directories to include in context
sandbox: Whether to run in sandbox mode
"""
super().__init__(workspace=workspace, timeout=timeout)
self.output_format = output_format
self.model = model
self.include_directories = include_directories
self.sandbox = sandbox
# Store last stats for retrieval
self._last_stats: Optional[Dict[str, Any]] = None
@property
def cli_command(self) -> str:
"""Return the CLI command name."""
return "gemini"
def _build_command(self, prompt: str, **options) -> List[str]:
"""
Build the Gemini CLI command.
Args:
prompt: The prompt to send
**options: Additional options
Returns:
List of command arguments
"""
cmd = ["gemini"]
# Add print mode flag
cmd.append("-p")
# Add prompt
cmd.append(prompt)
# Add model
cmd.extend(["-m", self.model])
# Add output format
cmd.extend(["--output-format", self.output_format])
# Add include directories if specified
if self.include_directories:
cmd.extend(["--include-directories", ",".join(self.include_directories)])
# Add sandbox flag if enabled
if self.sandbox:
cmd.append("--sandbox")
return cmd
async def execute(self, prompt: str, **options) -> str:
"""
Execute Gemini CLI and return the result.
Args:
prompt: The prompt/query to send
**options: Additional options
Returns:
str: The CLI output (parsed from JSON if output_format is "json")
"""
cmd = self._build_command(prompt, **options)
output = await self.execute_async(cmd)
# Parse JSON output if applicable
if self.output_format == "json":
try:
data = json.loads(output)
# Store stats for later retrieval
self._last_stats = data.get("stats")
# Extract the main response
return data.get("response", str(data))
except json.JSONDecodeError:
return output
return output
async def execute_with_stats(self, prompt: str, **options) -> Tuple[str, Optional[Dict[str, Any]]]:
"""
Execute Gemini CLI and return both result and usage stats.
Args:
prompt: The prompt/query to send
**options: Additional options
Returns:
Tuple[str, dict]: (result, stats) where stats contains usage information
"""
# Ensure JSON format for stats
original_format = self.output_format
self.output_format = "json"
try:
cmd = self._build_command(prompt, **options)
output = await self.execute_async(cmd)
try:
data = json.loads(output)
response = data.get("response", str(data))
stats = data.get("stats")
return response, stats
except json.JSONDecodeError:
return output, None
finally:
self.output_format = original_format
async def stream(self, prompt: str, **options) -> AsyncIterator[Dict[str, Any]]:
"""
Stream output from Gemini CLI.
Args:
prompt: The prompt/query to send
**options: Additional options
Yields:
dict: Parsed JSON events from the CLI
"""
# Use stream-json format for streaming
original_format = self.output_format
self.output_format = "stream-json"
try:
cmd = self._build_command(prompt, **options)
async for line in self.stream_async(cmd):
if line.strip():
try:
event = json.loads(line)
yield event
except json.JSONDecodeError:
yield {"type": "text", "content": line}
finally:
self.output_format = original_format
def get_last_stats(self) -> Optional[Dict[str, Any]]:
"""
Get the stats from the last execution.
Returns:
dict: Usage statistics or None if not available
"""
return self._last_stats
def get_env(self) -> Dict[str, str]:
"""Get environment variables for CLI execution."""
env = super().get_env()
# Add Google API key if available
if "GOOGLE_API_KEY" in os.environ:
env["GOOGLE_API_KEY"] = os.environ["GOOGLE_API_KEY"]
elif "GEMINI_API_KEY" in os.environ:
# Map GEMINI_API_KEY to GOOGLE_API_KEY
env["GOOGLE_API_KEY"] = os.environ["GEMINI_API_KEY"]
return env
+1 -1
View File
@@ -1 +1 @@
__version__ = "2.3.38"
__version__ = "2.3.47"
+1 -1
View File
@@ -12,7 +12,7 @@ dependencies = [
"rich>=13.7",
"markdown>=3.5",
"pyparsing>=3.0.0",
"praisonaiagents>=0.1.16",
"praisonaiagents>=0.1.17",
"python-dotenv>=0.19.0",
"instructor>=1.3.3",
"PyYAML>=6.0",
+21
View File
@@ -0,0 +1,21 @@
{
"$schema": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json",
"name": "io.github.MervinPraison/praisonai",
"title": "PraisonAI",
"description": "AI Agents Framework with Self Reflection and MCP support",
"repository": {
"url": "https://github.com/MervinPraison/PraisonAI",
"source": "github"
},
"version": "2.3.42",
"packages": [
{
"registryType": "pypi",
"identifier": "praisonai",
"version": "2.3.42",
"transport": {
"type": "stdio"
}
}
]
}
@@ -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"])
@@ -0,0 +1,182 @@
"""
Integration tests for AutoGenerator and WorkflowAutoGenerator.
These tests require a real OpenAI API key and make actual API calls.
Run with: OPENAI_API_KEY=your-key pytest tests/integration/test_auto_generator_integration.py -v
"""
import pytest
import os
import sys
import tempfile
import yaml
# Add the source path for imports
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
# Skip if no API key
if not os.environ.get('OPENAI_API_KEY'):
pytest.skip("OPENAI_API_KEY not set", allow_module_level=True)
try:
from praisonai.auto import AutoGenerator, WorkflowAutoGenerator
except ImportError as e:
pytest.skip(f"Could not import required modules: {e}", allow_module_level=True)
class TestAutoGeneratorIntegration:
"""Integration tests for AutoGenerator with real API calls."""
@pytest.mark.integration
def test_generate_agents_yaml(self):
"""Test generating agents.yaml with real API call."""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
agent_file = f.name
try:
generator = AutoGenerator(
topic="Create a simple greeting bot",
agent_file=agent_file,
framework="praisonai"
)
result_path = generator.generate()
# Verify file was created
assert os.path.exists(result_path)
# Verify file content
with open(result_path, 'r') as f:
data = yaml.safe_load(f)
# Check structure
assert 'framework' in data
assert 'topic' in data
assert 'roles' in data
assert len(data['roles']) > 0
# Check that tools are preserved (not empty list)
for role_id, role_data in data['roles'].items():
assert 'tools' in role_data
# Tools should be a list (may be empty if LLM didn't suggest any)
assert isinstance(role_data['tools'], list)
print(f"✓ Generated agents.yaml successfully")
print(f" Roles: {list(data['roles'].keys())}")
finally:
if os.path.exists(agent_file):
os.unlink(agent_file)
@pytest.mark.integration
def test_generate_with_merge(self):
"""Test generating and merging agents."""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
agent_file = f.name
try:
# First generation
generator1 = AutoGenerator(
topic="Create a researcher",
agent_file=agent_file,
framework="praisonai"
)
generator1.generate()
# Second generation with merge
generator2 = AutoGenerator(
topic="Create a writer",
agent_file=agent_file,
framework="praisonai"
)
generator2.generate(merge=True)
# Verify merged content
with open(agent_file, 'r') as f:
data = yaml.safe_load(f)
# Should have roles from both generations
assert len(data['roles']) >= 2
print(f"✓ Merged agents.yaml successfully")
print(f" Total roles: {len(data['roles'])}")
finally:
if os.path.exists(agent_file):
os.unlink(agent_file)
class TestWorkflowAutoGeneratorIntegration:
"""Integration tests for WorkflowAutoGenerator with real API calls."""
@pytest.mark.integration
def test_generate_sequential_workflow(self):
"""Test generating sequential workflow."""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
workflow_file = f.name
try:
generator = WorkflowAutoGenerator(
topic="Research and summarize AI news",
workflow_file=workflow_file
)
result_path = generator.generate(pattern="sequential")
# Verify file was created
assert os.path.exists(result_path)
# Verify file content
with open(result_path, 'r') as f:
data = yaml.safe_load(f)
# Check structure
assert 'name' in data
assert 'agents' in data
assert 'steps' in data
assert len(data['agents']) > 0
assert len(data['steps']) > 0
print(f"✓ Generated sequential workflow successfully")
print(f" Agents: {list(data['agents'].keys())}")
print(f" Steps: {len(data['steps'])}")
finally:
if os.path.exists(workflow_file):
os.unlink(workflow_file)
@pytest.mark.integration
def test_generate_parallel_workflow(self):
"""Test generating parallel workflow."""
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
workflow_file = f.name
try:
generator = WorkflowAutoGenerator(
topic="Multi-source research on climate change",
workflow_file=workflow_file
)
result_path = generator.generate(pattern="parallel")
# Verify file was created
assert os.path.exists(result_path)
# Verify file content
with open(result_path, 'r') as f:
data = yaml.safe_load(f)
# Check structure
assert 'agents' in data
assert 'steps' in data
print(f"✓ Generated parallel workflow successfully")
print(f" Agents: {list(data['agents'].keys())}")
finally:
if os.path.exists(workflow_file):
os.unlink(workflow_file)
if __name__ == '__main__':
pytest.main([__file__, '-v', '-m', 'integration'])
@@ -0,0 +1,271 @@
"""
Integration tests for CLI tool integrations.
These tests require the actual CLI tools to be installed and API keys to be set.
Run with: pytest tests/integration/test_cli_integrations.py -v
Environment variables required:
- OPENAI_API_KEY: For general testing
- ANTHROPIC_API_KEY or CLAUDE_API_KEY: For Claude Code
- GEMINI_API_KEY or GOOGLE_API_KEY: For Gemini CLI
- CURSOR_API_KEY: For Cursor CLI
"""
import pytest
import os
import asyncio
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))
# Skip all tests if no API keys are set
pytestmark = pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY"),
reason="OPENAI_API_KEY not set"
)
class TestIntegrationAvailability:
"""Test that integrations can be imported and checked."""
def test_import_all_integrations(self):
"""Test importing all integrations."""
from praisonai.integrations import (
BaseCLIIntegration,
ClaudeCodeIntegration,
GeminiCLIIntegration,
CodexCLIIntegration,
CursorCLIIntegration,
)
assert BaseCLIIntegration is not None
assert ClaudeCodeIntegration is not None
assert GeminiCLIIntegration is not None
assert CodexCLIIntegration is not None
assert CursorCLIIntegration is not None
def test_get_available_integrations(self):
"""Test getting available integrations."""
from praisonai.integrations import get_available_integrations
availability = get_available_integrations()
assert isinstance(availability, dict)
assert "claude" in availability
assert "gemini" in availability
assert "codex" in availability
assert "cursor" in availability
print("\n📊 Integration Availability:")
for name, available in availability.items():
status = "✅" if available else "❌"
print(f" {status} {name}")
class TestClaudeCodeIntegrationReal:
"""Real integration tests for Claude Code CLI."""
@pytest.mark.skipif(
not os.environ.get("ANTHROPIC_API_KEY") and not os.environ.get("CLAUDE_API_KEY"),
reason="ANTHROPIC_API_KEY or CLAUDE_API_KEY not set"
)
def test_claude_integration_creation(self):
"""Test creating Claude Code integration."""
from praisonai.integrations import ClaudeCodeIntegration
integration = ClaudeCodeIntegration(
workspace=".",
output_format="json",
skip_permissions=True
)
assert integration.cli_command == "claude"
print(f"\n🔧 Claude CLI available: {integration.is_available}")
@pytest.mark.skipif(
not os.environ.get("ANTHROPIC_API_KEY") and not os.environ.get("CLAUDE_API_KEY"),
reason="ANTHROPIC_API_KEY or CLAUDE_API_KEY not set"
)
def test_claude_as_tool(self):
"""Test Claude Code as agent tool."""
from praisonai.integrations import ClaudeCodeIntegration
integration = ClaudeCodeIntegration()
tool = integration.as_tool()
assert callable(tool)
assert tool.__name__ == "claude_tool"
print(f"\n🔧 Claude tool created: {tool.__name__}")
class TestGeminiCLIIntegrationReal:
"""Real integration tests for Gemini CLI."""
@pytest.mark.skipif(
not os.environ.get("GEMINI_API_KEY") and not os.environ.get("GOOGLE_API_KEY"),
reason="GEMINI_API_KEY or GOOGLE_API_KEY not set"
)
def test_gemini_integration_creation(self):
"""Test creating Gemini CLI integration."""
from praisonai.integrations import GeminiCLIIntegration
integration = GeminiCLIIntegration(
workspace=".",
model="gemini-2.5-pro",
output_format="json"
)
assert integration.cli_command == "gemini"
assert integration.model == "gemini-2.5-pro"
print(f"\n🔧 Gemini CLI available: {integration.is_available}")
@pytest.mark.skipif(
not os.environ.get("GEMINI_API_KEY") and not os.environ.get("GOOGLE_API_KEY"),
reason="GEMINI_API_KEY or GOOGLE_API_KEY not set"
)
def test_gemini_as_tool(self):
"""Test Gemini CLI as agent tool."""
from praisonai.integrations import GeminiCLIIntegration
integration = GeminiCLIIntegration()
tool = integration.as_tool()
assert callable(tool)
assert tool.__name__ == "gemini_tool"
print(f"\n🔧 Gemini tool created: {tool.__name__}")
class TestCursorCLIIntegrationReal:
"""Real integration tests for Cursor CLI."""
@pytest.mark.skipif(
not os.environ.get("CURSOR_API_KEY"),
reason="CURSOR_API_KEY not set"
)
def test_cursor_integration_creation(self):
"""Test creating Cursor CLI integration."""
from praisonai.integrations import CursorCLIIntegration
integration = CursorCLIIntegration(
workspace=".",
output_format="json",
force=False
)
assert integration.cli_command == "cursor-agent"
print(f"\n🔧 Cursor CLI available: {integration.is_available}")
@pytest.mark.skipif(
not os.environ.get("CURSOR_API_KEY"),
reason="CURSOR_API_KEY not set"
)
def test_cursor_as_tool(self):
"""Test Cursor CLI as agent tool."""
from praisonai.integrations import CursorCLIIntegration
integration = CursorCLIIntegration()
tool = integration.as_tool()
assert callable(tool)
assert tool.__name__ == "cursor-agent_tool"
print(f"\n🔧 Cursor tool created: {tool.__name__}")
class TestExternalAgentsHandler:
"""Test the ExternalAgentsHandler."""
def test_handler_creation(self):
"""Test creating the handler."""
from praisonai.cli.features import ExternalAgentsHandler
handler = ExternalAgentsHandler(verbose=True)
assert handler.feature_name == "external_agents"
assert handler.flag_name == "external-agent"
def test_handler_list_integrations(self):
"""Test listing integrations."""
from praisonai.cli.features import ExternalAgentsHandler
handler = ExternalAgentsHandler()
integrations = handler.list_integrations()
assert "claude" in integrations
assert "gemini" in integrations
assert "codex" in integrations
assert "cursor" in integrations
print("\n📋 Available integrations:", integrations)
def test_handler_check_availability(self):
"""Test checking availability."""
from praisonai.cli.features import ExternalAgentsHandler
handler = ExternalAgentsHandler()
availability = handler.check_availability()
print("\n📊 Handler Availability Check:")
for name, available in availability.items():
status = "✅" if available else "❌"
print(f" {status} {name}")
class TestPerformanceImpact:
"""Test that integrations don't impact performance."""
def test_lazy_loading_integrations(self):
"""Test that integrations are lazy loaded."""
import time
# Measure import time
start = time.time()
from praisonai.integrations import BaseCLIIntegration
import_time = time.time() - start
print(f"\n⏱️ Import time: {import_time*1000:.2f}ms")
# Should be fast (< 100ms)
assert import_time < 0.1, f"Import took too long: {import_time}s"
def test_lazy_loading_handler(self):
"""Test that handler is lazy loaded."""
import time
# Measure import time
start = time.time()
from praisonai.cli.features import ExternalAgentsHandler
import_time = time.time() - start
print(f"\n⏱️ Handler import time: {import_time*1000:.2f}ms")
# Should be fast (< 100ms)
assert import_time < 0.1, f"Import took too long: {import_time}s"
def test_availability_caching(self):
"""Test that availability checks are cached."""
import time
from praisonai.integrations import ClaudeCodeIntegration
integration = ClaudeCodeIntegration()
# First check
start = time.time()
_ = integration.is_available
first_time = time.time() - start
# Second check (should be cached)
start = time.time()
_ = integration.is_available
second_time = time.time() - start
print(f"\n⏱️ First availability check: {first_time*1000:.4f}ms")
print(f"⏱️ Second availability check: {second_time*1000:.4f}ms")
# Both should be very fast (< 10ms) due to caching
assert first_time < 0.01, f"First check too slow: {first_time}s"
assert second_time < 0.01, f"Second check too slow: {second_time}s"
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
@@ -0,0 +1,324 @@
"""
Real integration tests for Model Routing with actual API calls.
Tests the RouterHandler with real LLM providers.
Requires API keys: OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY
"""
import pytest
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
from praisonai.cli.features.router import RouterHandler
class TestRouterHandlerReal:
"""Real integration tests for RouterHandler."""
def test_router_handler_initialization(self):
"""Test RouterHandler initializes correctly."""
handler = RouterHandler(verbose=True)
assert handler is not None
assert handler.model_router is not None
assert len(handler.model_router) > 0
print(f"✅ RouterHandler initialized with {len(handler.model_router)} default models")
def test_complexity_analysis(self):
"""Test complexity analysis for different prompts."""
handler = RouterHandler()
# Simple prompts
simple_prompts = [
"What is 2+2?",
"Define photosynthesis",
"List the planets",
]
for prompt in simple_prompts:
complexity = handler.analyze_complexity(prompt)
print(f" '{prompt[:30]}...' -> {complexity}")
assert complexity == 'simple', f"Expected 'simple' for: {prompt}"
print("✅ Simple prompts correctly classified")
# Complex prompts
complex_prompts = [
"Analyze and research the comprehensive impact of AI on healthcare systems globally",
"Compare and evaluate different machine learning architectures for natural language processing",
]
for prompt in complex_prompts:
complexity = handler.analyze_complexity(prompt)
print(f" '{prompt[:40]}...' -> {complexity}")
assert complexity == 'complex', f"Expected 'complex' for: {prompt}"
print("✅ Complex prompts correctly classified")
def test_model_selection_by_complexity(self):
"""Test model selection based on complexity."""
handler = RouterHandler(verbose=True)
# Simple task - should select cheap model
simple_model = handler.select_model("What is the capital of France?")
print(f"Simple task model: {simple_model}")
assert simple_model in ['gpt-4o-mini', 'claude-3-haiku', 'gemini-1.5-flash']
# Complex task - should select powerful model
complex_model = handler.select_model(
"Analyze and research the comprehensive impact of quantum computing on cybersecurity"
)
print(f"Complex task model: {complex_model}")
assert complex_model in ['gpt-4-turbo', 'claude-3-opus', 'o1-preview']
print("✅ Model selection by complexity works correctly")
def test_model_selection_with_provider_preference(self):
"""Test model selection with provider preference."""
handler = RouterHandler(verbose=True)
# Prefer OpenAI
openai_model = handler.select_model(
"Write a simple greeting",
preferred_provider='openai'
)
print(f"OpenAI preferred: {openai_model}")
assert 'gpt' in openai_model.lower() or 'o1' in openai_model.lower()
# Prefer Anthropic
anthropic_model = handler.select_model(
"Write a simple greeting",
preferred_provider='anthropic'
)
print(f"Anthropic preferred: {anthropic_model}")
assert 'claude' in anthropic_model.lower()
# Prefer Google
google_model = handler.select_model(
"Write a simple greeting",
preferred_provider='google'
)
print(f"Google preferred: {google_model}")
assert 'gemini' in google_model.lower()
print("✅ Provider preference works correctly")
def test_custom_models_from_yaml(self):
"""Test loading custom models from YAML-like config."""
handler = RouterHandler(verbose=True)
# Define custom models (as would be in YAML)
custom_models = {
'my-cheap-model': {
'provider': 'openai',
'complexity': ['simple'],
'cost_per_1k': 0.0001,
'capabilities': ['text']
},
'my-balanced-model': {
'provider': 'openai',
'complexity': ['moderate'],
'cost_per_1k': 0.001,
'capabilities': ['text', 'function-calling']
},
'my-premium-model': {
'provider': 'anthropic',
'complexity': ['complex', 'very_complex'],
'cost_per_1k': 0.015,
'capabilities': ['text', 'vision', 'function-calling']
}
}
# Load custom models
handler.load_models_from_config(custom_models)
# Verify models are loaded
available = handler.get_available_models()
print(f"Available models: {available}")
assert 'my-cheap-model' in available
assert 'my-balanced-model' in available
assert 'my-premium-model' in available
# Test selection with custom models
simple_model = handler.select_model("What is 2+2?")
print(f"Simple task -> {simple_model}")
assert simple_model == 'my-cheap-model'
complex_model = handler.select_model(
"Analyze and research the comprehensive impact of AI"
)
print(f"Complex task -> {complex_model}")
assert complex_model == 'my-premium-model'
print("✅ Custom models from YAML config work correctly")
def test_cost_threshold(self):
"""Test cost threshold filtering."""
handler = RouterHandler(verbose=True)
# Load models with different costs
custom_models = {
'cheap': {
'provider': 'openai',
'complexity': ['simple', 'moderate', 'complex'],
'cost_per_1k': 0.0001,
'capabilities': ['text']
},
'expensive': {
'provider': 'anthropic',
'complexity': ['simple', 'moderate', 'complex'],
'cost_per_1k': 0.05,
'capabilities': ['text']
}
}
handler.load_models_from_config(custom_models)
# Set cost threshold
handler.set_cost_threshold(0.001)
# Should only select cheap model
model = handler.select_model("Any task here")
print(f"With cost threshold 0.001: {model}")
assert model == 'cheap'
# Verify cost
cost = handler.get_model_cost(model)
assert cost <= 0.001
print("✅ Cost threshold filtering works correctly")
def test_merge_with_defaults(self):
"""Test merging custom models with defaults."""
handler = RouterHandler(verbose=True)
custom_models = {
'my-custom-model': {
'provider': 'custom',
'complexity': ['moderate'],
'cost_per_1k': 0.002,
'capabilities': ['text']
}
}
handler.load_models_from_config(custom_models, merge_with_defaults=True)
available = handler.get_available_models()
print(f"Available models (merged): {available}")
# Should have both custom and default models
assert 'my-custom-model' in available
assert 'gpt-4o-mini' in available
print("✅ Merge with defaults works correctly")
class TestRealAPIIntegration:
"""Tests that make real API calls - requires API keys."""
@pytest.mark.skipif(
not os.getenv('OPENAI_API_KEY'),
reason="OPENAI_API_KEY not set"
)
def test_real_openai_call(self):
"""Test real OpenAI API call with routed model."""
from praisonaiagents import Agent
handler = RouterHandler(verbose=True)
# Select model for simple task
model = handler.select_model(
"What is 2+2?",
preferred_provider='openai'
)
print(f"Selected model: {model}")
# Create agent with selected model
agent = Agent(
name="TestAgent",
role="Calculator",
goal="Answer simple math questions",
llm=model,
verbose=True
)
# Run simple task
result = agent.start("What is 2+2? Answer with just the number.")
print(f"Result: {result}")
assert '4' in str(result)
print("✅ Real OpenAI API call successful")
@pytest.mark.skipif(
not os.getenv('ANTHROPIC_API_KEY'),
reason="ANTHROPIC_API_KEY not set"
)
def test_real_anthropic_call(self):
"""Test real Anthropic API call with routed model."""
from praisonaiagents import Agent
handler = RouterHandler(verbose=True)
# Select Anthropic model
model = handler.select_model(
"What is the capital of France?",
preferred_provider='anthropic'
)
print(f"Selected model: {model}")
# Create agent with selected model
agent = Agent(
name="TestAgent",
role="Geography Expert",
goal="Answer geography questions",
llm=model,
verbose=True
)
# Run task
result = agent.start("What is the capital of France? Answer with just the city name.")
print(f"Result: {result}")
assert 'paris' in str(result).lower()
print("✅ Real Anthropic API call successful")
@pytest.mark.skipif(
not os.getenv('GEMINI_API_KEY'),
reason="GEMINI_API_KEY not set"
)
def test_real_gemini_call(self):
"""Test real Gemini API call with routed model."""
from praisonaiagents import Agent
handler = RouterHandler(verbose=True)
# Select Google model
model = handler.select_model(
"What is 3 times 7?",
preferred_provider='google'
)
print(f"Selected model: {model}")
# Create agent with selected model
agent = Agent(
name="TestAgent",
role="Calculator",
goal="Answer math questions",
llm=f"gemini/{model}" if not model.startswith("gemini/") else model,
verbose=True
)
# Run task
result = agent.start("What is 3 times 7? Answer with just the number.")
print(f"Result: {result}")
assert '21' in str(result)
print("✅ Real Gemini API call successful")
if __name__ == '__main__':
# Run tests
pytest.main([__file__, '-v', '-s'])
@@ -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
@@ -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"
@@ -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", "[email protected]"],
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", "[email protected]"],
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", "[email protected]"],
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", "[email protected]"],
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", "[email protected]"],
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
@@ -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
@@ -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
@@ -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
@@ -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"]
@@ -0,0 +1 @@
"""Tests for praisonai.code module."""
@@ -0,0 +1,235 @@
"""
Integration tests for praisonai.code agent tools.
"""
import os
import tempfile
class TestAgentTools:
"""Tests for agent-compatible tool wrappers."""
def test_set_and_get_workspace(self):
"""Test setting and getting workspace."""
from praisonai.code import set_workspace, get_workspace
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
assert get_workspace() == temp_dir
def test_code_read_file(self):
"""Test code_read_file agent tool."""
from praisonai.code import code_read_file, set_workspace
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
# Create test file
file_path = os.path.join(temp_dir, "test.py")
with open(file_path, 'w') as f:
f.write("def hello():\n print('hello')\n")
result = code_read_file("test.py")
assert "File: test.py" in result
assert "def hello():" in result
assert "1 |" in result # Line numbers
def test_code_write_file(self):
"""Test code_write_file agent tool."""
from praisonai.code import code_write_file, set_workspace
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
result = code_write_file("new_file.py", "print('hello')")
assert "Created" in result
assert os.path.exists(os.path.join(temp_dir, "new_file.py"))
def test_code_list_files(self):
"""Test code_list_files agent tool."""
from praisonai.code import code_list_files, set_workspace
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
# Create some files
open(os.path.join(temp_dir, "file1.py"), 'w').close()
open(os.path.join(temp_dir, "file2.py"), 'w').close()
os.makedirs(os.path.join(temp_dir, "subdir"))
result = code_list_files(".")
assert "file1.py" in result
assert "file2.py" in result
assert "subdir" in result
def test_code_apply_diff(self):
"""Test code_apply_diff agent tool."""
from praisonai.code import code_apply_diff, set_workspace
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
# Create test file
file_path = os.path.join(temp_dir, "test.py")
with open(file_path, 'w') as f:
f.write("def old_function():\n return 'old'\n")
diff = """<<<<<<< SEARCH
:start_line:1
-------
def old_function():
return 'old'
=======
def new_function():
return 'new'
>>>>>>> REPLACE"""
result = code_apply_diff("test.py", diff)
assert "Successfully applied" in result
# Verify file was modified
with open(file_path, 'r') as f:
content = f.read()
assert "def new_function():" in content
def test_code_search_replace(self):
"""Test code_search_replace agent tool."""
from praisonai.code import code_search_replace, set_workspace
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
# Create test file
file_path = os.path.join(temp_dir, "test.py")
with open(file_path, 'w') as f:
f.write("old_name = 'value'\nprint(old_name)\n")
result = code_search_replace("test.py", "old_name", "new_name")
assert "Replaced" in result
assert "2" in result # 2 occurrences
def test_code_execute_command(self):
"""Test code_execute_command agent tool."""
from praisonai.code import code_execute_command, set_workspace
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
result = code_execute_command("echo 'Hello from test'")
assert "completed successfully" in result
assert "Hello" in result
def test_code_tools_list(self):
"""Test CODE_TOOLS list contains all tools."""
from praisonai.code import CODE_TOOLS
assert len(CODE_TOOLS) == 6
tool_names = [t.__name__ for t in CODE_TOOLS]
assert "code_read_file" in tool_names
assert "code_write_file" in tool_names
assert "code_list_files" in tool_names
assert "code_apply_diff" in tool_names
assert "code_search_replace" in tool_names
assert "code_execute_command" in tool_names
class TestDiffHelpers:
"""Tests for diff helper functions."""
def test_create_diff_block(self):
"""Test creating a diff block."""
from praisonai.code import create_diff_block
diff = create_diff_block(
search_content="old content",
replace_content="new content",
start_line=10,
)
assert "<<<<<<< SEARCH" in diff
assert ":start_line:10" in diff
assert "old content" in diff
assert "=======" in diff
assert "new content" in diff
assert ">>>>>>> REPLACE" in diff
def test_create_multi_diff(self):
"""Test creating multiple diff blocks."""
from praisonai.code import create_multi_diff
blocks = [
{'search': 'old1', 'replace': 'new1', 'start_line': 5},
{'search': 'old2', 'replace': 'new2', 'start_line': 20},
]
diff = create_multi_diff(blocks)
# Should contain two SEARCH/REPLACE blocks
assert diff.count("<<<<<<< SEARCH") == 2
assert diff.count(">>>>>>> REPLACE") == 2
assert "old1" in diff
assert "new1" in diff
assert "old2" in diff
assert "new2" in diff
class TestEndToEndWorkflow:
"""End-to-end workflow tests."""
def test_read_modify_write_workflow(self):
"""Test a complete read-modify-write workflow."""
from praisonai.code import (
set_workspace, code_read_file, code_apply_diff, code_execute_command
)
with tempfile.TemporaryDirectory() as temp_dir:
set_workspace(temp_dir)
# Create a Python file
file_path = os.path.join(temp_dir, "calculator.py")
with open(file_path, 'w') as f:
f.write('''def add(a, b):
return a + b
def subtract(a, b):
return a - b
if __name__ == "__main__":
print(add(2, 3))
''')
# Read the file
content = code_read_file("calculator.py")
assert "def add(a, b):" in content
# Apply a diff to add a multiply function
diff = '''<<<<<<< SEARCH
:start_line:4
-------
def subtract(a, b):
return a - b
=======
def subtract(a, b):
return a - b
def multiply(a, b):
return a * b
>>>>>>> REPLACE'''
result = code_apply_diff("calculator.py", diff)
assert "Successfully applied" in result
# Verify the change
content = code_read_file("calculator.py")
assert "def multiply(a, b):" in content
# Run the file to verify it's valid Python
result = code_execute_command("python calculator.py")
assert "5" in result # add(2, 3) = 5
Loaded 100 of 113 files, more files were not shown because too many files have changed in this diff. Show more