mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-10 03:31:28 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6d3ba7cdf8 | ||
|
|
96616498d2 | ||
|
|
1cad345415 | ||
|
|
802e7e52f9 | ||
|
|
53704f234c | ||
|
|
b575994d91 | ||
|
|
7ed8596261 | ||
|
|
deef65c4bb | ||
|
|
d51327effd | ||
|
|
30ce8857e8 | ||
|
|
397874ccfc | ||
|
|
2016902e9a | ||
|
|
d2d19511da | ||
|
|
bc9579b497 | ||
|
|
3161389341 | ||
|
|
b01eb2b8fa | ||
|
|
e609fb7425 | ||
|
|
168ef98473 | ||
|
|
ce2c2926f9 | ||
|
|
0a569e2b46 | ||
|
|
054aea8f66 | ||
|
|
34f1293cfc | ||
|
|
7c1e32afc0 | ||
|
|
1125c835ec | ||
|
|
889d04a843 | ||
|
|
fba811c477 |
No files matched your search
+4
-1
@@ -77,4 +77,7 @@ agents/praisonaiagents/praisonaiagents.egg-info
|
||||
.qodo
|
||||
CopilotKit*
|
||||
.gemini
|
||||
issue_details.json
|
||||
issue_details.json
|
||||
|
||||
# MCP Registry tokens
|
||||
.mcpregistry_*
|
||||
@@ -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) |
|
||||
@@ -240,12 +252,20 @@ npm install praisonai
|
||||
| 🪝 Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
|
||||
| 📈 Telemetry | [Example](examples/python/telemetry/production-telemetry-example.py) | [📖](https://docs.praison.ai/features/telemetry) |
|
||||
| 📹 Camera Integration | [Example](examples/python/camera/) | [📖](https://docs.praison.ai/features/camera-integration) |
|
||||
| 📄 Project Docs (.praison/docs/) | [Example](#docs-cli) | [📖](https://docs.praison.ai/cli/docs) |
|
||||
| 🔌 MCP Config Management | [Example](#mcp-config-cli) | [📖](https://docs.praison.ai/cli/mcp) |
|
||||
| 💬 AI Commit Messages | [Example](#ai-commit-cli) | [📖](https://docs.praison.ai/cli/commit) |
|
||||
| @ @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) |
|
||||
| 📄 Project Docs (.praison/docs/) | [Example](#docs-cli) | [📖](https://docs.praison.ai/docs/cli/docs) |
|
||||
| 🔌 MCP Config Management | [Example](#mcp-config-cli) | [📖](https://docs.praison.ai/docs/cli/mcp) |
|
||||
| 💬 AI Commit Messages | [Example](#ai-commit-cli) | [📖](https://docs.praison.ai/docs/cli/commit) |
|
||||
| @ @Mentions in Prompts | [Example](#mentions-in-prompts) | [📖](https://docs.praison.ai/docs/cli/mentions) |
|
||||
| 💾 Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| 📜 History in Context | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| **🖥️ CLI Features** | | |
|
||||
| ↳ Slash Commands | [Example](examples/python/cli/slash_commands_example.py) | [📖](https://docs.praison.ai/docs/cli/slash-commands) |
|
||||
| ↳ Autonomy Modes | [Example](examples/python/cli/autonomy_modes_example.py) | [📖](https://docs.praison.ai/docs/cli/autonomy-modes) |
|
||||
| ↳ Cost Tracking | [Example](examples/python/cli/cost_tracking_example.py) | [📖](https://docs.praison.ai/docs/cli/cost-tracking) |
|
||||
| ↳ Repository Map | [Example](examples/python/cli/repo_map_example.py) | [📖](https://docs.praison.ai/docs/cli/repo-map) |
|
||||
| ↳ Interactive TUI | [Example](examples/python/cli/interactive_tui_example.py) | [📖](https://docs.praison.ai/docs/cli/interactive-tui) |
|
||||
| ↳ Git Integration | [Example](examples/python/cli/git_integration_example.py) | [📖](https://docs.praison.ai/docs/cli/git-integration) |
|
||||
| ↳ Sandbox Execution | [Example](examples/python/cli/sandbox_execution_example.py) | [📖](https://docs.praison.ai/docs/cli/sandbox-execution) |
|
||||
|
||||
---
|
||||
|
||||
@@ -952,6 +972,32 @@ export OPENAI_API_KEY=xxxxxxxxxxxxxxxxxxxxxx
|
||||
praisonai --auto create a movie script about Robots in Mars
|
||||
```
|
||||
|
||||
### Interactive Mode CLI:
|
||||
```bash
|
||||
# Start interactive terminal mode (inspired by Gemini CLI, Codex CLI, Claude Code)
|
||||
praisonai --interactive
|
||||
praisonai -i
|
||||
|
||||
# Features:
|
||||
# - Streaming responses (no boxes)
|
||||
# - Built-in tools: read_file, write_file, list_files, execute_command, internet_search
|
||||
# - Slash commands: /help, /exit, /tools, /clear
|
||||
|
||||
# Chat mode - single prompt with interactive style (for testing/scripting)
|
||||
# Use --chat (or --chat-mode for backward compatibility)
|
||||
praisonai "list files in current folder" --chat
|
||||
praisonai "search the web for AI news" --chat
|
||||
praisonai "read README.md" --chat
|
||||
```
|
||||
|
||||
### Chat UI (Web Interface):
|
||||
```bash
|
||||
# Start web-based Chainlit chat interface (requires praisonai[chat])
|
||||
pip install "praisonai[chat]"
|
||||
praisonai chat
|
||||
# Opens browser at http://localhost:8084
|
||||
```
|
||||
|
||||
### Query Rewriting (works with any command):
|
||||
```bash
|
||||
# Rewrite query for better results (uses QueryRewriterAgent)
|
||||
@@ -1118,6 +1164,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 +1182,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 +1414,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
|
||||
@@ -1412,13 +1472,31 @@ praisonai mcp help
|
||||
|
||||
### AI Commit CLI:
|
||||
```bash
|
||||
# Generate AI commit message for staged changes
|
||||
# Full auto mode: stage all, security check, commit, and push
|
||||
praisonai commit -a
|
||||
|
||||
# Interactive mode (requires git add first)
|
||||
praisonai commit
|
||||
|
||||
# Generate, commit, and push
|
||||
# Interactive with auto-push
|
||||
praisonai commit --push
|
||||
|
||||
# Skip security check (not recommended)
|
||||
praisonai commit -a --no-verify
|
||||
```
|
||||
|
||||
**Features:**
|
||||
- 🤖 AI-generated conventional commit messages
|
||||
- 🔒 Built-in security scanning (API keys, passwords, secrets, sensitive files)
|
||||
- 📦 Auto-staging with `-a` flag
|
||||
- 🚀 Auto-push in full auto mode
|
||||
- ✏️ Edit message before commit in interactive mode
|
||||
|
||||
**Security Detection:**
|
||||
- API keys, secrets, tokens (AWS, GitHub, GitLab, Slack)
|
||||
- Passwords and private keys
|
||||
- Sensitive files (`.env`, `id_rsa`, `.pem`, `.key`, etc.)
|
||||
|
||||
### Serve CLI (API Server):
|
||||
```bash
|
||||
# Start API server for agents defined in YAML
|
||||
@@ -1449,6 +1527,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
|
||||
@@ -1674,32 +1801,32 @@ app.include_router(a2a.get_router())
|
||||
|
||||
| Feature | Docs |
|
||||
|---------|:----:|
|
||||
| 🔄 Query Rewrite - RAG optimization | [📖](https://docs.praison.ai/cli/query-rewrite) |
|
||||
| 🔬 Deep Research - Automated research | [📖](https://docs.praison.ai/cli/deep-research) |
|
||||
| 📋 Planning - Step-by-step execution | [📖](https://docs.praison.ai/cli/planning) |
|
||||
| 💾 Memory - Persistent agent memory | [📖](https://docs.praison.ai/cli/memory) |
|
||||
| 📜 Rules - Auto-discovered instructions | [📖](https://docs.praison.ai/cli/rules) |
|
||||
| 🔄 Workflow - Multi-step workflows | [📖](https://docs.praison.ai/cli/workflow) |
|
||||
| 🪝 Hooks - Event-driven actions | [📖](https://docs.praison.ai/cli/hooks) |
|
||||
| 🧠 Claude Memory - Anthropic memory tool | [📖](https://docs.praison.ai/cli/claude-memory) |
|
||||
| 🛡️ Guardrail - Output validation | [📖](https://docs.praison.ai/cli/guardrail) |
|
||||
| 📊 Metrics - Token usage tracking | [📖](https://docs.praison.ai/cli/metrics) |
|
||||
| 🖼️ Image - Vision processing | [📖](https://docs.praison.ai/cli/image) |
|
||||
| 📡 Telemetry - Usage monitoring | [📖](https://docs.praison.ai/cli/telemetry) |
|
||||
| 🔌 MCP - Model Context Protocol | [📖](https://docs.praison.ai/cli/mcp) |
|
||||
| ⚡ Fast Context - Codebase search | [📖](https://docs.praison.ai/cli/fast-context) |
|
||||
| 📚 Knowledge - RAG management | [📖](https://docs.praison.ai/cli/knowledge) |
|
||||
| 💬 Session - Conversation management | [📖](https://docs.praison.ai/cli/session) |
|
||||
| 🔧 Tools - Tool discovery | [📖](https://docs.praison.ai/cli/tools) |
|
||||
| 🤝 Handoff - Agent delegation | [📖](https://docs.praison.ai/cli/handoff) |
|
||||
| 🧠 Auto Memory - Memory extraction | [📖](https://docs.praison.ai/cli/auto-memory) |
|
||||
| 📋 Todo - Task management | [📖](https://docs.praison.ai/cli/todo) |
|
||||
| 🎯 Router - Smart model selection | [📖](https://docs.praison.ai/cli/router) |
|
||||
| 📈 Flow Display - Visual workflow | [📖](https://docs.praison.ai/cli/flow-display) |
|
||||
| ✨ Prompt Expansion - Detailed prompts | [📖](https://docs.praison.ai/cli/prompt-expansion) |
|
||||
| 🌐 Web Search - Real-time search | [📖](https://docs.praison.ai/cli/web-search) |
|
||||
| 📥 Web Fetch - URL content retrieval | [📖](https://docs.praison.ai/cli/web-fetch) |
|
||||
| 💾 Prompt Caching - Cost reduction | [📖](https://docs.praison.ai/cli/prompt-caching) |
|
||||
| 🔄 Query Rewrite - RAG optimization | [📖](https://docs.praison.ai/docs/cli/query-rewrite) |
|
||||
| 🔬 Deep Research - Automated research | [📖](https://docs.praison.ai/docs/cli/deep-research) |
|
||||
| 📋 Planning - Step-by-step execution | [📖](https://docs.praison.ai/docs/cli/planning) |
|
||||
| 💾 Memory - Persistent agent memory | [📖](https://docs.praison.ai/docs/cli/memory) |
|
||||
| 📜 Rules - Auto-discovered instructions | [📖](https://docs.praison.ai/docs/cli/rules) |
|
||||
| 🔄 Workflow - Multi-step workflows | [📖](https://docs.praison.ai/docs/cli/workflow) |
|
||||
| 🪝 Hooks - Event-driven actions | [📖](https://docs.praison.ai/docs/cli/hooks) |
|
||||
| 🧠 Claude Memory - Anthropic memory tool | [📖](https://docs.praison.ai/docs/cli/claude-memory) |
|
||||
| 🛡️ Guardrail - Output validation | [📖](https://docs.praison.ai/docs/cli/guardrail) |
|
||||
| 📊 Metrics - Token usage tracking | [📖](https://docs.praison.ai/docs/cli/metrics) |
|
||||
| 🖼️ Image - Vision processing | [📖](https://docs.praison.ai/docs/cli/image) |
|
||||
| 📡 Telemetry - Usage monitoring | [📖](https://docs.praison.ai/docs/cli/telemetry) |
|
||||
| 🔌 MCP - Model Context Protocol | [📖](https://docs.praison.ai/docs/cli/mcp) |
|
||||
| ⚡ Fast Context - Codebase search | [📖](https://docs.praison.ai/docs/cli/fast-context) |
|
||||
| 📚 Knowledge - RAG management | [📖](https://docs.praison.ai/docs/cli/knowledge) |
|
||||
| 💬 Session - Conversation management | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| 🔧 Tools - Tool discovery | [📖](https://docs.praison.ai/docs/cli/tools) |
|
||||
| 🤝 Handoff - Agent delegation | [📖](https://docs.praison.ai/docs/cli/handoff) |
|
||||
| 🧠 Auto Memory - Memory extraction | [📖](https://docs.praison.ai/docs/cli/auto-memory) |
|
||||
| 📋 Todo - Task management | [📖](https://docs.praison.ai/docs/cli/todo) |
|
||||
| 🎯 Router - Smart model selection | [📖](https://docs.praison.ai/docs/cli/router) |
|
||||
| 📈 Flow Display - Visual workflow | [📖](https://docs.praison.ai/docs/cli/flow-display) |
|
||||
| ✨ Prompt Expansion - Detailed prompts | [📖](https://docs.praison.ai/docs/cli/prompt-expansion) |
|
||||
| 🌐 Web Search - Real-time search | [📖](https://docs.praison.ai/docs/cli/web-search) |
|
||||
| 📥 Web Fetch - URL content retrieval | [📖](https://docs.praison.ai/docs/cli/web-fetch) |
|
||||
| 💾 Prompt Caching - Cost reduction | [📖](https://docs.praison.ai/docs/cli/prompt-caching) |
|
||||
|
||||
---
|
||||
|
||||
@@ -2251,7 +2378,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 +2387,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
|
||||
|
||||
@@ -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.52" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -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.52" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -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.52" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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.")
|
||||
@@ -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())
|
||||
@@ -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
|
||||
@@ -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"])
|
||||
Generated
+1
-1
@@ -2422,7 +2422,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.1.16"
|
||||
version = "0.1.17"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
+172
-38
@@ -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) |
|
||||
@@ -240,12 +252,20 @@ npm install praisonai
|
||||
| 🪝 Hooks | [Example](#9-hooks) | [📖](https://docs.praison.ai/features/hooks) |
|
||||
| 📈 Telemetry | [Example](examples/python/telemetry/production-telemetry-example.py) | [📖](https://docs.praison.ai/features/telemetry) |
|
||||
| 📹 Camera Integration | [Example](examples/python/camera/) | [📖](https://docs.praison.ai/features/camera-integration) |
|
||||
| 📄 Project Docs (.praison/docs/) | [Example](#docs-cli) | [📖](https://docs.praison.ai/cli/docs) |
|
||||
| 🔌 MCP Config Management | [Example](#mcp-config-cli) | [📖](https://docs.praison.ai/cli/mcp) |
|
||||
| 💬 AI Commit Messages | [Example](#ai-commit-cli) | [📖](https://docs.praison.ai/cli/commit) |
|
||||
| @ @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) |
|
||||
| 📄 Project Docs (.praison/docs/) | [Example](#docs-cli) | [📖](https://docs.praison.ai/docs/cli/docs) |
|
||||
| 🔌 MCP Config Management | [Example](#mcp-config-cli) | [📖](https://docs.praison.ai/docs/cli/mcp) |
|
||||
| 💬 AI Commit Messages | [Example](#ai-commit-cli) | [📖](https://docs.praison.ai/docs/cli/commit) |
|
||||
| @ @Mentions in Prompts | [Example](#mentions-in-prompts) | [📖](https://docs.praison.ai/docs/cli/mentions) |
|
||||
| 💾 Auto-Save Sessions | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| 📜 History in Context | [Example](#session-management-python) | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| **🖥️ CLI Features** | | |
|
||||
| ↳ Slash Commands | [Example](examples/python/cli/slash_commands_example.py) | [📖](https://docs.praison.ai/docs/cli/slash-commands) |
|
||||
| ↳ Autonomy Modes | [Example](examples/python/cli/autonomy_modes_example.py) | [📖](https://docs.praison.ai/docs/cli/autonomy-modes) |
|
||||
| ↳ Cost Tracking | [Example](examples/python/cli/cost_tracking_example.py) | [📖](https://docs.praison.ai/docs/cli/cost-tracking) |
|
||||
| ↳ Repository Map | [Example](examples/python/cli/repo_map_example.py) | [📖](https://docs.praison.ai/docs/cli/repo-map) |
|
||||
| ↳ Interactive TUI | [Example](examples/python/cli/interactive_tui_example.py) | [📖](https://docs.praison.ai/docs/cli/interactive-tui) |
|
||||
| ↳ Git Integration | [Example](examples/python/cli/git_integration_example.py) | [📖](https://docs.praison.ai/docs/cli/git-integration) |
|
||||
| ↳ Sandbox Execution | [Example](examples/python/cli/sandbox_execution_example.py) | [📖](https://docs.praison.ai/docs/cli/sandbox-execution) |
|
||||
|
||||
---
|
||||
|
||||
@@ -952,6 +972,32 @@ export OPENAI_API_KEY=xxxxxxxxxxxxxxxxxxxxxx
|
||||
praisonai --auto create a movie script about Robots in Mars
|
||||
```
|
||||
|
||||
### Interactive Mode CLI:
|
||||
```bash
|
||||
# Start interactive terminal mode (inspired by Gemini CLI, Codex CLI, Claude Code)
|
||||
praisonai --interactive
|
||||
praisonai -i
|
||||
|
||||
# Features:
|
||||
# - Streaming responses (no boxes)
|
||||
# - Built-in tools: read_file, write_file, list_files, execute_command, internet_search
|
||||
# - Slash commands: /help, /exit, /tools, /clear
|
||||
|
||||
# Chat mode - single prompt with interactive style (for testing/scripting)
|
||||
# Use --chat (or --chat-mode for backward compatibility)
|
||||
praisonai "list files in current folder" --chat
|
||||
praisonai "search the web for AI news" --chat
|
||||
praisonai "read README.md" --chat
|
||||
```
|
||||
|
||||
### Chat UI (Web Interface):
|
||||
```bash
|
||||
# Start web-based Chainlit chat interface (requires praisonai[chat])
|
||||
pip install "praisonai[chat]"
|
||||
praisonai chat
|
||||
# Opens browser at http://localhost:8084
|
||||
```
|
||||
|
||||
### Query Rewriting (works with any command):
|
||||
```bash
|
||||
# Rewrite query for better results (uses QueryRewriterAgent)
|
||||
@@ -1118,6 +1164,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 +1182,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 +1414,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
|
||||
@@ -1412,13 +1472,31 @@ praisonai mcp help
|
||||
|
||||
### AI Commit CLI:
|
||||
```bash
|
||||
# Generate AI commit message for staged changes
|
||||
# Full auto mode: stage all, security check, commit, and push
|
||||
praisonai commit -a
|
||||
|
||||
# Interactive mode (requires git add first)
|
||||
praisonai commit
|
||||
|
||||
# Generate, commit, and push
|
||||
# Interactive with auto-push
|
||||
praisonai commit --push
|
||||
|
||||
# Skip security check (not recommended)
|
||||
praisonai commit -a --no-verify
|
||||
```
|
||||
|
||||
**Features:**
|
||||
- 🤖 AI-generated conventional commit messages
|
||||
- 🔒 Built-in security scanning (API keys, passwords, secrets, sensitive files)
|
||||
- 📦 Auto-staging with `-a` flag
|
||||
- 🚀 Auto-push in full auto mode
|
||||
- ✏️ Edit message before commit in interactive mode
|
||||
|
||||
**Security Detection:**
|
||||
- API keys, secrets, tokens (AWS, GitHub, GitLab, Slack)
|
||||
- Passwords and private keys
|
||||
- Sensitive files (`.env`, `id_rsa`, `.pem`, `.key`, etc.)
|
||||
|
||||
### Serve CLI (API Server):
|
||||
```bash
|
||||
# Start API server for agents defined in YAML
|
||||
@@ -1449,6 +1527,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
|
||||
@@ -1674,32 +1801,32 @@ app.include_router(a2a.get_router())
|
||||
|
||||
| Feature | Docs |
|
||||
|---------|:----:|
|
||||
| 🔄 Query Rewrite - RAG optimization | [📖](https://docs.praison.ai/cli/query-rewrite) |
|
||||
| 🔬 Deep Research - Automated research | [📖](https://docs.praison.ai/cli/deep-research) |
|
||||
| 📋 Planning - Step-by-step execution | [📖](https://docs.praison.ai/cli/planning) |
|
||||
| 💾 Memory - Persistent agent memory | [📖](https://docs.praison.ai/cli/memory) |
|
||||
| 📜 Rules - Auto-discovered instructions | [📖](https://docs.praison.ai/cli/rules) |
|
||||
| 🔄 Workflow - Multi-step workflows | [📖](https://docs.praison.ai/cli/workflow) |
|
||||
| 🪝 Hooks - Event-driven actions | [📖](https://docs.praison.ai/cli/hooks) |
|
||||
| 🧠 Claude Memory - Anthropic memory tool | [📖](https://docs.praison.ai/cli/claude-memory) |
|
||||
| 🛡️ Guardrail - Output validation | [📖](https://docs.praison.ai/cli/guardrail) |
|
||||
| 📊 Metrics - Token usage tracking | [📖](https://docs.praison.ai/cli/metrics) |
|
||||
| 🖼️ Image - Vision processing | [📖](https://docs.praison.ai/cli/image) |
|
||||
| 📡 Telemetry - Usage monitoring | [📖](https://docs.praison.ai/cli/telemetry) |
|
||||
| 🔌 MCP - Model Context Protocol | [📖](https://docs.praison.ai/cli/mcp) |
|
||||
| ⚡ Fast Context - Codebase search | [📖](https://docs.praison.ai/cli/fast-context) |
|
||||
| 📚 Knowledge - RAG management | [📖](https://docs.praison.ai/cli/knowledge) |
|
||||
| 💬 Session - Conversation management | [📖](https://docs.praison.ai/cli/session) |
|
||||
| 🔧 Tools - Tool discovery | [📖](https://docs.praison.ai/cli/tools) |
|
||||
| 🤝 Handoff - Agent delegation | [📖](https://docs.praison.ai/cli/handoff) |
|
||||
| 🧠 Auto Memory - Memory extraction | [📖](https://docs.praison.ai/cli/auto-memory) |
|
||||
| 📋 Todo - Task management | [📖](https://docs.praison.ai/cli/todo) |
|
||||
| 🎯 Router - Smart model selection | [📖](https://docs.praison.ai/cli/router) |
|
||||
| 📈 Flow Display - Visual workflow | [📖](https://docs.praison.ai/cli/flow-display) |
|
||||
| ✨ Prompt Expansion - Detailed prompts | [📖](https://docs.praison.ai/cli/prompt-expansion) |
|
||||
| 🌐 Web Search - Real-time search | [📖](https://docs.praison.ai/cli/web-search) |
|
||||
| 📥 Web Fetch - URL content retrieval | [📖](https://docs.praison.ai/cli/web-fetch) |
|
||||
| 💾 Prompt Caching - Cost reduction | [📖](https://docs.praison.ai/cli/prompt-caching) |
|
||||
| 🔄 Query Rewrite - RAG optimization | [📖](https://docs.praison.ai/docs/cli/query-rewrite) |
|
||||
| 🔬 Deep Research - Automated research | [📖](https://docs.praison.ai/docs/cli/deep-research) |
|
||||
| 📋 Planning - Step-by-step execution | [📖](https://docs.praison.ai/docs/cli/planning) |
|
||||
| 💾 Memory - Persistent agent memory | [📖](https://docs.praison.ai/docs/cli/memory) |
|
||||
| 📜 Rules - Auto-discovered instructions | [📖](https://docs.praison.ai/docs/cli/rules) |
|
||||
| 🔄 Workflow - Multi-step workflows | [📖](https://docs.praison.ai/docs/cli/workflow) |
|
||||
| 🪝 Hooks - Event-driven actions | [📖](https://docs.praison.ai/docs/cli/hooks) |
|
||||
| 🧠 Claude Memory - Anthropic memory tool | [📖](https://docs.praison.ai/docs/cli/claude-memory) |
|
||||
| 🛡️ Guardrail - Output validation | [📖](https://docs.praison.ai/docs/cli/guardrail) |
|
||||
| 📊 Metrics - Token usage tracking | [📖](https://docs.praison.ai/docs/cli/metrics) |
|
||||
| 🖼️ Image - Vision processing | [📖](https://docs.praison.ai/docs/cli/image) |
|
||||
| 📡 Telemetry - Usage monitoring | [📖](https://docs.praison.ai/docs/cli/telemetry) |
|
||||
| 🔌 MCP - Model Context Protocol | [📖](https://docs.praison.ai/docs/cli/mcp) |
|
||||
| ⚡ Fast Context - Codebase search | [📖](https://docs.praison.ai/docs/cli/fast-context) |
|
||||
| 📚 Knowledge - RAG management | [📖](https://docs.praison.ai/docs/cli/knowledge) |
|
||||
| 💬 Session - Conversation management | [📖](https://docs.praison.ai/docs/cli/session) |
|
||||
| 🔧 Tools - Tool discovery | [📖](https://docs.praison.ai/docs/cli/tools) |
|
||||
| 🤝 Handoff - Agent delegation | [📖](https://docs.praison.ai/docs/cli/handoff) |
|
||||
| 🧠 Auto Memory - Memory extraction | [📖](https://docs.praison.ai/docs/cli/auto-memory) |
|
||||
| 📋 Todo - Task management | [📖](https://docs.praison.ai/docs/cli/todo) |
|
||||
| 🎯 Router - Smart model selection | [📖](https://docs.praison.ai/docs/cli/router) |
|
||||
| 📈 Flow Display - Visual workflow | [📖](https://docs.praison.ai/docs/cli/flow-display) |
|
||||
| ✨ Prompt Expansion - Detailed prompts | [📖](https://docs.praison.ai/docs/cli/prompt-expansion) |
|
||||
| 🌐 Web Search - Real-time search | [📖](https://docs.praison.ai/docs/cli/web-search) |
|
||||
| 📥 Web Fetch - URL content retrieval | [📖](https://docs.praison.ai/docs/cli/web-fetch) |
|
||||
| 💾 Prompt Caching - Cost reduction | [📖](https://docs.praison.ai/docs/cli/prompt-caching) |
|
||||
|
||||
---
|
||||
|
||||
@@ -2251,7 +2378,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 +2387,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
|
||||
|
||||
@@ -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.52.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.52.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -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
@@ -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,11 @@ 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
|
||||
- message_queue: Message queuing while agent is processing
|
||||
"""
|
||||
|
||||
# Type hints available for IDE support
|
||||
@@ -43,6 +48,16 @@ __all__ = [
|
||||
'FlowDisplayHandler',
|
||||
'WorkflowHandler',
|
||||
'N8nHandler',
|
||||
'ExternalAgentsHandler',
|
||||
# New CLI enhancement features
|
||||
'SlashCommandHandler',
|
||||
'AutonomyModeHandler',
|
||||
'CostTrackerHandler',
|
||||
'RepoMapHandler',
|
||||
'InteractiveTUIHandler',
|
||||
'GitIntegrationHandler',
|
||||
'SandboxExecutorHandler',
|
||||
'MessageQueueHandler',
|
||||
]
|
||||
|
||||
def __getattr__(name):
|
||||
@@ -95,4 +110,32 @@ 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
|
||||
elif name == 'MessageQueueHandler':
|
||||
from .message_queue import MessageQueueHandler
|
||||
return MessageQueueHandler
|
||||
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,587 @@
|
||||
"""
|
||||
Message Queue for Interactive Mode.
|
||||
|
||||
Provides FIFO message queuing while the AI agent is processing,
|
||||
inspired by Claude Code, Windsurf Cascade, Cursor, and Gemini CLI.
|
||||
|
||||
Features:
|
||||
- Queue messages while agent is processing
|
||||
- Auto-process queue when agent becomes idle
|
||||
- Visual indicators for queue status
|
||||
- Thread-safe operations
|
||||
"""
|
||||
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Enums
|
||||
# ============================================================================
|
||||
|
||||
class ProcessingState(Enum):
|
||||
"""State of the agent processing."""
|
||||
IDLE = "idle"
|
||||
PROCESSING = "processing"
|
||||
WAITING_APPROVAL = "waiting_approval"
|
||||
|
||||
|
||||
class MessagePriority(Enum):
|
||||
"""Priority levels for queued messages."""
|
||||
LOW = "low"
|
||||
NORMAL = "normal"
|
||||
HIGH = "high"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Data Classes
|
||||
# ============================================================================
|
||||
|
||||
@dataclass
|
||||
class QueuedMessage:
|
||||
"""A message in the queue with metadata."""
|
||||
content: str
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
priority: MessagePriority = MessagePriority.NORMAL
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# MessageQueue - Thread-safe FIFO queue
|
||||
# ============================================================================
|
||||
|
||||
class MessageQueue:
|
||||
"""
|
||||
Thread-safe FIFO message queue.
|
||||
|
||||
Provides basic queue operations with thread safety for concurrent access.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._messages: List[str] = []
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def is_empty(self) -> bool:
|
||||
"""Check if queue is empty."""
|
||||
with self._lock:
|
||||
return len(self._messages) == 0
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
"""Get number of messages in queue."""
|
||||
with self._lock:
|
||||
return len(self._messages)
|
||||
|
||||
def add(self, message: str) -> bool:
|
||||
"""
|
||||
Add a message to the end of the queue.
|
||||
|
||||
Args:
|
||||
message: The message to add
|
||||
|
||||
Returns:
|
||||
True if message was added, False if empty/whitespace
|
||||
"""
|
||||
trimmed = message.strip() if message else ""
|
||||
if not trimmed:
|
||||
return False
|
||||
|
||||
with self._lock:
|
||||
self._messages.append(trimmed)
|
||||
return True
|
||||
|
||||
def pop(self) -> Optional[str]:
|
||||
"""
|
||||
Remove and return the first message (FIFO).
|
||||
|
||||
Returns:
|
||||
The first message, or None if queue is empty
|
||||
"""
|
||||
with self._lock:
|
||||
if self._messages:
|
||||
return self._messages.pop(0)
|
||||
return None
|
||||
|
||||
def peek(self) -> Optional[str]:
|
||||
"""
|
||||
View the first message without removing it.
|
||||
|
||||
Returns:
|
||||
The first message, or None if queue is empty
|
||||
"""
|
||||
with self._lock:
|
||||
if self._messages:
|
||||
return self._messages[0]
|
||||
return None
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Remove all messages from the queue."""
|
||||
with self._lock:
|
||||
self._messages.clear()
|
||||
|
||||
def get_all(self) -> List[str]:
|
||||
"""
|
||||
Get all messages as a list (does not remove them).
|
||||
|
||||
Returns:
|
||||
Copy of all messages in queue order
|
||||
"""
|
||||
with self._lock:
|
||||
return self._messages.copy()
|
||||
|
||||
def remove_at(self, index: int) -> Optional[str]:
|
||||
"""
|
||||
Remove message at specific index.
|
||||
|
||||
Args:
|
||||
index: The index to remove (0-based)
|
||||
|
||||
Returns:
|
||||
The removed message, or None if index is invalid
|
||||
"""
|
||||
with self._lock:
|
||||
if 0 <= index < len(self._messages):
|
||||
return self._messages.pop(index)
|
||||
return None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# StateManager - Processing state management
|
||||
# ============================================================================
|
||||
|
||||
class StateManager:
|
||||
"""
|
||||
Manages the processing state with optional callbacks.
|
||||
|
||||
Tracks whether the agent is idle, processing, or waiting for approval.
|
||||
"""
|
||||
|
||||
def __init__(self, on_state_change: Optional[Callable[[ProcessingState, ProcessingState], None]] = None):
|
||||
self._state = ProcessingState.IDLE
|
||||
self._lock = threading.Lock()
|
||||
self._on_state_change = on_state_change
|
||||
|
||||
@property
|
||||
def current_state(self) -> ProcessingState:
|
||||
"""Get current processing state."""
|
||||
with self._lock:
|
||||
return self._state
|
||||
|
||||
@property
|
||||
def is_idle(self) -> bool:
|
||||
"""Check if currently idle."""
|
||||
return self.current_state == ProcessingState.IDLE
|
||||
|
||||
@property
|
||||
def is_processing(self) -> bool:
|
||||
"""Check if currently processing."""
|
||||
return self.current_state == ProcessingState.PROCESSING
|
||||
|
||||
def set_state(self, new_state: ProcessingState) -> None:
|
||||
"""
|
||||
Set the processing state.
|
||||
|
||||
Args:
|
||||
new_state: The new state to set
|
||||
"""
|
||||
with self._lock:
|
||||
old_state = self._state
|
||||
self._state = new_state
|
||||
|
||||
# Call callback outside lock to prevent deadlock
|
||||
if self._on_state_change and old_state != new_state:
|
||||
self._on_state_change(old_state, new_state)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# QueueDisplay - Visual formatting
|
||||
# ============================================================================
|
||||
|
||||
class QueueDisplay:
|
||||
"""
|
||||
Formats queue and status for display.
|
||||
|
||||
Provides visual indicators for queue status, inspired by Codex CLI.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
queue: MessageQueue,
|
||||
state_manager: Optional[StateManager] = None,
|
||||
max_message_length: int = 50
|
||||
):
|
||||
self._queue = queue
|
||||
self._state_manager = state_manager
|
||||
self._max_length = max_message_length
|
||||
|
||||
def format_queue(self) -> str:
|
||||
"""
|
||||
Format queued messages for display.
|
||||
|
||||
Returns:
|
||||
Formatted string with queued messages
|
||||
"""
|
||||
messages = self._queue.get_all()
|
||||
if not messages:
|
||||
return ""
|
||||
|
||||
lines = []
|
||||
for msg in messages:
|
||||
# Truncate long messages
|
||||
display_msg = msg
|
||||
if len(msg) > self._max_length:
|
||||
display_msg = msg[:self._max_length - 3] + "..."
|
||||
lines.append(f" ↳ {display_msg}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def format_status(self) -> str:
|
||||
"""
|
||||
Format current processing status.
|
||||
|
||||
Returns:
|
||||
Status string (e.g., "⏳ Processing..." or "")
|
||||
"""
|
||||
if not self._state_manager:
|
||||
return ""
|
||||
|
||||
state = self._state_manager.current_state
|
||||
if state == ProcessingState.PROCESSING:
|
||||
return "⏳ Processing..."
|
||||
elif state == ProcessingState.WAITING_APPROVAL:
|
||||
return "🔒 Waiting for approval..."
|
||||
return ""
|
||||
|
||||
def format_queue_count(self) -> str:
|
||||
"""
|
||||
Format queue count indicator.
|
||||
|
||||
Returns:
|
||||
Count string (e.g., "📋 Queued (2)")
|
||||
"""
|
||||
count = self._queue.count
|
||||
if count == 0:
|
||||
return ""
|
||||
return f"📋 Queued ({count})"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# MessageQueueHandler - Main handler class
|
||||
# ============================================================================
|
||||
|
||||
class MessageQueueHandler:
|
||||
"""
|
||||
Main handler for message queue functionality.
|
||||
|
||||
Coordinates queue, state, and processing.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
processor: Optional[Callable[[str], str]] = None,
|
||||
on_response: Optional[Callable[[str], None]] = None
|
||||
):
|
||||
self._processor = processor
|
||||
self._on_response = on_response
|
||||
self._queue = MessageQueue()
|
||||
self._state_manager = StateManager(on_state_change=self._on_state_change)
|
||||
self._processing_lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def queue(self) -> MessageQueue:
|
||||
"""Get the message queue."""
|
||||
return self._queue
|
||||
|
||||
@property
|
||||
def state_manager(self) -> StateManager:
|
||||
"""Get the state manager."""
|
||||
return self._state_manager
|
||||
|
||||
def submit(self, message: str) -> bool:
|
||||
"""
|
||||
Submit a message for processing.
|
||||
|
||||
If idle, processes immediately. If processing, queues the message.
|
||||
|
||||
Args:
|
||||
message: The message to submit
|
||||
|
||||
Returns:
|
||||
True if processed/queued, False if empty message
|
||||
"""
|
||||
trimmed = message.strip() if message else ""
|
||||
if not trimmed:
|
||||
return False
|
||||
|
||||
if self._state_manager.is_idle:
|
||||
# Process immediately
|
||||
self._process_message(trimmed)
|
||||
else:
|
||||
# Queue for later
|
||||
self._queue.add(trimmed)
|
||||
|
||||
return True
|
||||
|
||||
def on_processing_complete(self) -> None:
|
||||
"""
|
||||
Called when current processing completes.
|
||||
|
||||
Sets state to idle and processes next queued message if any.
|
||||
"""
|
||||
self._state_manager.set_state(ProcessingState.IDLE)
|
||||
self._process_next_in_queue()
|
||||
|
||||
def get_status(self) -> dict:
|
||||
"""
|
||||
Get current status.
|
||||
|
||||
Returns:
|
||||
Dict with queue_count and state
|
||||
"""
|
||||
return {
|
||||
'queue_count': self._queue.count,
|
||||
'state': self._state_manager.current_state.value,
|
||||
'messages': self._queue.get_all()
|
||||
}
|
||||
|
||||
def clear_queue(self) -> None:
|
||||
"""Clear all queued messages."""
|
||||
self._queue.clear()
|
||||
|
||||
def _process_message(self, message: str) -> None:
|
||||
"""Process a single message."""
|
||||
if not self._processor:
|
||||
return
|
||||
|
||||
with self._processing_lock:
|
||||
self._state_manager.set_state(ProcessingState.PROCESSING)
|
||||
try:
|
||||
response = self._processor(message)
|
||||
if self._on_response and response:
|
||||
self._on_response(response)
|
||||
finally:
|
||||
self._state_manager.set_state(ProcessingState.IDLE)
|
||||
|
||||
def _process_next_in_queue(self) -> None:
|
||||
"""Process the next message in queue if any."""
|
||||
next_msg = self._queue.pop()
|
||||
if next_msg:
|
||||
self._process_message(next_msg)
|
||||
|
||||
def _on_state_change(self, old_state: ProcessingState, new_state: ProcessingState) -> None:
|
||||
"""Handle state changes."""
|
||||
# Auto-process queue when becoming idle
|
||||
if new_state == ProcessingState.IDLE and old_state == ProcessingState.PROCESSING:
|
||||
# Don't auto-process here to avoid recursion
|
||||
# The on_processing_complete method handles this
|
||||
pass
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# AsyncProcessor - Background processing
|
||||
# ============================================================================
|
||||
|
||||
class AsyncProcessor:
|
||||
"""
|
||||
Runs work functions in background threads.
|
||||
|
||||
Allows agent processing to happen asynchronously while
|
||||
the main thread remains responsive for user input.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
on_complete: Optional[Callable[[str], None]] = None,
|
||||
on_status: Optional[Callable[[str], None]] = None,
|
||||
on_error: Optional[Callable[[Exception], None]] = None
|
||||
):
|
||||
self._on_complete = on_complete
|
||||
self._on_status = on_status
|
||||
self._on_error = on_error
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._running = False
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if currently processing."""
|
||||
with self._lock:
|
||||
return self._running
|
||||
|
||||
def start(self, work_fn: Callable[[], str]) -> None:
|
||||
"""
|
||||
Start processing work in background thread.
|
||||
|
||||
Args:
|
||||
work_fn: Function to execute (should return result string)
|
||||
"""
|
||||
def _worker():
|
||||
try:
|
||||
with self._lock:
|
||||
self._running = True
|
||||
|
||||
result = work_fn()
|
||||
|
||||
with self._lock:
|
||||
self._running = False
|
||||
|
||||
if self._on_complete:
|
||||
self._on_complete(result)
|
||||
except Exception as e:
|
||||
with self._lock:
|
||||
self._running = False
|
||||
if self._on_error:
|
||||
self._on_error(e)
|
||||
|
||||
self._thread = threading.Thread(target=_worker, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def update_status(self, status: str) -> None:
|
||||
"""Update status (can be called from work function)."""
|
||||
if self._on_status:
|
||||
self._on_status(status)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# LiveStatusDisplay - Real-time status display
|
||||
# ============================================================================
|
||||
|
||||
class LiveStatusDisplay:
|
||||
"""
|
||||
Tracks and displays real-time status during processing.
|
||||
|
||||
Shows tool calls, command executions, and current status.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._status = ""
|
||||
self._tool_calls: List[dict] = []
|
||||
self._commands: List[str] = []
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def current_status(self) -> str:
|
||||
"""Get current status message."""
|
||||
with self._lock:
|
||||
return self._status
|
||||
|
||||
@property
|
||||
def tool_calls(self) -> List[dict]:
|
||||
"""Get list of tool calls."""
|
||||
with self._lock:
|
||||
return self._tool_calls.copy()
|
||||
|
||||
@property
|
||||
def commands(self) -> List[str]:
|
||||
"""Get list of commands executed."""
|
||||
with self._lock:
|
||||
return self._commands.copy()
|
||||
|
||||
def update_status(self, status: str) -> None:
|
||||
"""Update current status message."""
|
||||
with self._lock:
|
||||
self._status = status
|
||||
|
||||
def add_tool_call(self, name: str, args: dict) -> None:
|
||||
"""Add a tool call to the display."""
|
||||
with self._lock:
|
||||
self._tool_calls.append({
|
||||
'name': name,
|
||||
'args': args,
|
||||
'timestamp': time.time()
|
||||
})
|
||||
|
||||
def add_command(self, command: str) -> None:
|
||||
"""Add a command execution to the display."""
|
||||
with self._lock:
|
||||
self._commands.append(command)
|
||||
|
||||
def format(self) -> str:
|
||||
"""Format current status for display."""
|
||||
lines = []
|
||||
|
||||
with self._lock:
|
||||
if self._status:
|
||||
lines.append(f"⏳ {self._status}")
|
||||
|
||||
for tool in self._tool_calls[-3:]: # Show last 3 tools
|
||||
lines.append(f" 🔧 {tool['name']}")
|
||||
|
||||
for cmd in self._commands[-2:]: # Show last 2 commands
|
||||
display_cmd = cmd[:40] + "..." if len(cmd) > 40 else cmd
|
||||
lines.append(f" 💻 {display_cmd}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all status."""
|
||||
with self._lock:
|
||||
self._status = ""
|
||||
self._tool_calls.clear()
|
||||
self._commands.clear()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# NonBlockingInput - Async input handling
|
||||
# ============================================================================
|
||||
|
||||
class NonBlockingInput:
|
||||
"""
|
||||
Handles user input asynchronously.
|
||||
|
||||
Allows users to type new messages while processing is ongoing.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._queue = MessageQueue()
|
||||
|
||||
@property
|
||||
def has_pending(self) -> bool:
|
||||
"""Check if there are pending inputs."""
|
||||
return not self._queue.is_empty
|
||||
|
||||
@property
|
||||
def pending_count(self) -> int:
|
||||
"""Get number of pending inputs."""
|
||||
return self._queue.count
|
||||
|
||||
def submit(self, message: str) -> bool:
|
||||
"""Submit a new input message."""
|
||||
return self._queue.add(message)
|
||||
|
||||
def pop(self) -> Optional[str]:
|
||||
"""Get next pending input."""
|
||||
return self._queue.pop()
|
||||
|
||||
def peek(self) -> Optional[str]:
|
||||
"""View next pending input without removing."""
|
||||
return self._queue.peek()
|
||||
|
||||
def get_all(self) -> List[str]:
|
||||
"""Get all pending inputs."""
|
||||
return self._queue.get_all()
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all pending inputs."""
|
||||
self._queue.clear()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Exports
|
||||
# ============================================================================
|
||||
|
||||
__all__ = [
|
||||
'ProcessingState',
|
||||
'MessagePriority',
|
||||
'QueuedMessage',
|
||||
'MessageQueue',
|
||||
'StateManager',
|
||||
'QueueDisplay',
|
||||
'MessageQueueHandler',
|
||||
'AsyncProcessor',
|
||||
'LiveStatusDisplay',
|
||||
'NonBlockingInput',
|
||||
]
|
||||
@@ -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"
|
||||
))
|
||||
@@ -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")
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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"
|
||||
@@ -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("<div>")
|
||||
'<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
|
||||
@@ -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.52 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 @@
|
||||
__version__ = "2.3.38"
|
||||
__version__ = "2.3.52"
|
||||
@@ -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",
|
||||
|
||||
@@ -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,151 @@
|
||||
"""
|
||||
Unit tests for --chat flag (alias for --chat-mode).
|
||||
|
||||
Tests verify that:
|
||||
1. --chat flag works as an alias for --chat-mode
|
||||
2. --chat flag does not conflict with 'chat' subcommand
|
||||
3. Both flags set chat_mode=True correctly
|
||||
4. Backward compatibility with --chat-mode is maintained
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import argparse
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))))
|
||||
|
||||
|
||||
class TestChatFlagParsing:
|
||||
"""Test argument parsing for --chat and --chat-mode flags."""
|
||||
|
||||
def test_chat_mode_flag_sets_chat_mode_true(self):
|
||||
"""--chat-mode should set chat_mode=True."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["--chat-mode", "test prompt"])
|
||||
assert args.chat_mode is True
|
||||
assert args.command == "test prompt"
|
||||
|
||||
def test_chat_flag_sets_chat_mode_true(self):
|
||||
"""--chat should set chat_mode=True (alias for --chat-mode)."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["--chat", "test prompt"])
|
||||
assert args.chat_mode is True
|
||||
assert args.command == "test prompt"
|
||||
|
||||
def test_chat_subcommand_does_not_set_chat_mode(self):
|
||||
"""'chat' as positional command should NOT set chat_mode."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["chat"])
|
||||
assert args.chat_mode is False
|
||||
assert args.command == "chat"
|
||||
|
||||
def test_chat_flag_after_prompt(self):
|
||||
"""--chat after prompt should work."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["What is 2+2?", "--chat"])
|
||||
assert args.chat_mode is True
|
||||
assert args.command == "What is 2+2?"
|
||||
|
||||
def test_chat_mode_flag_after_prompt(self):
|
||||
"""--chat-mode after prompt should work (backward compat)."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["What is 2+2?", "--chat-mode"])
|
||||
assert args.chat_mode is True
|
||||
assert args.command == "What is 2+2?"
|
||||
|
||||
def test_no_flag_defaults_to_false(self):
|
||||
"""Without --chat or --chat-mode, chat_mode should be False."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["some prompt"])
|
||||
assert args.chat_mode is False
|
||||
assert args.command == "some prompt"
|
||||
|
||||
|
||||
class TestChatFlagNoConflict:
|
||||
"""Test that --chat flag does not conflict with other features."""
|
||||
|
||||
def test_chat_flag_with_other_flags(self):
|
||||
"""--chat should work alongside other flags."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("--verbose", "-v", action="store_true")
|
||||
parser.add_argument("--model", type=str)
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["--chat", "--verbose", "--model", "gpt-4o", "test"])
|
||||
assert args.chat_mode is True
|
||||
assert args.verbose is True
|
||||
assert args.model == "gpt-4o"
|
||||
assert args.command == "test"
|
||||
|
||||
def test_chat_subcommand_with_other_flags(self):
|
||||
"""'chat' subcommand should work with other flags."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("--verbose", "-v", action="store_true")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
args, _ = parser.parse_known_args(["chat", "--verbose"])
|
||||
assert args.chat_mode is False # subcommand, not flag
|
||||
assert args.command == "chat"
|
||||
assert args.verbose is True
|
||||
|
||||
|
||||
class TestChatFlagIntegration:
|
||||
"""Integration tests with actual PraisonAI CLI parser."""
|
||||
|
||||
def test_praisonai_parser_has_chat_flag(self):
|
||||
"""PraisonAI CLI should have --chat as alias for --chat-mode."""
|
||||
# This test will fail until we implement the change
|
||||
from praisonai.cli.main import PraisonAI
|
||||
|
||||
# Create instance and get parser
|
||||
pai = PraisonAI()
|
||||
|
||||
# Mock sys.argv to avoid test environment issues
|
||||
import sys
|
||||
original_argv = sys.argv
|
||||
try:
|
||||
sys.argv = ["praisonai", "--chat", "test prompt"]
|
||||
args = pai.parse_args()
|
||||
assert hasattr(args, 'chat_mode')
|
||||
# After implementation, this should be True
|
||||
# For now, it may work due to argparse prefix matching
|
||||
finally:
|
||||
sys.argv = original_argv
|
||||
|
||||
def test_praisonai_chat_subcommand_still_works(self):
|
||||
"""'praisonai chat' subcommand should still launch Chainlit.
|
||||
|
||||
Note: In test environment, parse_args returns default args to avoid
|
||||
pytest interference. This test verifies the parser structure instead.
|
||||
"""
|
||||
# Test the parser structure directly without going through PraisonAI
|
||||
# which has test environment detection
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--chat-mode", "--chat", action="store_true", dest="chat_mode")
|
||||
parser.add_argument("command", nargs="?")
|
||||
|
||||
# Simulate 'praisonai chat' - should set command='chat', not chat_mode
|
||||
args, _ = parser.parse_known_args(["chat"])
|
||||
assert args.command == "chat"
|
||||
assert args.chat_mode is False # subcommand, not flag
|
||||
@@ -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,661 @@
|
||||
"""
|
||||
Unit tests for Message Queue feature.
|
||||
|
||||
Tests for:
|
||||
1. MessageQueue - FIFO queue for messages
|
||||
2. ProcessingState - State management
|
||||
3. QueuedMessage - Message with metadata
|
||||
4. MessageQueueHandler - Main handler class
|
||||
"""
|
||||
|
||||
import time
|
||||
import threading
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 1. MessageQueue Tests - FIFO operations
|
||||
# ============================================================================
|
||||
|
||||
class TestMessageQueue:
|
||||
"""Test MessageQueue FIFO operations."""
|
||||
|
||||
def test_create_empty_queue(self):
|
||||
"""Queue should start empty."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
assert queue.is_empty
|
||||
assert queue.count == 0
|
||||
|
||||
def test_add_message(self):
|
||||
"""Should add message to queue."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("Hello")
|
||||
assert not queue.is_empty
|
||||
assert queue.count == 1
|
||||
|
||||
def test_add_multiple_messages(self):
|
||||
"""Should add multiple messages in order."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
queue.add("Third")
|
||||
assert queue.count == 3
|
||||
|
||||
def test_pop_message_fifo(self):
|
||||
"""Should pop messages in FIFO order."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
msg = queue.pop()
|
||||
assert msg == "First"
|
||||
assert queue.count == 1
|
||||
|
||||
def test_pop_empty_queue(self):
|
||||
"""Should return None when popping empty queue."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
msg = queue.pop()
|
||||
assert msg is None
|
||||
|
||||
def test_peek_message(self):
|
||||
"""Should peek without removing."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
msg = queue.peek()
|
||||
assert msg == "First"
|
||||
assert queue.count == 2 # Not removed
|
||||
|
||||
def test_peek_empty_queue(self):
|
||||
"""Should return None when peeking empty queue."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
msg = queue.peek()
|
||||
assert msg is None
|
||||
|
||||
def test_clear_queue(self):
|
||||
"""Should clear all messages."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
queue.clear()
|
||||
assert queue.is_empty
|
||||
assert queue.count == 0
|
||||
|
||||
def test_get_all_messages(self):
|
||||
"""Should get all messages as list."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
messages = queue.get_all()
|
||||
assert messages == ["First", "Second"]
|
||||
assert queue.count == 2 # Not removed
|
||||
|
||||
def test_remove_at_index(self):
|
||||
"""Should remove message at specific index."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
queue.add("Third")
|
||||
removed = queue.remove_at(1)
|
||||
assert removed == "Second"
|
||||
assert queue.count == 2
|
||||
assert queue.get_all() == ["First", "Third"]
|
||||
|
||||
def test_remove_at_invalid_index(self):
|
||||
"""Should return None for invalid index."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
removed = queue.remove_at(5)
|
||||
assert removed is None
|
||||
assert queue.count == 1
|
||||
|
||||
def test_add_empty_message_ignored(self):
|
||||
"""Should ignore empty messages."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("")
|
||||
queue.add(" ")
|
||||
assert queue.is_empty
|
||||
|
||||
def test_thread_safety(self):
|
||||
"""Queue should be thread-safe."""
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
queue = MessageQueue()
|
||||
results = []
|
||||
|
||||
def add_messages():
|
||||
for i in range(100):
|
||||
queue.add(f"msg-{i}")
|
||||
|
||||
def pop_messages():
|
||||
for _ in range(100):
|
||||
msg = queue.pop()
|
||||
if msg:
|
||||
results.append(msg)
|
||||
|
||||
t1 = threading.Thread(target=add_messages)
|
||||
t2 = threading.Thread(target=pop_messages)
|
||||
t1.start()
|
||||
t2.start()
|
||||
t1.join()
|
||||
t2.join()
|
||||
|
||||
# All messages should be accounted for
|
||||
total = len(results) + queue.count
|
||||
assert total == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 2. ProcessingState Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestProcessingState:
|
||||
"""Test ProcessingState enum and state management."""
|
||||
|
||||
def test_state_enum_values(self):
|
||||
"""Should have correct state values."""
|
||||
from praisonai.cli.features.message_queue import ProcessingState
|
||||
assert ProcessingState.IDLE.value == "idle"
|
||||
assert ProcessingState.PROCESSING.value == "processing"
|
||||
assert ProcessingState.WAITING_APPROVAL.value == "waiting_approval"
|
||||
|
||||
def test_state_manager_initial_state(self):
|
||||
"""Should start in IDLE state."""
|
||||
from praisonai.cli.features.message_queue import StateManager, ProcessingState
|
||||
manager = StateManager()
|
||||
assert manager.current_state == ProcessingState.IDLE
|
||||
|
||||
def test_state_transition(self):
|
||||
"""Should transition between states."""
|
||||
from praisonai.cli.features.message_queue import StateManager, ProcessingState
|
||||
manager = StateManager()
|
||||
manager.set_state(ProcessingState.PROCESSING)
|
||||
assert manager.current_state == ProcessingState.PROCESSING
|
||||
|
||||
def test_is_idle(self):
|
||||
"""Should correctly report idle state."""
|
||||
from praisonai.cli.features.message_queue import StateManager, ProcessingState
|
||||
manager = StateManager()
|
||||
assert manager.is_idle
|
||||
manager.set_state(ProcessingState.PROCESSING)
|
||||
assert not manager.is_idle
|
||||
|
||||
def test_is_processing(self):
|
||||
"""Should correctly report processing state."""
|
||||
from praisonai.cli.features.message_queue import StateManager, ProcessingState
|
||||
manager = StateManager()
|
||||
assert not manager.is_processing
|
||||
manager.set_state(ProcessingState.PROCESSING)
|
||||
assert manager.is_processing
|
||||
|
||||
def test_state_callback(self):
|
||||
"""Should call callback on state change."""
|
||||
from praisonai.cli.features.message_queue import StateManager, ProcessingState
|
||||
callback_called = []
|
||||
|
||||
def on_state_change(old_state, new_state):
|
||||
callback_called.append((old_state, new_state))
|
||||
|
||||
manager = StateManager(on_state_change=on_state_change)
|
||||
manager.set_state(ProcessingState.PROCESSING)
|
||||
|
||||
assert len(callback_called) == 1
|
||||
assert callback_called[0] == (ProcessingState.IDLE, ProcessingState.PROCESSING)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 3. QueuedMessage Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestQueuedMessage:
|
||||
"""Test QueuedMessage dataclass."""
|
||||
|
||||
def test_create_queued_message(self):
|
||||
"""Should create message with content."""
|
||||
from praisonai.cli.features.message_queue import QueuedMessage
|
||||
msg = QueuedMessage(content="Hello")
|
||||
assert msg.content == "Hello"
|
||||
assert msg.timestamp is not None
|
||||
|
||||
def test_message_has_timestamp(self):
|
||||
"""Should have creation timestamp."""
|
||||
from praisonai.cli.features.message_queue import QueuedMessage
|
||||
before = time.time()
|
||||
msg = QueuedMessage(content="Hello")
|
||||
after = time.time()
|
||||
assert before <= msg.timestamp <= after
|
||||
|
||||
def test_message_priority(self):
|
||||
"""Should support priority levels."""
|
||||
from praisonai.cli.features.message_queue import QueuedMessage, MessagePriority
|
||||
msg = QueuedMessage(content="Hello", priority=MessagePriority.HIGH)
|
||||
assert msg.priority == MessagePriority.HIGH
|
||||
|
||||
def test_default_priority_normal(self):
|
||||
"""Should default to NORMAL priority."""
|
||||
from praisonai.cli.features.message_queue import QueuedMessage, MessagePriority
|
||||
msg = QueuedMessage(content="Hello")
|
||||
assert msg.priority == MessagePriority.NORMAL
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 4. MessageQueueHandler Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestMessageQueueHandler:
|
||||
"""Test MessageQueueHandler main class."""
|
||||
|
||||
def test_handler_creation(self):
|
||||
"""Should create handler with default settings."""
|
||||
from praisonai.cli.features.message_queue import MessageQueueHandler
|
||||
handler = MessageQueueHandler()
|
||||
assert handler.queue.is_empty
|
||||
assert handler.state_manager.is_idle
|
||||
|
||||
def test_submit_when_idle(self):
|
||||
"""Should process immediately when idle."""
|
||||
from praisonai.cli.features.message_queue import MessageQueueHandler
|
||||
processed = []
|
||||
|
||||
def processor(msg):
|
||||
processed.append(msg)
|
||||
return f"Response to: {msg}"
|
||||
|
||||
handler = MessageQueueHandler(processor=processor)
|
||||
handler.submit("Hello")
|
||||
|
||||
# Should process immediately
|
||||
assert "Hello" in processed
|
||||
|
||||
def test_queue_when_processing(self):
|
||||
"""Should queue message when already processing."""
|
||||
from praisonai.cli.features.message_queue import MessageQueueHandler, ProcessingState
|
||||
|
||||
handler = MessageQueueHandler()
|
||||
handler.state_manager.set_state(ProcessingState.PROCESSING)
|
||||
handler.submit("Hello")
|
||||
|
||||
# Should be queued, not processed
|
||||
assert handler.queue.count == 1
|
||||
|
||||
def test_auto_process_queue_on_idle(self):
|
||||
"""Should auto-process queue when becoming idle."""
|
||||
from praisonai.cli.features.message_queue import MessageQueueHandler, ProcessingState
|
||||
processed = []
|
||||
|
||||
def processor(msg):
|
||||
processed.append(msg)
|
||||
return f"Response: {msg}"
|
||||
|
||||
handler = MessageQueueHandler(processor=processor)
|
||||
|
||||
# Simulate processing state
|
||||
handler.state_manager.set_state(ProcessingState.PROCESSING)
|
||||
handler.submit("Queued1")
|
||||
handler.submit("Queued2")
|
||||
|
||||
assert handler.queue.count == 2
|
||||
|
||||
# Simulate becoming idle - should trigger queue processing
|
||||
handler.on_processing_complete()
|
||||
|
||||
# Queue should be processed
|
||||
assert "Queued1" in processed or handler.queue.count < 2
|
||||
|
||||
def test_get_queue_status(self):
|
||||
"""Should return queue status."""
|
||||
from praisonai.cli.features.message_queue import MessageQueueHandler, ProcessingState
|
||||
handler = MessageQueueHandler()
|
||||
handler.state_manager.set_state(ProcessingState.PROCESSING)
|
||||
handler.submit("Hello")
|
||||
handler.submit("World")
|
||||
|
||||
status = handler.get_status()
|
||||
assert status['queue_count'] == 2
|
||||
assert status['state'] == 'processing'
|
||||
|
||||
def test_clear_queue(self):
|
||||
"""Should clear queue via handler."""
|
||||
from praisonai.cli.features.message_queue import MessageQueueHandler, ProcessingState
|
||||
handler = MessageQueueHandler()
|
||||
handler.state_manager.set_state(ProcessingState.PROCESSING)
|
||||
handler.submit("Hello")
|
||||
handler.clear_queue()
|
||||
assert handler.queue.is_empty
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 5. QueueDisplay Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestQueueDisplay:
|
||||
"""Test QueueDisplay for visual indicators."""
|
||||
|
||||
def test_format_empty_queue(self):
|
||||
"""Should return empty string for empty queue."""
|
||||
from praisonai.cli.features.message_queue import QueueDisplay, MessageQueue
|
||||
queue = MessageQueue()
|
||||
display = QueueDisplay(queue)
|
||||
assert display.format_queue() == ""
|
||||
|
||||
def test_format_single_message(self):
|
||||
"""Should format single message."""
|
||||
from praisonai.cli.features.message_queue import QueueDisplay, MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("Hello world")
|
||||
display = QueueDisplay(queue)
|
||||
formatted = display.format_queue()
|
||||
assert "↳" in formatted
|
||||
assert "Hello world" in formatted
|
||||
|
||||
def test_format_multiple_messages(self):
|
||||
"""Should format multiple messages."""
|
||||
from praisonai.cli.features.message_queue import QueueDisplay, MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
display = QueueDisplay(queue)
|
||||
formatted = display.format_queue()
|
||||
assert "First" in formatted
|
||||
assert "Second" in formatted
|
||||
|
||||
def test_truncate_long_message(self):
|
||||
"""Should truncate long messages."""
|
||||
from praisonai.cli.features.message_queue import QueueDisplay, MessageQueue
|
||||
queue = MessageQueue()
|
||||
queue.add("A" * 100)
|
||||
display = QueueDisplay(queue, max_message_length=20)
|
||||
formatted = display.format_queue()
|
||||
assert "..." in formatted
|
||||
|
||||
def test_format_status_idle(self):
|
||||
"""Should show idle status."""
|
||||
from praisonai.cli.features.message_queue import QueueDisplay, MessageQueue, StateManager
|
||||
queue = MessageQueue()
|
||||
state = StateManager()
|
||||
display = QueueDisplay(queue, state_manager=state)
|
||||
status = display.format_status()
|
||||
assert "idle" in status.lower() or status == ""
|
||||
|
||||
def test_format_status_processing(self):
|
||||
"""Should show processing status."""
|
||||
from praisonai.cli.features.message_queue import (
|
||||
QueueDisplay, MessageQueue, StateManager, ProcessingState
|
||||
)
|
||||
queue = MessageQueue()
|
||||
state = StateManager()
|
||||
state.set_state(ProcessingState.PROCESSING)
|
||||
display = QueueDisplay(queue, state_manager=state)
|
||||
status = display.format_status()
|
||||
assert "processing" in status.lower() or "⏳" in status
|
||||
|
||||
def test_format_queue_count(self):
|
||||
"""Should show queue count."""
|
||||
from praisonai.cli.features.message_queue import QueueDisplay, MessageQueue, ProcessingState, StateManager
|
||||
queue = MessageQueue()
|
||||
queue.add("First")
|
||||
queue.add("Second")
|
||||
state = StateManager()
|
||||
state.set_state(ProcessingState.PROCESSING)
|
||||
display = QueueDisplay(queue, state_manager=state)
|
||||
count_str = display.format_queue_count()
|
||||
assert "2" in count_str
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 6. Integration Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestMessageQueueIntegration:
|
||||
"""Integration tests for message queue system."""
|
||||
|
||||
def test_full_workflow(self):
|
||||
"""Test complete queue workflow."""
|
||||
from praisonai.cli.features.message_queue import (
|
||||
MessageQueueHandler, ProcessingState
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
def processor(msg):
|
||||
results.append(msg)
|
||||
return f"Processed: {msg}"
|
||||
|
||||
handler = MessageQueueHandler(processor=processor)
|
||||
|
||||
# 1. Submit when idle - processes immediately
|
||||
handler.submit("First")
|
||||
assert "First" in results
|
||||
|
||||
# 2. Simulate processing
|
||||
handler.state_manager.set_state(ProcessingState.PROCESSING)
|
||||
|
||||
# 3. Queue messages while processing
|
||||
handler.submit("Second")
|
||||
handler.submit("Third")
|
||||
assert handler.queue.count == 2
|
||||
|
||||
# 4. Complete processing - should auto-process queue
|
||||
handler.on_processing_complete()
|
||||
|
||||
# Queue should be processed
|
||||
assert handler.queue.count < 2 or "Second" in results
|
||||
|
||||
def test_queue_with_display(self):
|
||||
"""Test queue with display formatting."""
|
||||
from praisonai.cli.features.message_queue import (
|
||||
MessageQueueHandler, QueueDisplay, ProcessingState
|
||||
)
|
||||
|
||||
handler = MessageQueueHandler()
|
||||
handler.state_manager.set_state(ProcessingState.PROCESSING)
|
||||
handler.submit("Task 1")
|
||||
handler.submit("Task 2")
|
||||
|
||||
display = QueueDisplay(handler.queue, state_manager=handler.state_manager)
|
||||
|
||||
# Should show processing status
|
||||
status = display.format_status()
|
||||
assert "processing" in status.lower() or "⏳" in status
|
||||
|
||||
# Should show queued messages
|
||||
queue_str = display.format_queue()
|
||||
assert "Task 1" in queue_str
|
||||
assert "Task 2" in queue_str
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 9. AsyncProcessor Tests - Background processing
|
||||
# ============================================================================
|
||||
|
||||
class TestAsyncProcessor:
|
||||
"""Test AsyncProcessor for background agent execution."""
|
||||
|
||||
def test_processor_creation(self):
|
||||
"""Should create processor with callback."""
|
||||
from praisonai.cli.features.message_queue import AsyncProcessor
|
||||
|
||||
results = []
|
||||
def on_complete(result):
|
||||
results.append(result)
|
||||
|
||||
processor = AsyncProcessor(on_complete=on_complete)
|
||||
assert processor is not None
|
||||
|
||||
def test_start_processing(self):
|
||||
"""Should start processing in background thread."""
|
||||
from praisonai.cli.features.message_queue import AsyncProcessor
|
||||
import time
|
||||
|
||||
results = []
|
||||
def work_fn():
|
||||
time.sleep(0.05)
|
||||
return "done"
|
||||
|
||||
def on_complete(result):
|
||||
results.append(result)
|
||||
|
||||
processor = AsyncProcessor(on_complete=on_complete)
|
||||
processor.start(work_fn)
|
||||
|
||||
# Should be processing
|
||||
assert processor.is_running
|
||||
|
||||
# Wait for completion
|
||||
time.sleep(0.15)
|
||||
assert not processor.is_running
|
||||
assert "done" in results
|
||||
|
||||
def test_is_running_property(self):
|
||||
"""Should track running state."""
|
||||
from praisonai.cli.features.message_queue import AsyncProcessor
|
||||
import time
|
||||
|
||||
def slow_work():
|
||||
time.sleep(0.1)
|
||||
return "result"
|
||||
|
||||
processor = AsyncProcessor()
|
||||
assert not processor.is_running
|
||||
|
||||
processor.start(slow_work)
|
||||
assert processor.is_running
|
||||
|
||||
time.sleep(0.15)
|
||||
assert not processor.is_running
|
||||
|
||||
def test_on_status_callback(self):
|
||||
"""Should call status callback during processing."""
|
||||
from praisonai.cli.features.message_queue import AsyncProcessor
|
||||
import time
|
||||
|
||||
statuses = []
|
||||
def on_status(status):
|
||||
statuses.append(status)
|
||||
|
||||
def work_with_status(status_callback):
|
||||
status_callback("Starting...")
|
||||
time.sleep(0.02)
|
||||
status_callback("Working...")
|
||||
return "done"
|
||||
|
||||
processor = AsyncProcessor(on_status=on_status)
|
||||
processor.start(lambda: work_with_status(on_status))
|
||||
|
||||
time.sleep(0.1)
|
||||
assert len(statuses) >= 1
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 10. LiveStatusDisplay Tests - Real-time status
|
||||
# ============================================================================
|
||||
|
||||
class TestLiveStatusDisplay:
|
||||
"""Test LiveStatusDisplay for real-time tool/command status."""
|
||||
|
||||
def test_display_creation(self):
|
||||
"""Should create display."""
|
||||
from praisonai.cli.features.message_queue import LiveStatusDisplay
|
||||
|
||||
display = LiveStatusDisplay()
|
||||
assert display is not None
|
||||
|
||||
def test_update_status(self):
|
||||
"""Should update current status."""
|
||||
from praisonai.cli.features.message_queue import LiveStatusDisplay
|
||||
|
||||
display = LiveStatusDisplay()
|
||||
display.update_status("Processing...")
|
||||
assert display.current_status == "Processing..."
|
||||
|
||||
def test_add_tool_call(self):
|
||||
"""Should track tool calls."""
|
||||
from praisonai.cli.features.message_queue import LiveStatusDisplay
|
||||
|
||||
display = LiveStatusDisplay()
|
||||
display.add_tool_call("read_file", {"path": "test.py"})
|
||||
|
||||
assert len(display.tool_calls) == 1
|
||||
assert display.tool_calls[0]['name'] == "read_file"
|
||||
|
||||
def test_add_command_execution(self):
|
||||
"""Should track command executions."""
|
||||
from praisonai.cli.features.message_queue import LiveStatusDisplay
|
||||
|
||||
display = LiveStatusDisplay()
|
||||
display.add_command("ls -la")
|
||||
|
||||
assert len(display.commands) == 1
|
||||
assert display.commands[0] == "ls -la"
|
||||
|
||||
def test_format_live_status(self):
|
||||
"""Should format live status for display."""
|
||||
from praisonai.cli.features.message_queue import LiveStatusDisplay
|
||||
|
||||
display = LiveStatusDisplay()
|
||||
display.update_status("Thinking...")
|
||||
display.add_tool_call("read_file", {"path": "main.py"})
|
||||
|
||||
formatted = display.format()
|
||||
assert "Thinking" in formatted or "read_file" in formatted
|
||||
|
||||
def test_clear_status(self):
|
||||
"""Should clear all status."""
|
||||
from praisonai.cli.features.message_queue import LiveStatusDisplay
|
||||
|
||||
display = LiveStatusDisplay()
|
||||
display.update_status("Working")
|
||||
display.add_tool_call("test", {})
|
||||
display.add_command("echo hi")
|
||||
|
||||
display.clear()
|
||||
|
||||
assert display.current_status == ""
|
||||
assert len(display.tool_calls) == 0
|
||||
assert len(display.commands) == 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 11. NonBlockingInput Tests - Async input handling
|
||||
# ============================================================================
|
||||
|
||||
class TestNonBlockingInput:
|
||||
"""Test NonBlockingInput for async user input."""
|
||||
|
||||
def test_input_handler_creation(self):
|
||||
"""Should create input handler."""
|
||||
from praisonai.cli.features.message_queue import NonBlockingInput
|
||||
|
||||
handler = NonBlockingInput()
|
||||
assert handler is not None
|
||||
|
||||
def test_submit_input(self):
|
||||
"""Should accept input while processing."""
|
||||
from praisonai.cli.features.message_queue import NonBlockingInput
|
||||
|
||||
handler = NonBlockingInput()
|
||||
handler.submit("new message")
|
||||
|
||||
assert handler.has_pending
|
||||
assert handler.pop() == "new message"
|
||||
|
||||
def test_multiple_inputs(self):
|
||||
"""Should queue multiple inputs."""
|
||||
from praisonai.cli.features.message_queue import NonBlockingInput
|
||||
|
||||
handler = NonBlockingInput()
|
||||
handler.submit("first")
|
||||
handler.submit("second")
|
||||
|
||||
assert handler.pending_count == 2
|
||||
assert handler.pop() == "first"
|
||||
assert handler.pop() == "second"
|
||||
@@ -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
|
||||
Loaded 100 of 116 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user