mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-10 03:31:28 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
073062cd96 | ||
|
|
936b422876 | ||
|
|
99ab421ce2 | ||
|
|
c1f334905e | ||
|
|
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 | ||
|
|
df687d86c3 | ||
|
|
afa37b0782 | ||
|
|
43a15e6049 |
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
|
||||
@@ -1628,36 +1755,78 @@ server.run() # stdio for Claude Desktop
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
## A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
## CLI Features
|
||||
|
||||
| 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) |
|
||||
|
||||
---
|
||||
|
||||
@@ -2209,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
|
||||
|
||||
@@ -2218,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.35" \
|
||||
"praisonai>=2.3.56" \
|
||||
"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.35" \
|
||||
"praisonai>=2.3.56" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
# Dockerfile for PraisonAI Serve API
|
||||
# Deploy to Railway, Render, Fly.io, or any container platform
|
||||
#
|
||||
# Build: docker build -f Dockerfile.serve -t praisonai-serve .
|
||||
# Run: docker run -p 8005:8005 -e OPENAI_API_KEY=sk-xxx praisonai-serve
|
||||
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install dependencies
|
||||
RUN pip install --no-cache-dir praisonai praisonaiagents
|
||||
|
||||
# Copy agents configuration
|
||||
COPY agents.yaml .
|
||||
|
||||
# Expose the API port
|
||||
EXPOSE 8005
|
||||
|
||||
# Health check
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
|
||||
CMD curl -f http://localhost:8005/health || exit 1
|
||||
|
||||
# Start the API server
|
||||
CMD ["praisonai", "serve", "agents.yaml", "--host", "0.0.0.0", "--port", "8005"]
|
||||
@@ -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.35" \
|
||||
"praisonai>=2.3.56" \
|
||||
"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,53 @@
|
||||
"""
|
||||
PraisonAI A2A Server Example
|
||||
|
||||
Expose a PraisonAI Agent as an A2A (Agent2Agent) Server.
|
||||
This enables other AI agents to discover and communicate with your agent.
|
||||
|
||||
Run:
|
||||
uvicorn a2a-server:app --reload
|
||||
|
||||
Endpoints:
|
||||
GET /.well-known/agent.json - Agent Card for discovery
|
||||
GET /status - Server status
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Search results for: {query}"
|
||||
|
||||
def calculate(expression: str) -> str:
|
||||
"""Calculate a mathematical expression."""
|
||||
try:
|
||||
return f"Result: {eval(expression)}"
|
||||
except Exception:
|
||||
return "Invalid expression"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics and answer questions",
|
||||
tools=[search_web, calculate]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(
|
||||
agent=agent,
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0"
|
||||
)
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI(
|
||||
title="PraisonAI A2A Server",
|
||||
description="A2A-compatible agent server"
|
||||
)
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
@@ -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"
|
||||
@@ -9,6 +9,8 @@ Usage:
|
||||
export OPENAI_API_KEY=your_key
|
||||
python benchmarks/execution_benchmark.py
|
||||
python benchmarks/execution_benchmark.py --no-save
|
||||
python benchmarks/execution_benchmark.py --iterations 5
|
||||
python benchmarks/execution_benchmark.py --model gpt-4o --iterations 10
|
||||
"""
|
||||
|
||||
import time
|
||||
@@ -16,14 +18,10 @@ import os
|
||||
import argparse
|
||||
from importlib.metadata import version as get_version
|
||||
|
||||
# Check API key
|
||||
if not os.environ.get('OPENAI_API_KEY'):
|
||||
print("ERROR: Set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
MODEL = "gpt-4o-mini"
|
||||
ITERATIONS = 3
|
||||
PROMPT = "What is 2+2? Reply with just the number."
|
||||
# Defaults (can be overridden via CLI)
|
||||
DEFAULT_MODEL = "gpt-4o-mini"
|
||||
DEFAULT_ITERATIONS = 3
|
||||
DEFAULT_PROMPT = "What is 2+2? Reply with just the number."
|
||||
|
||||
|
||||
def get_package_versions():
|
||||
@@ -42,24 +40,24 @@ def get_package_versions():
|
||||
return versions
|
||||
|
||||
|
||||
def benchmark_praisonai():
|
||||
def benchmark_praisonai(model, iterations, prompt):
|
||||
"""Benchmark PraisonAI agent execution using agent.start()"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("\n--- PraisonAI ---")
|
||||
print(f"Method: agent.start()")
|
||||
print("Method: agent.start()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
llm=MODEL,
|
||||
llm=model,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.start(PROMPT)
|
||||
response = agent.start(prompt)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
@@ -70,24 +68,24 @@ def benchmark_praisonai():
|
||||
return avg
|
||||
|
||||
|
||||
def benchmark_praisonai_litellm():
|
||||
def benchmark_praisonai_litellm(model, iterations, prompt):
|
||||
"""Benchmark PraisonAI with LiteLLM backend using agent.start()"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("\n--- PraisonAI (LiteLLM) ---")
|
||||
print(f"Method: agent.start()")
|
||||
print("Method: agent.start()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
llm=f"openai/{MODEL}",
|
||||
llm=f"openai/{model}",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.start(PROMPT)
|
||||
response = agent.start(prompt)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
@@ -98,23 +96,23 @@ def benchmark_praisonai_litellm():
|
||||
return avg
|
||||
|
||||
|
||||
def benchmark_agno():
|
||||
def benchmark_agno(model, iterations, prompt):
|
||||
"""Benchmark Agno agent execution using agent.run()"""
|
||||
from agno.agent import Agent
|
||||
from agno.models.openai import OpenAIChat
|
||||
|
||||
print("\n--- Agno ---")
|
||||
print(f"Method: agent.run()")
|
||||
print("Method: agent.run()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
model=OpenAIChat(id=MODEL),
|
||||
model=OpenAIChat(id=model),
|
||||
instructions=["You are a helpful assistant. Be very brief."],
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.run(PROMPT)
|
||||
response = agent.run(prompt)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
@@ -127,25 +125,25 @@ def benchmark_agno():
|
||||
return avg
|
||||
|
||||
|
||||
def benchmark_crewai():
|
||||
def benchmark_crewai(model, iterations, prompt):
|
||||
"""Benchmark CrewAI agent execution using crew.kickoff()"""
|
||||
from crewai import Agent, Task, Crew
|
||||
|
||||
print("\n--- CrewAI ---")
|
||||
print(f"Method: crew.kickoff()")
|
||||
print("Method: crew.kickoff()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
role="Calculator",
|
||||
goal="Answer math questions",
|
||||
backstory="You are a helpful assistant.",
|
||||
llm=MODEL,
|
||||
llm=model,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description=PROMPT,
|
||||
description=prompt,
|
||||
expected_output="A number",
|
||||
agent=agent
|
||||
)
|
||||
@@ -168,7 +166,7 @@ def benchmark_crewai():
|
||||
return avg
|
||||
|
||||
|
||||
def save_results(results: dict):
|
||||
def save_results(results: dict, model: str, iterations: int, prompt: str):
|
||||
"""Save results to markdown file."""
|
||||
import os
|
||||
from datetime import datetime
|
||||
@@ -183,9 +181,9 @@ def save_results(results: dict):
|
||||
with open(filepath, 'w') as f:
|
||||
f.write('# PraisonAI Agents - Real Execution Benchmark\n\n')
|
||||
f.write(f'**Generated:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n')
|
||||
f.write(f'**Model:** {MODEL}\n')
|
||||
f.write(f'**Iterations:** {ITERATIONS}\n')
|
||||
f.write(f'**Prompt:** "{PROMPT}"\n\n')
|
||||
f.write(f'**Model:** {model}\n')
|
||||
f.write(f'**Iterations:** {iterations}\n')
|
||||
f.write(f'**Prompt:** "{prompt}"\n\n')
|
||||
f.write('## Results\n\n')
|
||||
f.write('| Framework | Method | Avg Time | Relative |\n')
|
||||
f.write('|-----------|--------|----------|----------|\n')
|
||||
@@ -224,14 +222,26 @@ def save_results(results: dict):
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='PraisonAI Agents - Real Execution Benchmark')
|
||||
parser.add_argument('--no-save', action='store_true', help='Do not save results to file')
|
||||
parser.add_argument('--model', '-m', type=str, default=DEFAULT_MODEL, help=f'Model to use (default: {DEFAULT_MODEL})')
|
||||
parser.add_argument('--iterations', '-i', type=int, default=DEFAULT_ITERATIONS, help=f'Number of iterations (default: {DEFAULT_ITERATIONS})')
|
||||
parser.add_argument('--prompt', '-p', type=str, default=DEFAULT_PROMPT, help=f'Prompt to use (default: "{DEFAULT_PROMPT}")')
|
||||
args = parser.parse_args()
|
||||
|
||||
# Check API key
|
||||
if not os.environ.get('OPENAI_API_KEY'):
|
||||
print("ERROR: Set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
model = args.model
|
||||
iterations = args.iterations
|
||||
prompt = args.prompt
|
||||
|
||||
print('=' * 60)
|
||||
print('PraisonAI Agents - Real Execution Benchmark')
|
||||
print('=' * 60)
|
||||
print(f'\nModel: {MODEL}')
|
||||
print(f'Iterations: {ITERATIONS}')
|
||||
print(f'Prompt: "{PROMPT}"')
|
||||
print(f'\nModel: {model}')
|
||||
print(f'Iterations: {iterations}')
|
||||
print(f'Prompt: "{prompt}"')
|
||||
|
||||
versions = get_package_versions()
|
||||
print('\nPackage versions:')
|
||||
@@ -242,22 +252,22 @@ if __name__ == '__main__':
|
||||
|
||||
# Run benchmarks
|
||||
try:
|
||||
results['PraisonAI'] = benchmark_praisonai()
|
||||
results['PraisonAI'] = benchmark_praisonai(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"PraisonAI error: {e}")
|
||||
|
||||
try:
|
||||
results['PraisonAI (LiteLLM)'] = benchmark_praisonai_litellm()
|
||||
results['PraisonAI (LiteLLM)'] = benchmark_praisonai_litellm(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"PraisonAI (LiteLLM) error: {e}")
|
||||
|
||||
try:
|
||||
results['Agno'] = benchmark_agno()
|
||||
results['Agno'] = benchmark_agno(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"Agno error: {e}")
|
||||
|
||||
try:
|
||||
results['CrewAI'] = benchmark_crewai()
|
||||
results['CrewAI'] = benchmark_crewai(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"CrewAI error: {e}")
|
||||
|
||||
@@ -281,6 +291,6 @@ if __name__ == '__main__':
|
||||
|
||||
# Save results (unless --no-save)
|
||||
if results and not args.no_save:
|
||||
save_results(results)
|
||||
save_results(results, model, iterations, prompt)
|
||||
elif args.no_save:
|
||||
print('\nResults not saved (--no-save flag used)')
|
||||
@@ -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()
|
||||
@@ -305,3 +305,21 @@ __all__.extend([
|
||||
'LineRange'
|
||||
])
|
||||
|
||||
# AG-UI support (optional, lazy loaded)
|
||||
try:
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
__all__.append('AGUI')
|
||||
_agui_available = True
|
||||
except ImportError:
|
||||
_agui_available = False
|
||||
AGUI = None
|
||||
|
||||
# A2A support (optional, lazy loaded)
|
||||
try:
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
__all__.append('A2A')
|
||||
_a2a_available = True
|
||||
except ImportError:
|
||||
_a2a_available = False
|
||||
A2A = None
|
||||
|
||||
@@ -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,13 @@
|
||||
"""
|
||||
UI Integrations for PraisonAI Agents
|
||||
|
||||
This module provides UI protocol integrations for exposing PraisonAI Agents
|
||||
via various frontend protocols.
|
||||
|
||||
Available integrations:
|
||||
- agui: AG-UI (Agent-User Interface) protocol for CopilotKit and compatible frontends
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
|
||||
__all__ = ["AGUI"]
|
||||
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
A2A Protocol Integration for PraisonAI Agents
|
||||
|
||||
This module provides A2A (Agent2Agent) protocol support,
|
||||
enabling PraisonAI Agents to be exposed as A2A Servers
|
||||
for agent-to-agent communication.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Agent Card available at: GET /.well-known/agent.json
|
||||
# A2A endpoint at: POST /a2a (JSON-RPC)
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.a2a.a2a import A2A
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
# Enums
|
||||
TaskState,
|
||||
Role,
|
||||
# Part types
|
||||
TextPart,
|
||||
FilePart,
|
||||
DataPart,
|
||||
# Core types
|
||||
Message,
|
||||
TaskStatus,
|
||||
Task,
|
||||
Artifact,
|
||||
# Agent Card types
|
||||
AgentSkill,
|
||||
AgentCapabilities,
|
||||
AgentCard,
|
||||
# Request/Response types
|
||||
MessageConfig,
|
||||
SendMessageRequest,
|
||||
# Streaming events
|
||||
TaskStatusUpdateEvent,
|
||||
TaskArtifactUpdateEvent,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Main class
|
||||
"A2A",
|
||||
# Enums
|
||||
"TaskState",
|
||||
"Role",
|
||||
# Part types
|
||||
"TextPart",
|
||||
"FilePart",
|
||||
"DataPart",
|
||||
# Core types
|
||||
"Message",
|
||||
"TaskStatus",
|
||||
"Task",
|
||||
"Artifact",
|
||||
# Agent Card types
|
||||
"AgentSkill",
|
||||
"AgentCapabilities",
|
||||
"AgentCard",
|
||||
# Request/Response types
|
||||
"MessageConfig",
|
||||
"SendMessageRequest",
|
||||
# Streaming events
|
||||
"TaskStatusUpdateEvent",
|
||||
"TaskArtifactUpdateEvent",
|
||||
]
|
||||
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
A2A - Main A2A Interface Class
|
||||
|
||||
Exposes PraisonAI Agents via the A2A (Agent2Agent) protocol.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional, TYPE_CHECKING
|
||||
|
||||
from praisonaiagents.ui.a2a.types import AgentCard
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent, PraisonAIAgents
|
||||
from fastapi import APIRouter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class A2A:
|
||||
"""
|
||||
A2A Interface for PraisonAI Agents.
|
||||
|
||||
Exposes a PraisonAI Agent or PraisonAIAgents workflow via the A2A protocol,
|
||||
enabling agent-to-agent communication with other A2A-compatible systems.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Agent Card at: GET /.well-known/agent.json
|
||||
# A2A endpoint at: POST /a2a
|
||||
|
||||
Args:
|
||||
agent: Single PraisonAI Agent instance
|
||||
agents: PraisonAIAgents instance for multi-agent workflows
|
||||
name: Name for the A2A endpoint (defaults to agent name)
|
||||
description: Description of the agent
|
||||
url: URL where the A2A endpoint is hosted
|
||||
version: Version string for the agent
|
||||
prefix: URL prefix for the router
|
||||
tags: OpenAPI tags for the router
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: Optional["Agent"] = None,
|
||||
agents: Optional["PraisonAIAgents"] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
url: str = "http://localhost:8000/a2a",
|
||||
version: str = "1.0.0",
|
||||
prefix: str = "",
|
||||
tags: Optional[List[str]] = None,
|
||||
):
|
||||
if agent is None and agents is None:
|
||||
raise ValueError("A2A requires an agent or agents instance")
|
||||
|
||||
self.agent = agent
|
||||
self.agents = agents
|
||||
self.url = url
|
||||
self.version = version
|
||||
self.prefix = prefix
|
||||
self.tags = tags or ["A2A"]
|
||||
|
||||
# Set name from agent if not provided
|
||||
if name:
|
||||
self.name = name
|
||||
elif agent and hasattr(agent, 'name'):
|
||||
self.name = agent.name
|
||||
elif agents and hasattr(agents, 'name'):
|
||||
self.name = agents.name
|
||||
else:
|
||||
self.name = "PraisonAI Agent"
|
||||
|
||||
# Set description
|
||||
if description:
|
||||
self.description = description
|
||||
elif agent and hasattr(agent, 'role'):
|
||||
self.description = agent.role
|
||||
else:
|
||||
self.description = "PraisonAI Agent via A2A"
|
||||
|
||||
# Initialize task store
|
||||
self.task_store = TaskStore()
|
||||
|
||||
# Router cache
|
||||
self._router: Optional["APIRouter"] = None
|
||||
self._agent_card: Optional[AgentCard] = None
|
||||
|
||||
def get_agent_card(self) -> AgentCard:
|
||||
"""
|
||||
Get the Agent Card for this A2A instance.
|
||||
|
||||
Returns:
|
||||
AgentCard object for A2A discovery
|
||||
"""
|
||||
if self._agent_card is None:
|
||||
self._agent_card = generate_agent_card(
|
||||
agent=self.agent,
|
||||
url=self.url,
|
||||
version=self.version,
|
||||
streaming=True,
|
||||
)
|
||||
return self._agent_card
|
||||
|
||||
def get_router(self) -> "APIRouter":
|
||||
"""
|
||||
Get the FastAPI router for this A2A instance.
|
||||
|
||||
Returns:
|
||||
FastAPI APIRouter with A2A endpoints
|
||||
"""
|
||||
if self._router is None:
|
||||
self._router = self._create_router()
|
||||
return self._router
|
||||
|
||||
def _create_router(self) -> "APIRouter":
|
||||
"""Create FastAPI router with A2A endpoints."""
|
||||
from fastapi import APIRouter
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
router = APIRouter(prefix=self.prefix, tags=self.tags)
|
||||
|
||||
# Agent Card endpoint (well-known URI)
|
||||
@router.get("/.well-known/agent.json")
|
||||
async def get_agent_card():
|
||||
"""Return the Agent Card for discovery."""
|
||||
card = self.get_agent_card()
|
||||
return JSONResponse(content=card.model_dump(by_alias=True, exclude_none=True))
|
||||
|
||||
# Status endpoint
|
||||
@router.get("/status")
|
||||
async def get_status():
|
||||
"""Return server status."""
|
||||
return {
|
||||
"status": "ok",
|
||||
"name": self.name,
|
||||
"version": self.version,
|
||||
}
|
||||
|
||||
# TODO: Add JSON-RPC endpoint for message/send
|
||||
# TODO: Add streaming endpoint for message/stream
|
||||
# TODO: Add task management endpoints
|
||||
|
||||
return router
|
||||
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
A2A Agent Card Generation
|
||||
|
||||
Generate Agent Cards from PraisonAI Agents for A2A discovery.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
AgentCard,
|
||||
AgentCapabilities,
|
||||
AgentSkill,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
def extract_skills_from_tools(tools: Optional[List]) -> List[AgentSkill]:
|
||||
"""
|
||||
Extract AgentSkill objects from a list of tools.
|
||||
|
||||
Args:
|
||||
tools: List of tool functions or objects
|
||||
|
||||
Returns:
|
||||
List of AgentSkill objects
|
||||
"""
|
||||
if not tools:
|
||||
return []
|
||||
|
||||
skills = []
|
||||
|
||||
for tool in tools:
|
||||
# Skip non-callable items
|
||||
if not callable(tool):
|
||||
continue
|
||||
|
||||
# Get function name
|
||||
name = getattr(tool, '__name__', None)
|
||||
if not name:
|
||||
continue
|
||||
|
||||
# Get description from docstring
|
||||
docstring = inspect.getdoc(tool)
|
||||
description = docstring.split('\n')[0] if docstring else f"Tool: {name}"
|
||||
|
||||
# Create skill
|
||||
skill = AgentSkill(
|
||||
id=name,
|
||||
name=name,
|
||||
description=description,
|
||||
tags=["tool"],
|
||||
)
|
||||
skills.append(skill)
|
||||
|
||||
return skills
|
||||
|
||||
|
||||
def generate_agent_card(
|
||||
agent: "Agent",
|
||||
url: str,
|
||||
version: str = "1.0.0",
|
||||
streaming: bool = False,
|
||||
push_notifications: bool = False,
|
||||
) -> AgentCard:
|
||||
"""
|
||||
Generate an A2A Agent Card from a PraisonAI Agent.
|
||||
|
||||
Args:
|
||||
agent: PraisonAI Agent instance
|
||||
url: URL where the A2A endpoint is hosted
|
||||
version: Version string for the agent
|
||||
streaming: Whether streaming is supported
|
||||
push_notifications: Whether push notifications are supported
|
||||
|
||||
Returns:
|
||||
AgentCard object for A2A discovery
|
||||
"""
|
||||
# Get agent name
|
||||
name = getattr(agent, 'name', None) or "PraisonAI Agent"
|
||||
|
||||
# Build description from role, goal, or instructions
|
||||
description_parts = []
|
||||
|
||||
role = getattr(agent, 'role', None)
|
||||
if role:
|
||||
description_parts.append(role)
|
||||
|
||||
goal = getattr(agent, 'goal', None)
|
||||
if goal:
|
||||
description_parts.append(goal)
|
||||
|
||||
instructions = getattr(agent, 'instructions', None)
|
||||
if instructions and not description_parts:
|
||||
# Use first line of instructions if no role/goal
|
||||
description_parts.append(instructions.split('\n')[0][:200])
|
||||
|
||||
description = ". ".join(description_parts) if description_parts else "PraisonAI Agent"
|
||||
|
||||
# Extract skills from tools
|
||||
tools = getattr(agent, 'tools', None)
|
||||
skills = extract_skills_from_tools(tools)
|
||||
|
||||
# Build capabilities
|
||||
capabilities = AgentCapabilities(
|
||||
streaming=streaming,
|
||||
push_notifications=push_notifications,
|
||||
)
|
||||
|
||||
# Create Agent Card
|
||||
card = AgentCard(
|
||||
name=name,
|
||||
url=url,
|
||||
version=version,
|
||||
description=description,
|
||||
capabilities=capabilities,
|
||||
skills=skills if skills else None,
|
||||
provider={"name": "PraisonAI"},
|
||||
)
|
||||
|
||||
return card
|
||||
@@ -0,0 +1,131 @@
|
||||
"""
|
||||
A2A Message Conversion
|
||||
|
||||
Convert between A2A protocol messages and PraisonAI internal format.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
Message,
|
||||
Role,
|
||||
TextPart,
|
||||
Artifact,
|
||||
)
|
||||
|
||||
|
||||
def a2a_to_praisonai_messages(messages: List[Message]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert A2A messages to PraisonAI/OpenAI chat format.
|
||||
|
||||
Args:
|
||||
messages: List of A2A Message objects
|
||||
|
||||
Returns:
|
||||
List of dicts in OpenAI chat format {"role": ..., "content": ...}
|
||||
"""
|
||||
result = []
|
||||
|
||||
for msg in messages:
|
||||
# Map A2A role to OpenAI role
|
||||
role = "user" if msg.role == Role.USER else "assistant"
|
||||
|
||||
# Combine all text parts
|
||||
content_parts = []
|
||||
for part in msg.parts:
|
||||
if isinstance(part, TextPart):
|
||||
content_parts.append(part.text)
|
||||
elif hasattr(part, 'text'):
|
||||
content_parts.append(part.text)
|
||||
|
||||
content = " ".join(content_parts)
|
||||
|
||||
result.append({
|
||||
"role": role,
|
||||
"content": content,
|
||||
})
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def praisonai_to_a2a_message(
|
||||
response: str,
|
||||
message_id: Optional[str] = None,
|
||||
context_id: Optional[str] = None,
|
||||
task_id: Optional[str] = None,
|
||||
) -> Message:
|
||||
"""
|
||||
Convert PraisonAI response to A2A Message.
|
||||
|
||||
Args:
|
||||
response: Text response from agent
|
||||
message_id: Optional message ID (auto-generated if not provided)
|
||||
context_id: Optional context ID
|
||||
task_id: Optional task ID
|
||||
|
||||
Returns:
|
||||
A2A Message object
|
||||
"""
|
||||
msg_id = message_id or f"msg-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
return Message(
|
||||
message_id=msg_id,
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text=response)],
|
||||
context_id=context_id,
|
||||
task_id=task_id,
|
||||
)
|
||||
|
||||
|
||||
def create_artifact(
|
||||
content: str,
|
||||
artifact_id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> Artifact:
|
||||
"""
|
||||
Create an A2A Artifact from response content.
|
||||
|
||||
Args:
|
||||
content: Text content for the artifact
|
||||
artifact_id: Optional artifact ID (auto-generated if not provided)
|
||||
name: Optional artifact name
|
||||
description: Optional artifact description
|
||||
|
||||
Returns:
|
||||
A2A Artifact object
|
||||
"""
|
||||
art_id = artifact_id or f"art-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
return Artifact(
|
||||
artifact_id=art_id,
|
||||
name=name,
|
||||
description=description,
|
||||
parts=[TextPart(text=content)],
|
||||
)
|
||||
|
||||
|
||||
def extract_user_input(messages: List[Message]) -> str:
|
||||
"""
|
||||
Extract the last user input from a list of messages.
|
||||
|
||||
Args:
|
||||
messages: List of A2A Message objects
|
||||
|
||||
Returns:
|
||||
Text content of the last user message, or empty string if none
|
||||
"""
|
||||
# Find the last user message
|
||||
for msg in reversed(messages):
|
||||
if msg.role == Role.USER:
|
||||
# Combine text parts
|
||||
content_parts = []
|
||||
for part in msg.parts:
|
||||
if isinstance(part, TextPart):
|
||||
content_parts.append(part.text)
|
||||
elif hasattr(part, 'text'):
|
||||
content_parts.append(part.text)
|
||||
return " ".join(content_parts)
|
||||
|
||||
return ""
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
A2A Streaming Events
|
||||
|
||||
SSE streaming support for A2A task updates.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import AsyncIterator, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
Task,
|
||||
TaskState,
|
||||
TaskStatus,
|
||||
TaskStatusUpdateEvent,
|
||||
TaskArtifactUpdateEvent,
|
||||
Artifact,
|
||||
TextPart,
|
||||
)
|
||||
|
||||
|
||||
class A2AEventEncoder:
|
||||
"""Encode A2A events for SSE streaming."""
|
||||
|
||||
@staticmethod
|
||||
def encode_event(event_type: str, data: dict) -> str:
|
||||
"""
|
||||
Encode an event for SSE streaming.
|
||||
|
||||
Args:
|
||||
event_type: Type of event (e.g., "task.status", "task.artifact")
|
||||
data: Event data as dict
|
||||
|
||||
Returns:
|
||||
SSE-formatted string
|
||||
"""
|
||||
json_data = json.dumps(data, default=str)
|
||||
return f"event: {event_type}\ndata: {json_data}\n\n"
|
||||
|
||||
@staticmethod
|
||||
def encode_task_status(event: TaskStatusUpdateEvent) -> str:
|
||||
"""Encode TaskStatusUpdateEvent for SSE."""
|
||||
data = event.model_dump(by_alias=True, exclude_none=True)
|
||||
return A2AEventEncoder.encode_event("task.status", data)
|
||||
|
||||
@staticmethod
|
||||
def encode_task_artifact(event: TaskArtifactUpdateEvent) -> str:
|
||||
"""Encode TaskArtifactUpdateEvent for SSE."""
|
||||
data = event.model_dump(by_alias=True, exclude_none=True)
|
||||
return A2AEventEncoder.encode_event("task.artifact", data)
|
||||
|
||||
@staticmethod
|
||||
def encode_task(task: Task) -> str:
|
||||
"""Encode Task object for SSE."""
|
||||
data = task.model_dump(by_alias=True, exclude_none=True)
|
||||
return A2AEventEncoder.encode_event("task", data)
|
||||
|
||||
@staticmethod
|
||||
def encode_done() -> str:
|
||||
"""Encode stream completion event."""
|
||||
return "event: done\ndata: {}\n\n"
|
||||
|
||||
|
||||
def create_status_event(
|
||||
task_id: str,
|
||||
state: TaskState,
|
||||
context_id: Optional[str] = None,
|
||||
final: bool = False,
|
||||
) -> TaskStatusUpdateEvent:
|
||||
"""
|
||||
Create a TaskStatusUpdateEvent.
|
||||
|
||||
Args:
|
||||
task_id: Task ID
|
||||
state: New task state
|
||||
context_id: Optional context ID
|
||||
final: Whether this is the final event
|
||||
|
||||
Returns:
|
||||
TaskStatusUpdateEvent
|
||||
"""
|
||||
return TaskStatusUpdateEvent(
|
||||
task_id=task_id,
|
||||
context_id=context_id,
|
||||
status=TaskStatus(state=state),
|
||||
final=final,
|
||||
)
|
||||
|
||||
|
||||
def create_artifact_event(
|
||||
task_id: str,
|
||||
content: str,
|
||||
artifact_id: Optional[str] = None,
|
||||
context_id: Optional[str] = None,
|
||||
append: bool = False,
|
||||
last_chunk: bool = False,
|
||||
) -> TaskArtifactUpdateEvent:
|
||||
"""
|
||||
Create a TaskArtifactUpdateEvent.
|
||||
|
||||
Args:
|
||||
task_id: Task ID
|
||||
content: Artifact content
|
||||
artifact_id: Optional artifact ID
|
||||
context_id: Optional context ID
|
||||
append: Whether to append to existing artifact
|
||||
last_chunk: Whether this is the last chunk
|
||||
|
||||
Returns:
|
||||
TaskArtifactUpdateEvent
|
||||
"""
|
||||
import uuid
|
||||
art_id = artifact_id or f"art-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
artifact = Artifact(
|
||||
artifact_id=art_id,
|
||||
parts=[TextPart(text=content)],
|
||||
)
|
||||
|
||||
return TaskArtifactUpdateEvent(
|
||||
task_id=task_id,
|
||||
context_id=context_id,
|
||||
artifact=artifact,
|
||||
append=append,
|
||||
last_chunk=last_chunk,
|
||||
)
|
||||
|
||||
|
||||
async def stream_agent_response(
|
||||
agent,
|
||||
user_input: str,
|
||||
task_id: str,
|
||||
context_id: Optional[str] = None,
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream agent response as A2A events.
|
||||
|
||||
Args:
|
||||
agent: PraisonAI Agent instance
|
||||
user_input: User input text
|
||||
task_id: Task ID
|
||||
context_id: Optional context ID
|
||||
|
||||
Yields:
|
||||
SSE-formatted event strings
|
||||
"""
|
||||
encoder = A2AEventEncoder()
|
||||
|
||||
# Emit working status
|
||||
working_event = create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.WORKING,
|
||||
context_id=context_id,
|
||||
)
|
||||
yield encoder.encode_task_status(working_event)
|
||||
|
||||
try:
|
||||
# Get response from agent
|
||||
response = agent.chat(user_input)
|
||||
|
||||
# Emit artifact with response
|
||||
artifact_event = create_artifact_event(
|
||||
task_id=task_id,
|
||||
content=str(response),
|
||||
context_id=context_id,
|
||||
last_chunk=True,
|
||||
)
|
||||
yield encoder.encode_task_artifact(artifact_event)
|
||||
|
||||
# Emit completed status
|
||||
completed_event = create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.COMPLETED,
|
||||
context_id=context_id,
|
||||
final=True,
|
||||
)
|
||||
yield encoder.encode_task_status(completed_event)
|
||||
|
||||
except Exception:
|
||||
# Emit failed status
|
||||
failed_event = create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.FAILED,
|
||||
context_id=context_id,
|
||||
final=True,
|
||||
)
|
||||
yield encoder.encode_task_status(failed_event)
|
||||
|
||||
# Emit done
|
||||
yield encoder.encode_done()
|
||||
@@ -0,0 +1,209 @@
|
||||
"""
|
||||
A2A Task Store
|
||||
|
||||
In-memory storage for A2A Task lifecycle management.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
Task,
|
||||
TaskState,
|
||||
TaskStatus,
|
||||
Message,
|
||||
Artifact,
|
||||
)
|
||||
|
||||
|
||||
class TaskStore:
|
||||
"""
|
||||
In-memory store for A2A Tasks.
|
||||
|
||||
Manages task lifecycle including creation, status updates,
|
||||
artifact storage, and message history.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize empty task store."""
|
||||
self._tasks: Dict[str, Task] = {}
|
||||
|
||||
def create_task(self, message: Message, context_id: Optional[str] = None) -> Task:
|
||||
"""
|
||||
Create a new task from an incoming message.
|
||||
|
||||
Args:
|
||||
message: The initial message that triggers the task
|
||||
context_id: Optional context ID (uses message's context_id if not provided)
|
||||
|
||||
Returns:
|
||||
Newly created Task in SUBMITTED state
|
||||
"""
|
||||
task_id = f"task-{uuid.uuid4().hex[:12]}"
|
||||
ctx_id = context_id or message.context_id
|
||||
|
||||
# Create initial status
|
||||
status = TaskStatus(
|
||||
state=TaskState.SUBMITTED,
|
||||
timestamp=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
# Create task with initial message in history
|
||||
task = Task(
|
||||
id=task_id,
|
||||
context_id=ctx_id,
|
||||
status=status,
|
||||
history=[message],
|
||||
artifacts=None,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = task
|
||||
return task
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[Task]:
|
||||
"""
|
||||
Retrieve a task by its ID.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to look up
|
||||
|
||||
Returns:
|
||||
Task if found, None otherwise
|
||||
"""
|
||||
return self._tasks.get(task_id)
|
||||
|
||||
def update_status(
|
||||
self,
|
||||
task_id: str,
|
||||
state: TaskState,
|
||||
message: Optional[Message] = None,
|
||||
) -> Optional[Task]:
|
||||
"""
|
||||
Update the status of a task.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to update
|
||||
state: New task state
|
||||
message: Optional status message
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
|
||||
# Create new status
|
||||
new_status = TaskStatus(
|
||||
state=state,
|
||||
message=message,
|
||||
timestamp=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
# Update task (create new instance for immutability)
|
||||
updated = Task(
|
||||
id=task.id,
|
||||
context_id=task.context_id,
|
||||
status=new_status,
|
||||
artifacts=task.artifacts,
|
||||
history=task.history,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = updated
|
||||
return updated
|
||||
|
||||
def add_artifact(self, task_id: str, artifact: Artifact) -> Optional[Task]:
|
||||
"""
|
||||
Add an artifact to a task.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to update
|
||||
artifact: Artifact to add
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
|
||||
# Get existing artifacts or create new list
|
||||
artifacts = list(task.artifacts) if task.artifacts else []
|
||||
artifacts.append(artifact)
|
||||
|
||||
# Update task
|
||||
updated = Task(
|
||||
id=task.id,
|
||||
context_id=task.context_id,
|
||||
status=task.status,
|
||||
artifacts=artifacts,
|
||||
history=task.history,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = updated
|
||||
return updated
|
||||
|
||||
def add_to_history(self, task_id: str, message: Message) -> Optional[Task]:
|
||||
"""
|
||||
Add a message to task history.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to update
|
||||
message: Message to add to history
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
|
||||
# Get existing history or create new list
|
||||
history = list(task.history) if task.history else []
|
||||
history.append(message)
|
||||
|
||||
# Update task
|
||||
updated = Task(
|
||||
id=task.id,
|
||||
context_id=task.context_id,
|
||||
status=task.status,
|
||||
artifacts=task.artifacts,
|
||||
history=history,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = updated
|
||||
return updated
|
||||
|
||||
def list_tasks(self, context_id: Optional[str] = None) -> List[Task]:
|
||||
"""
|
||||
List all tasks, optionally filtered by context.
|
||||
|
||||
Args:
|
||||
context_id: Optional context ID to filter by
|
||||
|
||||
Returns:
|
||||
List of tasks
|
||||
"""
|
||||
if context_id:
|
||||
return [t for t in self._tasks.values() if t.context_id == context_id]
|
||||
return list(self._tasks.values())
|
||||
|
||||
def cancel_task(self, task_id: str) -> Optional[Task]:
|
||||
"""
|
||||
Cancel a task.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to cancel
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
return self.update_status(task_id, TaskState.CANCELLED)
|
||||
|
||||
def clear(self):
|
||||
"""Clear all tasks from the store."""
|
||||
self._tasks.clear()
|
||||
@@ -0,0 +1,266 @@
|
||||
"""
|
||||
A2A Protocol Types
|
||||
|
||||
Pydantic models for the Agent2Agent (A2A) protocol.
|
||||
Based on the A2A Protocol Specification v1.0.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Enums
|
||||
# =============================================================================
|
||||
|
||||
class TaskState(str, Enum):
|
||||
"""
|
||||
Defines the possible lifecycle states of a Task.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.3.
|
||||
"""
|
||||
SUBMITTED = "submitted"
|
||||
WORKING = "working"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
INPUT_REQUIRED = "input_required"
|
||||
AUTH_REQUIRED = "auth_required"
|
||||
|
||||
|
||||
class Role(str, Enum):
|
||||
"""
|
||||
Defines the sender of a message in A2A protocol communication.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.5.
|
||||
"""
|
||||
USER = "user"
|
||||
AGENT = "agent"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Part Types
|
||||
# =============================================================================
|
||||
|
||||
class TextPart(BaseModel):
|
||||
"""
|
||||
A text content part.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.6.
|
||||
"""
|
||||
text: str
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class FilePart(BaseModel):
|
||||
"""
|
||||
A file content part.
|
||||
|
||||
Can contain either a URI reference or inline bytes.
|
||||
Based on A2A Protocol Specification Section 4.1.7.
|
||||
"""
|
||||
file_uri: Optional[str] = Field(default=None, alias="fileWithUri")
|
||||
file_bytes: Optional[bytes] = Field(default=None, alias="fileWithBytes")
|
||||
media_type: Optional[str] = Field(default=None, alias="mediaType")
|
||||
name: Optional[str] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class DataPart(BaseModel):
|
||||
"""
|
||||
A structured data part (JSON blob).
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.8.
|
||||
"""
|
||||
data: Dict[str, Any]
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# Union type for all Part types
|
||||
Part = Union[TextPart, FilePart, DataPart]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Core Types
|
||||
# =============================================================================
|
||||
|
||||
class Message(BaseModel):
|
||||
"""
|
||||
A communication turn between client and server.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.4.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
message_id: str = Field(alias="messageId")
|
||||
role: Role
|
||||
parts: List[Part]
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
task_id: Optional[str] = Field(default=None, alias="taskId")
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
extensions: Optional[Dict[str, str]] = None
|
||||
reference_task_ids: Optional[List[str]] = Field(default=None, alias="referenceTaskIds")
|
||||
|
||||
|
||||
class TaskStatus(BaseModel):
|
||||
"""
|
||||
Container for the status of a task.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.2.
|
||||
"""
|
||||
state: TaskState
|
||||
message: Optional[Message] = None
|
||||
timestamp: Optional[str] = None
|
||||
|
||||
|
||||
class Artifact(BaseModel):
|
||||
"""
|
||||
An output generated by the agent as a result of a task.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.9.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
artifact_id: str = Field(alias="artifactId")
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
parts: List[Part]
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
extensions: Optional[Dict[str, str]] = None
|
||||
|
||||
|
||||
class Task(BaseModel):
|
||||
"""
|
||||
The core unit of action for A2A.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.1.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
id: str
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
status: TaskStatus
|
||||
artifacts: Optional[List[Artifact]] = None
|
||||
history: Optional[List[Message]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Agent Card Types
|
||||
# =============================================================================
|
||||
|
||||
class AgentSkill(BaseModel):
|
||||
"""
|
||||
Describes a skill/capability of an agent.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.4.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
id: str
|
||||
name: str
|
||||
description: str
|
||||
tags: Optional[List[str]] = None
|
||||
examples: Optional[List[str]] = None
|
||||
input_modes: Optional[List[str]] = Field(default=None, alias="inputModes")
|
||||
output_modes: Optional[List[str]] = Field(default=None, alias="outputModes")
|
||||
|
||||
|
||||
class AgentCapabilities(BaseModel):
|
||||
"""
|
||||
Describes the capabilities of an agent.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.4.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
streaming: bool = False
|
||||
push_notifications: bool = Field(default=False, alias="pushNotifications")
|
||||
state_transition_history: bool = Field(default=False, alias="stateTransitionHistory")
|
||||
|
||||
|
||||
class AgentCard(BaseModel):
|
||||
"""
|
||||
JSON metadata document describing an agent.
|
||||
|
||||
Published by an A2A Server for discovery.
|
||||
Based on A2A Protocol Specification Section 8.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
name: str
|
||||
url: str
|
||||
version: str
|
||||
description: Optional[str] = None
|
||||
capabilities: AgentCapabilities
|
||||
skills: Optional[List[AgentSkill]] = None
|
||||
default_input_modes: Optional[List[str]] = Field(default=None, alias="defaultInputModes")
|
||||
default_output_modes: Optional[List[str]] = Field(default=None, alias="defaultOutputModes")
|
||||
provider: Optional[Dict[str, str]] = None
|
||||
documentation_url: Optional[str] = Field(default=None, alias="documentationUrl")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Request/Response Types
|
||||
# =============================================================================
|
||||
|
||||
class MessageConfig(BaseModel):
|
||||
"""
|
||||
Configuration for message sending.
|
||||
|
||||
Based on A2A Protocol Specification Section 3.2.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
accept_content_types: Optional[List[str]] = Field(default=None, alias="acceptContentTypes")
|
||||
blocking: Optional[bool] = None
|
||||
history_length: Optional[int] = Field(default=None, alias="historyLength")
|
||||
|
||||
|
||||
class SendMessageRequest(BaseModel):
|
||||
"""
|
||||
Request to send a message to an agent.
|
||||
|
||||
Based on A2A Protocol Specification Section 3.1.1.
|
||||
"""
|
||||
message: Message
|
||||
configuration: Optional[MessageConfig] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Streaming Event Types
|
||||
# =============================================================================
|
||||
|
||||
class TaskStatusUpdateEvent(BaseModel):
|
||||
"""
|
||||
Event sent by agent to notify client of task status change.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.2.1.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
task_id: str = Field(alias="taskId")
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
status: TaskStatus
|
||||
final: bool = False
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class TaskArtifactUpdateEvent(BaseModel):
|
||||
"""
|
||||
Event sent when an artifact is generated.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.2.2.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
task_id: str = Field(alias="taskId")
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
artifact: Artifact
|
||||
append: bool = False
|
||||
last_chunk: bool = Field(default=False, alias="lastChunk")
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
@@ -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
|
||||
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
AG-UI Protocol Integration for PraisonAI Agents
|
||||
|
||||
This module provides AG-UI (Agent-User Interface) protocol support,
|
||||
enabling PraisonAI Agents to be exposed via a standardized streaming API
|
||||
compatible with CopilotKit and other AG-UI frontends.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.agui.agui import AGUI
|
||||
|
||||
__all__ = ["AGUI"]
|
||||
@@ -0,0 +1,203 @@
|
||||
"""
|
||||
AGUI - Main AG-UI Interface Class
|
||||
|
||||
Exposes PraisonAI Agents via the AG-UI protocol.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from typing import AsyncIterator, List, Optional, TYPE_CHECKING
|
||||
|
||||
from praisonaiagents.ui.agui.types import (
|
||||
BaseEvent,
|
||||
RunAgentInput,
|
||||
)
|
||||
from praisonaiagents.ui.agui.conversion import (
|
||||
agui_messages_to_praisonai,
|
||||
validate_state,
|
||||
extract_user_input,
|
||||
)
|
||||
from praisonaiagents.ui.agui.streaming import (
|
||||
create_run_error_event,
|
||||
async_stream_agent_response,
|
||||
async_stream_agents_response,
|
||||
)
|
||||
from praisonaiagents.ui.agui.encoder import EventEncoder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent, PraisonAIAgents
|
||||
from fastapi import APIRouter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AGUI:
|
||||
"""
|
||||
AG-UI Interface for PraisonAI Agents.
|
||||
|
||||
Exposes a PraisonAI Agent or PraisonAIAgents workflow via the AG-UI protocol,
|
||||
enabling integration with CopilotKit and other AG-UI compatible frontends.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
Args:
|
||||
agent: Single PraisonAI Agent instance
|
||||
agents: PraisonAIAgents instance for multi-agent workflows
|
||||
name: Name for the AG-UI endpoint (defaults to agent name)
|
||||
description: Description of the agent
|
||||
prefix: URL prefix for the router (e.g., "/api/v1")
|
||||
tags: OpenAPI tags for the router
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: Optional["Agent"] = None,
|
||||
agents: Optional["PraisonAIAgents"] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
prefix: str = "",
|
||||
tags: Optional[List[str]] = None,
|
||||
):
|
||||
if agent is None and agents is None:
|
||||
raise ValueError("AGUI requires an agent or agents instance")
|
||||
|
||||
self.agent = agent
|
||||
self.agents = agents
|
||||
self.prefix = prefix
|
||||
self.tags = tags or ["AGUI"]
|
||||
|
||||
# Set name from agent if not provided
|
||||
if name:
|
||||
self.name = name
|
||||
elif agent and hasattr(agent, 'name'):
|
||||
self.name = agent.name
|
||||
elif agents and hasattr(agents, 'name'):
|
||||
self.name = agents.name
|
||||
else:
|
||||
self.name = "PraisonAI Agent"
|
||||
|
||||
# Set description
|
||||
if description:
|
||||
self.description = description
|
||||
elif agent and hasattr(agent, 'role'):
|
||||
self.description = agent.role
|
||||
else:
|
||||
self.description = "PraisonAI Agent via AG-UI"
|
||||
|
||||
self._router: Optional["APIRouter"] = None
|
||||
|
||||
def get_router(self) -> "APIRouter":
|
||||
"""
|
||||
Get the FastAPI router for this AGUI instance.
|
||||
|
||||
Returns:
|
||||
FastAPI APIRouter with AG-UI endpoints
|
||||
"""
|
||||
from fastapi import APIRouter
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
if self._router is not None:
|
||||
return self._router
|
||||
|
||||
self._router = APIRouter(prefix=self.prefix, tags=self.tags)
|
||||
|
||||
# Attach routes
|
||||
self._attach_routes(self._router)
|
||||
|
||||
return self._router
|
||||
|
||||
def _attach_routes(self, router: "APIRouter") -> None:
|
||||
"""Attach AG-UI routes to the router."""
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
encoder = EventEncoder()
|
||||
|
||||
@router.post("/agui")
|
||||
async def run_agent_agui(run_input: RunAgentInput):
|
||||
"""Run the agent via AG-UI protocol."""
|
||||
async def event_generator():
|
||||
async for event in self._run_agent(run_input):
|
||||
yield encoder.encode(event)
|
||||
|
||||
return StreamingResponse(
|
||||
event_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "POST, GET, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "*",
|
||||
},
|
||||
)
|
||||
|
||||
@router.get("/status")
|
||||
async def get_status():
|
||||
"""Get agent status."""
|
||||
return {"status": "available"}
|
||||
|
||||
async def _run_agent(self, run_input: RunAgentInput) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Run the agent with the given input.
|
||||
|
||||
Args:
|
||||
run_input: AG-UI run input
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
run_id = run_input.run_id or str(uuid.uuid4())
|
||||
thread_id = run_input.thread_id
|
||||
|
||||
try:
|
||||
# Convert messages
|
||||
messages = run_input.messages or []
|
||||
|
||||
# Extract user input
|
||||
user_input = extract_user_input(messages)
|
||||
|
||||
# Validate state
|
||||
session_state = validate_state(run_input.state, thread_id)
|
||||
|
||||
# Get user_id from forwarded_props if available
|
||||
user_id = None
|
||||
if run_input.forwarded_props and isinstance(run_input.forwarded_props, dict):
|
||||
user_id = run_input.forwarded_props.get("user_id")
|
||||
|
||||
# Run appropriate handler
|
||||
if self.agent:
|
||||
async for event in async_stream_agent_response(
|
||||
agent=self.agent,
|
||||
user_input=user_input,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
session_state=session_state,
|
||||
messages=agui_messages_to_praisonai(messages),
|
||||
):
|
||||
yield event
|
||||
|
||||
elif self.agents:
|
||||
async for event in async_stream_agents_response(
|
||||
agents=self.agents,
|
||||
user_input=user_input,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
session_state=session_state,
|
||||
):
|
||||
yield event
|
||||
|
||||
else:
|
||||
yield create_run_error_event("No agent or agents configured")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error running agent: {e}", exc_info=True)
|
||||
yield create_run_error_event(str(e))
|
||||
@@ -0,0 +1,211 @@
|
||||
"""
|
||||
AG-UI Message Conversion
|
||||
|
||||
Converts between AG-UI message format and PraisonAI message format.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Set
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from praisonaiagents.ui.agui.types import Message, ToolCall, FunctionCall
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def agui_messages_to_praisonai(messages: List[Message]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert AG-UI messages to PraisonAI format.
|
||||
|
||||
Args:
|
||||
messages: List of AG-UI Message objects
|
||||
|
||||
Returns:
|
||||
List of message dictionaries in PraisonAI format
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
|
||||
# Track tool_call_ids that have results to avoid duplicates
|
||||
tool_call_ids_with_results: Set[str] = set()
|
||||
for msg in messages:
|
||||
if msg.role == "tool" and msg.tool_call_id:
|
||||
tool_call_ids_with_results.add(msg.tool_call_id)
|
||||
|
||||
result: List[Dict[str, Any]] = []
|
||||
seen_tool_call_ids: Set[str] = set()
|
||||
|
||||
for msg in messages:
|
||||
# Skip system messages - agent builds its own
|
||||
if msg.role == "system":
|
||||
continue
|
||||
|
||||
if msg.role == "tool":
|
||||
# Deduplicate tool results
|
||||
if msg.tool_call_id in seen_tool_call_ids:
|
||||
logger.debug(f"Skipping duplicate tool result: {msg.tool_call_id}")
|
||||
continue
|
||||
if msg.tool_call_id:
|
||||
seen_tool_call_ids.add(msg.tool_call_id)
|
||||
|
||||
result.append({
|
||||
"role": "tool",
|
||||
"content": msg.content or "",
|
||||
"tool_call_id": msg.tool_call_id,
|
||||
})
|
||||
|
||||
elif msg.role == "assistant":
|
||||
msg_dict: Dict[str, Any] = {
|
||||
"role": "assistant",
|
||||
"content": msg.content or "",
|
||||
}
|
||||
|
||||
if msg.tool_calls:
|
||||
# Include all tool_calls (filtering only when there are results to match)
|
||||
tool_calls_list = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
"type": "function",
|
||||
}
|
||||
for tc in msg.tool_calls
|
||||
]
|
||||
if tool_calls_list:
|
||||
msg_dict["tool_calls"] = tool_calls_list
|
||||
|
||||
result.append(msg_dict)
|
||||
|
||||
elif msg.role == "user":
|
||||
result.append({
|
||||
"role": "user",
|
||||
"content": msg.content or "",
|
||||
})
|
||||
|
||||
else:
|
||||
logger.warning(f"Unknown message role: {msg.role}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def praisonai_messages_to_agui(messages: List[Dict[str, Any]]) -> List[Message]:
|
||||
"""
|
||||
Convert PraisonAI messages to AG-UI format.
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries in PraisonAI format
|
||||
|
||||
Returns:
|
||||
List of AG-UI Message objects
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
|
||||
result: List[Message] = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "tool":
|
||||
result.append(Message(
|
||||
role="tool",
|
||||
content=content,
|
||||
tool_call_id=msg.get("tool_call_id"),
|
||||
))
|
||||
|
||||
elif role == "assistant":
|
||||
tool_calls = None
|
||||
if "tool_calls" in msg and msg["tool_calls"]:
|
||||
tool_calls = [
|
||||
ToolCall(
|
||||
id=tc.get("id", ""),
|
||||
function=FunctionCall(
|
||||
name=tc.get("function", {}).get("name", ""),
|
||||
arguments=tc.get("function", {}).get("arguments", "{}"),
|
||||
),
|
||||
)
|
||||
for tc in msg["tool_calls"]
|
||||
]
|
||||
|
||||
result.append(Message(
|
||||
role="assistant",
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
))
|
||||
|
||||
elif role == "user":
|
||||
result.append(Message(
|
||||
role="user",
|
||||
content=content,
|
||||
))
|
||||
|
||||
elif role == "system":
|
||||
result.append(Message(
|
||||
role="system",
|
||||
content=content,
|
||||
))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def validate_state(state: Any, thread_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Validate and convert state to a dictionary.
|
||||
|
||||
Args:
|
||||
state: The state to validate (can be dict, Pydantic model, or other)
|
||||
thread_id: Thread ID for logging
|
||||
|
||||
Returns:
|
||||
State as a dictionary, or None if invalid
|
||||
"""
|
||||
if state is None:
|
||||
return None
|
||||
|
||||
if isinstance(state, dict):
|
||||
return state
|
||||
|
||||
if isinstance(state, BaseModel):
|
||||
try:
|
||||
return state.model_dump()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Try to_dict method
|
||||
if hasattr(state, "to_dict") and callable(getattr(state, "to_dict")):
|
||||
try:
|
||||
result = state.to_dict()
|
||||
if isinstance(result, dict):
|
||||
return result
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.warning(
|
||||
f"State must be a dict, got {type(state).__name__}. "
|
||||
f"State will be ignored. Thread: {thread_id}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def extract_user_input(messages: List[Message]) -> str:
|
||||
"""
|
||||
Extract the last user message content.
|
||||
|
||||
Args:
|
||||
messages: List of AG-UI messages
|
||||
|
||||
Returns:
|
||||
Content of the last user message, or empty string
|
||||
"""
|
||||
if not messages:
|
||||
return ""
|
||||
|
||||
for msg in reversed(messages):
|
||||
if msg.role == "user" and msg.content:
|
||||
return msg.content
|
||||
|
||||
return ""
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
AG-UI Event Encoder
|
||||
|
||||
Encodes AG-UI events to Server-Sent Events (SSE) format.
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.agui.types import BaseEvent
|
||||
|
||||
|
||||
class EventEncoder:
|
||||
"""Encodes AG-UI events to SSE format."""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def get_content_type(self) -> str:
|
||||
"""Returns the content type for SSE."""
|
||||
return "text/event-stream"
|
||||
|
||||
def encode(self, event: BaseEvent) -> str:
|
||||
"""
|
||||
Encode an event to SSE format.
|
||||
|
||||
Args:
|
||||
event: The event to encode
|
||||
|
||||
Returns:
|
||||
SSE formatted string
|
||||
"""
|
||||
return f"data: {event.model_dump_json(by_alias=True, exclude_none=True)}\n\n"
|
||||
@@ -0,0 +1,372 @@
|
||||
"""
|
||||
AG-UI Event Streaming
|
||||
|
||||
Provides utilities for streaming PraisonAI agent responses as AG-UI events.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Set
|
||||
|
||||
from praisonaiagents.ui.agui.types import (
|
||||
BaseEvent,
|
||||
TextMessageStartEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
ToolCallStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallEndEvent,
|
||||
ToolCallResultEvent,
|
||||
RunStartedEvent,
|
||||
RunFinishedEvent,
|
||||
RunErrorEvent,
|
||||
StepStartedEvent,
|
||||
StepFinishedEvent,
|
||||
StateSnapshotEvent,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EventBuffer:
|
||||
"""Buffer to manage event ordering constraints."""
|
||||
|
||||
active_tool_call_ids: Set[str] = field(default_factory=set)
|
||||
ended_tool_call_ids: Set[str] = field(default_factory=set)
|
||||
current_text_message_id: str = ""
|
||||
next_text_message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
pending_tool_calls_parent_id: str = ""
|
||||
|
||||
def start_tool_call(self, tool_call_id: str) -> None:
|
||||
"""Start a new tool call."""
|
||||
self.active_tool_call_ids.add(tool_call_id)
|
||||
|
||||
def end_tool_call(self, tool_call_id: str) -> None:
|
||||
"""End a tool call."""
|
||||
self.active_tool_call_ids.discard(tool_call_id)
|
||||
self.ended_tool_call_ids.add(tool_call_id)
|
||||
|
||||
def start_text_message(self) -> str:
|
||||
"""Start a new text message and return its ID."""
|
||||
self.current_text_message_id = self.next_text_message_id
|
||||
self.next_text_message_id = str(uuid.uuid4())
|
||||
return self.current_text_message_id
|
||||
|
||||
def get_parent_message_id_for_tool_call(self) -> str:
|
||||
"""Get the message ID to use as parent for tool calls."""
|
||||
if self.pending_tool_calls_parent_id:
|
||||
return self.pending_tool_calls_parent_id
|
||||
return self.current_text_message_id
|
||||
|
||||
|
||||
# Event Creation Functions
|
||||
|
||||
def create_text_message_events(
|
||||
content: str,
|
||||
message_id: Optional[str] = None,
|
||||
role: str = "assistant"
|
||||
) -> Iterator[BaseEvent]:
|
||||
"""
|
||||
Create text message events for a complete message.
|
||||
|
||||
Args:
|
||||
content: The message content
|
||||
message_id: Optional message ID (generated if not provided)
|
||||
role: Message role (default: assistant)
|
||||
|
||||
Yields:
|
||||
TextMessageStartEvent, TextMessageContentEvent(s), TextMessageEndEvent
|
||||
"""
|
||||
msg_id = message_id or str(uuid.uuid4())
|
||||
|
||||
yield TextMessageStartEvent(message_id=msg_id, role=role)
|
||||
|
||||
if content:
|
||||
yield TextMessageContentEvent(message_id=msg_id, delta=content)
|
||||
|
||||
yield TextMessageEndEvent(message_id=msg_id)
|
||||
|
||||
|
||||
def stream_text_chunks(
|
||||
chunks: Iterator[str],
|
||||
message_id: Optional[str] = None,
|
||||
role: str = "assistant"
|
||||
) -> Iterator[BaseEvent]:
|
||||
"""
|
||||
Stream text chunks as AG-UI events.
|
||||
|
||||
Args:
|
||||
chunks: Iterator of text chunks
|
||||
message_id: Optional message ID
|
||||
role: Message role
|
||||
|
||||
Yields:
|
||||
AG-UI events for the text stream
|
||||
"""
|
||||
msg_id = message_id or str(uuid.uuid4())
|
||||
|
||||
yield TextMessageStartEvent(message_id=msg_id, role=role)
|
||||
|
||||
for chunk in chunks:
|
||||
if chunk:
|
||||
yield TextMessageContentEvent(message_id=msg_id, delta=chunk)
|
||||
|
||||
yield TextMessageEndEvent(message_id=msg_id)
|
||||
|
||||
|
||||
def create_tool_call_events(
|
||||
tool_call_id: str,
|
||||
tool_name: str,
|
||||
arguments: str,
|
||||
parent_message_id: Optional[str] = None
|
||||
) -> Iterator[BaseEvent]:
|
||||
"""
|
||||
Create tool call events.
|
||||
|
||||
Args:
|
||||
tool_call_id: Unique ID for the tool call
|
||||
tool_name: Name of the tool being called
|
||||
arguments: JSON string of arguments
|
||||
parent_message_id: Optional parent message ID
|
||||
|
||||
Yields:
|
||||
ToolCallStartEvent, ToolCallArgsEvent, ToolCallEndEvent
|
||||
"""
|
||||
yield ToolCallStartEvent(
|
||||
tool_call_id=tool_call_id,
|
||||
tool_call_name=tool_name,
|
||||
parent_message_id=parent_message_id
|
||||
)
|
||||
|
||||
yield ToolCallArgsEvent(
|
||||
tool_call_id=tool_call_id,
|
||||
delta=arguments
|
||||
)
|
||||
|
||||
yield ToolCallEndEvent(tool_call_id=tool_call_id)
|
||||
|
||||
|
||||
def create_tool_result_event(
|
||||
tool_call_id: str,
|
||||
content: str,
|
||||
message_id: Optional[str] = None
|
||||
) -> ToolCallResultEvent:
|
||||
"""
|
||||
Create a tool result event.
|
||||
|
||||
Args:
|
||||
tool_call_id: ID of the tool call this is a result for
|
||||
content: Result content
|
||||
message_id: Optional message ID
|
||||
|
||||
Returns:
|
||||
ToolCallResultEvent
|
||||
"""
|
||||
return ToolCallResultEvent(
|
||||
message_id=message_id or str(uuid.uuid4()),
|
||||
tool_call_id=tool_call_id,
|
||||
content=content
|
||||
)
|
||||
|
||||
|
||||
# Run Lifecycle Events
|
||||
|
||||
def create_run_started_event(thread_id: str, run_id: str) -> RunStartedEvent:
|
||||
"""Create a run started event."""
|
||||
return RunStartedEvent(thread_id=thread_id, run_id=run_id)
|
||||
|
||||
|
||||
def create_run_finished_event(
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
result: Optional[Any] = None
|
||||
) -> RunFinishedEvent:
|
||||
"""Create a run finished event."""
|
||||
return RunFinishedEvent(thread_id=thread_id, run_id=run_id, result=result)
|
||||
|
||||
|
||||
def create_run_error_event(
|
||||
message: str,
|
||||
code: Optional[str] = None
|
||||
) -> RunErrorEvent:
|
||||
"""Create a run error event."""
|
||||
return RunErrorEvent(message=message, code=code)
|
||||
|
||||
|
||||
# Step Events
|
||||
|
||||
def create_step_started_event(step_name: str) -> StepStartedEvent:
|
||||
"""Create a step started event."""
|
||||
return StepStartedEvent(step_name=step_name)
|
||||
|
||||
|
||||
def create_step_finished_event(step_name: str) -> StepFinishedEvent:
|
||||
"""Create a step finished event."""
|
||||
return StepFinishedEvent(step_name=step_name)
|
||||
|
||||
|
||||
# State Events
|
||||
|
||||
def create_state_snapshot_event(state: Dict[str, Any]) -> StateSnapshotEvent:
|
||||
"""Create a state snapshot event."""
|
||||
return StateSnapshotEvent(snapshot=state)
|
||||
|
||||
|
||||
# Async Streaming
|
||||
|
||||
async def async_stream_response(
|
||||
response_stream: AsyncIterator[str],
|
||||
thread_id: str,
|
||||
run_id: str
|
||||
) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Stream an async response as AG-UI events.
|
||||
|
||||
Args:
|
||||
response_stream: Async iterator of response chunks
|
||||
thread_id: Thread ID
|
||||
run_id: Run ID
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
message_id = str(uuid.uuid4())
|
||||
message_started = False
|
||||
|
||||
try:
|
||||
async for chunk in response_stream:
|
||||
if not message_started:
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
message_started = True
|
||||
|
||||
if chunk:
|
||||
yield TextMessageContentEvent(message_id=message_id, delta=chunk)
|
||||
|
||||
if message_started:
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
# Emit run finished
|
||||
yield create_run_finished_event(thread_id, run_id)
|
||||
|
||||
except Exception as e:
|
||||
if message_started:
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
yield create_run_error_event(str(e))
|
||||
|
||||
|
||||
async def async_stream_agent_response(
|
||||
agent,
|
||||
user_input: str,
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
session_state: Optional[Dict[str, Any]] = None,
|
||||
messages: Optional[List[Dict[str, Any]]] = None
|
||||
) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Stream an agent's response as AG-UI events.
|
||||
|
||||
Args:
|
||||
agent: PraisonAI Agent instance
|
||||
user_input: User input string
|
||||
thread_id: Thread ID
|
||||
run_id: Run ID
|
||||
session_state: Optional session state
|
||||
messages: Optional message history
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
message_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# Check if agent has async chat method
|
||||
if hasattr(agent, 'achat'):
|
||||
response = await agent.achat(user_input)
|
||||
elif hasattr(agent, 'chat'):
|
||||
# Run sync chat in executor
|
||||
loop = asyncio.get_event_loop()
|
||||
response = await loop.run_in_executor(None, agent.chat, user_input)
|
||||
else:
|
||||
raise ValueError("Agent must have 'chat' or 'achat' method")
|
||||
|
||||
# Emit text message events
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
|
||||
if response:
|
||||
yield TextMessageContentEvent(message_id=message_id, delta=str(response))
|
||||
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
# Emit run finished
|
||||
yield create_run_finished_event(thread_id, run_id)
|
||||
|
||||
except Exception as e:
|
||||
yield create_run_error_event(str(e))
|
||||
|
||||
|
||||
async def async_stream_agents_response(
|
||||
agents,
|
||||
user_input: str,
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
session_state: Optional[Dict[str, Any]] = None
|
||||
) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Stream a PraisonAIAgents workflow response as AG-UI events.
|
||||
|
||||
Args:
|
||||
agents: PraisonAIAgents instance
|
||||
user_input: User input string
|
||||
thread_id: Thread ID
|
||||
run_id: Run ID
|
||||
session_state: Optional session state
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
message_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# Check if agents has async start method
|
||||
if hasattr(agents, 'astart'):
|
||||
result = await agents.astart(user_input)
|
||||
elif hasattr(agents, 'start'):
|
||||
# Run sync start in executor
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(None, agents.start, user_input)
|
||||
else:
|
||||
raise ValueError("Agents must have 'start' or 'astart' method")
|
||||
|
||||
# Emit text message events with result
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
|
||||
if result:
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
content = result.get("output", str(result))
|
||||
elif hasattr(result, "raw"):
|
||||
content = result.raw
|
||||
else:
|
||||
content = str(result)
|
||||
|
||||
yield TextMessageContentEvent(message_id=message_id, delta=content)
|
||||
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
# Emit run finished
|
||||
yield create_run_finished_event(thread_id, run_id)
|
||||
|
||||
except Exception as e:
|
||||
yield create_run_error_event(str(e))
|
||||
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
AG-UI Protocol Types
|
||||
|
||||
This module defines the event types and message types for the AG-UI protocol,
|
||||
following the ag-ui-protocol specification.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class EventType(str, Enum):
|
||||
"""AG-UI Event Types."""
|
||||
TEXT_MESSAGE_START = "TEXT_MESSAGE_START"
|
||||
TEXT_MESSAGE_CONTENT = "TEXT_MESSAGE_CONTENT"
|
||||
TEXT_MESSAGE_END = "TEXT_MESSAGE_END"
|
||||
TOOL_CALL_START = "TOOL_CALL_START"
|
||||
TOOL_CALL_ARGS = "TOOL_CALL_ARGS"
|
||||
TOOL_CALL_END = "TOOL_CALL_END"
|
||||
TOOL_CALL_RESULT = "TOOL_CALL_RESULT"
|
||||
STATE_SNAPSHOT = "STATE_SNAPSHOT"
|
||||
STATE_DELTA = "STATE_DELTA"
|
||||
RUN_STARTED = "RUN_STARTED"
|
||||
RUN_FINISHED = "RUN_FINISHED"
|
||||
RUN_ERROR = "RUN_ERROR"
|
||||
STEP_STARTED = "STEP_STARTED"
|
||||
STEP_FINISHED = "STEP_FINISHED"
|
||||
CUSTOM = "CUSTOM"
|
||||
|
||||
|
||||
class BaseEvent(BaseModel):
|
||||
"""Base event for all AG-UI events."""
|
||||
type: EventType
|
||||
timestamp: Optional[int] = None
|
||||
|
||||
model_config = {"use_enum_values": True}
|
||||
|
||||
|
||||
# Text Message Events
|
||||
class TextMessageStartEvent(BaseEvent):
|
||||
"""Event indicating the start of a text message."""
|
||||
type: Literal["TEXT_MESSAGE_START"] = "TEXT_MESSAGE_START"
|
||||
message_id: str
|
||||
role: str = "assistant"
|
||||
|
||||
|
||||
class TextMessageContentEvent(BaseEvent):
|
||||
"""Event containing text message content."""
|
||||
type: Literal["TEXT_MESSAGE_CONTENT"] = "TEXT_MESSAGE_CONTENT"
|
||||
message_id: str
|
||||
delta: str
|
||||
|
||||
|
||||
class TextMessageEndEvent(BaseEvent):
|
||||
"""Event indicating the end of a text message."""
|
||||
type: Literal["TEXT_MESSAGE_END"] = "TEXT_MESSAGE_END"
|
||||
message_id: str
|
||||
|
||||
|
||||
# Tool Call Events
|
||||
class ToolCallStartEvent(BaseEvent):
|
||||
"""Event indicating the start of a tool call."""
|
||||
type: Literal["TOOL_CALL_START"] = "TOOL_CALL_START"
|
||||
tool_call_id: str
|
||||
tool_call_name: str
|
||||
parent_message_id: Optional[str] = None
|
||||
|
||||
|
||||
class ToolCallArgsEvent(BaseEvent):
|
||||
"""Event containing tool call arguments."""
|
||||
type: Literal["TOOL_CALL_ARGS"] = "TOOL_CALL_ARGS"
|
||||
tool_call_id: str
|
||||
delta: str
|
||||
|
||||
|
||||
class ToolCallEndEvent(BaseEvent):
|
||||
"""Event indicating the end of a tool call."""
|
||||
type: Literal["TOOL_CALL_END"] = "TOOL_CALL_END"
|
||||
tool_call_id: str
|
||||
|
||||
|
||||
class ToolCallResultEvent(BaseEvent):
|
||||
"""Event containing tool call result."""
|
||||
type: Literal["TOOL_CALL_RESULT"] = "TOOL_CALL_RESULT"
|
||||
message_id: str
|
||||
tool_call_id: str
|
||||
content: str
|
||||
role: Optional[Literal["tool"]] = "tool"
|
||||
|
||||
|
||||
# Run Lifecycle Events
|
||||
class RunStartedEvent(BaseEvent):
|
||||
"""Event indicating a run has started."""
|
||||
type: Literal["RUN_STARTED"] = "RUN_STARTED"
|
||||
thread_id: str
|
||||
run_id: str
|
||||
parent_run_id: Optional[str] = None
|
||||
|
||||
|
||||
class RunFinishedEvent(BaseEvent):
|
||||
"""Event indicating a run has finished."""
|
||||
type: Literal["RUN_FINISHED"] = "RUN_FINISHED"
|
||||
thread_id: str
|
||||
run_id: str
|
||||
result: Optional[Any] = None
|
||||
|
||||
|
||||
class RunErrorEvent(BaseEvent):
|
||||
"""Event indicating a run error."""
|
||||
type: Literal["RUN_ERROR"] = "RUN_ERROR"
|
||||
message: str
|
||||
code: Optional[str] = None
|
||||
|
||||
|
||||
# Step Events
|
||||
class StepStartedEvent(BaseEvent):
|
||||
"""Event indicating a step has started."""
|
||||
type: Literal["STEP_STARTED"] = "STEP_STARTED"
|
||||
step_name: str
|
||||
|
||||
|
||||
class StepFinishedEvent(BaseEvent):
|
||||
"""Event indicating a step has finished."""
|
||||
type: Literal["STEP_FINISHED"] = "STEP_FINISHED"
|
||||
step_name: str
|
||||
|
||||
|
||||
# State Events
|
||||
class StateSnapshotEvent(BaseEvent):
|
||||
"""Event containing a state snapshot."""
|
||||
type: Literal["STATE_SNAPSHOT"] = "STATE_SNAPSHOT"
|
||||
snapshot: Dict[str, Any]
|
||||
|
||||
|
||||
# Message Types
|
||||
class FunctionCall(BaseModel):
|
||||
"""Function call in a tool call."""
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
class ToolCall(BaseModel):
|
||||
"""Tool call in a message."""
|
||||
id: str
|
||||
function: FunctionCall
|
||||
type: str = "function"
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
"""AG-UI Message type."""
|
||||
role: Literal["user", "assistant", "system", "tool"]
|
||||
content: Optional[str] = None
|
||||
tool_calls: Optional[List[ToolCall]] = None
|
||||
tool_call_id: Optional[str] = None
|
||||
name: Optional[str] = None
|
||||
|
||||
|
||||
class RunAgentInput(BaseModel):
|
||||
"""Input for running an agent."""
|
||||
thread_id: str
|
||||
run_id: Optional[str] = None
|
||||
messages: Optional[List[Message]] = Field(default_factory=list)
|
||||
state: Optional[Dict[str, Any]] = None
|
||||
forwarded_props: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# Union type for all events
|
||||
Event = Union[
|
||||
TextMessageStartEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
ToolCallStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallEndEvent,
|
||||
ToolCallResultEvent,
|
||||
StateSnapshotEvent,
|
||||
RunStartedEvent,
|
||||
RunFinishedEvent,
|
||||
RunErrorEvent,
|
||||
StepStartedEvent,
|
||||
StepFinishedEvent,
|
||||
]
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.1.14"
|
||||
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 @@
|
||||
"""A2A Protocol Tests"""
|
||||
@@ -0,0 +1,135 @@
|
||||
"""
|
||||
Tests for Main A2A Class
|
||||
|
||||
TDD: Write tests first, then implement A2A class.
|
||||
"""
|
||||
|
||||
|
||||
class TestA2AClass:
|
||||
"""Tests for A2A main class."""
|
||||
|
||||
def test_a2a_creation_with_agent(self):
|
||||
"""Test creating A2A with single agent."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
assert a2a.agent is not None
|
||||
assert a2a.name == "Test"
|
||||
|
||||
def test_a2a_creation_with_custom_name(self):
|
||||
"""Test creating A2A with custom name."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent, name="Custom Name")
|
||||
|
||||
assert a2a.name == "Custom Name"
|
||||
|
||||
def test_a2a_requires_agent(self):
|
||||
"""Test A2A requires agent or agents."""
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
import pytest
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
A2A()
|
||||
|
||||
def test_a2a_get_agent_card(self):
|
||||
"""Test getting Agent Card from A2A."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Card Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
card = a2a.get_agent_card()
|
||||
|
||||
assert card.name == "Card Test"
|
||||
assert card.url == "http://localhost:8000/a2a"
|
||||
|
||||
def test_a2a_get_router(self):
|
||||
"""Test getting FastAPI router from A2A."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Router Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
router = a2a.get_router()
|
||||
|
||||
assert router is not None
|
||||
# Check router has routes
|
||||
assert len(router.routes) > 0
|
||||
|
||||
|
||||
class TestA2ARouter:
|
||||
"""Tests for A2A FastAPI router endpoints."""
|
||||
|
||||
def test_agent_card_endpoint(self):
|
||||
"""Test /.well-known/agent.json endpoint."""
|
||||
from fastapi.testclient import TestClient
|
||||
from fastapi import FastAPI
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Endpoint Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/.well-known/agent.json")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "Endpoint Test"
|
||||
|
||||
def test_status_endpoint(self):
|
||||
"""Test /status endpoint."""
|
||||
from fastapi.testclient import TestClient
|
||||
from fastapi import FastAPI
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Status Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "ok"
|
||||
|
||||
|
||||
class TestA2AWithTools:
|
||||
"""Tests for A2A with agent tools."""
|
||||
|
||||
def test_a2a_agent_card_includes_skills(self):
|
||||
"""Test Agent Card includes skills from tools."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
def my_tool(x: str) -> str:
|
||||
"""A test tool."""
|
||||
return x
|
||||
|
||||
agent = Agent(
|
||||
name="Tool Agent",
|
||||
role="Helper",
|
||||
goal="Help",
|
||||
tools=[my_tool]
|
||||
)
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
card = a2a.get_agent_card()
|
||||
|
||||
assert card.skills is not None
|
||||
assert len(card.skills) >= 1
|
||||
@@ -0,0 +1,200 @@
|
||||
"""
|
||||
Tests for A2A Agent Card Generation
|
||||
|
||||
TDD: Write tests first, then implement agent_card module.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class TestAgentCardGeneration:
|
||||
"""Tests for generating Agent Card from PraisonAI Agent."""
|
||||
|
||||
def test_generate_agent_card_basic(self):
|
||||
"""Test generating Agent Card from basic Agent."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
name="Test Agent",
|
||||
role="Test Role",
|
||||
goal="Test Goal"
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.name == "Test Agent"
|
||||
assert card.url == "http://localhost:8000/a2a"
|
||||
assert card.version == "1.0.0"
|
||||
|
||||
def test_generate_agent_card_with_description(self):
|
||||
"""Test Agent Card includes description from role."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
name="Research Agent",
|
||||
role="Research Analyst",
|
||||
goal="Research topics thoroughly"
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.description is not None
|
||||
assert "Research" in card.description
|
||||
|
||||
def test_generate_agent_card_with_instructions(self):
|
||||
"""Test Agent Card from Agent with instructions."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant that answers questions."
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.name is not None
|
||||
assert card.description is not None
|
||||
|
||||
def test_generate_agent_card_capabilities(self):
|
||||
"""Test Agent Card has correct capabilities."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
|
||||
card = generate_agent_card(
|
||||
agent,
|
||||
url="http://localhost:8000/a2a",
|
||||
streaming=True
|
||||
)
|
||||
|
||||
assert card.capabilities.streaming is True
|
||||
|
||||
|
||||
class TestSkillsExtraction:
|
||||
"""Tests for extracting skills from Agent tools."""
|
||||
|
||||
def test_extract_skills_from_function(self):
|
||||
"""Test extracting skill from a function tool."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Search Agent",
|
||||
role="Searcher",
|
||||
goal="Search the web",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.skills is not None
|
||||
assert len(card.skills) >= 1
|
||||
|
||||
skill_names = [s.name for s in card.skills]
|
||||
assert "search_web" in skill_names
|
||||
|
||||
def test_extract_skills_with_description(self):
|
||||
"""Test skill description comes from function docstring."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
def calculate(expression: str) -> str:
|
||||
"""Calculate a mathematical expression."""
|
||||
return str(eval(expression))
|
||||
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
role="Math Helper",
|
||||
goal="Do math",
|
||||
tools=[calculate]
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
calc_skill = next((s for s in card.skills if s.name == "calculate"), None)
|
||||
assert calc_skill is not None
|
||||
assert "mathematical" in calc_skill.description.lower()
|
||||
|
||||
def test_extract_skills_multiple_tools(self):
|
||||
"""Test extracting skills from multiple tools."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
def tool_one(x: str) -> str:
|
||||
"""First tool."""
|
||||
return x
|
||||
|
||||
def tool_two(y: str) -> str:
|
||||
"""Second tool."""
|
||||
return y
|
||||
|
||||
agent = Agent(
|
||||
name="Multi Tool Agent",
|
||||
role="Helper",
|
||||
goal="Help",
|
||||
tools=[tool_one, tool_two]
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert len(card.skills) >= 2
|
||||
|
||||
def test_no_skills_when_no_tools(self):
|
||||
"""Test Agent Card with no tools has empty skills."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
name="Simple Agent",
|
||||
role="Helper",
|
||||
goal="Help"
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
# Skills should be None or empty list
|
||||
assert card.skills is None or len(card.skills) == 0
|
||||
|
||||
|
||||
class TestAgentCardSerialization:
|
||||
"""Tests for Agent Card JSON serialization."""
|
||||
|
||||
def test_agent_card_to_json(self):
|
||||
"""Test Agent Card can be serialized to JSON."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
json_str = card.model_dump_json()
|
||||
data = json.loads(json_str)
|
||||
|
||||
assert "name" in data
|
||||
assert "url" in data
|
||||
assert "capabilities" in data
|
||||
|
||||
def test_agent_card_camel_case_keys(self):
|
||||
"""Test Agent Card JSON uses camelCase keys."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
card = generate_agent_card(
|
||||
agent,
|
||||
url="http://localhost:8000/a2a",
|
||||
streaming=True
|
||||
)
|
||||
|
||||
# Use by_alias to get camelCase
|
||||
data = card.model_dump(by_alias=True)
|
||||
|
||||
# Check capabilities uses camelCase
|
||||
assert "pushNotifications" in data["capabilities"]
|
||||
@@ -0,0 +1,190 @@
|
||||
"""
|
||||
Tests for A2A Message Conversion
|
||||
|
||||
TDD: Write tests first, then implement conversion module.
|
||||
"""
|
||||
|
||||
|
||||
class TestA2AToPraisonAI:
|
||||
"""Tests for converting A2A messages to PraisonAI format."""
|
||||
|
||||
def test_convert_user_message(self):
|
||||
"""Test converting A2A user message to PraisonAI format."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
a2a_msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello, how are you?")]
|
||||
)
|
||||
|
||||
result = a2a_to_praisonai_messages([a2a_msg])
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[0]["content"] == "Hello, how are you?"
|
||||
|
||||
def test_convert_agent_message(self):
|
||||
"""Test converting A2A agent message to PraisonAI format."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
a2a_msg = Message(
|
||||
message_id="msg-2",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="I'm doing well, thanks!")]
|
||||
)
|
||||
|
||||
result = a2a_to_praisonai_messages([a2a_msg])
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "assistant"
|
||||
assert result[0]["content"] == "I'm doing well, thanks!"
|
||||
|
||||
def test_convert_multiple_parts(self):
|
||||
"""Test converting message with multiple text parts."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
a2a_msg = Message(
|
||||
message_id="msg-3",
|
||||
role=Role.USER,
|
||||
parts=[
|
||||
TextPart(text="First part."),
|
||||
TextPart(text="Second part.")
|
||||
]
|
||||
)
|
||||
|
||||
result = a2a_to_praisonai_messages([a2a_msg])
|
||||
|
||||
assert "First part" in result[0]["content"]
|
||||
assert "Second part" in result[0]["content"]
|
||||
|
||||
def test_convert_conversation_history(self):
|
||||
"""Test converting full conversation history."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
messages = [
|
||||
Message(message_id="m1", role=Role.USER, parts=[TextPart(text="Hi")]),
|
||||
Message(message_id="m2", role=Role.AGENT, parts=[TextPart(text="Hello!")]),
|
||||
Message(message_id="m3", role=Role.USER, parts=[TextPart(text="How are you?")]),
|
||||
]
|
||||
|
||||
result = a2a_to_praisonai_messages(messages)
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[1]["role"] == "assistant"
|
||||
assert result[2]["role"] == "user"
|
||||
|
||||
|
||||
class TestPraisonAIToA2A:
|
||||
"""Tests for converting PraisonAI responses to A2A format."""
|
||||
|
||||
def test_convert_string_response(self):
|
||||
"""Test converting string response to A2A Message."""
|
||||
from praisonaiagents.ui.a2a.conversion import praisonai_to_a2a_message
|
||||
|
||||
response = "This is the agent's response."
|
||||
|
||||
result = praisonai_to_a2a_message(response)
|
||||
|
||||
assert result.role.value == "agent"
|
||||
assert len(result.parts) == 1
|
||||
assert result.parts[0].text == response
|
||||
|
||||
def test_convert_response_with_context(self):
|
||||
"""Test converting response with context_id."""
|
||||
from praisonaiagents.ui.a2a.conversion import praisonai_to_a2a_message
|
||||
|
||||
response = "Response text"
|
||||
|
||||
result = praisonai_to_a2a_message(response, context_id="ctx-123")
|
||||
|
||||
assert result.context_id == "ctx-123"
|
||||
|
||||
def test_convert_response_with_task(self):
|
||||
"""Test converting response with task_id."""
|
||||
from praisonaiagents.ui.a2a.conversion import praisonai_to_a2a_message
|
||||
|
||||
response = "Response text"
|
||||
|
||||
result = praisonai_to_a2a_message(response, task_id="task-456")
|
||||
|
||||
assert result.task_id == "task-456"
|
||||
|
||||
|
||||
class TestArtifactConversion:
|
||||
"""Tests for converting responses to A2A Artifacts."""
|
||||
|
||||
def test_create_artifact_from_response(self):
|
||||
"""Test creating Artifact from response."""
|
||||
from praisonaiagents.ui.a2a.conversion import create_artifact
|
||||
|
||||
response = "Generated content here."
|
||||
|
||||
artifact = create_artifact(response)
|
||||
|
||||
assert artifact.artifact_id is not None
|
||||
assert len(artifact.parts) == 1
|
||||
assert artifact.parts[0].text == response
|
||||
|
||||
def test_create_artifact_with_name(self):
|
||||
"""Test creating Artifact with name."""
|
||||
from praisonaiagents.ui.a2a.conversion import create_artifact
|
||||
|
||||
response = "Report content"
|
||||
|
||||
artifact = create_artifact(response, name="Report", description="Generated report")
|
||||
|
||||
assert artifact.name == "Report"
|
||||
assert artifact.description == "Generated report"
|
||||
|
||||
|
||||
class TestExtractUserInput:
|
||||
"""Tests for extracting user input from A2A messages."""
|
||||
|
||||
def test_extract_from_single_message(self):
|
||||
"""Test extracting user input from single message."""
|
||||
from praisonaiagents.ui.a2a.conversion import extract_user_input
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
msg = Message(
|
||||
message_id="m1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="What is AI?")]
|
||||
)
|
||||
|
||||
result = extract_user_input([msg])
|
||||
|
||||
assert result == "What is AI?"
|
||||
|
||||
def test_extract_from_conversation(self):
|
||||
"""Test extracting last user input from conversation."""
|
||||
from praisonaiagents.ui.a2a.conversion import extract_user_input
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
messages = [
|
||||
Message(message_id="m1", role=Role.USER, parts=[TextPart(text="First question")]),
|
||||
Message(message_id="m2", role=Role.AGENT, parts=[TextPart(text="Answer")]),
|
||||
Message(message_id="m3", role=Role.USER, parts=[TextPart(text="Follow up question")]),
|
||||
]
|
||||
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == "Follow up question"
|
||||
|
||||
def test_extract_empty_when_no_user_message(self):
|
||||
"""Test extracting returns empty when no user message."""
|
||||
from praisonaiagents.ui.a2a.conversion import extract_user_input
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
messages = [
|
||||
Message(message_id="m1", role=Role.AGENT, parts=[TextPart(text="Hello")]),
|
||||
]
|
||||
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == ""
|
||||
@@ -0,0 +1,116 @@
|
||||
"""
|
||||
Tests for A2A Streaming Events
|
||||
"""
|
||||
|
||||
|
||||
class TestA2AEventEncoder:
|
||||
"""Tests for A2AEventEncoder."""
|
||||
|
||||
def test_encode_event(self):
|
||||
"""Test encoding a generic event."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
|
||||
result = A2AEventEncoder.encode_event("test", {"key": "value"})
|
||||
|
||||
assert "event: test" in result
|
||||
assert "data:" in result
|
||||
assert "key" in result
|
||||
|
||||
def test_encode_task_status(self):
|
||||
"""Test encoding TaskStatusUpdateEvent."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
from praisonaiagents.ui.a2a.types import TaskStatusUpdateEvent, TaskStatus, TaskState
|
||||
|
||||
event = TaskStatusUpdateEvent(
|
||||
task_id="task-123",
|
||||
status=TaskStatus(state=TaskState.WORKING),
|
||||
final=False,
|
||||
)
|
||||
|
||||
result = A2AEventEncoder.encode_task_status(event)
|
||||
|
||||
assert "event: task.status" in result
|
||||
assert "task-123" in result
|
||||
|
||||
def test_encode_task_artifact(self):
|
||||
"""Test encoding TaskArtifactUpdateEvent."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
from praisonaiagents.ui.a2a.types import TaskArtifactUpdateEvent, Artifact, TextPart
|
||||
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Result")]
|
||||
)
|
||||
event = TaskArtifactUpdateEvent(
|
||||
task_id="task-456",
|
||||
artifact=artifact,
|
||||
last_chunk=True,
|
||||
)
|
||||
|
||||
result = A2AEventEncoder.encode_task_artifact(event)
|
||||
|
||||
assert "event: task.artifact" in result
|
||||
assert "art-1" in result
|
||||
|
||||
def test_encode_done(self):
|
||||
"""Test encoding done event."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
|
||||
result = A2AEventEncoder.encode_done()
|
||||
|
||||
assert "event: done" in result
|
||||
|
||||
|
||||
class TestCreateEvents:
|
||||
"""Tests for event creation helpers."""
|
||||
|
||||
def test_create_status_event(self):
|
||||
"""Test creating status event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_status_event
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
|
||||
event = create_status_event(
|
||||
task_id="task-1",
|
||||
state=TaskState.WORKING,
|
||||
)
|
||||
|
||||
assert event.task_id == "task-1"
|
||||
assert event.status.state == TaskState.WORKING
|
||||
|
||||
def test_create_status_event_final(self):
|
||||
"""Test creating final status event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_status_event
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
|
||||
event = create_status_event(
|
||||
task_id="task-2",
|
||||
state=TaskState.COMPLETED,
|
||||
final=True,
|
||||
)
|
||||
|
||||
assert event.final is True
|
||||
|
||||
def test_create_artifact_event(self):
|
||||
"""Test creating artifact event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_artifact_event
|
||||
|
||||
event = create_artifact_event(
|
||||
task_id="task-3",
|
||||
content="Generated content",
|
||||
)
|
||||
|
||||
assert event.task_id == "task-3"
|
||||
assert event.artifact is not None
|
||||
assert len(event.artifact.parts) == 1
|
||||
|
||||
def test_create_artifact_event_last_chunk(self):
|
||||
"""Test creating last chunk artifact event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_artifact_event
|
||||
|
||||
event = create_artifact_event(
|
||||
task_id="task-4",
|
||||
content="Final content",
|
||||
last_chunk=True,
|
||||
)
|
||||
|
||||
assert event.last_chunk is True
|
||||
@@ -0,0 +1,225 @@
|
||||
"""
|
||||
Tests for A2A Task Store
|
||||
|
||||
TDD: Write tests first, then implement task_store module.
|
||||
"""
|
||||
|
||||
|
||||
class TestTaskStore:
|
||||
"""Tests for TaskStore class."""
|
||||
|
||||
def test_create_task(self):
|
||||
"""Test creating a new task."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
|
||||
assert task.id is not None
|
||||
assert task.status.state == TaskState.SUBMITTED
|
||||
|
||||
def test_create_task_with_context(self):
|
||||
"""Test creating task with context_id."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-2",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")],
|
||||
context_id="ctx-123"
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
|
||||
assert task.context_id == "ctx-123"
|
||||
|
||||
def test_get_task(self):
|
||||
"""Test retrieving a task by ID."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-3",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
created = store.create_task(msg)
|
||||
retrieved = store.get_task(created.id)
|
||||
|
||||
assert retrieved is not None
|
||||
assert retrieved.id == created.id
|
||||
|
||||
def test_get_task_not_found(self):
|
||||
"""Test getting non-existent task returns None."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
|
||||
store = TaskStore()
|
||||
result = store.get_task("non-existent-id")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_update_status(self):
|
||||
"""Test updating task status."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-4",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
updated = store.update_status(task.id, TaskState.WORKING)
|
||||
|
||||
assert updated.status.state == TaskState.WORKING
|
||||
|
||||
def test_update_status_to_completed(self):
|
||||
"""Test updating task to completed state."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-5",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
store.update_status(task.id, TaskState.WORKING)
|
||||
updated = store.update_status(task.id, TaskState.COMPLETED)
|
||||
|
||||
assert updated.status.state == TaskState.COMPLETED
|
||||
|
||||
def test_add_artifact(self):
|
||||
"""Test adding artifact to task."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, Artifact
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-6",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Result")]
|
||||
)
|
||||
|
||||
updated = store.add_artifact(task.id, artifact)
|
||||
|
||||
assert updated.artifacts is not None
|
||||
assert len(updated.artifacts) == 1
|
||||
|
||||
def test_add_message_to_history(self):
|
||||
"""Test adding message to task history."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
msg1 = Message(
|
||||
message_id="msg-7",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Question")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg1)
|
||||
|
||||
msg2 = Message(
|
||||
message_id="msg-8",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Answer")]
|
||||
)
|
||||
|
||||
updated = store.add_to_history(task.id, msg2)
|
||||
|
||||
assert updated.history is not None
|
||||
assert len(updated.history) >= 1
|
||||
|
||||
def test_list_tasks(self):
|
||||
"""Test listing all tasks."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
|
||||
for i in range(3):
|
||||
msg = Message(
|
||||
message_id=f"msg-{i}",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text=f"Message {i}")]
|
||||
)
|
||||
store.create_task(msg)
|
||||
|
||||
tasks = store.list_tasks()
|
||||
|
||||
assert len(tasks) == 3
|
||||
|
||||
def test_cancel_task(self):
|
||||
"""Test cancelling a task."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-cancel",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
cancelled = store.cancel_task(task.id)
|
||||
|
||||
assert cancelled.status.state == TaskState.CANCELLED
|
||||
|
||||
|
||||
class TestTaskStoreContextManagement:
|
||||
"""Tests for context management in TaskStore."""
|
||||
|
||||
def test_tasks_grouped_by_context(self):
|
||||
"""Test getting tasks by context_id."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
|
||||
# Create tasks in context A
|
||||
for i in range(2):
|
||||
msg = Message(
|
||||
message_id=f"msg-a-{i}",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text=f"A{i}")],
|
||||
context_id="context-a"
|
||||
)
|
||||
store.create_task(msg)
|
||||
|
||||
# Create task in context B
|
||||
msg_b = Message(
|
||||
message_id="msg-b-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="B1")],
|
||||
context_id="context-b"
|
||||
)
|
||||
store.create_task(msg_b)
|
||||
|
||||
tasks_a = store.list_tasks(context_id="context-a")
|
||||
tasks_b = store.list_tasks(context_id="context-b")
|
||||
|
||||
assert len(tasks_a) == 2
|
||||
assert len(tasks_b) == 1
|
||||
@@ -0,0 +1,395 @@
|
||||
"""
|
||||
Tests for A2A Protocol Types
|
||||
|
||||
TDD: Write tests first, then implement types.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class TestTaskState:
|
||||
"""Tests for TaskState enum."""
|
||||
|
||||
def test_task_state_submitted(self):
|
||||
"""Test SUBMITTED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.SUBMITTED == "submitted"
|
||||
|
||||
def test_task_state_working(self):
|
||||
"""Test WORKING state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.WORKING == "working"
|
||||
|
||||
def test_task_state_completed(self):
|
||||
"""Test COMPLETED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.COMPLETED == "completed"
|
||||
|
||||
def test_task_state_failed(self):
|
||||
"""Test FAILED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.FAILED == "failed"
|
||||
|
||||
def test_task_state_cancelled(self):
|
||||
"""Test CANCELLED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.CANCELLED == "cancelled"
|
||||
|
||||
def test_task_state_input_required(self):
|
||||
"""Test INPUT_REQUIRED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.INPUT_REQUIRED == "input_required"
|
||||
|
||||
def test_task_state_auth_required(self):
|
||||
"""Test AUTH_REQUIRED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.AUTH_REQUIRED == "auth_required"
|
||||
|
||||
|
||||
class TestRole:
|
||||
"""Tests for Role enum."""
|
||||
|
||||
def test_role_user(self):
|
||||
"""Test USER role exists."""
|
||||
from praisonaiagents.ui.a2a.types import Role
|
||||
assert Role.USER == "user"
|
||||
|
||||
def test_role_agent(self):
|
||||
"""Test AGENT role exists."""
|
||||
from praisonaiagents.ui.a2a.types import Role
|
||||
assert Role.AGENT == "agent"
|
||||
|
||||
|
||||
class TestPart:
|
||||
"""Tests for Part models."""
|
||||
|
||||
def test_text_part_creation(self):
|
||||
"""Test TextPart can be created."""
|
||||
from praisonaiagents.ui.a2a.types import TextPart
|
||||
part = TextPart(text="Hello, world!")
|
||||
assert part.text == "Hello, world!"
|
||||
|
||||
def test_file_part_with_uri(self):
|
||||
"""Test FilePart with URI."""
|
||||
from praisonaiagents.ui.a2a.types import FilePart
|
||||
part = FilePart(
|
||||
file_uri="https://example.com/file.pdf",
|
||||
media_type="application/pdf",
|
||||
name="document.pdf"
|
||||
)
|
||||
assert part.file_uri == "https://example.com/file.pdf"
|
||||
assert part.media_type == "application/pdf"
|
||||
|
||||
def test_file_part_with_bytes(self):
|
||||
"""Test FilePart with bytes."""
|
||||
from praisonaiagents.ui.a2a.types import FilePart
|
||||
part = FilePart(
|
||||
file_bytes=b"file content",
|
||||
media_type="text/plain",
|
||||
name="test.txt"
|
||||
)
|
||||
assert part.file_bytes == b"file content"
|
||||
|
||||
def test_data_part_creation(self):
|
||||
"""Test DataPart can be created."""
|
||||
from praisonaiagents.ui.a2a.types import DataPart
|
||||
part = DataPart(data={"key": "value", "count": 42})
|
||||
assert part.data == {"key": "value", "count": 42}
|
||||
|
||||
|
||||
class TestMessage:
|
||||
"""Tests for Message model."""
|
||||
|
||||
def test_message_creation(self):
|
||||
"""Test Message can be created."""
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-123",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
assert msg.message_id == "msg-123"
|
||||
assert msg.role == Role.USER
|
||||
assert len(msg.parts) == 1
|
||||
|
||||
def test_message_with_context(self):
|
||||
"""Test Message with context_id."""
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-456",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Response")],
|
||||
context_id="ctx-789"
|
||||
)
|
||||
assert msg.context_id == "ctx-789"
|
||||
|
||||
def test_message_with_task(self):
|
||||
"""Test Message with task_id."""
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-789",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Working on it")],
|
||||
task_id="task-123"
|
||||
)
|
||||
assert msg.task_id == "task-123"
|
||||
|
||||
|
||||
class TestTaskStatus:
|
||||
"""Tests for TaskStatus model."""
|
||||
|
||||
def test_task_status_creation(self):
|
||||
"""Test TaskStatus can be created."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatus, TaskState
|
||||
status = TaskStatus(state=TaskState.WORKING)
|
||||
assert status.state == TaskState.WORKING
|
||||
|
||||
def test_task_status_with_message(self):
|
||||
"""Test TaskStatus with message."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatus, TaskState, Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Processing...")]
|
||||
)
|
||||
status = TaskStatus(state=TaskState.WORKING, message=msg)
|
||||
assert status.message is not None
|
||||
|
||||
def test_task_status_with_timestamp(self):
|
||||
"""Test TaskStatus with timestamp."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatus, TaskState
|
||||
ts = datetime.utcnow().isoformat()
|
||||
status = TaskStatus(state=TaskState.COMPLETED, timestamp=ts)
|
||||
assert status.timestamp == ts
|
||||
|
||||
|
||||
class TestTask:
|
||||
"""Tests for Task model."""
|
||||
|
||||
def test_task_creation(self):
|
||||
"""Test Task can be created."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState
|
||||
task = Task(
|
||||
id="task-123",
|
||||
status=TaskStatus(state=TaskState.SUBMITTED)
|
||||
)
|
||||
assert task.id == "task-123"
|
||||
assert task.status.state == TaskState.SUBMITTED
|
||||
|
||||
def test_task_with_context(self):
|
||||
"""Test Task with context_id."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState
|
||||
task = Task(
|
||||
id="task-456",
|
||||
context_id="ctx-789",
|
||||
status=TaskStatus(state=TaskState.WORKING)
|
||||
)
|
||||
assert task.context_id == "ctx-789"
|
||||
|
||||
def test_task_with_artifacts(self):
|
||||
"""Test Task with artifacts."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState, Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Result")]
|
||||
)
|
||||
task = Task(
|
||||
id="task-789",
|
||||
status=TaskStatus(state=TaskState.COMPLETED),
|
||||
artifacts=[artifact]
|
||||
)
|
||||
assert len(task.artifacts) == 1
|
||||
|
||||
def test_task_with_history(self):
|
||||
"""Test Task with message history."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState, Message, Role, TextPart
|
||||
msg1 = Message(message_id="m1", role=Role.USER, parts=[TextPart(text="Q")])
|
||||
msg2 = Message(message_id="m2", role=Role.AGENT, parts=[TextPart(text="A")])
|
||||
task = Task(
|
||||
id="task-abc",
|
||||
status=TaskStatus(state=TaskState.COMPLETED),
|
||||
history=[msg1, msg2]
|
||||
)
|
||||
assert len(task.history) == 2
|
||||
|
||||
|
||||
class TestArtifact:
|
||||
"""Tests for Artifact model."""
|
||||
|
||||
def test_artifact_creation(self):
|
||||
"""Test Artifact can be created."""
|
||||
from praisonaiagents.ui.a2a.types import Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-123",
|
||||
parts=[TextPart(text="Generated content")]
|
||||
)
|
||||
assert artifact.artifact_id == "art-123"
|
||||
|
||||
def test_artifact_with_name(self):
|
||||
"""Test Artifact with name and description."""
|
||||
from praisonaiagents.ui.a2a.types import Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-456",
|
||||
name="Report",
|
||||
description="Generated report",
|
||||
parts=[TextPart(text="Report content")]
|
||||
)
|
||||
assert artifact.name == "Report"
|
||||
assert artifact.description == "Generated report"
|
||||
|
||||
|
||||
class TestAgentSkill:
|
||||
"""Tests for AgentSkill model."""
|
||||
|
||||
def test_agent_skill_creation(self):
|
||||
"""Test AgentSkill can be created."""
|
||||
from praisonaiagents.ui.a2a.types import AgentSkill
|
||||
skill = AgentSkill(
|
||||
id="search",
|
||||
name="Web Search",
|
||||
description="Search the web for information"
|
||||
)
|
||||
assert skill.id == "search"
|
||||
assert skill.name == "Web Search"
|
||||
|
||||
def test_agent_skill_with_tags(self):
|
||||
"""Test AgentSkill with tags."""
|
||||
from praisonaiagents.ui.a2a.types import AgentSkill
|
||||
skill = AgentSkill(
|
||||
id="calc",
|
||||
name="Calculator",
|
||||
description="Perform calculations",
|
||||
tags=["math", "utility"]
|
||||
)
|
||||
assert skill.tags == ["math", "utility"]
|
||||
|
||||
def test_agent_skill_with_examples(self):
|
||||
"""Test AgentSkill with examples."""
|
||||
from praisonaiagents.ui.a2a.types import AgentSkill
|
||||
skill = AgentSkill(
|
||||
id="translate",
|
||||
name="Translator",
|
||||
description="Translate text",
|
||||
examples=["Translate hello to Spanish", "What is bonjour in English?"]
|
||||
)
|
||||
assert len(skill.examples) == 2
|
||||
|
||||
|
||||
class TestAgentCapabilities:
|
||||
"""Tests for AgentCapabilities model."""
|
||||
|
||||
def test_capabilities_creation(self):
|
||||
"""Test AgentCapabilities can be created."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCapabilities
|
||||
caps = AgentCapabilities(streaming=True)
|
||||
assert caps.streaming is True
|
||||
|
||||
def test_capabilities_defaults(self):
|
||||
"""Test AgentCapabilities defaults."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCapabilities
|
||||
caps = AgentCapabilities()
|
||||
assert caps.streaming is False
|
||||
assert caps.push_notifications is False
|
||||
|
||||
|
||||
class TestAgentCard:
|
||||
"""Tests for AgentCard model."""
|
||||
|
||||
def test_agent_card_creation(self):
|
||||
"""Test AgentCard can be created."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCard, AgentCapabilities
|
||||
card = AgentCard(
|
||||
name="Research Agent",
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0",
|
||||
capabilities=AgentCapabilities(streaming=True)
|
||||
)
|
||||
assert card.name == "Research Agent"
|
||||
assert card.url == "http://localhost:8000/a2a"
|
||||
|
||||
def test_agent_card_with_skills(self):
|
||||
"""Test AgentCard with skills."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCard, AgentCapabilities, AgentSkill
|
||||
skill = AgentSkill(id="search", name="Search", description="Search web")
|
||||
card = AgentCard(
|
||||
name="Search Agent",
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0",
|
||||
capabilities=AgentCapabilities(),
|
||||
skills=[skill]
|
||||
)
|
||||
assert len(card.skills) == 1
|
||||
|
||||
def test_agent_card_json_serialization(self):
|
||||
"""Test AgentCard can be serialized to JSON."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCard, AgentCapabilities
|
||||
card = AgentCard(
|
||||
name="Test Agent",
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0",
|
||||
capabilities=AgentCapabilities()
|
||||
)
|
||||
json_str = card.model_dump_json()
|
||||
data = json.loads(json_str)
|
||||
assert data["name"] == "Test Agent"
|
||||
|
||||
|
||||
class TestSendMessageRequest:
|
||||
"""Tests for SendMessageRequest model."""
|
||||
|
||||
def test_send_message_request_creation(self):
|
||||
"""Test SendMessageRequest can be created."""
|
||||
from praisonaiagents.ui.a2a.types import SendMessageRequest, Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
req = SendMessageRequest(message=msg)
|
||||
assert req.message.message_id == "msg-1"
|
||||
|
||||
def test_send_message_request_with_config(self):
|
||||
"""Test SendMessageRequest with configuration."""
|
||||
from praisonaiagents.ui.a2a.types import SendMessageRequest, Message, Role, TextPart, MessageConfig
|
||||
msg = Message(
|
||||
message_id="msg-2",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Search for AI")]
|
||||
)
|
||||
config = MessageConfig(
|
||||
accept_content_types=["text/plain", "application/json"]
|
||||
)
|
||||
req = SendMessageRequest(message=msg, configuration=config)
|
||||
assert req.configuration is not None
|
||||
|
||||
|
||||
class TestStreamingEvents:
|
||||
"""Tests for streaming event models."""
|
||||
|
||||
def test_task_status_update_event(self):
|
||||
"""Test TaskStatusUpdateEvent creation."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatusUpdateEvent, TaskStatus, TaskState
|
||||
event = TaskStatusUpdateEvent(
|
||||
task_id="task-123",
|
||||
status=TaskStatus(state=TaskState.WORKING),
|
||||
final=False
|
||||
)
|
||||
assert event.task_id == "task-123"
|
||||
assert event.final is False
|
||||
|
||||
def test_task_artifact_update_event(self):
|
||||
"""Test TaskArtifactUpdateEvent creation."""
|
||||
from praisonaiagents.ui.a2a.types import TaskArtifactUpdateEvent, Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Partial result")]
|
||||
)
|
||||
event = TaskArtifactUpdateEvent(
|
||||
task_id="task-456",
|
||||
artifact=artifact,
|
||||
append=True,
|
||||
last_chunk=False
|
||||
)
|
||||
assert event.append is True
|
||||
assert event.last_chunk is False
|
||||
@@ -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 @@
|
||||
"""Tests for AG-UI integration."""
|
||||
@@ -0,0 +1,188 @@
|
||||
"""
|
||||
Test AGUI Class - TDD Tests for Main Interface Class
|
||||
|
||||
Phase 5: AGUI Class Tests
|
||||
- Test AGUI initialization with Agent
|
||||
- Test AGUI initialization with PraisonAIAgents
|
||||
- Test get_router returns FastAPI router
|
||||
- Test validation errors
|
||||
"""
|
||||
|
||||
|
||||
class TestAGUIInitialization:
|
||||
"""Test AGUI class initialization."""
|
||||
|
||||
def test_agui_init_with_agent(self):
|
||||
"""Test AGUI initialization with single Agent."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
assert agui.agent == agent
|
||||
assert agui.agents is None
|
||||
|
||||
def test_agui_init_with_agents(self):
|
||||
"""Test AGUI initialization with PraisonAIAgents."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
task = Task(description="Test task", expected_output="Result", agent=agent)
|
||||
agents = PraisonAIAgents(agents=[agent], tasks=[task])
|
||||
|
||||
agui = AGUI(agents=agents)
|
||||
|
||||
assert agui.agents == agents
|
||||
assert agui.agent is None
|
||||
|
||||
def test_agui_requires_agent_or_agents(self):
|
||||
"""Test AGUI raises error if neither agent nor agents provided."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
import pytest
|
||||
|
||||
with pytest.raises(ValueError, match="requires an agent or agents"):
|
||||
AGUI()
|
||||
|
||||
def test_agui_custom_prefix(self):
|
||||
"""Test AGUI with custom prefix."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, prefix="/api/v1")
|
||||
|
||||
assert agui.prefix == "/api/v1"
|
||||
|
||||
def test_agui_custom_tags(self):
|
||||
"""Test AGUI with custom tags."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, tags=["Custom", "Tags"])
|
||||
|
||||
assert agui.tags == ["Custom", "Tags"]
|
||||
|
||||
def test_agui_default_tags(self):
|
||||
"""Test AGUI default tags."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
assert agui.tags == ["AGUI"]
|
||||
|
||||
|
||||
class TestAGUIRouter:
|
||||
"""Test AGUI router generation."""
|
||||
|
||||
def test_get_router_returns_api_router(self):
|
||||
"""Test get_router returns FastAPI APIRouter."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import APIRouter
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
assert isinstance(router, APIRouter)
|
||||
|
||||
def test_router_has_agui_endpoint(self):
|
||||
"""Test router has /agui POST endpoint."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
# Check routes
|
||||
route_paths = [route.path for route in router.routes]
|
||||
assert "/agui" in route_paths
|
||||
|
||||
def test_router_has_status_endpoint(self):
|
||||
"""Test router has /status GET endpoint."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
# Check routes
|
||||
route_paths = [route.path for route in router.routes]
|
||||
assert "/status" in route_paths
|
||||
|
||||
def test_router_with_prefix(self):
|
||||
"""Test router respects prefix."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, prefix="/api/v1")
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
assert router.prefix == "/api/v1"
|
||||
|
||||
|
||||
class TestAGUIIntegration:
|
||||
"""Test AGUI integration with FastAPI."""
|
||||
|
||||
def test_include_router_in_fastapi(self):
|
||||
"""Test including AGUI router in FastAPI app."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
# Check app has routes
|
||||
route_paths = [route.path for route in app.routes]
|
||||
assert "/agui" in route_paths
|
||||
assert "/status" in route_paths
|
||||
|
||||
|
||||
class TestAGUIName:
|
||||
"""Test AGUI name and description."""
|
||||
|
||||
def test_agui_name(self):
|
||||
"""Test AGUI name property."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, name="My Agent")
|
||||
|
||||
assert agui.name == "My Agent"
|
||||
|
||||
def test_agui_description(self):
|
||||
"""Test AGUI description property."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, description="A test agent")
|
||||
|
||||
assert agui.description == "A test agent"
|
||||
|
||||
def test_agui_default_name_from_agent(self):
|
||||
"""Test AGUI uses agent name by default."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="TestAgent", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
assert agui.name == "TestAgent"
|
||||
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
Test AG-UI Message Conversion - TDD Tests for Message Format Conversion
|
||||
|
||||
Phase 3: Message Conversion Tests
|
||||
- Test converting AG-UI messages to PraisonAI format
|
||||
- Test converting PraisonAI messages to AG-UI format
|
||||
- Test handling of tool calls
|
||||
- Test handling of tool results
|
||||
- Test edge cases (empty messages, missing fields)
|
||||
"""
|
||||
|
||||
class TestAGUIToPraisonAIConversion:
|
||||
"""Test converting AG-UI messages to PraisonAI format."""
|
||||
|
||||
def test_convert_user_message(self):
|
||||
"""Test converting a user message."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [Message(role="user", content="Hello")]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[0]["content"] == "Hello"
|
||||
|
||||
def test_convert_assistant_message(self):
|
||||
"""Test converting an assistant message."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [Message(role="assistant", content="Hi there!")]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "assistant"
|
||||
assert result[0]["content"] == "Hi there!"
|
||||
|
||||
def test_convert_tool_message(self):
|
||||
"""Test converting a tool result message."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [Message(role="tool", content="Search results", tool_call_id="tc-123")]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "tool"
|
||||
assert result[0]["content"] == "Search results"
|
||||
assert result[0]["tool_call_id"] == "tc-123"
|
||||
|
||||
def test_convert_assistant_with_tool_calls(self):
|
||||
"""Test converting assistant message with tool calls."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message, ToolCall, FunctionCall
|
||||
|
||||
agui_messages = [
|
||||
Message(
|
||||
role="assistant",
|
||||
content="",
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id="tc-123",
|
||||
function=FunctionCall(name="search", arguments='{"query": "test"}')
|
||||
)
|
||||
]
|
||||
)
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "assistant"
|
||||
assert "tool_calls" in result[0]
|
||||
assert len(result[0]["tool_calls"]) == 1
|
||||
assert result[0]["tool_calls"][0]["function"]["name"] == "search"
|
||||
|
||||
def test_convert_multiple_messages(self):
|
||||
"""Test converting multiple messages."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [
|
||||
Message(role="user", content="Hello"),
|
||||
Message(role="assistant", content="Hi!"),
|
||||
Message(role="user", content="How are you?"),
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[1]["role"] == "assistant"
|
||||
assert result[2]["role"] == "user"
|
||||
|
||||
def test_convert_empty_messages(self):
|
||||
"""Test converting empty message list."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
|
||||
result = agui_messages_to_praisonai([])
|
||||
assert result == []
|
||||
|
||||
def test_skip_system_messages(self):
|
||||
"""Test that system messages are skipped (agent builds its own)."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [
|
||||
Message(role="system", content="You are helpful"),
|
||||
Message(role="user", content="Hello"),
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
# System message should be skipped
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
|
||||
|
||||
class TestPraisonAIToAGUIConversion:
|
||||
"""Test converting PraisonAI messages to AG-UI format."""
|
||||
|
||||
def test_convert_user_message(self):
|
||||
"""Test converting a user message."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [{"role": "user", "content": "Hello"}]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "user"
|
||||
assert result[0].content == "Hello"
|
||||
|
||||
def test_convert_assistant_message(self):
|
||||
"""Test converting an assistant message."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [{"role": "assistant", "content": "Hi there!"}]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "assistant"
|
||||
assert result[0].content == "Hi there!"
|
||||
|
||||
def test_convert_tool_result(self):
|
||||
"""Test converting a tool result message."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [
|
||||
{"role": "tool", "content": "Search results", "tool_call_id": "tc-123"}
|
||||
]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "tool"
|
||||
assert result[0].content == "Search results"
|
||||
assert result[0].tool_call_id == "tc-123"
|
||||
|
||||
def test_convert_assistant_with_tool_calls(self):
|
||||
"""Test converting assistant message with tool calls."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "tc-123",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"arguments": '{"query": "test"}'
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "assistant"
|
||||
assert result[0].tool_calls is not None
|
||||
assert len(result[0].tool_calls) == 1
|
||||
assert result[0].tool_calls[0].function.name == "search"
|
||||
|
||||
def test_convert_empty_messages(self):
|
||||
"""Test converting empty message list."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
result = praisonai_messages_to_agui([])
|
||||
assert result == []
|
||||
|
||||
|
||||
class TestStateConversion:
|
||||
"""Test state conversion utilities."""
|
||||
|
||||
def test_validate_dict_state(self):
|
||||
"""Test validating dict state."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
|
||||
state = {"key": "value"}
|
||||
result = validate_state(state, "thread-123")
|
||||
|
||||
assert result == {"key": "value"}
|
||||
|
||||
def test_validate_none_state(self):
|
||||
"""Test validating None state."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
|
||||
result = validate_state(None, "thread-123")
|
||||
assert result is None
|
||||
|
||||
def test_validate_pydantic_state(self):
|
||||
"""Test validating Pydantic model state."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
from pydantic import BaseModel
|
||||
|
||||
class TestState(BaseModel):
|
||||
key: str = "value"
|
||||
|
||||
state = TestState()
|
||||
result = validate_state(state, "thread-123")
|
||||
|
||||
assert result == {"key": "value"}
|
||||
|
||||
def test_validate_invalid_state_returns_none(self):
|
||||
"""Test that invalid state returns None."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
|
||||
# A string is not a valid state
|
||||
result = validate_state("invalid", "thread-123")
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestToolCallDeduplication:
|
||||
"""Test tool call deduplication logic."""
|
||||
|
||||
def test_deduplicate_tool_results(self):
|
||||
"""Test that duplicate tool results are deduplicated."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
# Same tool_call_id appearing twice
|
||||
agui_messages = [
|
||||
Message(role="tool", content="Result 1", tool_call_id="tc-123"),
|
||||
Message(role="tool", content="Result 2", tool_call_id="tc-123"), # Duplicate
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
# Should only have one result
|
||||
assert len(result) == 1
|
||||
assert result[0]["content"] == "Result 1" # First one wins
|
||||
|
||||
|
||||
class TestExtractUserInput:
|
||||
"""Test extracting user input from messages."""
|
||||
|
||||
def test_extract_last_user_message(self):
|
||||
"""Test extracting the last user message."""
|
||||
from praisonaiagents.ui.agui.conversion import extract_user_input
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
messages = [
|
||||
Message(role="user", content="First message"),
|
||||
Message(role="assistant", content="Response"),
|
||||
Message(role="user", content="Last message"),
|
||||
]
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == "Last message"
|
||||
|
||||
def test_extract_from_empty_messages(self):
|
||||
"""Test extracting from empty messages."""
|
||||
from praisonaiagents.ui.agui.conversion import extract_user_input
|
||||
|
||||
result = extract_user_input([])
|
||||
assert result == ""
|
||||
|
||||
def test_extract_no_user_messages(self):
|
||||
"""Test extracting when no user messages exist."""
|
||||
from praisonaiagents.ui.agui.conversion import extract_user_input
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
messages = [
|
||||
Message(role="assistant", content="Hello"),
|
||||
]
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == ""
|
||||
@@ -0,0 +1,311 @@
|
||||
"""
|
||||
Test AG-UI Integration - TDD Tests for End-to-End Integration
|
||||
|
||||
Phase 7: Integration Tests
|
||||
- Test full flow with single Agent
|
||||
- Test full flow with PraisonAIAgents
|
||||
- Test tool calling flow
|
||||
- Test error recovery
|
||||
- Test state persistence
|
||||
"""
|
||||
|
||||
class TestSingleAgentIntegration:
|
||||
"""Test integration with single Agent."""
|
||||
|
||||
def test_single_agent_chat_flow(self):
|
||||
"""Test complete chat flow with single agent."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
role="Helpful Assistant",
|
||||
goal="Help users with their questions",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Send a message
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Say hello"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
# Parse SSE events
|
||||
events = []
|
||||
for line in response.text.split("\n"):
|
||||
if line.startswith("data: "):
|
||||
event_data = json.loads(line[6:])
|
||||
events.append(event_data)
|
||||
|
||||
# Should have run started event
|
||||
event_types = [e.get("type") for e in events]
|
||||
assert "RUN_STARTED" in event_types
|
||||
assert "RUN_FINISHED" in event_types or "RUN_ERROR" in event_types
|
||||
|
||||
|
||||
class TestMultiAgentIntegration:
|
||||
"""Test integration with PraisonAIAgents."""
|
||||
|
||||
def test_multi_agent_workflow_flow(self):
|
||||
"""Test complete workflow with multiple agents."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
# Create agents
|
||||
researcher = Agent(
|
||||
name="Researcher",
|
||||
role="Research Analyst",
|
||||
goal="Research topics thoroughly",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
writer = Agent(
|
||||
name="Writer",
|
||||
role="Content Writer",
|
||||
goal="Write engaging content",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create tasks
|
||||
research_task = Task(
|
||||
description="Research the topic: {topic}",
|
||||
expected_output="Research findings",
|
||||
agent=researcher
|
||||
)
|
||||
write_task = Task(
|
||||
description="Write about the research",
|
||||
expected_output="Written article",
|
||||
agent=writer
|
||||
)
|
||||
|
||||
# Create PraisonAIAgents
|
||||
agents = PraisonAIAgents(
|
||||
agents=[researcher, writer],
|
||||
tasks=[research_task, write_task]
|
||||
)
|
||||
|
||||
agui = AGUI(agents=agents)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Research AI"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestToolCallingIntegration:
|
||||
"""Test integration with tool calling."""
|
||||
|
||||
def test_agent_with_tools(self):
|
||||
"""Test agent with tools integration."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
def search_tool(query: str) -> str:
|
||||
"""Search for information."""
|
||||
return f"Search results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="SearchAgent",
|
||||
role="Search Assistant",
|
||||
goal="Help users search for information",
|
||||
tools=[search_tool],
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Search for Python tutorials"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
# Parse events
|
||||
events = []
|
||||
for line in response.text.split("\n"):
|
||||
if line.startswith("data: "):
|
||||
event_data = json.loads(line[6:])
|
||||
events.append(event_data)
|
||||
|
||||
# Check for tool call events (may or may not have them depending on LLM response)
|
||||
event_types = [e.get("type") for e in events]
|
||||
assert "RUN_STARTED" in event_types
|
||||
|
||||
|
||||
class TestStateManagement:
|
||||
"""Test state management in AG-UI."""
|
||||
|
||||
def test_state_passed_to_agent(self):
|
||||
"""Test that state is passed to agent."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(
|
||||
name="StatefulAgent",
|
||||
role="Stateful Assistant",
|
||||
goal="Remember context",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"state": {"user_name": "John", "preferences": {"language": "en"}}
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestConversationHistory:
|
||||
"""Test conversation history handling."""
|
||||
|
||||
def test_multiple_messages_in_history(self):
|
||||
"""Test handling multiple messages in conversation history."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(
|
||||
name="ChatAgent",
|
||||
role="Chat Assistant",
|
||||
goal="Have conversations",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [
|
||||
{"role": "user", "content": "My name is John"},
|
||||
{"role": "assistant", "content": "Hello John!"},
|
||||
{"role": "user", "content": "What is my name?"}
|
||||
]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestErrorRecovery:
|
||||
"""Test error recovery scenarios."""
|
||||
|
||||
def test_graceful_error_handling(self):
|
||||
"""Test graceful error handling."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
# Agent with invalid LLM should handle error gracefully
|
||||
agent = Agent(
|
||||
name="ErrorAgent",
|
||||
role="Error Test",
|
||||
goal="Test error handling",
|
||||
llm="invalid-model-name"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should not crash, should return error event
|
||||
assert response.status_code == 200
|
||||
|
||||
# Parse events
|
||||
events = []
|
||||
for line in response.text.split("\n"):
|
||||
if line.startswith("data: "):
|
||||
try:
|
||||
event_data = json.loads(line[6:])
|
||||
events.append(event_data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Should have error event
|
||||
event_types = [e.get("type") for e in events]
|
||||
assert "RUN_ERROR" in event_types or "RUN_FINISHED" in event_types
|
||||
|
||||
|
||||
class TestPerformance:
|
||||
"""Test performance - ensure no impact on existing functionality."""
|
||||
|
||||
def test_agui_import_does_not_slow_package(self):
|
||||
"""Test that importing AGUI doesn't slow down package import."""
|
||||
import time
|
||||
|
||||
start = time.time()
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
import_time = time.time() - start
|
||||
|
||||
# Import should be fast (< 1 second)
|
||||
assert import_time < 1.0
|
||||
|
||||
def test_agui_lazy_loading(self):
|
||||
"""Test that AGUI uses lazy loading where possible."""
|
||||
# This test ensures we don't import heavy dependencies until needed
|
||||
import sys
|
||||
|
||||
# Remove agui from cache if present
|
||||
modules_to_remove = [k for k in sys.modules.keys() if 'agui' in k]
|
||||
for mod in modules_to_remove:
|
||||
del sys.modules[mod]
|
||||
|
||||
# Import should not load FastAPI until get_router is called
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
# FastAPI should be loaded when get_router is called
|
||||
router = agui.get_router()
|
||||
assert router is not None
|
||||
@@ -0,0 +1,266 @@
|
||||
"""
|
||||
Test AG-UI Router - TDD Tests for FastAPI HTTP Endpoints
|
||||
|
||||
Phase 6: FastAPI Router Tests
|
||||
- Test /agui POST endpoint
|
||||
- Test /status GET endpoint
|
||||
- Test request validation
|
||||
- Test streaming response
|
||||
- Test error handling
|
||||
"""
|
||||
|
||||
class TestStatusEndpoint:
|
||||
"""Test /status endpoint."""
|
||||
|
||||
def test_status_endpoint_returns_available(self):
|
||||
"""Test status endpoint returns available status."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"status": "available"}
|
||||
|
||||
|
||||
class TestAGUIEndpoint:
|
||||
"""Test /agui POST endpoint."""
|
||||
|
||||
def test_agui_endpoint_accepts_post(self):
|
||||
"""Test /agui endpoint accepts POST requests."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Minimal valid request
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should return streaming response (200)
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_agui_endpoint_returns_streaming_response(self):
|
||||
"""Test /agui endpoint returns streaming response."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Check content type is event-stream
|
||||
assert "text/event-stream" in response.headers.get("content-type", "")
|
||||
|
||||
def test_agui_endpoint_requires_thread_id(self):
|
||||
"""Test /agui endpoint requires thread_id."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Missing thread_id
|
||||
response = client.post("/agui", json={
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should return validation error
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
class TestAGUIEndpointWithState:
|
||||
"""Test /agui endpoint with state."""
|
||||
|
||||
def test_agui_endpoint_accepts_state(self):
|
||||
"""Test /agui endpoint accepts state parameter."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"state": {"custom_key": "custom_value"}
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestAGUIEndpointWithRunId:
|
||||
"""Test /agui endpoint with run_id."""
|
||||
|
||||
def test_agui_endpoint_accepts_run_id(self):
|
||||
"""Test /agui endpoint accepts run_id parameter."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"run_id": "run-456",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_agui_endpoint_generates_run_id_if_missing(self):
|
||||
"""Test /agui endpoint generates run_id if not provided."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should succeed even without run_id
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestAGUIEndpointCORS:
|
||||
"""Test /agui endpoint CORS headers."""
|
||||
|
||||
def test_agui_endpoint_cors_headers(self):
|
||||
"""Test /agui endpoint includes CORS headers."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Check CORS headers
|
||||
assert response.headers.get("access-control-allow-origin") == "*"
|
||||
|
||||
|
||||
class TestAGUIEndpointErrorHandling:
|
||||
"""Test /agui endpoint error handling."""
|
||||
|
||||
def test_agui_endpoint_handles_agent_error(self):
|
||||
"""Test /agui endpoint handles agent errors gracefully."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create agent that will fail
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# This should not crash the server
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should return 200 with error event in stream
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestAGUIEndpointWithForwardedProps:
|
||||
"""Test /agui endpoint with forwarded_props."""
|
||||
|
||||
def test_agui_endpoint_accepts_forwarded_props(self):
|
||||
"""Test /agui endpoint accepts forwarded_props."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"forwarded_props": {"user_id": "user-123"}
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
@@ -0,0 +1,240 @@
|
||||
"""
|
||||
Test AG-UI Event Streaming - TDD Tests for Streaming PraisonAI Responses
|
||||
|
||||
Phase 4: Event Streaming Tests
|
||||
- Test streaming text content as AG-UI events
|
||||
- Test streaming tool calls as AG-UI events
|
||||
- Test run lifecycle events (start, finish, error)
|
||||
- Test step events
|
||||
- Test event ordering
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
|
||||
class TestTextMessageStreaming:
|
||||
"""Test streaming text messages."""
|
||||
|
||||
def test_stream_simple_text(self):
|
||||
"""Test streaming simple text content."""
|
||||
from praisonaiagents.ui.agui.streaming import create_text_message_events
|
||||
|
||||
events = list(create_text_message_events(
|
||||
content="Hello, world!",
|
||||
message_id="msg-123"
|
||||
))
|
||||
|
||||
# Should have start, content, end events
|
||||
assert len(events) == 3
|
||||
assert events[0].type == "TEXT_MESSAGE_START"
|
||||
assert events[0].message_id == "msg-123"
|
||||
assert events[1].type == "TEXT_MESSAGE_CONTENT"
|
||||
assert events[1].delta == "Hello, world!"
|
||||
assert events[2].type == "TEXT_MESSAGE_END"
|
||||
|
||||
def test_stream_chunked_text(self):
|
||||
"""Test streaming text in chunks."""
|
||||
from praisonaiagents.ui.agui.streaming import stream_text_chunks
|
||||
|
||||
chunks = ["Hello", ", ", "world", "!"]
|
||||
events = list(stream_text_chunks(chunks, message_id="msg-123"))
|
||||
|
||||
# Should have start, 4 content events, end
|
||||
assert len(events) == 6
|
||||
assert events[0].type == "TEXT_MESSAGE_START"
|
||||
assert events[1].delta == "Hello"
|
||||
assert events[2].delta == ", "
|
||||
assert events[3].delta == "world"
|
||||
assert events[4].delta == "!"
|
||||
assert events[5].type == "TEXT_MESSAGE_END"
|
||||
|
||||
def test_stream_empty_text(self):
|
||||
"""Test streaming empty text."""
|
||||
from praisonaiagents.ui.agui.streaming import create_text_message_events
|
||||
|
||||
events = list(create_text_message_events(content="", message_id="msg-123"))
|
||||
|
||||
# Should still have start and end, but no content
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "TEXT_MESSAGE_START"
|
||||
assert events[1].type == "TEXT_MESSAGE_END"
|
||||
|
||||
|
||||
class TestToolCallStreaming:
|
||||
"""Test streaming tool calls."""
|
||||
|
||||
def test_stream_tool_call(self):
|
||||
"""Test streaming a tool call."""
|
||||
from praisonaiagents.ui.agui.streaming import create_tool_call_events
|
||||
|
||||
events = list(create_tool_call_events(
|
||||
tool_call_id="tc-123",
|
||||
tool_name="search",
|
||||
arguments='{"query": "test"}',
|
||||
parent_message_id="msg-123"
|
||||
))
|
||||
|
||||
# Should have start, args, end events
|
||||
assert len(events) == 3
|
||||
assert events[0].type == "TOOL_CALL_START"
|
||||
assert events[0].tool_call_id == "tc-123"
|
||||
assert events[0].tool_call_name == "search"
|
||||
assert events[1].type == "TOOL_CALL_ARGS"
|
||||
assert events[1].delta == '{"query": "test"}'
|
||||
assert events[2].type == "TOOL_CALL_END"
|
||||
|
||||
def test_stream_tool_result(self):
|
||||
"""Test streaming a tool result."""
|
||||
from praisonaiagents.ui.agui.streaming import create_tool_result_event
|
||||
|
||||
event = create_tool_result_event(
|
||||
tool_call_id="tc-123",
|
||||
content="Search results...",
|
||||
message_id="msg-456"
|
||||
)
|
||||
|
||||
assert event.type == "TOOL_CALL_RESULT"
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.content == "Search results..."
|
||||
|
||||
|
||||
class TestRunLifecycleEvents:
|
||||
"""Test run lifecycle events."""
|
||||
|
||||
def test_create_run_started_event(self):
|
||||
"""Test creating run started event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_run_started_event
|
||||
|
||||
event = create_run_started_event(thread_id="thread-123", run_id="run-456")
|
||||
|
||||
assert event.type == "RUN_STARTED"
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
|
||||
def test_create_run_finished_event(self):
|
||||
"""Test creating run finished event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_run_finished_event
|
||||
|
||||
event = create_run_finished_event(thread_id="thread-123", run_id="run-456")
|
||||
|
||||
assert event.type == "RUN_FINISHED"
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
|
||||
def test_create_run_error_event(self):
|
||||
"""Test creating run error event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_run_error_event
|
||||
|
||||
event = create_run_error_event(message="Something went wrong", code="ERR001")
|
||||
|
||||
assert event.type == "RUN_ERROR"
|
||||
assert event.message == "Something went wrong"
|
||||
assert event.code == "ERR001"
|
||||
|
||||
|
||||
class TestStepEvents:
|
||||
"""Test step events."""
|
||||
|
||||
def test_create_step_started_event(self):
|
||||
"""Test creating step started event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_step_started_event
|
||||
|
||||
event = create_step_started_event(step_name="research")
|
||||
|
||||
assert event.type == "STEP_STARTED"
|
||||
assert event.step_name == "research"
|
||||
|
||||
def test_create_step_finished_event(self):
|
||||
"""Test creating step finished event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_step_finished_event
|
||||
|
||||
event = create_step_finished_event(step_name="research")
|
||||
|
||||
assert event.type == "STEP_FINISHED"
|
||||
assert event.step_name == "research"
|
||||
|
||||
|
||||
class TestEventBuffer:
|
||||
"""Test EventBuffer for managing event ordering."""
|
||||
|
||||
def test_event_buffer_creation(self):
|
||||
"""Test EventBuffer creation."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
assert buffer.active_tool_call_ids == set()
|
||||
assert buffer.current_text_message_id == ""
|
||||
|
||||
def test_event_buffer_start_tool_call(self):
|
||||
"""Test starting a tool call in buffer."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
buffer.start_tool_call("tc-123")
|
||||
|
||||
assert "tc-123" in buffer.active_tool_call_ids
|
||||
|
||||
def test_event_buffer_end_tool_call(self):
|
||||
"""Test ending a tool call in buffer."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
buffer.start_tool_call("tc-123")
|
||||
buffer.end_tool_call("tc-123")
|
||||
|
||||
assert "tc-123" not in buffer.active_tool_call_ids
|
||||
assert "tc-123" in buffer.ended_tool_call_ids
|
||||
|
||||
def test_event_buffer_start_text_message(self):
|
||||
"""Test starting a text message in buffer."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
msg_id = buffer.start_text_message()
|
||||
|
||||
assert msg_id != ""
|
||||
assert buffer.current_text_message_id == msg_id
|
||||
|
||||
|
||||
class TestAsyncStreaming:
|
||||
"""Test async streaming functionality."""
|
||||
|
||||
def test_async_stream_response(self):
|
||||
"""Test async streaming of agent response."""
|
||||
async def run_test():
|
||||
from praisonaiagents.ui.agui.streaming import async_stream_response
|
||||
|
||||
# Mock response chunks
|
||||
async def mock_chunks():
|
||||
yield "Hello"
|
||||
yield ", "
|
||||
yield "world!"
|
||||
|
||||
events = []
|
||||
async for event in async_stream_response(
|
||||
response_stream=mock_chunks(),
|
||||
thread_id="thread-123",
|
||||
run_id="run-456"
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
# Should have run_started, text events, run_finished
|
||||
assert len(events) > 0
|
||||
assert events[0].type == "RUN_STARTED"
|
||||
assert events[-1].type == "RUN_FINISHED"
|
||||
|
||||
asyncio.run(run_test())
|
||||
|
||||
|
||||
class TestStateSnapshotEvents:
|
||||
"""Test state snapshot events."""
|
||||
|
||||
def test_create_state_snapshot_event(self):
|
||||
"""Test creating state snapshot event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_state_snapshot_event
|
||||
|
||||
state = {"task": "research", "progress": 50}
|
||||
event = create_state_snapshot_event(state)
|
||||
|
||||
assert event.type == "STATE_SNAPSHOT"
|
||||
assert event.snapshot == state
|
||||
@@ -0,0 +1,327 @@
|
||||
"""
|
||||
Test AG-UI Types - TDD Tests for Event and Message Types
|
||||
|
||||
Phase 2: Core Types Tests
|
||||
- Test EventType enum values
|
||||
- Test BaseEvent creation
|
||||
- Test all event type classes
|
||||
- Test message type classes
|
||||
- Test RunAgentInput validation
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class TestEventType:
|
||||
"""Test EventType enum."""
|
||||
|
||||
def test_event_type_text_message_start(self):
|
||||
"""Test TEXT_MESSAGE_START event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TEXT_MESSAGE_START == "TEXT_MESSAGE_START"
|
||||
|
||||
def test_event_type_text_message_content(self):
|
||||
"""Test TEXT_MESSAGE_CONTENT event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TEXT_MESSAGE_CONTENT == "TEXT_MESSAGE_CONTENT"
|
||||
|
||||
def test_event_type_text_message_end(self):
|
||||
"""Test TEXT_MESSAGE_END event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TEXT_MESSAGE_END == "TEXT_MESSAGE_END"
|
||||
|
||||
def test_event_type_tool_call_start(self):
|
||||
"""Test TOOL_CALL_START event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_START == "TOOL_CALL_START"
|
||||
|
||||
def test_event_type_tool_call_args(self):
|
||||
"""Test TOOL_CALL_ARGS event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_ARGS == "TOOL_CALL_ARGS"
|
||||
|
||||
def test_event_type_tool_call_end(self):
|
||||
"""Test TOOL_CALL_END event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_END == "TOOL_CALL_END"
|
||||
|
||||
def test_event_type_tool_call_result(self):
|
||||
"""Test TOOL_CALL_RESULT event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_RESULT == "TOOL_CALL_RESULT"
|
||||
|
||||
def test_event_type_run_started(self):
|
||||
"""Test RUN_STARTED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.RUN_STARTED == "RUN_STARTED"
|
||||
|
||||
def test_event_type_run_finished(self):
|
||||
"""Test RUN_FINISHED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.RUN_FINISHED == "RUN_FINISHED"
|
||||
|
||||
def test_event_type_run_error(self):
|
||||
"""Test RUN_ERROR event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.RUN_ERROR == "RUN_ERROR"
|
||||
|
||||
def test_event_type_step_started(self):
|
||||
"""Test STEP_STARTED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.STEP_STARTED == "STEP_STARTED"
|
||||
|
||||
def test_event_type_step_finished(self):
|
||||
"""Test STEP_FINISHED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.STEP_FINISHED == "STEP_FINISHED"
|
||||
|
||||
def test_event_type_state_snapshot(self):
|
||||
"""Test STATE_SNAPSHOT event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.STATE_SNAPSHOT == "STATE_SNAPSHOT"
|
||||
|
||||
|
||||
class TestBaseEvent:
|
||||
"""Test BaseEvent class."""
|
||||
|
||||
def test_base_event_creation(self):
|
||||
"""Test BaseEvent can be created."""
|
||||
from praisonaiagents.ui.agui.types import BaseEvent, EventType
|
||||
event = BaseEvent(type=EventType.RUN_STARTED)
|
||||
assert event.type == EventType.RUN_STARTED
|
||||
|
||||
def test_base_event_timestamp_optional(self):
|
||||
"""Test BaseEvent timestamp is optional."""
|
||||
from praisonaiagents.ui.agui.types import BaseEvent, EventType
|
||||
event = BaseEvent(type=EventType.RUN_STARTED)
|
||||
assert event.timestamp is None
|
||||
|
||||
def test_base_event_to_dict(self):
|
||||
"""Test BaseEvent can be converted to dict."""
|
||||
from praisonaiagents.ui.agui.types import BaseEvent, EventType
|
||||
event = BaseEvent(type=EventType.RUN_STARTED)
|
||||
data = event.model_dump()
|
||||
assert data["type"] == "RUN_STARTED"
|
||||
|
||||
|
||||
class TestTextMessageEvents:
|
||||
"""Test text message event classes."""
|
||||
|
||||
def test_text_message_start_event(self):
|
||||
"""Test TextMessageStartEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageStartEvent
|
||||
event = TextMessageStartEvent(message_id="msg-123", role="assistant")
|
||||
assert event.message_id == "msg-123"
|
||||
assert event.role == "assistant"
|
||||
assert event.type == "TEXT_MESSAGE_START"
|
||||
|
||||
def test_text_message_content_event(self):
|
||||
"""Test TextMessageContentEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageContentEvent
|
||||
event = TextMessageContentEvent(message_id="msg-123", delta="Hello")
|
||||
assert event.message_id == "msg-123"
|
||||
assert event.delta == "Hello"
|
||||
assert event.type == "TEXT_MESSAGE_CONTENT"
|
||||
|
||||
def test_text_message_end_event(self):
|
||||
"""Test TextMessageEndEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageEndEvent
|
||||
event = TextMessageEndEvent(message_id="msg-123")
|
||||
assert event.message_id == "msg-123"
|
||||
assert event.type == "TEXT_MESSAGE_END"
|
||||
|
||||
|
||||
class TestToolCallEvents:
|
||||
"""Test tool call event classes."""
|
||||
|
||||
def test_tool_call_start_event(self):
|
||||
"""Test ToolCallStartEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallStartEvent
|
||||
event = ToolCallStartEvent(
|
||||
tool_call_id="tc-123",
|
||||
tool_call_name="search",
|
||||
parent_message_id="msg-123"
|
||||
)
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.tool_call_name == "search"
|
||||
assert event.parent_message_id == "msg-123"
|
||||
assert event.type == "TOOL_CALL_START"
|
||||
|
||||
def test_tool_call_args_event(self):
|
||||
"""Test ToolCallArgsEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallArgsEvent
|
||||
event = ToolCallArgsEvent(tool_call_id="tc-123", delta='{"query": "test"}')
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.delta == '{"query": "test"}'
|
||||
assert event.type == "TOOL_CALL_ARGS"
|
||||
|
||||
def test_tool_call_end_event(self):
|
||||
"""Test ToolCallEndEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallEndEvent
|
||||
event = ToolCallEndEvent(tool_call_id="tc-123")
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.type == "TOOL_CALL_END"
|
||||
|
||||
def test_tool_call_result_event(self):
|
||||
"""Test ToolCallResultEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallResultEvent
|
||||
event = ToolCallResultEvent(
|
||||
message_id="msg-456",
|
||||
tool_call_id="tc-123",
|
||||
content="Search results..."
|
||||
)
|
||||
assert event.message_id == "msg-456"
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.content == "Search results..."
|
||||
assert event.type == "TOOL_CALL_RESULT"
|
||||
|
||||
|
||||
class TestRunEvents:
|
||||
"""Test run lifecycle event classes."""
|
||||
|
||||
def test_run_started_event(self):
|
||||
"""Test RunStartedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import RunStartedEvent
|
||||
event = RunStartedEvent(thread_id="thread-123", run_id="run-456")
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
assert event.type == "RUN_STARTED"
|
||||
|
||||
def test_run_finished_event(self):
|
||||
"""Test RunFinishedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import RunFinishedEvent
|
||||
event = RunFinishedEvent(thread_id="thread-123", run_id="run-456")
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
assert event.type == "RUN_FINISHED"
|
||||
|
||||
def test_run_error_event(self):
|
||||
"""Test RunErrorEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import RunErrorEvent
|
||||
event = RunErrorEvent(message="Something went wrong")
|
||||
assert event.message == "Something went wrong"
|
||||
assert event.type == "RUN_ERROR"
|
||||
|
||||
|
||||
class TestStepEvents:
|
||||
"""Test step event classes."""
|
||||
|
||||
def test_step_started_event(self):
|
||||
"""Test StepStartedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import StepStartedEvent
|
||||
event = StepStartedEvent(step_name="research")
|
||||
assert event.step_name == "research"
|
||||
assert event.type == "STEP_STARTED"
|
||||
|
||||
def test_step_finished_event(self):
|
||||
"""Test StepFinishedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import StepFinishedEvent
|
||||
event = StepFinishedEvent(step_name="research")
|
||||
assert event.step_name == "research"
|
||||
assert event.type == "STEP_FINISHED"
|
||||
|
||||
|
||||
class TestStateEvents:
|
||||
"""Test state event classes."""
|
||||
|
||||
def test_state_snapshot_event(self):
|
||||
"""Test StateSnapshotEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import StateSnapshotEvent
|
||||
event = StateSnapshotEvent(snapshot={"key": "value"})
|
||||
assert event.snapshot == {"key": "value"}
|
||||
assert event.type == "STATE_SNAPSHOT"
|
||||
|
||||
|
||||
class TestMessageTypes:
|
||||
"""Test message type classes."""
|
||||
|
||||
def test_user_message(self):
|
||||
"""Test UserMessage creation."""
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
msg = Message(role="user", content="Hello")
|
||||
assert msg.role == "user"
|
||||
assert msg.content == "Hello"
|
||||
|
||||
def test_assistant_message(self):
|
||||
"""Test AssistantMessage creation."""
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
msg = Message(role="assistant", content="Hi there!")
|
||||
assert msg.role == "assistant"
|
||||
assert msg.content == "Hi there!"
|
||||
|
||||
def test_tool_message(self):
|
||||
"""Test ToolMessage creation."""
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
msg = Message(role="tool", content="Result", tool_call_id="tc-123")
|
||||
assert msg.role == "tool"
|
||||
assert msg.content == "Result"
|
||||
assert msg.tool_call_id == "tc-123"
|
||||
|
||||
|
||||
class TestToolCall:
|
||||
"""Test ToolCall type."""
|
||||
|
||||
def test_tool_call_creation(self):
|
||||
"""Test ToolCall creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCall, FunctionCall
|
||||
tc = ToolCall(
|
||||
id="tc-123",
|
||||
function=FunctionCall(name="search", arguments='{"query": "test"}')
|
||||
)
|
||||
assert tc.id == "tc-123"
|
||||
assert tc.function.name == "search"
|
||||
assert tc.function.arguments == '{"query": "test"}'
|
||||
|
||||
|
||||
class TestRunAgentInput:
|
||||
"""Test RunAgentInput type."""
|
||||
|
||||
def test_run_agent_input_creation(self):
|
||||
"""Test RunAgentInput creation."""
|
||||
from praisonaiagents.ui.agui.types import RunAgentInput, Message
|
||||
input_data = RunAgentInput(
|
||||
thread_id="thread-123",
|
||||
run_id="run-456",
|
||||
messages=[Message(role="user", content="Hello")]
|
||||
)
|
||||
assert input_data.thread_id == "thread-123"
|
||||
assert input_data.run_id == "run-456"
|
||||
assert len(input_data.messages) == 1
|
||||
|
||||
def test_run_agent_input_optional_fields(self):
|
||||
"""Test RunAgentInput optional fields."""
|
||||
from praisonaiagents.ui.agui.types import RunAgentInput
|
||||
input_data = RunAgentInput(thread_id="thread-123")
|
||||
assert input_data.thread_id == "thread-123"
|
||||
assert input_data.run_id is None
|
||||
assert input_data.messages is None or len(input_data.messages) == 0
|
||||
assert input_data.state is None
|
||||
|
||||
|
||||
class TestEventSerialization:
|
||||
"""Test event serialization to JSON."""
|
||||
|
||||
def test_event_json_serialization(self):
|
||||
"""Test events can be serialized to JSON."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageStartEvent
|
||||
event = TextMessageStartEvent(message_id="msg-123", role="assistant")
|
||||
json_str = event.model_dump_json()
|
||||
data = json.loads(json_str)
|
||||
assert data["type"] == "TEXT_MESSAGE_START"
|
||||
assert data["message_id"] == "msg-123"
|
||||
|
||||
def test_event_sse_format(self):
|
||||
"""Test events can be formatted as SSE."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageStartEvent
|
||||
from praisonaiagents.ui.agui.encoder import EventEncoder
|
||||
|
||||
encoder = EventEncoder()
|
||||
event = TextMessageStartEvent(message_id="msg-123", role="assistant")
|
||||
sse = encoder.encode(event)
|
||||
|
||||
assert sse.startswith("data: ")
|
||||
assert sse.endswith("\n\n")
|
||||
# Parse the JSON part
|
||||
json_part = sse[6:-2] # Remove "data: " prefix and "\n\n" suffix
|
||||
data = json.loads(json_part)
|
||||
assert data["type"] == "TEXT_MESSAGE_START"
|
||||
@@ -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.14"
|
||||
version = "0.1.17"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
+214
-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
|
||||
@@ -1628,36 +1755,78 @@ server.run() # stdio for Claude Desktop
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
## A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
## CLI Features
|
||||
|
||||
| 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) |
|
||||
|
||||
---
|
||||
|
||||
@@ -2209,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
|
||||
|
||||
@@ -2218,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
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
# PraisonAI CLI - Benchmark Results
|
||||
|
||||
**Generated:** 2025-12-19 16:04:40
|
||||
|
||||
## Import Times
|
||||
|
||||
| Module | Time (ms) | Target | Status |
|
||||
|--------|-----------|--------|--------|
|
||||
| praisonai | 4.11 | <10 | ✅ |
|
||||
| message_queue | 6.87 | <50 | ✅ |
|
||||
| at_mentions | 0.58 | <50 | ✅ |
|
||||
| profiler | 0.90 | <10 | ✅ |
|
||||
|
||||
## Instantiation Times
|
||||
|
||||
| Class | Time (μs) | Target | Status |
|
||||
|-------|-----------|--------|--------|
|
||||
| MessageQueue | 0.26 | <100 | ✅ |
|
||||
| StateManager | 0.30 | <100 | ✅ |
|
||||
| FileSearchService | 8.42 | <100 | ✅ |
|
||||
|
||||
## Lazy Loading
|
||||
|
||||
| Operation | Time (ms) |
|
||||
|-----------|----------|
|
||||
| PraisonAI lazy load | 3886.47 |
|
||||
|
||||
## How to Reproduce
|
||||
|
||||
```bash
|
||||
cd praisonai
|
||||
python benchmarks/simple_benchmark.py
|
||||
```
|
||||
@@ -0,0 +1,231 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PraisonAI CLI - Performance Benchmark
|
||||
|
||||
Measures import times and module loading performance.
|
||||
|
||||
Usage:
|
||||
python benchmarks/simple_benchmark.py
|
||||
python benchmarks/simple_benchmark.py --no-save
|
||||
"""
|
||||
|
||||
import time
|
||||
import sys
|
||||
import os
|
||||
import argparse
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
ITERATIONS = 100
|
||||
|
||||
|
||||
def measure_import(module_name: str) -> float:
|
||||
"""Measure import time in milliseconds."""
|
||||
# Clear from cache if present
|
||||
modules_to_remove = [k for k in sys.modules.keys() if k.startswith(module_name)]
|
||||
for mod in modules_to_remove:
|
||||
del sys.modules[mod]
|
||||
|
||||
start = time.perf_counter()
|
||||
__import__(module_name)
|
||||
return (time.perf_counter() - start) * 1000
|
||||
|
||||
|
||||
def measure_instantiation(create_fn, iterations=ITERATIONS) -> float:
|
||||
"""Measure average instantiation time in microseconds."""
|
||||
times = []
|
||||
for _ in range(iterations):
|
||||
start = time.perf_counter()
|
||||
create_fn()
|
||||
times.append((time.perf_counter() - start) * 1_000_000)
|
||||
return sum(times) / len(times)
|
||||
|
||||
|
||||
def run_benchmark():
|
||||
"""Run the benchmark."""
|
||||
results = {}
|
||||
|
||||
print('=' * 60)
|
||||
print('PraisonAI CLI - Performance Benchmark')
|
||||
print('=' * 60)
|
||||
print(f'\nIterations: {ITERATIONS}')
|
||||
print('Metrics: Import time (ms), Instantiation time (μs)\n')
|
||||
|
||||
# 1. Package import time
|
||||
print("1. Testing package imports...")
|
||||
|
||||
# Clear cache
|
||||
modules_to_remove = [k for k in sys.modules.keys() if 'praisonai' in k]
|
||||
for mod in modules_to_remove:
|
||||
del sys.modules[mod]
|
||||
|
||||
start = time.perf_counter()
|
||||
import praisonai
|
||||
results['praisonai_import_ms'] = (time.perf_counter() - start) * 1000
|
||||
print(f" import praisonai: {results['praisonai_import_ms']:.2f}ms")
|
||||
|
||||
# 2. CLI features import
|
||||
print("\n2. Testing CLI features imports...")
|
||||
|
||||
start = time.perf_counter()
|
||||
from praisonai.cli.features.message_queue import MessageQueue
|
||||
results['message_queue_import_ms'] = (time.perf_counter() - start) * 1000
|
||||
print(f" message_queue: {results['message_queue_import_ms']:.2f}ms")
|
||||
|
||||
start = time.perf_counter()
|
||||
from praisonai.cli.features.at_mentions import CombinedCompleter
|
||||
results['at_mentions_import_ms'] = (time.perf_counter() - start) * 1000
|
||||
print(f" at_mentions: {results['at_mentions_import_ms']:.2f}ms")
|
||||
|
||||
# 3. Profiler import (should be fast, standalone)
|
||||
print("\n3. Testing profiler import...")
|
||||
start = time.perf_counter()
|
||||
from praisonai.profiler import Profiler
|
||||
results['profiler_import_ms'] = (time.perf_counter() - start) * 1000
|
||||
print(f" profiler: {results['profiler_import_ms']:.2f}ms")
|
||||
|
||||
# 4. Instantiation benchmarks
|
||||
print(f"\n4. Testing instantiation ({ITERATIONS} iterations)...")
|
||||
|
||||
from praisonai.cli.features.message_queue import MessageQueue, StateManager
|
||||
from praisonai.cli.features.at_mentions import FileSearchService
|
||||
|
||||
results['MessageQueue_us'] = measure_instantiation(lambda: MessageQueue())
|
||||
print(f" MessageQueue: {results['MessageQueue_us']:.2f}μs")
|
||||
|
||||
results['StateManager_us'] = measure_instantiation(lambda: StateManager())
|
||||
print(f" StateManager: {results['StateManager_us']:.2f}μs")
|
||||
|
||||
results['FileSearchService_us'] = measure_instantiation(
|
||||
lambda: FileSearchService(root_dir=".")
|
||||
)
|
||||
print(f" FileSearchService: {results['FileSearchService_us']:.2f}μs")
|
||||
|
||||
# 5. PraisonAI class (lazy load)
|
||||
print("\n5. Testing PraisonAI class access (lazy load)...")
|
||||
start = time.perf_counter()
|
||||
PraisonAI = praisonai.PraisonAI
|
||||
results['PraisonAI_lazy_ms'] = (time.perf_counter() - start) * 1000
|
||||
print(f" PraisonAI lazy load: {results['PraisonAI_lazy_ms']:.2f}ms")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def print_summary(results: dict):
|
||||
"""Print benchmark summary."""
|
||||
print('\n' + '=' * 60)
|
||||
print('SUMMARY')
|
||||
print('=' * 60)
|
||||
|
||||
print('\nImport Times (target: <50ms for features, <10ms for package):')
|
||||
print(f" {'Module':<25} {'Time':>10} {'Status':>10}")
|
||||
print(f" {'-'*25} {'-'*10} {'-'*10}")
|
||||
|
||||
import_checks = [
|
||||
('praisonai', results['praisonai_import_ms'], 10),
|
||||
('message_queue', results['message_queue_import_ms'], 50),
|
||||
('at_mentions', results['at_mentions_import_ms'], 50),
|
||||
('profiler', results['profiler_import_ms'], 10),
|
||||
]
|
||||
|
||||
all_pass = True
|
||||
for name, time_ms, threshold in import_checks:
|
||||
status = '✅ PASS' if time_ms < threshold else '❌ SLOW'
|
||||
if time_ms >= threshold:
|
||||
all_pass = False
|
||||
print(f" {name:<25} {time_ms:>8.2f}ms {status:>10}")
|
||||
|
||||
print('\nInstantiation Times (target: <100μs):')
|
||||
print(f" {'Class':<25} {'Time':>10} {'Status':>10}")
|
||||
print(f" {'-'*25} {'-'*10} {'-'*10}")
|
||||
|
||||
inst_checks = [
|
||||
('MessageQueue', results['MessageQueue_us'], 100),
|
||||
('StateManager', results['StateManager_us'], 100),
|
||||
('FileSearchService', results['FileSearchService_us'], 100),
|
||||
]
|
||||
|
||||
for name, time_us, threshold in inst_checks:
|
||||
status = '✅ PASS' if time_us < threshold else '❌ SLOW'
|
||||
if time_us >= threshold:
|
||||
all_pass = False
|
||||
print(f" {name:<25} {time_us:>8.2f}μs {status:>10}")
|
||||
|
||||
print('\n' + '=' * 60)
|
||||
if all_pass:
|
||||
print('✅ All benchmarks PASSED')
|
||||
else:
|
||||
print('❌ Some benchmarks FAILED')
|
||||
print('=' * 60)
|
||||
|
||||
return all_pass
|
||||
|
||||
|
||||
def save_results(results: dict):
|
||||
"""Save benchmark results to markdown file."""
|
||||
filepath = os.path.join(os.path.dirname(__file__), 'BENCHMARK_RESULTS.md')
|
||||
|
||||
with open(filepath, 'w') as f:
|
||||
f.write('# PraisonAI CLI - Benchmark Results\n\n')
|
||||
f.write(f'**Generated:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n\n')
|
||||
|
||||
f.write('## Import Times\n\n')
|
||||
f.write('| Module | Time (ms) | Target | Status |\n')
|
||||
f.write('|--------|-----------|--------|--------|\n')
|
||||
|
||||
import_data = [
|
||||
('praisonai', results['praisonai_import_ms'], 10),
|
||||
('message_queue', results['message_queue_import_ms'], 50),
|
||||
('at_mentions', results['at_mentions_import_ms'], 50),
|
||||
('profiler', results['profiler_import_ms'], 10),
|
||||
]
|
||||
|
||||
for name, time_ms, threshold in import_data:
|
||||
status = '✅' if time_ms < threshold else '❌'
|
||||
f.write(f'| {name} | {time_ms:.2f} | <{threshold} | {status} |\n')
|
||||
|
||||
f.write('\n## Instantiation Times\n\n')
|
||||
f.write('| Class | Time (μs) | Target | Status |\n')
|
||||
f.write('|-------|-----------|--------|--------|\n')
|
||||
|
||||
inst_data = [
|
||||
('MessageQueue', results['MessageQueue_us'], 100),
|
||||
('StateManager', results['StateManager_us'], 100),
|
||||
('FileSearchService', results['FileSearchService_us'], 100),
|
||||
]
|
||||
|
||||
for name, time_us, threshold in inst_data:
|
||||
status = '✅' if time_us < threshold else '❌'
|
||||
f.write(f'| {name} | {time_us:.2f} | <{threshold} | {status} |\n')
|
||||
|
||||
f.write('\n## Lazy Loading\n\n')
|
||||
f.write(f'| Operation | Time (ms) |\n')
|
||||
f.write(f'|-----------|----------|\n')
|
||||
f.write(f'| PraisonAI lazy load | {results["PraisonAI_lazy_ms"]:.2f} |\n')
|
||||
|
||||
f.write('\n## How to Reproduce\n\n')
|
||||
f.write('```bash\n')
|
||||
f.write('cd praisonai\n')
|
||||
f.write('python benchmarks/simple_benchmark.py\n')
|
||||
f.write('```\n')
|
||||
|
||||
print(f'\nResults saved to: {filepath}')
|
||||
return filepath
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='PraisonAI CLI Performance Benchmark')
|
||||
parser.add_argument('--no-save', action='store_true', help='Do not save results to file')
|
||||
args = parser.parse_args()
|
||||
|
||||
results = run_benchmark()
|
||||
all_pass = print_summary(results)
|
||||
|
||||
if not args.no_save:
|
||||
save_results(results)
|
||||
|
||||
return 0 if all_pass else 1
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
@@ -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.35.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.35.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.56.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.56.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -1,31 +1,34 @@
|
||||
# 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"
|
||||
os.environ["EC_TELEMETRY"] = "false"
|
||||
from .cli import PraisonAI
|
||||
|
||||
# Version is lightweight, import directly
|
||||
from .version import __version__
|
||||
|
||||
# Re-export all classes from praisonaiagents to enable:
|
||||
# from praisonai import Agent, Task, PraisonAIAgents
|
||||
try:
|
||||
import praisonaiagents
|
||||
# Import all symbols from praisonaiagents using * import
|
||||
from praisonaiagents import *
|
||||
except ImportError:
|
||||
# If praisonaiagents is not available, these imports will fail gracefully
|
||||
pass
|
||||
|
||||
# Define __all__ to include both PraisonAI core classes and praisonaiagents exports
|
||||
# Define __all__ for lazy loading
|
||||
__all__ = [
|
||||
# Core PraisonAI classes
|
||||
'PraisonAI',
|
||||
'__version__',
|
||||
]
|
||||
|
||||
# Dynamically extend __all__ with praisonaiagents exports
|
||||
try:
|
||||
import praisonaiagents
|
||||
__all__.extend(praisonaiagents.__all__)
|
||||
except (ImportError, AttributeError):
|
||||
# If praisonaiagents is not available or doesn't have __all__, fail gracefully
|
||||
pass
|
||||
# Lazy loading for heavy imports
|
||||
def __getattr__(name):
|
||||
"""Lazy load heavy modules to improve import time."""
|
||||
if name == 'PraisonAI':
|
||||
from .cli import PraisonAI
|
||||
return PraisonAI
|
||||
|
||||
# Try praisonaiagents exports
|
||||
try:
|
||||
import praisonaiagents
|
||||
if hasattr(praisonaiagents, name):
|
||||
return getattr(praisonaiagents, name)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
+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 += """
|
||||
|
||||
@@ -8,12 +8,19 @@ Structure:
|
||||
- features/: Feature handlers for CLI flags and commands
|
||||
"""
|
||||
|
||||
from .main import PraisonAI
|
||||
|
||||
__all__ = ["PraisonAI"]
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
"""Lazy load PraisonAI to avoid slow imports."""
|
||||
if name == "PraisonAI":
|
||||
from .main import PraisonAI
|
||||
return PraisonAI
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def main():
|
||||
"""CLI entry point function."""
|
||||
from .main import PraisonAI
|
||||
praison_ai = PraisonAI()
|
||||
praison_ai.main()
|
||||
@@ -21,6 +21,12 @@ 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
|
||||
- at_mentions: @ mention autocomplete for files and directories
|
||||
"""
|
||||
|
||||
# Type hints available for IDE support
|
||||
@@ -43,6 +49,20 @@ __all__ = [
|
||||
'FlowDisplayHandler',
|
||||
'WorkflowHandler',
|
||||
'N8nHandler',
|
||||
'ExternalAgentsHandler',
|
||||
# New CLI enhancement features
|
||||
'SlashCommandHandler',
|
||||
'AutonomyModeHandler',
|
||||
'CostTrackerHandler',
|
||||
'RepoMapHandler',
|
||||
'InteractiveTUIHandler',
|
||||
'GitIntegrationHandler',
|
||||
'SandboxExecutorHandler',
|
||||
'MessageQueueHandler',
|
||||
# @ mention autocomplete
|
||||
'AtMentionCompleter',
|
||||
'CombinedCompleter',
|
||||
'FileSearchService',
|
||||
]
|
||||
|
||||
def __getattr__(name):
|
||||
@@ -95,4 +115,42 @@ 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
|
||||
# @ mention autocomplete
|
||||
elif name == 'AtMentionCompleter':
|
||||
from .at_mentions import AtMentionCompleter
|
||||
return AtMentionCompleter
|
||||
elif name == 'CombinedCompleter':
|
||||
from .at_mentions import CombinedCompleter
|
||||
return CombinedCompleter
|
||||
elif name == 'FileSearchService':
|
||||
from .at_mentions import FileSearchService
|
||||
return FileSearchService
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,416 @@
|
||||
"""
|
||||
@ Mention Autocomplete for PraisonAI CLI.
|
||||
|
||||
Provides file/directory autocomplete when user types @.
|
||||
Uses only stdlib imports for zero performance impact.
|
||||
|
||||
Features:
|
||||
- Detect @ trigger in input
|
||||
- Fuzzy search files/directories
|
||||
- Cache results for performance
|
||||
- Integrate with prompt_toolkit
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import fnmatch
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Dict, Tuple
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Data Classes
|
||||
# ============================================================================
|
||||
|
||||
@dataclass
|
||||
class AtMentionContext:
|
||||
"""Context for an @ mention in user input."""
|
||||
is_active: bool
|
||||
query: str
|
||||
start_pos: int
|
||||
mention_type: str = "file" # file, directory, web, etc.
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileSuggestion:
|
||||
"""A file or directory suggestion."""
|
||||
path: str
|
||||
file_type: str # "file" or "directory"
|
||||
score: int = 100
|
||||
display: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.display:
|
||||
self.display = self.path
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# @ Mention Detection
|
||||
# ============================================================================
|
||||
|
||||
def detect_at_mention(text: str, cursor_pos: int) -> Optional[AtMentionContext]:
|
||||
"""
|
||||
Detect if cursor is within an @ mention context.
|
||||
|
||||
Scans backward from cursor to find @ symbol.
|
||||
Returns None if not in @ context.
|
||||
|
||||
Args:
|
||||
text: Full input text
|
||||
cursor_pos: Current cursor position
|
||||
|
||||
Returns:
|
||||
AtMentionContext if in @ mention, None otherwise
|
||||
"""
|
||||
if cursor_pos <= 0 or cursor_pos > len(text):
|
||||
return None
|
||||
|
||||
# Scan backward from cursor to find @
|
||||
at_pos = -1
|
||||
for i in range(cursor_pos - 1, -1, -1):
|
||||
char = text[i]
|
||||
|
||||
# Found @
|
||||
if char == '@':
|
||||
at_pos = i
|
||||
break
|
||||
|
||||
# Space/newline breaks @ context
|
||||
if char in ' \t\n\r':
|
||||
return None
|
||||
|
||||
if at_pos == -1:
|
||||
return None
|
||||
|
||||
# Extract query between @ and cursor
|
||||
query = text[at_pos + 1:cursor_pos]
|
||||
|
||||
return AtMentionContext(
|
||||
is_active=True,
|
||||
query=query,
|
||||
start_pos=at_pos
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# File Search Service
|
||||
# ============================================================================
|
||||
|
||||
class FileSearchService:
|
||||
"""
|
||||
Service for searching files in a directory.
|
||||
|
||||
Features:
|
||||
- Walk directory tree
|
||||
- Fuzzy match file paths
|
||||
- Cache results with TTL
|
||||
- Respect .gitignore (basic)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root_dir: str,
|
||||
cache_ttl: int = 30,
|
||||
max_depth: int = 5,
|
||||
ignore_patterns: Optional[List[str]] = None
|
||||
):
|
||||
self._root = os.path.abspath(root_dir)
|
||||
self._cache_ttl = cache_ttl
|
||||
self._max_depth = max_depth
|
||||
self._ignore_patterns = ignore_patterns or [
|
||||
'.git', '__pycache__', 'node_modules', '.venv', 'venv',
|
||||
'*.pyc', '*.pyo', '.DS_Store', '*.egg-info'
|
||||
]
|
||||
|
||||
# Cache: query -> (results, timestamp)
|
||||
self._cache: Dict[str, Tuple[List[FileSuggestion], float]] = {}
|
||||
|
||||
# File list cache
|
||||
self._file_list: Optional[List[str]] = None
|
||||
self._file_list_time: float = 0
|
||||
|
||||
def _should_ignore(self, name: str) -> bool:
|
||||
"""Check if file/dir should be ignored."""
|
||||
for pattern in self._ignore_patterns:
|
||||
if fnmatch.fnmatch(name, pattern):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _get_file_list(self) -> List[str]:
|
||||
"""Get list of all files (cached)."""
|
||||
now = time.time()
|
||||
|
||||
# Return cached if fresh
|
||||
if self._file_list is not None and (now - self._file_list_time) < self._cache_ttl:
|
||||
return self._file_list
|
||||
|
||||
files = []
|
||||
|
||||
try:
|
||||
for root, dirs, filenames in os.walk(self._root):
|
||||
# Calculate depth
|
||||
rel_root = os.path.relpath(root, self._root)
|
||||
depth = 0 if rel_root == '.' else rel_root.count(os.sep) + 1
|
||||
|
||||
if depth > self._max_depth:
|
||||
dirs.clear() # Don't descend further
|
||||
continue
|
||||
|
||||
# Filter ignored directories
|
||||
dirs[:] = [d for d in dirs if not self._should_ignore(d)]
|
||||
|
||||
# Add directories
|
||||
for d in dirs:
|
||||
rel_path = os.path.relpath(os.path.join(root, d), self._root)
|
||||
files.append(rel_path + os.sep)
|
||||
|
||||
# Add files
|
||||
for f in filenames:
|
||||
if not self._should_ignore(f):
|
||||
rel_path = os.path.relpath(os.path.join(root, f), self._root)
|
||||
files.append(rel_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
self._file_list = files
|
||||
self._file_list_time = now
|
||||
return files
|
||||
|
||||
def _fuzzy_match(self, text: str, query: str) -> int:
|
||||
"""
|
||||
Simple fuzzy match scoring.
|
||||
|
||||
Returns score 0-100, higher is better.
|
||||
"""
|
||||
if not query:
|
||||
return 50 # Default score for empty query
|
||||
|
||||
text_lower = text.lower()
|
||||
query_lower = query.lower()
|
||||
|
||||
# Exact match
|
||||
if text_lower == query_lower:
|
||||
return 100
|
||||
|
||||
# Starts with
|
||||
if text_lower.startswith(query_lower):
|
||||
return 90
|
||||
|
||||
# Contains
|
||||
if query_lower in text_lower:
|
||||
return 80
|
||||
|
||||
# Fuzzy: all query chars appear in order
|
||||
text_idx = 0
|
||||
query_idx = 0
|
||||
matches = 0
|
||||
|
||||
while text_idx < len(text_lower) and query_idx < len(query_lower):
|
||||
if text_lower[text_idx] == query_lower[query_idx]:
|
||||
matches += 1
|
||||
query_idx += 1
|
||||
text_idx += 1
|
||||
|
||||
if query_idx == len(query_lower):
|
||||
# All query chars found
|
||||
return 50 + int(30 * matches / len(text))
|
||||
|
||||
return 0
|
||||
|
||||
def search(self, query: str, max_results: int = 20) -> List[FileSuggestion]:
|
||||
"""
|
||||
Search for files matching query.
|
||||
|
||||
Args:
|
||||
query: Search query (fuzzy matched)
|
||||
max_results: Maximum results to return
|
||||
|
||||
Returns:
|
||||
List of FileSuggestion sorted by score
|
||||
"""
|
||||
# Check cache
|
||||
cache_key = f"{query}:{max_results}"
|
||||
now = time.time()
|
||||
|
||||
if cache_key in self._cache:
|
||||
results, cached_time = self._cache[cache_key]
|
||||
if (now - cached_time) < self._cache_ttl:
|
||||
return results
|
||||
|
||||
# Get all files
|
||||
all_files = self._get_file_list()
|
||||
|
||||
# Score and filter
|
||||
scored: List[Tuple[int, str]] = []
|
||||
for path in all_files:
|
||||
score = self._fuzzy_match(path, query)
|
||||
if score > 0:
|
||||
scored.append((score, path))
|
||||
|
||||
# Sort by score descending
|
||||
scored.sort(key=lambda x: (-x[0], x[1]))
|
||||
|
||||
# Build results
|
||||
results = []
|
||||
for score, path in scored[:max_results]:
|
||||
file_type = "directory" if path.endswith(os.sep) else "file"
|
||||
results.append(FileSuggestion(
|
||||
path=path,
|
||||
file_type=file_type,
|
||||
score=score
|
||||
))
|
||||
|
||||
# Cache results
|
||||
self._cache[cache_key] = (results, now)
|
||||
|
||||
return results
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""Clear all caches."""
|
||||
self._cache.clear()
|
||||
self._file_list = None
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Completers for prompt_toolkit
|
||||
# ============================================================================
|
||||
|
||||
# Lazy import to avoid performance impact
|
||||
_Completer = None
|
||||
_Completion = None
|
||||
|
||||
def _get_prompt_toolkit():
|
||||
"""Lazy import prompt_toolkit."""
|
||||
global _Completer, _Completion
|
||||
if _Completer is None:
|
||||
from prompt_toolkit.completion import Completer, Completion
|
||||
_Completer = Completer
|
||||
_Completion = Completion
|
||||
return _Completer, _Completion
|
||||
|
||||
|
||||
class AtMentionCompleter:
|
||||
"""
|
||||
Completer for @ mentions.
|
||||
|
||||
Shows file suggestions when user types @.
|
||||
Implements prompt_toolkit Completer interface.
|
||||
"""
|
||||
|
||||
def __init__(self, root_dir: str):
|
||||
self._file_service = FileSearchService(root_dir)
|
||||
|
||||
def get_completions(self, document, complete_event):
|
||||
"""Get completions for @ mentions."""
|
||||
_, Completion = _get_prompt_toolkit()
|
||||
|
||||
text = document.text_before_cursor
|
||||
cursor_pos = len(text)
|
||||
|
||||
# Detect @ context
|
||||
context = detect_at_mention(text, cursor_pos)
|
||||
if context is None or not context.is_active:
|
||||
return
|
||||
|
||||
# Search files
|
||||
results = self._file_service.search(context.query, max_results=15)
|
||||
|
||||
# Yield completions
|
||||
for suggestion in results:
|
||||
# Icon based on type
|
||||
icon = "📁 " if suggestion.file_type == "directory" else "📄 "
|
||||
display = f"{icon}{suggestion.path}"
|
||||
|
||||
# Calculate replacement position
|
||||
# Replace from @ to cursor
|
||||
start_position = -(cursor_pos - context.start_pos)
|
||||
|
||||
yield Completion(
|
||||
text=f"@{suggestion.path}",
|
||||
start_position=start_position,
|
||||
display=display
|
||||
)
|
||||
|
||||
|
||||
def create_combined_completer(commands: List[str], root_dir: str):
|
||||
"""
|
||||
Factory function to create a CombinedCompleter.
|
||||
|
||||
Returns a proper Completer subclass that handles both / and @.
|
||||
"""
|
||||
Completer, Completion = _get_prompt_toolkit()
|
||||
file_service = FileSearchService(root_dir)
|
||||
|
||||
class CombinedCompleter(Completer):
|
||||
"""
|
||||
Combined completer for both / commands and @ mentions.
|
||||
|
||||
Properly inherits from prompt_toolkit Completer.
|
||||
"""
|
||||
|
||||
def get_completions(self, document, complete_event):
|
||||
"""Get completions for / or @."""
|
||||
text = document.text_before_cursor.lstrip()
|
||||
|
||||
# Check for slash command
|
||||
if text.startswith('/'):
|
||||
cmd_text = text[1:].lower()
|
||||
for cmd in commands:
|
||||
if cmd.lower().startswith(cmd_text):
|
||||
yield Completion(
|
||||
f'/{cmd}',
|
||||
start_position=-len(text),
|
||||
display=f'/{cmd}'
|
||||
)
|
||||
return
|
||||
|
||||
# Check for @ mention
|
||||
full_text = document.text_before_cursor
|
||||
cursor_pos = len(full_text)
|
||||
|
||||
context = detect_at_mention(full_text, cursor_pos)
|
||||
if context is None or not context.is_active:
|
||||
return
|
||||
|
||||
# Search files
|
||||
results = file_service.search(context.query, max_results=15)
|
||||
|
||||
# Yield completions
|
||||
for suggestion in results:
|
||||
icon = "📁 " if suggestion.file_type == "directory" else "📄 "
|
||||
display = f"{icon}{suggestion.path}"
|
||||
start_position = -(cursor_pos - context.start_pos)
|
||||
|
||||
yield Completion(
|
||||
text=f"@{suggestion.path}",
|
||||
start_position=start_position,
|
||||
display=display
|
||||
)
|
||||
|
||||
return CombinedCompleter()
|
||||
|
||||
|
||||
# Backward compatible alias
|
||||
class CombinedCompleter:
|
||||
"""
|
||||
Wrapper class for backward compatibility.
|
||||
|
||||
Use create_combined_completer() for proper Completer subclass.
|
||||
"""
|
||||
|
||||
def __new__(cls, commands: List[str], root_dir: str):
|
||||
return create_combined_completer(commands, root_dir)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Exports
|
||||
# ============================================================================
|
||||
|
||||
__all__ = [
|
||||
'AtMentionContext',
|
||||
'FileSuggestion',
|
||||
'detect_at_mention',
|
||||
'FileSearchService',
|
||||
'AtMentionCompleter',
|
||||
'CombinedCompleter',
|
||||
]
|
||||
@@ -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',
|
||||
]
|
||||
@@ -35,7 +35,7 @@ class N8nHandler(FlagHandler):
|
||||
"""
|
||||
|
||||
def __init__(self, verbose: bool = False, n8n_url: str = "http://localhost:5678",
|
||||
use_execute_command: bool = False):
|
||||
use_execute_command: bool = False, api_url: str = "http://127.0.0.1:8005"):
|
||||
"""
|
||||
Initialize the n8n handler.
|
||||
|
||||
@@ -44,10 +44,11 @@ class N8nHandler(FlagHandler):
|
||||
n8n_url: Base URL of the n8n instance
|
||||
use_execute_command: Use Execute Command nodes (runs praisonai directly)
|
||||
instead of HTTP Request nodes
|
||||
api_url: PraisonAI API URL that n8n will call (for tunnel/cloud deployments)
|
||||
"""
|
||||
super().__init__(verbose=verbose)
|
||||
self.n8n_url = n8n_url
|
||||
self.praisonai_api_url = "http://127.0.0.1:8005"
|
||||
self.praisonai_api_url = api_url
|
||||
self.use_execute_command = use_execute_command
|
||||
|
||||
@property
|
||||
|
||||
@@ -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")
|
||||
|
||||
+1154
-41
File diff suppressed because it is too large.
Load diff
Loaded 100 of 156 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user